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