Add the test of JwtVendor

Signed-off-by: Ryan Liang <[email protected]>
DarshitChanpura · Sep 29, 2023 · 29bb23b · 29bb23b
1 parent 79e84a7
commit 29bb23b
Show file tree

Hide file tree

Showing 4 changed files with 189 additions and 8 deletions.
diff --git a/src/main/java/org/opensearch/security/action/onbehalf/CreateOnBehalfOfTokenAction.java b/src/main/java/org/opensearch/security/action/onbehalf/CreateOnBehalfOfTokenAction.java
@@ -12,11 +12,11 @@
 package org.opensearch.security.action.onbehalf;
 
 import java.io.IOException;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import com.google.common.collect.ImmutableList;
 import org.greenrobot.eventbus.Subscribe;
@@ -152,8 +152,8 @@ public void accept(RestChannel channel) throws Exception {
                         user.getName(),
                         service,
                         tokenDuration,
-                        new HashSet<>(mappedRoles),
-                        new HashSet<>(user.getRoles()),
+                        mappedRoles.stream().collect(Collectors.toList()),
+                        user.getRoles().stream().collect(Collectors.toList()),
                         isRoleEncrypted
                     );
                     builder.field("authenticationToken", token);

diff --git a/src/main/java/org/opensearch/security/authtoken/jwt/JwtVendor.java b/src/main/java/org/opensearch/security/authtoken/jwt/JwtVendor.java
@@ -12,8 +12,8 @@
 package org.opensearch.security.authtoken.jwt;
 
 import java.time.Instant;
+import java.util.List;
 import java.util.Optional;
-import java.util.Set;
 import java.util.function.LongSupplier;
 
 import com.google.common.base.Strings;
