cache API auth keys to offload an app and avoid frequent HTTP requests towards DynamoDB

Author: sergejsvisockis

service/pom.xml

@@ -60,6 +60,11 @@
             <version>3.0.5</version>
         </dependency>
 
+        <dependency>
+            <groupId>com.github.ben-manes.caffeine</groupId>
+            <artifactId>caffeine</artifactId>
+        </dependency>
+
         <!-- Test dependencies -->
         <dependency>
             <groupId>org.junit.jupiter</groupId>

service/src/main/java/io/github/sergejsvisockis/documentservice/auth/AuthenticationService.java

@@ -1,41 +1,48 @@
 package io.github.sergejsvisockis.documentservice.auth;
 
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
 import io.github.sergejsvisockis.documentservice.repository.ApiKey;
 import io.github.sergejsvisockis.documentservice.repository.ApiKeyRepository;
 import jakarta.servlet.http.HttpServletRequest;
-import lombok.RequiredArgsConstructor;
 import org.springframework.security.authentication.BadCredentialsException;
 import org.springframework.security.core.Authentication;
 import org.springframework.security.core.authority.AuthorityUtils;
 import org.springframework.stereotype.Service;
 
 import java.time.LocalDateTime;
 import java.util.Base64;
-import java.util.Optional;
+import java.util.concurrent.TimeUnit;
 
 @Service
