[Security/Extension] JWT Vendor for extensions (opensearch-project#2567)

* JWT Vendor for extensions
Signed-off-by: Ryan Liang <[email protected]>

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

@@ -0,0 +1,174 @@
+/*
+ * 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;
+
+import java.time.Instant;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.LongSupplier;
+
+import com.google.common.base.Strings;
+import org.apache.cxf.jaxrs.json.basic.JsonMapObjectReaderWriter;
+import org.apache.cxf.rs.security.jose.jwk.JsonWebKey;
+import org.apache.cxf.rs.security.jose.jwk.KeyType;
+import org.apache.cxf.rs.security.jose.jwk.PublicKeyUse;
+import org.apache.cxf.rs.security.jose.jws.JwsUtils;
+import org.apache.cxf.rs.security.jose.jwt.JoseJwtProducer;
+import org.apache.cxf.rs.security.jose.jwt.JwtClaims;
+import org.apache.cxf.rs.security.jose.jwt.JwtToken;
+import org.apache.cxf.rs.security.jose.jwt.JwtUtils;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.transport.TransportAddress;
+import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.security.securityconf.ConfigModel;
+import org.opensearch.security.support.ConfigConstants;
+import org.opensearch.security.user.User;
+import org.opensearch.threadpool.ThreadPool;
+
+public class JwtVendor {
+    private static final Logger logger = LogManager.getLogger(JwtVendor.class);
+
+    private static JsonMapObjectReaderWriter jsonMapReaderWriter = new JsonMapObjectReaderWriter();
+
+    private JsonWebKey signingKey;
+    private JoseJwtProducer jwtProducer;
+    private final LongSupplier timeProvider;
+
+    //TODO: Relocate/Remove them at once we make the descisions about the `roles`
+    private ConfigModel configModel;
+    private ThreadContext threadContext;
+
+    public JwtVendor(Settings settings) {
+        JoseJwtProducer jwtProducer = new JoseJwtProducer();
+        try {
+            this.signingKey = createJwkFromSettings(settings);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+        this.jwtProducer = jwtProducer;
+        timeProvider = System::currentTimeMillis;
+    }
+
+    //For testing the expiration in the future
+    public JwtVendor(Settings settings, final LongSupplier timeProvider) {
+        JoseJwtProducer jwtProducer = new JoseJwtProducer();
+        try {
+            this.signingKey = createJwkFromSettings(settings);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+        this.jwtProducer = jwtProducer;
+        this.timeProvider = timeProvider;
+    }
+
+    /*
+    * The default configuration of this web key should be:
+    *   KeyType: OCTET
+    *   PublicKeyUse: SIGN
+    *   Encryption Algorithm: HS512
+    * */
+    static JsonWebKey createJwkFromSettings(Settings settings) throws Exception {
+        String signingKey = settings.get("signing_key");
+
+        if (!Strings.isNullOrEmpty(signingKey)) {
+
+            JsonWebKey jwk = new JsonWebKey();
+
+            jwk.setKeyType(KeyType.OCTET);
+            jwk.setAlgorithm("HS512");
+            jwk.setPublicKeyUse(PublicKeyUse.SIGN);
+            jwk.setProperty("k", signingKey);
+
+            return jwk;
+        } else {
+            Settings jwkSettings = settings.getAsSettings("jwt").getAsSettings("key");
+
+            if (jwkSettings.isEmpty()) {
+                throw new Exception(
+                        "Settings for key is missing. Please specify at least the option signing_key with a shared secret.");
+            }
+
+            JsonWebKey jwk = new JsonWebKey();
+
+            for (String key : jwkSettings.keySet()) {
+                jwk.setProperty(key, jwkSettings.get(key));
+            }
+
+            return jwk;
+        }
+    }
+
+    //TODO:Getting roles from User
+    public Map<String, String> prepareClaimsForUser(User user, ThreadPool threadPool) {
+        Map<String, String> claims = new HashMap<>();
+        this.threadContext = threadPool.getThreadContext();
+        final TransportAddress caller = threadContext.getTransient(ConfigConstants.OPENDISTRO_SECURITY_REMOTE_ADDRESS);
+        Set<String> mappedRoles = mapRoles(user, caller);
+        claims.put("sub", user.getName());
+        claims.put("roles", String.join(",", mappedRoles));
+        return claims;
+    }
+
+    public Set<String> mapRoles(final User user, final TransportAddress caller) {
+        return this.configModel.mapSecurityRoles(user, caller);
+    }
+
+    public String createJwt(String issuer, String subject, String audience, Integer expirySeconds) throws Exception {
+        long timeMillis = timeProvider.getAsLong();
+        Instant now = Instant.ofEpochMilli(timeProvider.getAsLong());
+
+        jwtProducer.setSignatureProvider(JwsUtils.getSignatureProvider(signingKey));
+        JwtClaims jwtClaims = new JwtClaims();
+        JwtToken jwt = new JwtToken(jwtClaims);
+
+        jwtClaims.setIssuer(issuer);
+
+        jwtClaims.setIssuedAt(timeMillis);
+
+        jwtClaims.setSubject(subject);
+
+        jwtClaims.setAudience(audience);
+
+        jwtClaims.setNotBefore(timeMillis);
+
+        if (expirySeconds == null) {
+            long expiryTime = timeProvider.getAsLong() + (300 * 1000);
+            jwtClaims.setExpiryTime(expiryTime);
+        } else if (expirySeconds > 0) {
+            long expiryTime = timeProvider.getAsLong() + (expirySeconds * 1000);
+            jwtClaims.setExpiryTime(expiryTime);
+        } else {
+            throw new Exception("The expiration time should be a positive integer");
+        }
+
+        // TODO: Should call preparelaims() if we need roles in claim;
+
+        String encodedJwt = jwtProducer.processJwt(jwt);
+
+        if (logger.isDebugEnabled()) {
+            logger.debug(
+                    "Created JWT: "
+                            + encodedJwt
+                            + "\n"
+                            + jsonMapReaderWriter.toJson(jwt.getJwsHeaders())
+                            + "\n"
+                            + JwtUtils.claimsToJson(jwt.getClaims())
+            );
+        }
+
+        return encodedJwt;
+    }
+}

