Cache hit with service (#632)

Signed-off-by: Valentin Delaye <jonesbusy@users.noreply.github.com>

Author: jonesbusy

src/main/java/land/oras/Registry.java

@@ -906,7 +906,8 @@ private void handleError(HttpClient.ResponseWrapper<?> responseWrapper) {
                 LOG.debug("Response: {}", responseWrapper.response());
                 throw new OrasException((HttpClient.ResponseWrapper<String>) responseWrapper);
             }
-            throw new OrasException(new HttpClient.ResponseWrapper<>("", responseWrapper.statusCode(), Map.of()));
+            throw new OrasException(new HttpClient.ResponseWrapper<>(
+                    "", responseWrapper.statusCode(), Map.of(), responseWrapper.service()));
         }
     }
 
@@ -917,6 +918,7 @@ private void handleError(HttpClient.ResponseWrapper<?> responseWrapper) {
     private void logResponse(HttpClient.ResponseWrapper<?> response) {
         LOG.debug("Status Code: {}", response.statusCode());
         LOG.debug("Headers: {}", response.headers());
+        LOG.debug("Service: {}", response.service());
         String contentType = response.headers().get(Const.CONTENT_TYPE_HEADER.toLowerCase());
         boolean isBinaryResponse = contentType != null && contentType.contains("octet-stream");
         // Only log non-binary responses

src/main/java/land/oras/auth/HttpClient.java

@@ -438,7 +438,8 @@ public <T> TokenResponse refreshToken(
                                         : entry.getValue())));
 
         // Put in the cache
-        TokenResponse token = JsonUtils.fromJson(responseWrapper.response(), TokenResponse.class);
+        TokenResponse token = JsonUtils.fromJson(responseWrapper.response(), TokenResponse.class)
+                .forService(service);
         TokenCache.put(newScopes, token);
         return token;
     }
@@ -501,7 +502,7 @@ private <T> ResponseWrapper<T> executeRequest(
             if (cachedToken == null) {
                 LOG.trace("No cached token found for scopes: {}", newScopes);
             } else {
-                LOG.trace("Cached token for scopes: {}", newScopes);
+                LOG.trace("Found cached token for scopes: {}", newScopes.withService(cachedToken.service()));
             }
 
             // Add authentication header if any (from provider or cached token)
@@ -568,22 +569,17 @@ private <T> ResponseWrapper<T> redoRequest(
             AuthProvider authProvider) {
         if ((response.statusCode() == 401 || response.statusCode() == 403)) {
             LOG.debug("Requesting new token...");
-            HttpClient.TokenResponse token = refreshToken(toResponseWrapper(response), scopes, authProvider);
+            HttpClient.TokenResponse token =
+                    refreshToken(toResponseWrapper(response, scopes.getService()), scopes, authProvider);
             if (token.issued_at() != null && token.expires_in() != null) {
                 LOG.debug(
-                        "Found token issued_at {}, expire_id {} and expiring at {} ",
+                        "Received token issued_at {}, expire_id {} and expiring at {} ",
                         token.issued_at(),
                         token.expires_in(),
                         token.issued_at().plusSeconds(token.expires_in()));
             }
-            String bearerToken = token.token();
-            if (bearerToken == null) {
-                // Docker registry auth spec allows either token or auth_token (or both if they are the same)
-                bearerToken = token.access_token();
-            }
-            if (bearerToken == null) {
-                throw new OrasException("No Bearer token received");
-            }
+            String bearerToken = token.getEffectiveToken();
+            String service = token.service();
             try {
                 builder = builder.setHeader(Const.AUTHORIZATION_HEADER, "Bearer " + bearerToken);
                 HttpResponse<T> newResponse = client.send(builder.build(), handler);
@@ -601,25 +597,26 @@ private <T> ResponseWrapper<T> redoRequest(
                     }
 
                     return toResponseWrapper(
-                            client.send(builder.uri(URI.create(location)).build(), handler));
+                            client.send(builder.uri(URI.create(location)).build(), handler), service);
                 }
-                return toResponseWrapper(newResponse);
+                return toResponseWrapper(newResponse, service);
 
             } catch (Exception e) {
                 LOG.error("Failed to redo request", e);
                 throw new OrasException("Unable to redo HTTP request", e);
             }
         }