-@RequiredArgsConstructor
 public class AuthenticationService {
 
     static final String AUTH_TOKEN_HEADER_NAME = "x-api-key";
     static final String INCORRECT_API_KEY_MESSAGE = "Incorrect API key passed";
-    static final String NO_API_KEY_FOUND_MESSAGE = "NO associated API key found";
+    static final String NO_API_KEY_FOUND_MESSAGE = "No associated API key found";
     static final String API_KEY_EXPIRED_MESSAGE = "An API key has expired";
 
+    private final Cache<String, ApiKey> cache;
     private final ApiKeyRepository apiKeyRepository;
 
+    public AuthenticationService(ApiKeyRepository apiKeyRepository) {
+        this.apiKeyRepository = apiKeyRepository;
+        this.cache = Caffeine.newBuilder()
+                .expireAfterAccess(2, TimeUnit.HOURS)
+                .initialCapacity(10)
+                .maximumSize(70)
+                .build();
+    }
+
     public Authentication getAuthentication(HttpServletRequest request) {
         String apiKey = request.getHeader(AUTH_TOKEN_HEADER_NAME);
         if (apiKey == null) {
             throw new BadCredentialsException(INCORRECT_API_KEY_MESSAGE);
         }
 
-        Optional<ApiKey> apiAuthKey = findApiAuthKey(apiKey);
-        if (apiAuthKey.isEmpty()) {
-            throw new BadCredentialsException(NO_API_KEY_FOUND_MESSAGE);
-        }
+        ApiKey apiAuthKey = findApiAuthKey(apiKey);
 
-        if (!isValid(apiAuthKey.get())) {
+        if (!isValid(apiAuthKey)) {
             throw new BadCredentialsException(API_KEY_EXPIRED_MESSAGE);
         }
 
@@ -47,9 +54,21 @@ private boolean isValid(ApiKey apiKey) {
         return LocalDateTime.now().isBefore(validTo);
     }
 
-    private Optional<ApiKey> findApiAuthKey(String key) {
-        return apiKeyRepository.findApiKey(encodeApiKey(key));
+    private ApiKey findApiAuthKey(String key) {
+        ApiKey fromCache = cache.getIfPresent(key);
+
+        if (fromCache != null) {
+            return fromCache;
+        }
+
+        ApiKey apiKey = apiKeyRepository.findApiKey(encodeApiKey(key));
+
+        if (apiKey == null) {
+            throw new BadCredentialsException(NO_API_KEY_FOUND_MESSAGE);
+        }
 
+        cache.put(key, apiKey);
+        return apiKey;
     }
 
     private String encodeApiKey(String key) {

service/src/main/java/io/github/sergejsvisockis/documentservice/repository/ApiKeyRepository.java

@@ -8,15 +8,14 @@
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 
 import java.util.Map;
-import java.util.Optional;
 
 @Repository
 @RequiredArgsConstructor
 public class ApiKeyRepository {
 
     private final DynamoDbTemplate dynamoDbTemplate;
 
-    public Optional<ApiKey> findApiKey(String key) {
+    public ApiKey findApiKey(String key) {
         ScanEnhancedRequest request = ScanEnhancedRequest.builder()
                 .filterExpression(Expression.builder()
                         .expression("apiKey = :key")
@@ -26,7 +25,8 @@ public Optional<ApiKey> findApiKey(String key) {
         return dynamoDbTemplate.scan(request, ApiKey.class)
                 .stream()
                 .flatMap(p -> p.items().stream())
-                .findFirst();
+                .findFirst()
+                .orElse(null);
     }
 
 }

service/src/test/java/io/github/sergejsvisockis/documentservice/auth/AuthenticationServiceTest.java

@@ -7,17 +7,19 @@
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.InjectMocks;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.springframework.security.authentication.BadCredentialsException;
 import org.springframework.security.core.Authentication;
 
 import java.time.LocalDateTime;
 import java.util.Base64;
-import java.util.Optional;
 
+import static io.github.sergejsvisockis.documentservice.auth.AuthenticationService.API_KEY_EXPIRED_MESSAGE;
+import static io.github.sergejsvisockis.documentservice.auth.AuthenticationService.INCORRECT_API_KEY_MESSAGE;
+import static io.github.sergejsvisockis.documentservice.auth.AuthenticationService.NO_API_KEY_FOUND_MESSAGE;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.when;
 
 @ExtendWith(MockitoExtension.class)
@@ -27,10 +29,6 @@ class AuthenticationServiceTest {
     private static final String VALID_API_KEY = "valid-api-key";
     private static final String ENCODED_API_KEY = Base64.getEncoder().encodeToString(VALID_API_KEY.getBytes());
 
-    private static final String INCORRECT_API_KEY_MESSAGE = "Incorrect API key passed";
-    private static final String NO_API_KEY_FOUND_MESSAGE = "NO associated API key found";
-    private static final String API_KEY_EXPIRED_MESSAGE = "An API key has expired";
-
     @Mock
     private ApiKeyRepository apiKeyRepository;
 
@@ -48,7 +46,7 @@ void shouldReturnAuthenticationWhenApiKeyIsValid() {
         apiKey.setExpirationDate(LocalDateTime.now().plusDays(1).toString());
 
         when(request.getHeader(AUTH_TOKEN_HEADER_NAME)).thenReturn(VALID_API_KEY);
-        when(apiKeyRepository.findApiKey(ENCODED_API_KEY)).thenReturn(Optional.of(apiKey));
+        when(apiKeyRepository.findApiKey(ENCODED_API_KEY)).thenReturn(apiKey);
 
         // when
         Authentication authentication = authenticationService.getAuthentication(request);
@@ -61,7 +59,7 @@ void shouldReturnAuthenticationWhenApiKeyIsValid() {
     void shouldThrowExceptionWhenApiKeyIsNotFound() {
         // given
         when(request.getHeader(AUTH_TOKEN_HEADER_NAME)).thenReturn(VALID_API_KEY);
-        when(apiKeyRepository.findApiKey(anyString())).thenReturn(Optional.empty());
+        when(apiKeyRepository.findApiKey(ENCODED_API_KEY)).thenReturn(null);
 
         // when & then
         BadCredentialsException exception = assertThrows(BadCredentialsException.class,
@@ -83,12 +81,11 @@ void shouldThrowExceptionWhenApiKeyIsNull() {
     @Test
     void shouldThrowExceptionWhenApiKeyHasExpired() {
         // given
-        ApiKey apiKey = new ApiKey();
-        apiKey.setApiKey(ENCODED_API_KEY);
-        apiKey.setExpirationDate(LocalDateTime.now().minusDays(1).toString());
-
+        ApiKey apiKey = Mockito.mock(ApiKey.class);
+        when(apiKey.getExpirationDate()).thenReturn(LocalDateTime.now().minusDays(1).toString());
         when(request.getHeader(AUTH_TOKEN_HEADER_NAME)).thenReturn(VALID_API_KEY);
-        when(apiKeyRepository.findApiKey(ENCODED_API_KEY)).thenReturn(Optional.of(apiKey));
+        when(apiKeyRepository.findApiKey(ENCODED_API_KEY)).thenReturn(apiKey);
+
 
         // when & then
         BadCredentialsException exception = assertThrows(BadCredentialsException.class,