src/main/java/org/opensearch/security/filter/SecurityRestFilter.java

@@ -124,6 +124,7 @@ public void handleRequest(RestRequest request, RestChannel channel, NodeClient c
                 if (!checkAndAuthenticateRequest(request, channel, client)) {
                     User user = threadContext.getTransient(ConfigConstants.OPENDISTRO_SECURITY_USER);
                     if (userIsSuperAdmin(user, adminDNs) || (whitelistingSettings.checkRequestIsAllowed(request, channel, client) && allowlistingSettings.checkRequestIsAllowed(request, channel, client))) {
+                        //TODO: If the request is going to the extension, issue a JWT for authenticated user.
                         original.handleRequest(request, channel, client);
                     }
                 }

src/test/java/org/opensearch/security/authtoken/jwt/JwtVendorTest.java

@@ -0,0 +1,80 @@
+/*
+ * 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;
+
+import java.util.function.LongSupplier;
+
+import org.apache.cxf.rs.security.jose.jwk.JsonWebKey;
+import org.apache.cxf.rs.security.jose.jws.JwsJwtCompactConsumer;
+import org.apache.cxf.rs.security.jose.jwt.JwtToken;
+import org.junit.Assert;
+import org.junit.Test;
+
+import org.opensearch.common.settings.Settings;
+
+public class JwtVendorTest {
+
+    @Test
+    public void testCreateJwkFromSettings() throws Exception {
+        Settings settings = Settings.builder()
+                .put("signing_key", "abc123").build();
+
+        JsonWebKey jwk = JwtVendor.createJwkFromSettings(settings);
+        Assert.assertEquals("HS512", jwk.getAlgorithm());
+        Assert.assertEquals("sig", jwk.getPublicKeyUse().toString());
+        Assert.assertEquals("abc123", jwk.getProperty("k"));
+    }
+
+    @Test (expected = Exception.class)
+    public void testCreateJwkFromSettingsWithoutSigningKey() throws Exception{
+        Settings settings = Settings.builder()
+                .put("jwt", "").build();
+        JwtVendor.createJwkFromSettings(settings);
+    }
+
+    @Test
+    public void testCreateJwt() throws Exception {
+        String issuer = "cluster_0";
+        String subject = "admin";
+        String audience = "extension_0";
+        Integer expirySeconds = 300;
+        LongSupplier currentTime = () -> (int)100;
+        Settings settings =  Settings.builder().put("signing_key", "abc123").build();
+        Long expectedExp = currentTime.getAsLong() + (expirySeconds * 1000);
+
+        JwtVendor jwtVendor = new JwtVendor(settings, currentTime);
+        String encodedJwt = jwtVendor.createJwt(issuer, subject, audience, expirySeconds);
+
+        JwsJwtCompactConsumer jwtConsumer = new JwsJwtCompactConsumer(encodedJwt);
+        JwtToken jwt = jwtConsumer.getJwtToken();
+
+        Assert.assertEquals("cluster_0", jwt.getClaim("iss"));
+        Assert.assertEquals("admin", jwt.getClaim("sub"));
+        Assert.assertEquals("extension_0", jwt.getClaim("aud"));
+        Assert.assertNotNull(jwt.getClaim("iat"));
+        Assert.assertNotNull(jwt.getClaim("exp"));
+        Assert.assertEquals(expectedExp, jwt.getClaim("exp"));
+    }
+
+    @Test (expected = Exception.class)
+    public void testCreateJwtWithBadExpiry() throws Exception {
+        String issuer = "cluster_0";
+        String subject = "admin";
+        String audience = "extension_0";
+        Integer expirySeconds = -300;
+
+        Settings settings =  Settings.builder().put("signing_key", "abc123").build();
+        JwtVendor jwtVendor = new JwtVendor(settings);
+
+        jwtVendor.createJwt(issuer, subject, audience, expirySeconds);
+    }
+}