-        return toResponseWrapper(response);
+        return toResponseWrapper(response, scopes.getService());
     }
 
-    private <T> ResponseWrapper<T> toResponseWrapper(HttpResponse<T> response) {
+    private <T> ResponseWrapper<T> toResponseWrapper(HttpResponse<T> response, @Nullable String service) {
         return new ResponseWrapper<>(
                 response.body(),
                 response.statusCode(),
                 response.headers().map().entrySet().stream()
                         .collect(Collectors.toMap(
-                                Map.Entry::getKey, e -> e.getValue().get(0))));
+                                Map.Entry::getKey, e -> e.getValue().get(0))),
+                service);
     }
 
     /**
@@ -649,8 +646,10 @@ private void logRequest(HttpRequest request, byte[] body) {
      * @param response The response
      * @param statusCode The status code
      * @param headers The headers
+     * @param service The service (not on response but on HTTP headers)
      */
-    public record ResponseWrapper<T>(T response, int statusCode, Map<String, String> headers) {}
+    public record ResponseWrapper<T>(
+            T response, int statusCode, Map<String, String> headers, @Nullable String service) {}
 
     /**
      * Insecure trust manager when skipping TLS verification
@@ -697,6 +696,16 @@ public record TokenResponse(
             @Nullable ZonedDateTime issued_at) {
 
         /**
+         * Create a new token response with the service field set
+         * @param service The service
+         * @return A new token response with the service field set
+         */
+        public TokenResponse forService(String service) {
+            return new TokenResponse(token, access_token, service, expires_in, issued_at);
+        }
+
+        /**
+         * >>>>>>> 6379975 (Store token into caffeine cache (#631))
          * Get the effective token
          * @return The effective token, which is either the access_token or the token field depending on which one is present
          */

src/main/java/land/oras/auth/Scopes.java

