From f55e6b501e87458235a3c4cd515eb27363c9f91a Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Fri, 3 May 2024 11:01:55 +0300 Subject: [PATCH 1/5] chore: add logging --- .../aidial/core/controller/DeploymentPostController.java | 6 ++++++ 1 file changed, 6 insertions(+) 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..d110992d8 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,13 @@ private void setupProxyApiKeyData() { ApiKeyData proxyApiKeyData = new ApiKeyData(); context.setProxyApiKeyData(proxyApiKeyData); ApiKeyData.initFromContext(proxyApiKeyData, context); + log.info("Start assigning per request api key. Trace: {}. Span: {}. Key: {}. Deployment: {}", + context.getTraceId(), context.getSpanId(), + context.getProject(), context.getDeployment().getName()); proxy.getApiKeyStore().assignPerRequestApiKey(proxyApiKeyData); + log.info("Complete assigning per request api key. Trace: {}. Span: {}. Key: {}. Deployment: {}", + context.getTraceId(), context.getSpanId(), + context.getProject(), context.getDeployment().getName()); } private void handleRateLimitHit(RateLimitResult result) { From 5fc461c580869c8ebc5fbd1700b1c2bb4e26dfee Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Fri, 3 May 2024 11:33:25 +0300 Subject: [PATCH 2/5] chore: logging --- .../aidial/core/controller/DeploymentPostController.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) 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 d110992d8..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,11 +144,10 @@ private void setupProxyApiKeyData() { ApiKeyData proxyApiKeyData = new ApiKeyData(); context.setProxyApiKeyData(proxyApiKeyData); ApiKeyData.initFromContext(proxyApiKeyData, context); - log.info("Start assigning per request api key. Trace: {}. Span: {}. Key: {}. Deployment: {}", - context.getTraceId(), context.getSpanId(), - context.getProject(), context.getDeployment().getName()); + long start = System.currentTimeMillis(); proxy.getApiKeyStore().assignPerRequestApiKey(proxyApiKeyData); - log.info("Complete assigning per request api key. Trace: {}. Span: {}. Key: {}. Deployment: {}", + 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()); } From 357630341992d38b6ca3708e784fe252d050710f Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Fri, 3 May 2024 13:09:05 +0300 Subject: [PATCH 3/5] chore: disable bucket lock --- .../aidial/core/security/ApiKeyStore.java | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) 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..e6cfacb4c 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,16 @@ 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; - }); +// 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; +// }); } /** From 2f69d8d6c4fbe9d25698d29a3e429418eb6edce3 Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Fri, 3 May 2024 13:14:16 +0300 Subject: [PATCH 4/5] chore: fix --- src/main/java/com/epam/aidial/core/security/ApiKeyStore.java | 3 --- 1 file changed, 3 deletions(-) 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 e6cfacb4c..23927e467 100644 --- a/src/main/java/com/epam/aidial/core/security/ApiKeyStore.java +++ b/src/main/java/com/epam/aidial/core/security/ApiKeyStore.java @@ -50,7 +50,6 @@ 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); @@ -58,8 +57,6 @@ public synchronized void assignPerRequestApiKey(ApiKeyData 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; -// }); } /** From 5616cd5a18f7c754eb5f9a8102381ff630838288 Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Fri, 3 May 2024 14:36:45 +0300 Subject: [PATCH 5/5] chore: rollback assign per request key --- .../aidial/core/security/ApiKeyStore.java | 19 +++----- .../aidial/core/security/ApiKeyStoreTest.java | 47 ------------------- 2 files changed, 6 insertions(+), 60 deletions(-) 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 23927e467..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,13 +50,9 @@ public class ApiKeyStore { *

*/ public synchronized void assignPerRequestApiKey(ApiKeyData data) { - ResourceDescription resource = generateApiKey(); - String apiKey = resource.getName(); + String apiKey = generateApiKey(); + keys.put(apiKey, data); 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)); - } } /** @@ -106,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(); @@ -134,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())); - } }