feat: Tolerate trivial auth header when non-trivial api-key header is…

… provided #675 (#677)
epam · Feb 6, 2025 · 6f27b59 · 6f27b59
1 parent a05612d
commit 6f27b59
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/server/src/main/java/com/epam/aidial/core/server/Proxy.java b/server/src/main/java/com/epam/aidial/core/server/Proxy.java
@@ -223,6 +223,12 @@ private Future<AuthorizationResult> authorizeRequest(HttpServerRequest request)
                             });
         }
 
+        // see https://github.com/epam/ai-dial-core/issues/675
+        if ("Bearer".equalsIgnoreCase(authorization.trim())) {
+            return apiKeyStore.getApiKeyData(apiKey)
+                    .map(apiKeyData -> new AuthorizationResult(apiKeyData, null));
+        }
+
         if (apiKey.equals(AccessTokenValidator.extractTokenFromHeader(authorization))) {
             // we don't know exactly what kind of credentials a client provided to us.
             // we try if it's access token the first and then API key

diff --git a/server/src/test/java/com/epam/aidial/core/server/ProxyTest.java b/server/src/test/java/com/epam/aidial/core/server/ProxyTest.java
@@ -490,6 +490,39 @@ public void testHandle_SuccessApiKey() {
         verify(response).setStatusCode(OK.getCode());
     }
 
+    @Test
+    public void testHandle_AzureOpenAiRequest() {
+        when(request.version()).thenReturn(HttpVersion.HTTP_1_1);
+        when(request.method()).thenReturn(HttpMethod.GET);
+        MultiMap headers = mock(MultiMap.class);
+        when(request.headers()).thenReturn(headers);
+        when(request.getHeader(eq(HttpHeaders.CONTENT_TYPE))).thenReturn(null);
+        when(request.getHeader(eq(HttpHeaders.AUTHORIZATION))).thenReturn("bearer");
+        when(headers.get(eq(HEADER_API_KEY))).thenReturn("key1");
+        when(headers.get(eq(HttpHeaders.CONTENT_LENGTH))).thenReturn(Integer.toString(512));
+        when(request.path()).thenReturn("/foo");
+        when(request.uri()).thenReturn("/foo");
+
+        Config config = new Config();
+        Route route = new Route();
+        route.setMethods(Set.of("GET"));
+        route.setName("route");
+        route.setPaths(List.of(Pattern.compile("/foo")));
+        route.setResponse(new Route.Response());
+        LinkedHashMap<String, Route> routes = new LinkedHashMap<>();
+        routes.put("route", route);
+        config.setRoutes(routes);
+        when(configStore.load()).thenReturn(config);
+        ApiKeyData apiKeyData = new ApiKeyData();
+        Key originalKey = new Key();
+        apiKeyData.setOriginalKey(originalKey);
+        when(apiKeyStore.getApiKeyData("key1")).thenReturn(Future.succeededFuture(apiKeyData));
+
+        proxy.handle(request);
+
+        verify(response).setStatusCode(OK.getCode());
+    }
+
     @Test
     public void testHandle_SuccessAccessToken() {
         when(request.version()).thenReturn(HttpVersion.HTTP_1_1);