@@ -136,7 +136,7 @@ public Scopes withNewScope(String scope) {
      * @param service The service to set
      * @return A new Scopes object with the given service
      */
-    public Scopes withService(String service) {
+    public Scopes withService(@Nullable String service) {
         return new Scopes(containerRef, service, scopes);
     }
 

src/main/java/land/oras/auth/TokenCache.java

@@ -35,6 +35,11 @@
 @NullMarked
 public final class TokenCache {
 
+    /**
+     * Hard cache limit
+     */
+    public static final int MAX_CACHE_SIZE = 500;
+
     /**
      * Logger for this class
      */
@@ -47,11 +52,17 @@ private TokenCache() {
         // Private constructor to prevent instantiation
     }
 
+    /**
+     * Cache for storing service information based on the service URL. This is used to avoid redundant
+     */
+    private static final Cache<String, String> SERVICE_CACHE =
+            Caffeine.newBuilder().maximumSize(MAX_CACHE_SIZE).build();
+
     /**
      * The cache
      */
     private static final Cache<Scopes, HttpClient.TokenResponse> CACHE = Caffeine.newBuilder()
-            .maximumSize(500)
+            .maximumSize(MAX_CACHE_SIZE)
             .expireAfter(new Expiry<Scopes, HttpClient.TokenResponse>() {
                 @Override
                 public long expireAfterCreate(Scopes key, HttpClient.TokenResponse token, long currentTime) {
@@ -80,6 +91,10 @@ public long expireAfterRead(
     public static void put(Scopes scopes, HttpClient.TokenResponse token) {
         LOG.trace("Caching token for scopes: {}", scopes);
         CACHE.put(scopes, token);
+        if (scopes.getService() != null) {
+            LOG.trace("Caching service for registry: {}", scopes.getRegistry());
+            SERVICE_CACHE.put(scopes.getRegistry(), scopes.getService());
+        }
     }
 
     /**
@@ -88,7 +103,16 @@ public static void put(Scopes scopes, HttpClient.TokenResponse token) {
      * @return the token response associated with the scopes, or null if not found or expired
      */
     public static HttpClient.@Nullable TokenResponse get(Scopes scopes) {
-        return CACHE.getIfPresent(scopes);
+        HttpClient.TokenResponse token = CACHE.getIfPresent(scopes);
+        if (token == null) {
+            // Lookup for specific service
+            String service = SERVICE_CACHE.getIfPresent(scopes.getRegistry());
+            if (service == null) {
+                return null;
+            }
+            return CACHE.getIfPresent(scopes.withService(service));
+        }
+        return token;
     }
 
     /**

src/test/java/land/oras/DockerIoITCase.java

@@ -91,8 +91,10 @@ void shouldPullAnonymousIndexUnqualified() {
     void shouldPullOneBlob() {
         Registry registry = Registry.builder().build();
         ContainerRef containerRef1 = ContainerRef.parse("jbangdev/jbang-action");
-        Manifest manifest = registry.getManifest(containerRef1);
+        Index index = registry.getIndex(containerRef1);
         String effectiveRegistry = containerRef1.getEffectiveRegistry(registry);
+        Manifest manifest = registry.getManifest(
+                containerRef1.withDigest(index.getManifests().get(0).getDigest()));
         Layer oneLayer = manifest.getLayers().get(0);
         registry.fetchBlob(
                 containerRef1.forRegistry(effectiveRegistry).withDigest(oneLayer.getDigest()),

src/test/java/land/oras/RegistryWireMockTest.java

@@ -541,13 +541,16 @@ void shouldGetToken(WireMockRuntimeInfo wmRuntimeInfo) {
     void shouldGetAuthToken(WireMockRuntimeInfo wmRuntimeInfo) {
         byte[] blob = tokenScenario(wmRuntimeInfo, "get-auth-token", null, "access-token");
         assertEquals("blob-data", new String(blob));
-    }
+        blob = tokenScenario(wmRuntimeInfo, "get-auth-token", null, "access-token");
+        assertEquals("blob-data", new String(blob));
+        blob = tokenScenario(wmRuntimeInfo, "get-auth-token", null, "access-token");
+        assertEquals("blob-data", new String(blob));
 
-    @Test
-    void shouldThrowIfNoTokenFound(WireMockRuntimeInfo wmRuntimeInfo) {
-        assertThrows(OrasException.class, () -> {
-            tokenScenario(wmRuntimeInfo, "get-auth-token", null, null);
-        });
+        // Ensure only one request on token endpoint du to caching
+        WireMock.verify(
+                1,
+                WireMock.getRequestedFor(
+                        WireMock.urlEqualTo("/token?scope=repository:library/get-auth-token:pull&service=localhost")));
     }
 
     @Test

src/test/java/land/oras/exception/OrasExceptionTest.java

@@ -50,7 +50,8 @@ void shouldWrapResponse() {
         HttpClient.ResponseWrapper<String> response = new HttpClient.ResponseWrapper<>(
                 JsonUtils.toJson(new Error("5001", "foo", "the details")),
                 500,
-                Map.of(Const.CONTENT_TYPE_HEADER, Const.DEFAULT_JSON_MEDIA_TYPE));
+                Map.of(Const.CONTENT_TYPE_HEADER, Const.DEFAULT_JSON_MEDIA_TYPE),
+                null);
         OrasException orasException = new OrasException(response);
 
         // Getters
@@ -64,7 +65,8 @@ void shouldWrapResponse() {
 
     @Test
     void shouldWrapInvalidResponse() {
-        HttpClient.ResponseWrapper<String> response = new HttpClient.ResponseWrapper<>("corrupted", 500, Map.of());
+        HttpClient.ResponseWrapper<String> response =
+                new HttpClient.ResponseWrapper<>("corrupted", 500, Map.of(), null);
         OrasException orasException = new OrasException(response);
         assertEquals("Response code: 500", orasException.getMessage(), "Message should be correct");
         assertNull(orasException.getError(), "Error should be null");