diff --git a/src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java b/src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java index 55e2ec883..5b05febd6 100644 --- a/src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java +++ b/src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java @@ -144,7 +144,12 @@ private void setupProxyApiKeyData() { ApiKeyData proxyApiKeyData = new ApiKeyData(); context.setProxyApiKeyData(proxyApiKeyData); ApiKeyData.initFromContext(proxyApiKeyData, context); + long start = System.currentTimeMillis(); proxy.getApiKeyStore().assignPerRequestApiKey(proxyApiKeyData); + log.info("Complete assigning per request api key for {} ms. Trace: {}. Span: {}. Key: {}. Deployment: {}", + System.currentTimeMillis() - start, + context.getTraceId(), context.getSpanId(), + context.getProject(), context.getDeployment().getName()); } private void handleRateLimitHit(RateLimitResult result) { diff --git a/src/main/java/com/epam/aidial/core/security/ApiKeyStore.java b/src/main/java/com/epam/aidial/core/security/ApiKeyStore.java index a390cd110..1683f150b 100644 --- a/src/main/java/com/epam/aidial/core/security/ApiKeyStore.java +++ b/src/main/java/com/epam/aidial/core/security/ApiKeyStore.java @@ -50,16 +50,9 @@ public class ApiKeyStore { *

*/ public synchronized void assignPerRequestApiKey(ApiKeyData data) { - lockService.underBucketLock(API_KEY_DATA_LOCATION, () -> { - ResourceDescription resource = generateApiKey(); - String apiKey = resource.getName(); - data.setPerRequestKey(apiKey); - String json = ProxyUtil.convertToString(data); - if (resourceService.putResource(resource, json, false, false) == null) { - throw new IllegalStateException(String.format("API key %s already exists in the storage", apiKey)); - } - return apiKey; - }); + String apiKey = generateApiKey(); + keys.put(apiKey, data); + data.setPerRequestKey(apiKey); } /** @@ -109,8 +102,7 @@ public synchronized void addProjectKeys(Map projectKeys) { Key value = entry.getValue(); ResourceDescription resource = toResource(apiKey); if (resourceService.hasResource(resource)) { - resource = generateApiKey(); - apiKey = resource.getName(); + apiKey = generateApiKey(); } value.setKey(apiKey); ApiKeyData apiKeyData = new ApiKeyData(); @@ -137,15 +129,13 @@ public void updatePerRequestApiKeyData(ApiKeyData apiKeyData) { resourceService.putResource(resource, json, true, false); } - private ResourceDescription generateApiKey() { + private String generateApiKey() { String apiKey = generateKey(); - ResourceDescription resource = toResource(apiKey); - while (resourceService.hasResource(resource) || keys.containsKey(apiKey)) { + while (keys.containsKey(apiKey)) { log.warn("duplicate API key is found. Trying to generate a new one"); apiKey = generateKey(); - resource = toResource(apiKey); } - return resource; + return apiKey; } private static ResourceDescription toResource(String apiKey) { diff --git a/src/test/java/com/epam/aidial/core/security/ApiKeyStoreTest.java b/src/test/java/com/epam/aidial/core/security/ApiKeyStoreTest.java index 5391a9458..3496d11e2 100644 --- a/src/test/java/com/epam/aidial/core/security/ApiKeyStoreTest.java +++ b/src/test/java/com/epam/aidial/core/security/ApiKeyStoreTest.java @@ -105,41 +105,6 @@ public void testAssignApiKey() { assertNotNull(apiKeyData.getPerRequestKey()); } - @Test - public void testAddProjectKeys() { - when(vertx.executeBlocking(any(Callable.class))).thenAnswer(invocation -> { - Callable callable = invocation.getArgument(0); - return Future.succeededFuture(callable.call()); - }); - - Key key1 = new Key(); - key1.setProject("prj1"); - key1.setRole("role1"); - Map projectKeys1 = Map.of("key1", key1); - - store.addProjectKeys(projectKeys1); - - ApiKeyData apiKeyData = new ApiKeyData(); - store.assignPerRequestApiKey(apiKeyData); - - Key key2 = new Key(); - key1.setProject("prj1"); - key1.setRole("role1"); - Map projectKeys2 = Map.of("key2", key2); - - store.addProjectKeys(projectKeys2); - - // old key must be removed - assertNull(store.getApiKeyData("key1").result()); - // new key must be accessed - Future res1 = store.getApiKeyData("key2"); - assertNotNull(res1.result()); - assertEquals(key2, res1.result().getOriginalKey()); - // existing per request key must be accessed - assertNotNull(store.getApiKeyData(apiKeyData.getPerRequestKey()).result()); - - } - @Test public void testGetApiKeyData() { ApiKeyData apiKeyData = new ApiKeyData(); @@ -158,16 +123,4 @@ public void testGetApiKeyData() { assertNull(store.getApiKeyData("unknown-key").result()); } - - @Test - public void testInvalidateApiKey() { - ApiKeyData apiKeyData = new ApiKeyData(); - store.assignPerRequestApiKey(apiKeyData); - - assertNotNull(apiKeyData.getPerRequestKey()); - - store.invalidatePerRequestApiKey(apiKeyData); - - assertNull(store.getApiKeyData(apiKeyData.getPerRequestKey())); - } }