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()));
- }
}