Add dependencies injection

Signed-off-by: Ryan Liang <jiallian@amazon.com>
opensearch-project · Mar 24, 2023 · 7e73c87 · 7e73c87
1 parent 114cbde
commit 7e73c87
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 14 deletions.
diff --git a/src/main/java/org/opensearch/security/authtoken/jwt/JwtVendor.java b/src/main/java/org/opensearch/security/authtoken/jwt/JwtVendor.java
@@ -15,6 +15,8 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.LongSupplier;
+import java.time.Instant;
 
 import com.google.common.base.Strings;
 import org.apache.cxf.jaxrs.json.basic.JsonMapObjectReaderWriter;
@@ -44,6 +46,7 @@ public class JwtVendor {
 
     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;
@@ -57,6 +60,19 @@ public JwtVendor(Settings settings) {
             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;
     }
 
     static JsonWebKey createJwkFromSettings(Settings settings) throws Exception {
@@ -105,26 +121,29 @@ 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 expiryMin) throws Exception {
+    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(Instant.now().toEpochMilli());
+        jwtClaims.setIssuedAt(timeMillis);
 
         jwtClaims.setSubject(subject);
 
         jwtClaims.setAudience(audience);
 
-        jwtClaims.setNotBefore(System.currentTimeMillis() / 1000);
+        jwtClaims.setNotBefore(timeMillis);
 
-        if (expiryMin == null) {
-            long expiryTime = System.currentTimeMillis() / 1000 + (60 * 5);
+        if (expirySeconds == null) {
+            long expiryTime = timeProvider.getAsLong() + (300 * 1000);
             jwtClaims.setExpiryTime(expiryTime);
-        } else if (expiryMin > 0) {
-            long expiryTime = System.currentTimeMillis() / 1000 + (60 * expiryMin);
+        } 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");

diff --git a/src/test/java/org/opensearch/security/authtoken/jwt/JwtVendorTest.java b/src/test/java/org/opensearch/security/authtoken/jwt/JwtVendorTest.java
@@ -11,7 +11,7 @@
 
 package org.opensearch.security.authtoken.jwt;
 
-import java.util.Map;
+import java.util.function.LongSupplier;
 
 import org.apache.cxf.rs.security.jose.jwk.JsonWebKey;
 import org.apache.cxf.rs.security.jose.jws.JwsJwtCompactConsumer;
@@ -32,7 +32,6 @@ public void testCreateJwkFromSettings() throws Exception {
         Assert.assertEquals("HS512", jwk.getAlgorithm());
         Assert.assertEquals("sig", jwk.getPublicKeyUse().toString());
         Assert.assertEquals("abc123", jwk.getProperty("k"));
-        System.out.print(jwk.getPublicKeyUse());
     }
 
     @Test (expected = Exception.class)
@@ -47,30 +46,35 @@ public void testCreateJwt() throws Exception {
         String issuer = "cluster_0";
         String subject = "admin";
         String audience = "extension_0";
-        Integer expiryMin = 5;
+        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);
-        String encodedJwt = jwtVendor.createJwt(issuer, subject, audience, expiryMin);
+        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 expiryMin = -1;
+        Integer expirySeconds = -300;
 
         Settings settings =  Settings.builder().put("signing_key", "abc123").build();
         JwtVendor jwtVendor = new JwtVendor(settings);
 
-        jwtVendor.createJwt(issuer, subject, audience, expiryMin);
+        jwtVendor.createJwt(issuer, subject, audience, expirySeconds);
     }
 }