fix: handle caching when detecting a repo requires auth

jib-core/src/main/java/com/google/cloud/tools/jib/builder/steps/PullBaseImageStep.java

@@ -50,6 +50,7 @@
 import com.google.cloud.tools.jib.json.JsonTemplateMapper;
 import com.google.cloud.tools.jib.registry.ManifestAndDigest;
 import com.google.cloud.tools.jib.registry.RegistryClient;
+import com.google.cloud.tools.jib.registry.RegistryUnauthorizedExceptionHandler;
 import com.google.cloud.tools.jib.registry.credentials.CredentialRetrievalException;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
@@ -64,6 +65,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Callable;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /** Pulls the base image manifests for the specified platforms. */
@@ -132,12 +134,8 @@ public ImagesAndRegistryClient call()
       } else if (imageReference.getDigest().isPresent()) {
         List<Image> images = getCachedBaseImages();
         if (!images.isEmpty()) {
-          RegistryClient noAuthRegistryClient =
-              buildContext.newBaseImageRegistryClientFactory().newRegistryClient();
-          // TODO: passing noAuthRegistryClient may be problematic. It may return 401 unauthorized
-          // if layers have to be downloaded.
-          // https://github.com/GoogleContainerTools/jib/issues/2220
-          return new ImagesAndRegistryClient(images, noAuthRegistryClient);
+          RegistryClient onDemandAuthRegistryClient = createOnDemandAuthenticatingRegistryClient();
+          return new ImagesAndRegistryClient(images, onDemandAuthRegistryClient);
         }
       }
 
@@ -147,63 +145,62 @@ public ImagesAndRegistryClient call()
         return mirrorPull.get();
       }
 
-      try {
-        // First, try with no credentials. This works with public GCR images (but not Docker Hub).
-        // TODO: investigate if we should just pass credentials up front. However, this involves
-        // some risk. https://github.com/GoogleContainerTools/jib/pull/2200#discussion_r359069026
-        // contains some related discussions.
-        RegistryClient noAuthRegistryClient =
-            buildContext.newBaseImageRegistryClientFactory().newRegistryClient();
-        return new ImagesAndRegistryClient(
-            pullBaseImages(noAuthRegistryClient, progressDispatcher.newChildProducer()),
-            noAuthRegistryClient);
-
-      } catch (RegistryUnauthorizedException ex) {
-        eventHandlers.dispatch(
-            LogEvent.lifecycle(
-                "The base image requires auth. Trying again for " + imageReference + "..."));
-
-        Credential credential =
-            RegistryCredentialRetriever.getBaseImageCredential(buildContext).orElse(null);
-        RegistryClient registryClient =
-            buildContext
-                .newBaseImageRegistryClientFactory()
-                .setCredential(credential)
-                .newRegistryClient();
+      RegistryClient onDemandAuthRegistryClient = createOnDemandAuthenticatingRegistryClient();
+      return new ImagesAndRegistryClient(
+          pullBaseImages(onDemandAuthRegistryClient, progressDispatcher.newChildProducer()),
+          onDemandAuthRegistryClient);
+    }
+  }
+
+  private RegistryClient createOnDemandAuthenticatingRegistryClient()
+      throws CredentialRetrievalException {
+    Credential credential =
+        RegistryCredentialRetriever.getBaseImageCredential(buildContext).orElse(null);
+    return buildContext
+        .newBaseImageRegistryClientFactory()
+        .setCredential(credential)
+        .setUnauthorizedExceptionHandlerSupplier(
+            () -> PullBaseImageStep::handleRegistryUnauthorizedException)
+        .newRegistryClient();
+  }
 