@@ -109,8 +109,8 @@ public String createJwt(
         String subject,
         String audience,
         Integer expirySeconds,
-        Set<String> roles,
-        Set<String> backendRoles,
+        List<String> roles,
+        List<String> backendRoles,
         boolean isRoleEncrypted
     ) throws Exception {
         final long nowAsMillis = timeProvider.getAsLong();

diff --git a/src/main/java/org/opensearch/security/identity/SecurityTokenManager.java b/src/main/java/org/opensearch/security/identity/SecurityTokenManager.java
@@ -33,6 +33,7 @@
 import java.util.Base64;
 import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 public class SecurityTokenManager implements TokenManager {
 
@@ -83,8 +84,8 @@ public AuthToken issueOnBehalfOfToken(Subject subject, OnBehalfOfClaims claims)
                 user.getName(),
                 claims.getAudience(),
                 300,
-                mappedRoles,
-                user.getRoles(),
+                mappedRoles.stream().collect(Collectors.toList()),
+                user.getRoles().stream().collect(Collectors.toList()),
                 false
             );
         } catch (Exception e) {

diff --git a/src/test/java/org/opensearch/security/authtoken/jwt/JwtVendorTest.java b/src/test/java/org/opensearch/security/authtoken/jwt/JwtVendorTest.java
@@ -0,0 +1,180 @@
+/*
+ * 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 org.apache.commons.lang3.RandomStringUtils;
+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;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.function.LongSupplier;
+
+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
+    public void testCreateJwkFromSettingsWithoutSigningKey() {
+        Settings settings = Settings.builder().put("jwt", "").build();
+        Throwable exception = Assert.assertThrows(RuntimeException.class, () -> {
+            try {
+                JwtVendor.createJwkFromSettings(settings);
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        });
+        Assert.assertEquals(
+            "java.lang.Exception: Settings for signing key is missing. Please specify at least the option signing_key with a shared secret.",
+            exception.getMessage()
+        );
+    }
+
+    @Test
+    public void testCreateJwtWithRoles() throws Exception {
+        String issuer = "cluster_0";
+        String subject = "admin";
+        String audience = "audience_0";
+        List<String> roles = List.of("IT", "HR");
+        List<String> backendRoles = List.of("Sales", "Support");
+        String expectedRoles = "IT,HR";
+        int expirySeconds = 300;
+        LongSupplier currentTime = () -> (long) 100;
+        String claimsEncryptionKey = RandomStringUtils.randomAlphanumeric(16);
+        Settings settings = Settings.builder().put("signing_key", "abc123").put("encryption_key", claimsEncryptionKey).build();
+        Long expectedExp = currentTime.getAsLong() + expirySeconds;
+
+        JwtVendor jwtVendor = new JwtVendor(settings, Optional.of(currentTime));
+        String encodedJwt = jwtVendor.createJwt(issuer, subject, audience, expirySeconds, roles, backendRoles, true);
+
+        JwsJwtCompactConsumer jwtConsumer = new JwsJwtCompactConsumer(encodedJwt);
+        JwtToken jwt = jwtConsumer.getJwtToken();
+
+        Assert.assertEquals("cluster_0", jwt.getClaim("iss"));
+        Assert.assertEquals("admin", jwt.getClaim("sub"));
+        Assert.assertEquals("audience_0", jwt.getClaim("aud"));
+        Assert.assertNotNull(jwt.getClaim("iat"));
+        Assert.assertNotNull(jwt.getClaim("exp"));
+        Assert.assertEquals(expectedExp, jwt.getClaim("exp"));
+        EncryptionDecryptionUtil encryptionUtil = new EncryptionDecryptionUtil(claimsEncryptionKey);
+        Assert.assertEquals(expectedRoles, encryptionUtil.decrypt(jwt.getClaim("er").toString()));
+        Assert.assertNull(jwt.getClaim("br"));
+    }
+
+    @Test
+    public void testCreateJwtWithRoleSecurityMode() throws Exception {
+        String issuer = "cluster_0";
+        String subject = "admin";
+        String audience = "audience_0";
+        List<String> roles = List.of("IT", "HR");
+        List<String> backendRoles = List.of("Sales", "Support");
+        String expectedRoles = "IT,HR";
+        String expectedBackendRoles = "Sales,Support";
+
+        int expirySeconds = 300;
+        LongSupplier currentTime = () -> (long) 100;
+        String claimsEncryptionKey = RandomStringUtils.randomAlphanumeric(16);
+        Settings settings = Settings.builder().put("signing_key", "abc123").put("encryption_key", claimsEncryptionKey).build();
+        Long expectedExp = currentTime.getAsLong() + expirySeconds;
+
+        JwtVendor jwtVendor = new JwtVendor(settings, Optional.of(currentTime));
+        String encodedJwt = jwtVendor.createJwt(issuer, subject, audience, expirySeconds, roles, backendRoles, false);
+
+        JwsJwtCompactConsumer jwtConsumer = new JwsJwtCompactConsumer(encodedJwt);
+        JwtToken jwt = jwtConsumer.getJwtToken();
+
+        Assert.assertEquals("cluster_0", jwt.getClaim("iss"));
+        Assert.assertEquals("admin", jwt.getClaim("sub"));
+        Assert.assertEquals("audience_0", jwt.getClaim("aud"));
+        Assert.assertNotNull(jwt.getClaim("iat"));
+        Assert.assertNotNull(jwt.getClaim("exp"));
+        Assert.assertEquals(expectedExp, jwt.getClaim("exp"));
+        EncryptionDecryptionUtil encryptionUtil = new EncryptionDecryptionUtil(claimsEncryptionKey);
+        Assert.assertEquals(expectedRoles, encryptionUtil.decrypt(jwt.getClaim("er").toString()));
+        Assert.assertNotNull(jwt.getClaim("br"));
+        Assert.assertEquals(expectedBackendRoles, jwt.getClaim("br"));
+    }
+
+    @Test
+    public void testCreateJwtWithBadExpiry() {
+        String issuer = "cluster_0";
+        String subject = "admin";
+        String audience = "audience_0";
+        List<String> roles = List.of("admin");
+        Integer expirySeconds = -300;
+        String claimsEncryptionKey = RandomStringUtils.randomAlphanumeric(16);
+        Settings settings = Settings.builder().put("signing_key", "abc123").put("encryption_key", claimsEncryptionKey).build();
+        JwtVendor jwtVendor = new JwtVendor(settings, Optional.empty());
+
+        Throwable exception = Assert.assertThrows(RuntimeException.class, () -> {
+            try {
+                jwtVendor.createJwt(issuer, subject, audience, expirySeconds, roles, List.of(), true);
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        });
+        Assert.assertEquals("java.lang.Exception: The expiration time should be a positive integer", exception.getMessage());
+    }
+
+    @Test
+    public void testCreateJwtWithBadEncryptionKey() {
+        String issuer = "cluster_0";
+        String subject = "admin";
+        String audience = "audience_0";
+        List<String> roles = List.of("admin");
+        Integer expirySeconds = 300;
+
+        Settings settings = Settings.builder().put("signing_key", "abc123").build();
+
+        Throwable exception = Assert.assertThrows(RuntimeException.class, () -> {
+            try {
+                new JwtVendor(settings, Optional.empty()).createJwt(issuer, subject, audience, expirySeconds, roles, List.of(), true);
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        });
+        Assert.assertEquals("java.lang.IllegalArgumentException: encryption_key cannot be null", exception.getMessage());
+    }
+
+    @Test
+    public void testCreateJwtWithBadRoles() {
+        String issuer = "cluster_0";
+        String subject = "admin";
+        String audience = "audience_0";
+        List<String> roles = null;
+        Integer expirySeconds = 300;
+        String claimsEncryptionKey = RandomStringUtils.randomAlphanumeric(16);
+        Settings settings = Settings.builder().put("signing_key", "abc123").put("encryption_key", claimsEncryptionKey).build();
+        JwtVendor jwtVendor = new JwtVendor(settings, Optional.empty());
+
+        Throwable exception = Assert.assertThrows(RuntimeException.class, () -> {
+            try {
+                jwtVendor.createJwt(issuer, subject, audience, expirySeconds, roles, List.of(), true);
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        });
+        Assert.assertEquals("java.lang.Exception: Roles cannot be null", exception.getMessage());
+    }
+}