Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore: add logging #321

Closed
wants to merge 5 commits into from
Closed

chore: add logging #321

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f55e6b5
chore: add logging
astsiapanay May 3, 2024
5fc461c
chore: logging
astsiapanay May 3, 2024
3576303
chore: disable bucket lock
astsiapanay May 3, 2024
2f69d8d
chore: fix
astsiapanay May 3, 2024
5616cd5
chore: rollback assign per request key
astsiapanay May 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/main/java/com/epam/aidial/core/controller/DeploymentPostController.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
24 changes: 7 additions & 17 deletions src/main/java/com/epam/aidial/core/security/ApiKeyStore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,16 +50,9 @@ public class ApiKeyStore {
* </p>
*/
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);
}

/**
Expand Down Expand Up @@ -109,8 +102,7 @@ public synchronized void addProjectKeys(Map<String, Key> 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();
Expand All @@ -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) {
Expand Down
47 changes: 0 additions & 47 deletions src/test/java/com/epam/aidial/core/security/ApiKeyStoreTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, Key> 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<String, Key> projectKeys2 = Map.of("key2", key2);

store.addProjectKeys(projectKeys2);

// old key must be removed
assertNull(store.getApiKeyData("key1").result());
// new key must be accessed
Future<ApiKeyData> 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();
Expand All @@ -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()));
}
}
Loading