Add unit test for token revocation

wso2 · Nov 13, 2023 · bdfb360 · bdfb360
1 parent 723c9b4
commit bdfb360
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 2 deletions.
diff --git a/....wso2.apk.enforcer/src/main/java/org/wso2/apk/enforcer/security/jwt/JWTAuthenticator.java b/....wso2.apk.enforcer/src/main/java/org/wso2/apk/enforcer/security/jwt/JWTAuthenticator.java
@@ -137,9 +137,9 @@ public AuthenticationContext authenticate(RequestContext requestContext) throws
 
             JWTValidationInfo validationInfo = getJwtValidationInfo(jwtToken, organization, environment);
             if (RevokedTokenRedisClient.getRevokedTokens().contains(validationInfo.getIdentifier())) {
-                log.info("Expired JWT token. ", validationInfo.getIdentifier());
+                log.info("Revoked JWT token. ", validationInfo.getIdentifier());
                 throw new APISecurityException(APIConstants.StatusCodes.UNAUTHENTICATED.getCode(),
-                        APISecurityConstants.API_AUTH_INVALID_CREDENTIALS, "Invalid JWT token");
+                        APISecurityConstants.API_AUTH_INVALID_CREDENTIALS, APISecurityConstants.API_AUTH_INVALID_CREDENTIALS_MESSAGE);
             }
             if (validationInfo != null) {
                 if (validationInfo.isValid()) {

diff --git a/...orcer/org.wso2.apk.enforcer/src/test/java/org/wso2/apk/enforcer/jwt/JWTValidatorTest.java b/...orcer/org.wso2.apk.enforcer/src/test/java/org/wso2/apk/enforcer/jwt/JWTValidatorTest.java
@@ -158,6 +158,110 @@ public void testJWTValidator() throws APISecurityException, EnforcerException {
         }
     }
 
+    @Test
+    public void testRevokedToken() throws APISecurityException, EnforcerException {
+        HashSet<String> revokedTokens = new HashSet<>();
+        String revokedTokenJTI = "b8938768-23fd-4dec-8b70-bed45eb7c33d";
+        revokedTokens.add(revokedTokenJTI);
+        RevokedTokenRedisClient.setRevokedTokens(revokedTokens);
+        String organization = "org1";
+        String environment = "development";
+        String issuer = "https://localhost:9443/oauth2/token";
+        String signature = "sBgeoqJn0log5EZflj_G7ADvm6B3KQ9bdfFCEFVQS1U3oY9" +
+                "-cqPwAPyOLLh95pdfjYjakkf1UtjPZjeIupwXnzg0SffIc704RoVlZocAx9Ns2XihjU6Imx2MbXq9ARmQxQkyGVkJUMTwZ8" +
+                "-SfOnprfrhX2cMQQS8m2Lp7hcsvWFRGKxAKIeyUrbY4ihRIA5vOUrMBWYUx9Di1N7qdKA4S3e8O4KQX2VaZPBzN594c9TG" +
+                "riiH8AuuqnrftfvidSnlRLaFJmko8-QZo8jDepwacaFhtcaPVVJFG4uYP-_-N6sqfxLw3haazPN0_xU0T1zJLPRLC5HPfZMJDMGp" +
+                "EuSe9w";
+        String jwt = "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsIng1dCI6Ik5UZG1aak00WkRrM05qWTBZemM1T" +
+                "W1abU9EZ3dNVEUzTVdZd05ERTVNV1JsWkRnNE56YzRaQT09In0" +
+                ".eyJhdWQiOiJodHRwOlwvXC9vcmcud3NvMi5hcGltZ3RcL2dhdGV" +
+                "3YXkiLCJzdWIiOiJhZG1pbkBjYXJib24uc3VwZXIiLCJhcHBsaWNhdGlvbiI6eyJvd25lciI6ImFkbWluIiwidGllclF1b3RhVHlwZ" +
+                "SI6InJlcXVlc3RDb3VudCIsInRpZXIiOiJVbmxpbWl0ZWQiLCJuYW1lIjoiRGVmYXVsdEFwcGxpY2F0aW9uIiwiaWQiOjEsInV1aWQ" +
+                "iOm51bGx9LCJzY29wZSI6ImFtX2FwcGxpY2F0aW9uX3Njb3BlIGRlZmF1bHQiLCJpc3MiOiJodHRwczpcL1wvbG9jYWxob3N0Ojk0" +
+                "NDNcL29hdXRoMlwvdG9rZW4iLCJ0aWVySW5mbyI6e30sImtleXR5cGUiOiJQUk9EVUNUSU9OIiwic3Vic2NyaWJlZEFQSXMiOltdL" +
+                "CJjb25zdW1lcktleSI6IlhnTzM5NklIRks3ZUZZeWRycVFlNEhLR3oxa2EiLCJleHAiOjE1OTAzNDIzMTMsImlhdCI6MTU5MDMzO" +
+                "DcxMywianRpIjoiYjg5Mzg3NjgtMjNmZC00ZGVjLThiNzAtYmVkNDVlYjdjMzNkIn0." + signature;
+        JWTValidationInfo jwtValidationInfo = new JWTValidationInfo();
+        jwtValidationInfo.setValid(true);
+        jwtValidationInfo.setExpiryTime(System.currentTimeMillis() + 5000L);
+        jwtValidationInfo.setConsumerKey(UUID.randomUUID().toString());
+        jwtValidationInfo.setUser("user1");
+        jwtValidationInfo.setKeyManager("Default");
+        jwtValidationInfo.setIdentifier(revokedTokenJTI);
+
+        JWTConfigurationDto jwtConfigurationDto = new JWTConfigurationDto();
+        JWTAuthenticator jwtAuthenticator = new JWTAuthenticator(jwtConfigurationDto, true);
+        RequestContext requestContext = Mockito.mock(RequestContext.class);
+
+        ArrayList<ResourceConfig> resourceConfigs = new ArrayList<>();
+        ResourceConfig resourceConfig = Mockito.mock(ResourceConfig.class);
+        AuthenticationConfig authenticationConfig = new AuthenticationConfig();
+        JWTAuthenticationConfig jwtAuthenticationConfig = new JWTAuthenticationConfig();
+        jwtAuthenticationConfig.setHeader("Authorization");
+        authenticationConfig.setJwtAuthenticationConfig(jwtAuthenticationConfig);
+        Mockito.when(resourceConfig.getAuthenticationConfig()).thenReturn(authenticationConfig);
+        Mockito.when(resourceConfig.getMethod()).thenReturn(ResourceConfig.HttpMethods.GET);
+        resourceConfigs.add(resourceConfig);
+        Mockito.when(requestContext.getMatchedResourcePaths()).thenReturn(resourceConfigs);
+        Map<String, String> headers = new HashMap<>();
+        headers.put("Authorization", "Bearer " + jwt);
+        Mockito.when(requestContext.getHeaders()).thenReturn(headers);
+        Mockito.when(requestContext.getAuthenticationContext()).thenReturn(new AuthenticationContext());
+        APIConfig apiConfig = Mockito.mock(APIConfig.class);
+        Mockito.when(apiConfig.getName()).thenReturn("api1");
+        Mockito.when(apiConfig.getEnvironment()).thenReturn(environment);
+        Mockito.when(apiConfig.getOrganizationId()).thenReturn(organization);
+        Mockito.when(requestContext.getMatchedAPI()).thenReturn(apiConfig);
+        CacheProvider cacheProvider = Mockito.mock(CacheProvider.class);
+        LoadingCache gatewayKeyCache = Mockito.mock(LoadingCache.class);
+        LoadingCache invalidTokenCache = Mockito.mock(LoadingCache.class);
+        Mockito.when(cacheProvider.getGatewayKeyCache()).thenReturn(gatewayKeyCache);
+        Mockito.when(cacheProvider.getInvalidTokenCache()).thenReturn(invalidTokenCache);
+
+        try (MockedStatic<LogManager> logManagerDummy = Mockito.mockStatic(LogManager.class);
+             MockedStatic<LogFactory> logFactoryDummy = Mockito.mockStatic(LogFactory.class);
+                MockedStatic<CacheProviderUtil> cacheProviderUtilDummy = Mockito.mockStatic(CacheProviderUtil.class);
+                MockedStatic<ConfigHolder> configHolderDummy = Mockito.mockStatic(ConfigHolder.class);
+                MockedStatic<SubscriptionDataStoreImpl> subscriptionDataStoreImplDummy =
+                     Mockito.mockStatic(SubscriptionDataStoreImpl.class);
+                MockedStatic<KeyValidator> keyValidaterDummy = Mockito.mockStatic(KeyValidator.class)) {
+            Logger logger = Mockito.mock(Logger.class);
+            logManagerDummy.when(() -> LogManager.getLogger(JWTAuthenticator.class)).thenReturn(logger);
+            Log logger2 = Mockito.mock(Log.class);
+            logFactoryDummy.when(() -> LogFactory.getLog(AbstractAPIMgtGatewayJWTGenerator.class)).thenReturn(logger2);
+            cacheProviderUtilDummy.when(() -> CacheProviderUtil.getOrganizationCache(organization))
+                    .thenReturn(cacheProvider);
+            ExtendedTokenIssuerDto tokenIssuerDto = Mockito.mock(ExtendedTokenIssuerDto.class);
+            Mockito.when(tokenIssuerDto.getIssuer()).thenReturn(issuer);
+            JWKSConfigurationDTO jwksConfigurationDTO = new JWKSConfigurationDTO();
+            jwksConfigurationDTO.setEnabled(true);
+            Mockito.when(tokenIssuerDto.getJwksConfigurationDTO()).thenReturn(jwksConfigurationDTO);
+
+            EnforcerConfig enforcerConfig = Mockito.mock(EnforcerConfig.class);
+            ConfigHolder configHolder = Mockito.mock(ConfigHolder.class);
+            configHolderDummy.when(ConfigHolder::getInstance).thenReturn(configHolder);
+            Mockito.when(configHolder.getConfig()).thenReturn(enforcerConfig);
+            JWTTransformer jwtTransformer = new DefaultJWTTransformer();
+            Mockito.when(enforcerConfig.getJwtTransformer(issuer)).thenReturn(jwtTransformer);
+            JWTValidator jwtValidator = Mockito.mock(JWTValidator.class);
+
+            SubscriptionDataStoreImpl subscriptionDataStore = Mockito.mock(SubscriptionDataStoreImpl.class);
+            Mockito.when(subscriptionDataStore.getJWTValidatorByIssuer(issuer,
+                    organization, environment)).thenReturn(jwtValidator);
+            subscriptionDataStoreImplDummy.when(SubscriptionDataStoreImpl::getInstance).thenReturn(subscriptionDataStore);
+            Mockito.when(jwtValidator.validateToken(Mockito.eq(jwt), Mockito.any())).thenReturn(jwtValidationInfo);
+            keyValidaterDummy.when(()->KeyValidator.validateScopes(Mockito.any())).thenReturn(true);
+            try {
+                jwtAuthenticator.authenticate(requestContext);
+                Assert.fail("Authentication should fail for revoked tokens");
+            } catch (APISecurityException e) {
+                Assert.assertEquals(e.getMessage(), e.getMessage(), APISecurityConstants.API_AUTH_INVALID_CREDENTIALS_MESSAGE);
+            }
+        } finally {
+            RevokedTokenRedisClient.setRevokedTokens(new HashSet<>());
+        }
+    }
+
     @Test
     public void testCachedJWTValidator() throws APISecurityException, EnforcerException {
         String organization = "org1";