+  /**
+   * Handles an unauthorized exception by performing authentication on demand.
+   *
+   * @param registryClient the registry client to be reconfigured
+   * @param ex the exception that was caught
+   * @return a supplier of the next handler, used only if another exception is thrown
+   * @throws RegistryException if a registry error occurs
+   * @throws IOException if an I/O error occurs
+   */
+  static Supplier<RegistryUnauthorizedExceptionHandler> handleRegistryUnauthorizedException(
+      final RegistryClient registryClient, final RegistryUnauthorizedException ex)
+      throws RegistryException, IOException {
+    // Double indentation keeps code at same level as original code
+    {
+      {
+        final EventHandlers eventHandlers = registryClient.getEventHandlers();
+        final String imageReference = ex.getImageReference();
         String wwwAuthenticate = ex.getHttpResponseException().getHeaders().getAuthenticate();
         if (wwwAuthenticate != null) {
           eventHandlers.dispatch(
               LogEvent.debug("WWW-Authenticate for " + imageReference + ": " + wwwAuthenticate));
           registryClient.authPullByWwwAuthenticate(wwwAuthenticate);
-          return new ImagesAndRegistryClient(
-              pullBaseImages(registryClient, progressDispatcher.newChildProducer()),
-              registryClient);
-
         } else {
           // Not getting WWW-Authenticate is unexpected in practice, and we may just blame the
           // server and fail. However, to keep some old behavior, try a few things as a last resort.
           // TODO: consider removing this fallback branch.
-          if (credential != null && !credential.isOAuth2RefreshToken()) {
-            eventHandlers.dispatch(
-                LogEvent.debug("Trying basic auth as fallback for " + imageReference + "..."));
-            registryClient.configureBasicAuth();
-            try {
-              return new ImagesAndRegistryClient(
-                  pullBaseImages(registryClient, progressDispatcher.newChildProducer()),
-                  registryClient);
-            } catch (RegistryUnauthorizedException ignored) {
-              // Fall back to try bearer auth.
-            }
-          }
-
           eventHandlers.dispatch(
               LogEvent.debug("Trying bearer auth as fallback for " + imageReference + "..."));
-          registryClient.doPullBearerAuth();
-          return new ImagesAndRegistryClient(
-              pullBaseImages(registryClient, progressDispatcher.newChildProducer()),
-              registryClient);
+          if (!registryClient.doPullBearerAuth()) {
+            eventHandlers.dispatch(
+                LogEvent.error("Failed to bearer auth for pull of " + imageReference));
+            throw ex;
+          }
         }
+
+        // If another exception occurs, use the default behavior
+        return RegistryClient::defaultRegistryUnauthorizedExceptionHandler;
       }
     }
   }

jib-core/src/main/java/com/google/cloud/tools/jib/builder/steps/RegistryCredentialRetriever.java

@@ -39,7 +39,13 @@ static Optional<Credential> getBaseImageCredential(BuildContext buildContext)
   /** Retrieves credentials for the target image. */
   static Optional<Credential> getTargetImageCredential(BuildContext buildContext)
       throws CredentialRetrievalException {
-    return retrieve(buildContext.getTargetImageConfiguration(), buildContext.getEventHandlers());
+    final Optional<Credential> credential =
+        retrieve(buildContext.getTargetImageConfiguration(), buildContext.getEventHandlers());
+    if (!credential.isPresent()) {
+      logNoCredentialsRetrieved(
+          buildContext.getTargetImageConfiguration(), buildContext.getEventHandlers());
+    }
+    return credential;
   }
 
   private static Optional<Credential> retrieve(
@@ -52,10 +58,14 @@ private static Optional<Credential> retrieve(
       }
     }
 
+    return Optional.empty();
+  }
+
+  private static void logNoCredentialsRetrieved(
+      ImageConfiguration imageConfiguration, EventHandlers eventHandlers) {
     String registry = imageConfiguration.getImageRegistry();
     String repository = imageConfiguration.getImageRepository();
     eventHandlers.dispatch(
         LogEvent.info("No credentials could be retrieved for " + registry + "/" + repository));
-    return Optional.empty();
   }
 }

jib-core/src/main/java/com/google/cloud/tools/jib/registry/RegistryClient.java

@@ -48,6 +48,7 @@
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.ThreadSafe;
@@ -65,6 +66,8 @@ public static class Factory {
 
     @Nullable private String userAgent;
     @Nullable private Credential credential;
+    Supplier<RegistryUnauthorizedExceptionHandler> registryUnauthorizedExceptionHandlerSupplier =
+        RegistryClient::defaultRegistryUnauthorizedExceptionHandler;
 
     private Factory(
         EventHandlers eventHandlers,
@@ -86,6 +89,20 @@ public Factory setCredential(@Nullable Credential credential) {
       return this;
     }
 
+    /**
+     * Sets the supplier for the registry unauthorized exception handler.
+     *
+     * @param registryUnauthorizedExceptionHandlerSupplier the supplier for the exception handler
+     * @return this
+     */
+    public Factory setUnauthorizedExceptionHandlerSupplier(
+        Supplier<RegistryUnauthorizedExceptionHandler>
+            registryUnauthorizedExceptionHandlerSupplier) {
+      this.registryUnauthorizedExceptionHandlerSupplier =
+          registryUnauthorizedExceptionHandlerSupplier;
+      return this;
+    }
+
     /**
      * Sets the value of {@code User-Agent} in headers for registry requests.
      *
@@ -104,7 +121,12 @@ public Factory setUserAgent(@Nullable String userAgent) {
      */
     public RegistryClient newRegistryClient() {
       return new RegistryClient(
-          eventHandlers, credential, registryEndpointRequestProperties, userAgent, httpClient);
+          eventHandlers,
+          credential,
+          registryUnauthorizedExceptionHandlerSupplier,
+          registryEndpointRequestProperties,
+          userAgent,
+          httpClient);
     }
   }
 
