fix: Core shouldn't check user roles in case authorization by API key… (

#182) Co-authored-by: Aliaksandr Stsiapanay <[email protected]>
epam · Jan 29, 2024 · 5c1b5f4 · 5c1b5f4
1 parent a1c190d
commit 5c1b5f4
Show file tree

Hide file tree

Showing 7 changed files with 20 additions and 15 deletions.
diff --git a/src/main/java/com/epam/aidial/core/controller/DeploymentController.java b/src/main/java/com/epam/aidial/core/controller/DeploymentController.java
@@ -83,7 +83,7 @@ public static boolean hasAccessByUserRoles(ProxyContext context, Deployment depl
         Set<String> expectedUserRoles = deployment.getUserRoles();
         List<String> actualUserRoles = context.getUserRoles();
 
-        if (expectedUserRoles.isEmpty()) {
+        if (expectedUserRoles.isEmpty() || actualUserRoles == null) {
             return true;
         }
 

diff --git a/src/main/java/com/epam/aidial/core/security/AccessTokenValidator.java b/src/main/java/com/epam/aidial/core/security/AccessTokenValidator.java
@@ -13,8 +13,6 @@
 import java.util.List;
 import java.util.Objects;
 
-import static com.epam.aidial.core.security.IdentityProvider.CLAIMS_WITH_EMPTY_ROLES;
-
 public class AccessTokenValidator {
 
     private final List<IdentityProvider> providers = new ArrayList<>();
@@ -38,7 +36,7 @@ public AccessTokenValidator(JsonObject idpConfig, Vertx vertx) {
     public Future<ExtractedClaims> extractClaims(String authHeader) {
         try {
             if (authHeader == null) {
-                return Future.succeededFuture(CLAIMS_WITH_EMPTY_ROLES);
+                return Future.succeededFuture();
             }
             String encodedToken = Objects.requireNonNull(extractTokenFromHeader(authHeader), "Can't extract access token from header");
             DecodedJWT jwt = IdentityProvider.decodeJwtToken(encodedToken);

diff --git a/src/main/java/com/epam/aidial/core/security/IdentityProvider.java b/src/main/java/com/epam/aidial/core/security/IdentityProvider.java
@@ -15,7 +15,6 @@
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.security.interfaces.RSAPublicKey;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -29,8 +28,6 @@
 @Slf4j
 public class IdentityProvider {
 
-    public static final ExtractedClaims CLAIMS_WITH_EMPTY_ROLES = new ExtractedClaims(null, Collections.emptyList(), null);
-
     private final String[] rolePath;
 
     private final JwkProvider jwkProvider;

diff --git a/src/test/java/com/epam/aidial/core/ProxyTest.java b/src/test/java/com/epam/aidial/core/ProxyTest.java
@@ -221,7 +221,7 @@ public void testHandle_OpenAiRequestSuccess() {
         when(configStore.load()).thenReturn(config);
         when(apiKeyStore.getApiKeyData("key1")).thenReturn(new ApiKeyData());
 
-        when(accessTokenValidator.extractClaims(any())).thenReturn(Future.succeededFuture(IdentityProvider.CLAIMS_WITH_EMPTY_ROLES));
+        when(accessTokenValidator.extractClaims(any())).thenReturn(Future.succeededFuture());
 
         proxy.handle(request);
 
@@ -274,7 +274,7 @@ public void testHandle_SuccessApiKey() {
         when(configStore.load()).thenReturn(config);
         when(apiKeyStore.getApiKeyData("key1")).thenReturn(new ApiKeyData());
 
-        when(accessTokenValidator.extractClaims(any())).thenReturn(Future.succeededFuture(IdentityProvider.CLAIMS_WITH_EMPTY_ROLES));
+        when(accessTokenValidator.extractClaims(any())).thenReturn(Future.succeededFuture());
 
         proxy.handle(request);
 

diff --git a/src/test/java/com/epam/aidial/core/controller/DeploymentControllerTest.java b/src/test/java/com/epam/aidial/core/controller/DeploymentControllerTest.java
@@ -101,6 +101,16 @@ public void testHasAssessByRole_DeploymentRolesEmpty() {
         assertTrue(DeploymentController.hasAccessByUserRoles(proxyContext, deployment));
     }
 
+    @Test
+    public void testHasAssessByRole_DeploymentRolesIsNull() {
+        ProxyContext proxyContext = mock(ProxyContext.class);
+        Deployment deployment = mock(Deployment.class);
+
+        when(deployment.getUserRoles()).thenReturn(Collections.emptySet());
+
+        assertTrue(DeploymentController.hasAccessByUserRoles(proxyContext, deployment));
+    }
+
     @Test
     public void testHasAssessByRole_RoleMismatch() {
         ProxyContext proxyContext = mock(ProxyContext.class);

diff --git a/src/test/java/com/epam/aidial/core/limiter/RateLimiterTest.java b/src/test/java/com/epam/aidial/core/limiter/RateLimiterTest.java
@@ -41,7 +41,7 @@ public void beforeEach() {
     public void testLimit_EntityNotFound() {
         ApiKeyData apiKeyData = new ApiKeyData();
         apiKeyData.setOriginalKey(new Key());
-        ProxyContext proxyContext = new ProxyContext(new Config(), request, apiKeyData, IdentityProvider.CLAIMS_WITH_EMPTY_ROLES, "unknown-trace-id", "span-id");
+        ProxyContext proxyContext = new ProxyContext(new Config(), request, apiKeyData, null, "unknown-trace-id", "span-id");
         proxyContext.setDeployment(new Application());
 
         RateLimitResult result = rateLimiter.limit(proxyContext);
@@ -66,7 +66,7 @@ public void testLimit_ApiKeyLimitNotFound() {
         key.setRole("role");
         ApiKeyData apiKeyData = new ApiKeyData();
         apiKeyData.setOriginalKey(key);
-        ProxyContext proxyContext = new ProxyContext(new Config(), request, apiKeyData, IdentityProvider.CLAIMS_WITH_EMPTY_ROLES, "trace-id", "span-id");
+        ProxyContext proxyContext = new ProxyContext(new Config(), request, apiKeyData, null, "trace-id", "span-id");
         proxyContext.setDeployment(new Model());
 
         RateLimitResult result = rateLimiter.limit(proxyContext);
@@ -84,7 +84,7 @@ public void testLimit_ApiKeyLimitNotFoundWithNullRole() {
         key.setProject("project");
         ApiKeyData apiKeyData = new ApiKeyData();
         apiKeyData.setOriginalKey(key);
-        ProxyContext proxyContext = new ProxyContext(new Config(), request, apiKeyData, IdentityProvider.CLAIMS_WITH_EMPTY_ROLES, "trace-id", "span-id");
+        ProxyContext proxyContext = new ProxyContext(new Config(), request, apiKeyData, null, "trace-id", "span-id");
         proxyContext.setDeployment(new Model());
 
         RateLimitResult result = rateLimiter.limit(proxyContext);
@@ -107,7 +107,7 @@ public void testLimit_ApiKeyLimitNegative() {
         config.setRoles(Map.of("role", role));
         ApiKeyData apiKeyData = new ApiKeyData();
         apiKeyData.setOriginalKey(key);
-        ProxyContext proxyContext = new ProxyContext(config, request, apiKeyData, IdentityProvider.CLAIMS_WITH_EMPTY_ROLES, "trace-id", "span-id");
+        ProxyContext proxyContext = new ProxyContext(config, request, apiKeyData, null, "trace-id", "span-id");
         Model model = new Model();
         model.setName("model");
         proxyContext.setDeployment(model);
@@ -131,7 +131,7 @@ public void testLimit_ApiKeySuccess() {
         config.setRoles(Map.of("role", role));
         ApiKeyData apiKeyData = new ApiKeyData();
         apiKeyData.setOriginalKey(key);
-        ProxyContext proxyContext = new ProxyContext(config, request, apiKeyData, IdentityProvider.CLAIMS_WITH_EMPTY_ROLES, "trace-id", "span-id");
+        ProxyContext proxyContext = new ProxyContext(config, request, apiKeyData, null, "trace-id", "span-id");
         Model model = new Model();
         model.setName("model");
         proxyContext.setDeployment(model);

diff --git a/src/test/java/com/epam/aidial/core/security/AccessTokenValidatorTest.java b/src/test/java/com/epam/aidial/core/security/AccessTokenValidatorTest.java
@@ -52,7 +52,7 @@ public void testExtractClaims_01() {
         assertNotNull(future);
         future.onComplete(res -> {
             assertTrue(res.succeeded());
-            assertEquals(IdentityProvider.CLAIMS_WITH_EMPTY_ROLES, res.result());
+            assertNull(res.result());
         });
     }