Refactor JWT Vendor to take a claims builder and rename oboEnabled to enabled

CHANGELOG.md

@@ -4,8 +4,6 @@ All notable changes to this project are documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to the [Semantic Versioning](https://semver.org/spec/v2.0.0.html). See the [CONTRIBUTING guide](./CONTRIBUTING.md#Changelog) for instructions on how to add changelog entries.
 
 ## [Unreleased 3.x]
-### Added
-- Provide SecureHttpTransportParameters to complement SecureTransportParameters counterpart ([#5432](https://github.com/opensearch-project/security/pull/5432))
 
 ### Features
 
@@ -23,7 +21,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ### Refactoring
 
-
+* Refactor JWT Vender to take a claims builder and rename oboEnabled to enabled ([#5436](https://github.com/opensearch-project/security/pull/5436))
 
 ### Maintenance
 

src/integrationTest/java/org/opensearch/security/http/OnBehalfOfJwtAuthenticationTest.java

@@ -62,7 +62,6 @@ public class OnBehalfOfJwtAuthenticationTest {
     static final TestSecurityConfig.User ADMIN_USER = new TestSecurityConfig.User("admin").roles(ALL_ACCESS);
 
     private static final String CREATE_OBO_TOKEN_PATH = "_plugins/_security/api/generateonbehalfoftoken";
-    private static Boolean oboEnabled = true;
     private static final String signingKey = Base64.getEncoder()
         .encodeToString(
             "jwt signing key for an on behalf of token authentication backend for testing of OBO authentication".getBytes(
@@ -118,7 +117,7 @@ public class OnBehalfOfJwtAuthenticationTest {
         .roles(HOST_MAPPING_ROLE, ROLE_WITH_OBO_PERM);
 
     private static OnBehalfOfConfig defaultOnBehalfOfConfig() {
-        return new OnBehalfOfConfig().oboEnabled(oboEnabled).signingKey(signingKey).encryptionKey(encryptionKey);
+        return new OnBehalfOfConfig().enabled(true).signingKey(signingKey).encryptionKey(encryptionKey);
     }
 
     @ClassRule
@@ -155,7 +154,7 @@ public void shouldAuthenticateWithOBOTokenEndPoint() {
     }
 
     @Test
-    public void shouldNotAuthenticateWithATemperedOBOToken() {
+    public void shouldNotAuthenticateWithATamperedOBOToken() {
         String oboToken = generateOboToken(ADMIN_USER_NAME, DEFAULT_PASSWORD);
         oboToken = oboToken.substring(0, oboToken.length() - 1); // tampering the token
         Header adminOboAuthHeader = new BasicHeader("Authorization", "Bearer " + oboToken);
@@ -242,11 +241,11 @@ public void shouldNotAllowTokenWhenOboIsDisabled() {
         authenticateWithOboToken(oboHeader, OBO_USER_NAME_WITH_PERM, HttpStatus.SC_OK);
 
         // Disable OBO via config and see that the authenticator doesn't authorize
-        patchOnBehalfOfConfig(defaultOnBehalfOfConfig().oboEnabled(false));
+        patchOnBehalfOfConfig(defaultOnBehalfOfConfig().enabled(false));
         authenticateWithOboToken(oboHeader, OBO_USER_NAME_WITH_PERM, HttpStatus.SC_UNAUTHORIZED);
 
         // Reenable OBO via config and see that the authenticator is working again
-        patchOnBehalfOfConfig(defaultOnBehalfOfConfig().oboEnabled(true));
+        patchOnBehalfOfConfig(defaultOnBehalfOfConfig().enabled(true));
         authenticateWithOboToken(oboHeader, OBO_USER_NAME_WITH_PERM, HttpStatus.SC_OK);
     }
 

src/integrationTest/java/org/opensearch/test/framework/OnBehalfOfConfig.java

@@ -18,12 +18,12 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 
 public class OnBehalfOfConfig implements ToXContentObject {
-    private Boolean oboEnabled;
+    private Boolean enabled;
     private String signing_key;
     private String encryption_key;
 
-    public OnBehalfOfConfig oboEnabled(Boolean oboEnabled) {
-        this.oboEnabled = oboEnabled;
+    public OnBehalfOfConfig enabled(Boolean enabled) {
+        this.enabled = enabled;
         return this;
     }
 
@@ -40,7 +40,7 @@ public OnBehalfOfConfig encryptionKey(String encryption_key) {
     @Override
     public XContentBuilder toXContent(XContentBuilder xContentBuilder, ToXContent.Params params) throws IOException {
         xContentBuilder.startObject();
-        xContentBuilder.field("enabled", oboEnabled);
+        xContentBuilder.field("enabled", enabled);
         xContentBuilder.field("signing_key", signing_key);
         if (StringUtils.isNoneBlank(encryption_key)) {
             xContentBuilder.field("encryption_key", encryption_key);

src/main/java/org/opensearch/security/authtoken/jwt/JwtVendor.java

@@ -11,19 +11,19 @@
 
 package org.opensearch.security.authtoken.jwt;
 
+import java.security.AccessController;
+import java.security.PrivilegedAction;
 import java.text.ParseException;
 import java.util.Base64;
 import java.util.Date;
-import java.util.List;
-import java.util.Optional;
-import java.util.function.LongSupplier;
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
 import org.opensearch.OpenSearchException;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.security.authtoken.jwt.claims.JwtClaimsBuilder;
 
 import com.nimbusds.jose.JOSEException;
 import com.nimbusds.jose.JWSAlgorithm;
@@ -34,7 +34,6 @@
 import com.nimbusds.jose.jwk.JWK;
 import com.nimbusds.jose.jwk.KeyUse;
 import com.nimbusds.jose.jwk.OctetSequenceKey;
-import com.nimbusds.jwt.JWTClaimsSet;
 import com.nimbusds.jwt.SignedJWT;
 
 import static org.opensearch.security.util.AuthTokenUtils.isKeyNull;
@@ -44,21 +43,11 @@ public class JwtVendor {
 
     private final JWK signingKey;
     private final JWSSigner signer;
-    private final LongSupplier timeProvider;
-    private final EncryptionDecryptionUtil encryptionDecryptionUtil;
-    private static final Integer MAX_EXPIRY_SECONDS = 600;
 
-    public JwtVendor(final Settings settings, final Optional<LongSupplier> timeProvider) {
+    public JwtVendor(Settings settings) {
         final Tuple<JWK, JWSSigner> tuple = createJwkFromSettings(settings);
         signingKey = tuple.v1();
         signer = tuple.v2();
-
-        if (isKeyNull(settings, "encryption_key")) {
-            throw new IllegalArgumentException("encryption_key cannot be null");
-        } else {
-            this.encryptionDecryptionUtil = new EncryptionDecryptionUtil(settings.get("encryption_key"));
-        }
-        this.timeProvider = timeProvider.orElse(System::currentTimeMillis);
     }
 
     /*
@@ -96,46 +85,14 @@ static Tuple<JWK, JWSSigner> createJwkFromSettings(final Settings settings) {
         }
     }
 
-    public ExpiringBearerAuthToken createJwt(
-        final String issuer,
-        final String subject,
-        final String audience,
-        final long requestedExpirySeconds,
-        final List<String> roles,
-        final List<String> backendRoles,
-        final boolean includeBackendRoles
-    ) throws JOSEException, ParseException {
-        final long currentTimeMs = timeProvider.getAsLong();
-        final Date now = new Date(currentTimeMs);
-
-        final JWTClaimsSet.Builder claimsBuilder = new JWTClaimsSet.Builder();
-        claimsBuilder.issuer(issuer);
-        claimsBuilder.issueTime(now);
-        claimsBuilder.subject(subject);
-        claimsBuilder.audience(audience);
-        claimsBuilder.notBeforeTime(now);
-
-        final long expirySeconds = Math.min(requestedExpirySeconds, MAX_EXPIRY_SECONDS);
-        if (expirySeconds <= 0) {
-            throw new IllegalArgumentException("The expiration time should be a positive integer");
-        }
-        final Date expiryTime = new Date(currentTimeMs + expirySeconds * 1000);
-        claimsBuilder.expirationTime(expiryTime);
-
-        if (roles != null) {
-            final String listOfRoles = String.join(",", roles);
-            claimsBuilder.claim("er", encryptionDecryptionUtil.encrypt(listOfRoles));
-        } else {
-            throw new IllegalArgumentException("Roles cannot be null");
-        }
-
-        if (includeBackendRoles && backendRoles != null) {
-            final String listOfBackendRoles = String.join(",", backendRoles);
-            claimsBuilder.claim("br", listOfBackendRoles);
-        }
+    @SuppressWarnings("removal")
+    public ExpiringBearerAuthToken createJwt(JwtClaimsBuilder claimsBuilder, String subject, Date expiryTime, Long expirySeconds)
+        throws JOSEException, ParseException {
 
         final JWSHeader header = new JWSHeader.Builder(JWSAlgorithm.parse(signingKey.getAlgorithm().getName())).build();
-        final SignedJWT signedJwt = new SignedJWT(header, claimsBuilder.build());
+        final SignedJWT signedJwt = AccessController.doPrivileged(
+            (PrivilegedAction<SignedJWT>) () -> new SignedJWT(header, claimsBuilder.build())
+        );
 
         // Sign the JWT so it can be serialized
         signedJwt.sign(signer);

src/main/java/org/opensearch/security/authtoken/jwt/claims/JwtClaimsBuilder.java

@@ -0,0 +1,75 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.security.authtoken.jwt.claims;
+
+import java.util.Date;
+
+import com.nimbusds.jwt.JWTClaimsSet;
+
+/**
+ * Builder for creating JWT claims.
+ *
+ * This class builds a claims set with standard claims such as issued_at, not before time, subject, issuer, audience,
+ * expiration time, and custom claims.
+ */
+public class JwtClaimsBuilder {
+    private final JWTClaimsSet.Builder builder;
+
+    public JwtClaimsBuilder() {
+        this.builder = new JWTClaimsSet.Builder();
+    }
+
+    public JwtClaimsBuilder issueTime(Date issueTime) {
+        builder.issueTime(issueTime);
+        return this;
+    }
+
+    public JwtClaimsBuilder notBeforeTime(Date notBeforeTime) {
+        builder.notBeforeTime(notBeforeTime);
+        return this;
+    }
+
+    public JwtClaimsBuilder subject(String subject) {
+        builder.subject(subject);
+        return this;
+    }
+
+    public JwtClaimsBuilder issuer(String issuer) {
+        builder.issuer(issuer);
+        return this;
+    }
+
+    public JwtClaimsBuilder audience(String audience) {
+        builder.audience(audience);
+        return this;
+    }
+
+    public JwtClaimsBuilder issuedAt(Date issuedAt) {
+        builder.issueTime(issuedAt);
+        return this;
+    }
+
+    public JwtClaimsBuilder expirationTime(Date expirationTime) {
+        builder.expirationTime(expirationTime);
+        return this;
+    }
+
+    public JwtClaimsBuilder addCustomClaim(String claimName, String value) {
+        builder.claim(claimName, value);
+        return this;
+    }
+
+    public JWTClaimsSet build() {
+        return builder.build();
+    }
+
+}

src/main/java/org/opensearch/security/authtoken/jwt/claims/OBOJwtClaimsBuilder.java

@@ -0,0 +1,39 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+package org.opensearch.security.authtoken.jwt.claims;
+
+import java.util.List;
+
+import org.opensearch.security.authtoken.jwt.EncryptionDecryptionUtil;
+
+public class OBOJwtClaimsBuilder extends JwtClaimsBuilder {
+    private final EncryptionDecryptionUtil encryptionDecryptionUtil;
+
+    public OBOJwtClaimsBuilder(String encryptionKey) {
+        super();
+        this.encryptionDecryptionUtil = new EncryptionDecryptionUtil(encryptionKey);
+    }
+
+    public OBOJwtClaimsBuilder addRoles(List<String> roles) {
+        final String listOfRoles = String.join(",", roles);
+        this.addCustomClaim("er", encryptionDecryptionUtil.encrypt(listOfRoles));
+        return this;
+    }
+
+    public OBOJwtClaimsBuilder addBackendRoles(Boolean includeBackendRoles, List<String> backendRoles) {
+        if (includeBackendRoles && backendRoles != null) {
+            final String listOfBackendRoles = String.join(",", backendRoles);
+            this.addCustomClaim("br", listOfBackendRoles);
+        }
+        return this;
+    }
+}

src/main/java/org/opensearch/security/http/OnBehalfOfAuthenticator.java

@@ -63,15 +63,14 @@ public class OnBehalfOfAuthenticator implements HTTPAuthenticator {
 
     private final JwtParser jwtParser;
     private final String encryptionKey;
-    private final Boolean oboEnabled;
+    private final Boolean enabled;
     private final String clusterName;
 
     private final EncryptionDecryptionUtil encryptionUtil;
 
     @SuppressWarnings("removal")
     public OnBehalfOfAuthenticator(Settings settings, String clusterName) {
-        String oboEnabledSetting = settings.get("enabled", "true");
-        oboEnabled = Boolean.parseBoolean(oboEnabledSetting);
+        enabled = settings.getAsBoolean("enabled", Boolean.TRUE);
         encryptionKey = settings.get("encryption_key");
 
         final SecurityManager sm = System.getSecurityManager();
@@ -165,7 +164,7 @@ public AuthCredentials run() {
     }
 
     private AuthCredentials extractCredentials0(final SecurityRequest request) {
-        if (!oboEnabled) {
+        if (!enabled) {
             log.debug("On-behalf-of authentication is disabled");
             return null;
         }