@@ -231,6 +253,8 @@ static Multimap<String, String> decodeTokenRepositoryGrants(String token) {
 
   private final EventHandlers eventHandlers;
   @Nullable private final Credential credential;
+  private final Supplier<RegistryUnauthorizedExceptionHandler>
+      registryUnauthorizedExceptionHandlerSupplier;
   private final RegistryEndpointRequestProperties registryEndpointRequestProperties;
   @Nullable private final String userAgent;
   private final FailoverHttpClient httpClient;
@@ -254,11 +278,14 @@ static Multimap<String, String> decodeTokenRepositoryGrants(String token) {
   private RegistryClient(
       EventHandlers eventHandlers,
       @Nullable Credential credential,
+      Supplier<RegistryUnauthorizedExceptionHandler> registryUnauthorizedExceptionHandlerSupplier,
       RegistryEndpointRequestProperties registryEndpointRequestProperties,
       @Nullable String userAgent,
       FailoverHttpClient httpClient) {
     this.eventHandlers = eventHandlers;
     this.credential = credential;
+    this.registryUnauthorizedExceptionHandlerSupplier =
+        registryUnauthorizedExceptionHandlerSupplier;
     this.registryEndpointRequestProperties = registryEndpointRequestProperties;
     this.userAgent = userAgent;
     this.httpClient = httpClient;
@@ -595,6 +622,16 @@ private static boolean isBearerAuth(@Nullable Authorization authorization) {
     return authorization != null && "bearer".equalsIgnoreCase(authorization.getScheme());
   }
 
+  /**
+   * Obtains the event handlers. This is intended to be used by the {@link
+   * RegistryUnauthorizedExceptionHandler}.
+   *
+   * @return the event
+   */
+  public EventHandlers getEventHandlers() {
+    return eventHandlers;
+  }
+
   @Nullable
   @VisibleForTesting
   String getUserAgent() {
@@ -610,7 +647,8 @@ String getUserAgent() {
    */
   private <T> T callRegistryEndpoint(RegistryEndpointProvider<T> registryEndpointProvider)
       throws IOException, RegistryException {
-    int bearerTokenRefreshes = 0;
+    Supplier<RegistryUnauthorizedExceptionHandler> handlerSupplier =
+        this.registryUnauthorizedExceptionHandlerSupplier;
     while (true) {
       try {
         return new RegistryEndpointCaller<>(
@@ -623,18 +661,37 @@ private <T> T callRegistryEndpoint(RegistryEndpointProvider<T> registryEndpointP
             .call();
 
       } catch (RegistryUnauthorizedException ex) {
+        handlerSupplier = handlerSupplier.get().handle(this, ex);
+      }
+    }
+  }
+
+  static class DefaultRegistryUnauthorizedExceptionHandler
+      implements RegistryUnauthorizedExceptionHandler {
+    int bearerTokenRefreshes = 0;
+
+    public Supplier<RegistryUnauthorizedExceptionHandler> handle(
+        final RegistryClient registryClient, final RegistryUnauthorizedException ex)
+        throws RegistryException {
+      {
         if (ex.getHttpResponseException().getStatusCode()
                 != HttpStatusCodes.STATUS_CODE_UNAUTHORIZED
-            || !isBearerAuth(authorization.get())
+            || !isBearerAuth(registryClient.authorization.get())
             || ++bearerTokenRefreshes >= MAX_BEARER_TOKEN_REFRESH_TRIES) {
           throw ex;
         }
 
         // Because we successfully did bearer authentication initially, getting 401 here probably
         // means the token was expired.
         String wwwAuthenticate = ex.getHttpResponseException().getHeaders().getAuthenticate();
-        authorization.set(refreshBearerAuth(wwwAuthenticate));
+        registryClient.authorization.set(registryClient.refreshBearerAuth(wwwAuthenticate));
+
+        return () -> this;
       }
     }
   }
+
+  public static RegistryUnauthorizedExceptionHandler defaultRegistryUnauthorizedExceptionHandler() {
+    return new DefaultRegistryUnauthorizedExceptionHandler();
+  }
 }

jib-core/src/main/java/com/google/cloud/tools/jib/registry/RegistryUnauthorizedExceptionHandler.java

@@ -0,0 +1,26 @@
+package com.google.cloud.tools.jib.registry;
+
+import com.google.cloud.tools.jib.api.RegistryException;
+import com.google.cloud.tools.jib.api.RegistryUnauthorizedException;
+import java.io.IOException;
+import java.util.function.Supplier;
+
+public interface RegistryUnauthorizedExceptionHandler {
+
+  /**
+   * Handle the exception caught by the registry client when it attempted to communicate with the
+   * registry.
+   *
+   * <p>The most obvious action is to simply throw {@code ex}. Other possible actions are to
+   * reauthenticate the client.
+   *
+   * @param registryClient the registry client which may be reconfigured
+   * @param ex the exception being handled on behalf of the client
+   * @return a supplier to use if another exception occurs on the next retry
+   * @throws IOException if an I/O error occurs
+   * @throws RegistryException if a registry error occurs
+   */
+  Supplier<RegistryUnauthorizedExceptionHandler> handle(
+      RegistryClient registryClient, RegistryUnauthorizedException ex)
+      throws IOException, RegistryException;
+}

jib-core/src/test/java/com/google/cloud/tools/jib/builder/steps/PullBaseImageStepTest.java

@@ -54,6 +54,7 @@
 import com.google.cloud.tools.jib.registry.ManifestAndDigest;
 import com.google.cloud.tools.jib.registry.RegistryClient;
 import com.google.cloud.tools.jib.registry.credentials.CredentialRetrievalException;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableSet;
 import java.io.IOException;
@@ -94,13 +95,18 @@ public void setUp() {
     Mockito.when(buildContext.getBaseImageLayersCache()).thenReturn(cache);
     Mockito.when(buildContext.newBaseImageRegistryClientFactory())
         .thenReturn(registryClientFactory);
+    Mockito.when(registryClientFactory.setCredential(Mockito.any()))
+        .thenReturn(registryClientFactory);
+    Mockito.when(registryClientFactory.setUnauthorizedExceptionHandlerSupplier(Mockito.any()))
+        .thenReturn(registryClientFactory);
     Mockito.when(registryClientFactory.newRegistryClient()).thenReturn(registryClient);
     Mockito.when(buildContext.getContainerConfiguration()).thenReturn(containerConfig);
     Mockito.when(containerConfig.getPlatforms())
         .thenReturn(ImmutableSet.of(new Platform("slim arch", "fat system")));
     Mockito.when(progressDispatcherFactory.create(Mockito.anyString(), Mockito.anyLong()))
         .thenReturn(progressDispatcher);
     Mockito.when(progressDispatcher.newChildProducer()).thenReturn(progressDispatcherFactory);
+    Mockito.when(imageConfiguration.getCredentialRetrievers()).thenReturn(ImmutableList.of());
 
     pullBaseImageStep = new PullBaseImageStep(buildContext, progressDispatcherFactory);
   }
@@ -678,6 +684,10 @@ private static RegistryClient.Factory setUpWorkingRegistryClientFactoryWithV22Ma
     manifest.setContainerConfiguration(1234, digest);
 
     RegistryClient.Factory clientFactory = Mockito.mock(RegistryClient.Factory.class);
+    Mockito.doCallRealMethod().when(clientFactory).setCredential(Mockito.any());
+    Mockito.doCallRealMethod()
+        .when(clientFactory)
+        .setUnauthorizedExceptionHandlerSupplier(Mockito.any());
     RegistryClient client = Mockito.mock(RegistryClient.class);
     Mockito.when(clientFactory.newRegistryClient()).thenReturn(client);
     Mockito.when(client.pullManifest(Mockito.any()))
@@ -709,6 +719,10 @@ private static RegistryClient.Factory setUpWorkingRegistryClientFactoryWithV22Ma
     manifest.setContainerConfiguration(1234, digest);
 
     RegistryClient.Factory clientFactory = Mockito.mock(RegistryClient.Factory.class);
+    Mockito.doCallRealMethod().when(clientFactory).setCredential(Mockito.any());
+    Mockito.doCallRealMethod()
+        .when(clientFactory)
+        .setUnauthorizedExceptionHandlerSupplier(Mockito.any());
     RegistryClient client = Mockito.mock(RegistryClient.class);
     Mockito.when(clientFactory.newRegistryClient()).thenReturn(client);
     Mockito.when(client.pullManifest(eq("sha256:aaaaaaa")))