Add registry mirror support in Jib build plugins

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

@@ -56,8 +56,10 @@
 import com.google.common.collect.ImmutableList;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Callable;
@@ -96,7 +98,7 @@ public ImagesAndRegistryClient call()
           CacheCorruptedException, CredentialRetrievalException {
     EventHandlers eventHandlers = buildContext.getEventHandlers();
     try (ProgressEventDispatcher progressDispatcher =
-            progressDispatcherFactory.create("pulling base image manifest", 3);
+            progressDispatcherFactory.create("pulling base image manifest", 4);
         TimerEventDispatcher ignored1 = new TimerEventDispatcher(eventHandlers, DESCRIPTION)) {
 
       // Skip this step if this is a scratch image
@@ -138,6 +140,12 @@ public ImagesAndRegistryClient call()
         }
       }
 
+      Optional<ImagesAndRegistryClient> mirrorPull =
+          tryMirrors(buildContext, progressDispatcher.newChildProducer());
+      if (mirrorPull.isPresent()) {
+        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
@@ -199,6 +207,60 @@ public ImagesAndRegistryClient call()
     }
   }
 
+  @VisibleForTesting
+  Optional<ImagesAndRegistryClient> tryMirrors(
+      BuildContext buildContext, ProgressEventDispatcher.Factory progressDispatcherFactory)
+      throws LayerCountMismatchException, BadContainerConfigurationFormatException {
+    EventHandlers eventHandlers = buildContext.getEventHandlers();
+
+    Collection<Map.Entry<String, String>> mirrorEntries =
+        buildContext.getRegistryMirrors().entries();
+    try (ProgressEventDispatcher progressDispatcher1 =
+            progressDispatcherFactory.create("trying mirrors", mirrorEntries.size());
+        TimerEventDispatcher ignored1 = new TimerEventDispatcher(eventHandlers, "trying mirrors")) {
+      for (Map.Entry<String, String> entry : mirrorEntries) {
+        String registry = entry.getKey();
+        String mirror = entry.getValue();
+        eventHandlers.dispatch(LogEvent.debug("mirror config: " + registry + " --> " + mirror));
+
+        if (!buildContext.getBaseImageConfiguration().getImageRegistry().equals(registry)) {
+          progressDispatcher1.dispatchProgress(1);
+          continue;
+        }
+
+        eventHandlers.dispatch(LogEvent.info("trying mirror " + mirror + " for the base image"));
+        try (ProgressEventDispatcher progressDispatcher2 =
+            progressDispatcher1.newChildProducer().create("trying mirror " + mirror, 2)) {
+          // First, try with no credentials. This works with public GCR images.
+          RegistryClient registryClient =
+              buildContext.newBaseImageRegistryClientFactory(mirror).newRegistryClient();
+          try {
+            List<Image> images =
+                pullBaseImages(registryClient, progressDispatcher2.newChildProducer());
+            eventHandlers.dispatch(LogEvent.info("pulled manifest from mirror " + mirror));
+            return Optional.of(new ImagesAndRegistryClient(images, registryClient));
+
+          } catch (RegistryUnauthorizedException ex) {
+            // in case if a mirror requires bearer auth
+            eventHandlers.dispatch(LogEvent.debug("mirror " + mirror + " requires auth"));
+            registryClient.doPullBearerAuth();
+            List<Image> images =
+                pullBaseImages(registryClient, progressDispatcher2.newChildProducer());
+            eventHandlers.dispatch(LogEvent.info("pulled manifest from mirror " + mirror));
+            return Optional.of(new ImagesAndRegistryClient(images, registryClient));
+          }
+
+        } catch (IOException | RegistryException ex) {
+          // Ignore errors from this mirror and continue.
+          eventHandlers.dispatch(
+              LogEvent.debug(
+                  "failed to get manifest from mirror " + mirror + ": " + ex.getMessage()));
+        }
+      }
+      return Optional.empty();
+    }
+  }
+
   /**
    * Pulls the base images specified in the platforms list.
    *
@@ -350,6 +412,7 @@ private ContainerConfigurationTemplate pullContainerConfigJson(
         new ThrottledProgressEventDispatcherWrapper(
             progressDispatcherFactory,
             "pull container configuration " + manifest.getContainerConfiguration().getDigest())) {
+
       String containerConfigString =
           Blobs.writeToString(
               registryClient.pullBlob(

jib-core/src/main/java/com/google/cloud/tools/jib/configuration/BuildContext.java

@@ -527,12 +527,26 @@ public ImmutableListMultimap<String, String> getRegistryMirrors() {
 
   /**
    * Creates a new {@link RegistryClient.Factory} for the base image with fields from the build
-   * configuration.
+   * configuration. The server URL is derived from the base {@link ImageConfiguration}.
    *
    * @return a new {@link RegistryClient.Factory}
    */
   public RegistryClient.Factory newBaseImageRegistryClientFactory() {
-    return newRegistryClientFactory(baseImageConfiguration);
+    return newBaseImageRegistryClientFactory(baseImageConfiguration.getImageRegistry());
+  }
+
+  /**
+   * Creates a new {@link RegistryClient.Factory} for the base image repository on the registry
+   * {@code serverUrl}. Compared to @link #newBaseImageRegistryClientFactory()), this method is
+   * useful to try a mirror.
+   *
+   * @param serverUrl the server URL for the registry (for example, {@code gcr.io})
+   * @return a new {@link RegistryClient.Factory}
+   */
+  public RegistryClient.Factory newBaseImageRegistryClientFactory(String serverUrl) {
+    return RegistryClient.factory(
+            getEventHandlers(), serverUrl, baseImageConfiguration.getImageRepository(), httpClient)
+        .setUserAgent(makeUserAgent());
   }
 
   /**
@@ -554,14 +568,10 @@ public RegistryClient.Factory newTargetImageRegistryClientFactory() {
               httpClient)
           .setUserAgent(makeUserAgent());
     }
-    return newRegistryClientFactory(targetImageConfiguration);
-  }
-
-  private RegistryClient.Factory newRegistryClientFactory(ImageConfiguration imageConfiguration) {
     return RegistryClient.factory(
             getEventHandlers(),
-            imageConfiguration.getImageRegistry(),
-            imageConfiguration.getImageRepository(),
+            targetImageConfiguration.getImageRegistry(),
+            targetImageConfiguration.getImageRepository(),
             httpClient)
         .setUserAgent(makeUserAgent());
   }

jib-core/src/main/java/com/google/cloud/tools/jib/image/json/BuildableManifestTemplate.java

@@ -40,7 +40,6 @@ public interface BuildableManifestTemplate extends ManifestTemplate {
    * @see <a href="https://github.com/opencontainers/image-spec/blob/master/descriptor.md">OCI
    *     Content Descriptors</a>
    */
-  @VisibleForTesting
   class ContentDescriptorTemplate implements JsonTemplate {
 
     @SuppressWarnings("unused")

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

@@ -19,8 +19,10 @@
 import com.google.cloud.tools.jib.api.DescriptorDigest;
 import com.google.cloud.tools.jib.api.ImageReference;
 import com.google.cloud.tools.jib.api.InvalidImageReferenceException;
+import com.google.cloud.tools.jib.api.LogEvent;
 import com.google.cloud.tools.jib.api.RegistryException;
 import com.google.cloud.tools.jib.api.buildplan.Platform;
+import com.google.cloud.tools.jib.blob.Blobs;
 import com.google.cloud.tools.jib.builder.ProgressEventDispatcher;
 import com.google.cloud.tools.jib.builder.steps.PullBaseImageStep.ImagesAndRegistryClient;
 import com.google.cloud.tools.jib.cache.Cache;
@@ -41,18 +43,22 @@
 import com.google.cloud.tools.jib.image.json.V22ManifestListTemplate;
 import com.google.cloud.tools.jib.image.json.V22ManifestTemplate;
 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.credentials.CredentialRetrievalException;
+import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableSet;
 import java.io.IOException;
 import java.security.DigestException;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.Consumer;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.InOrder;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.junit.MockitoJUnitRunner;
@@ -62,8 +68,10 @@
 public class PullBaseImageStepTest {
 
   @Mock private ProgressEventDispatcher.Factory progressDispatcherFactory;
+  @Mock private ProgressEventDispatcher progressDispatcher;
   @Mock private BuildContext buildContext;
   @Mock private RegistryClient registryClient;
+  @Mock private RegistryClient.Factory registryClientFactory;
   @Mock private ImageConfiguration imageConfiguration;
   @Mock private ContainerConfiguration containerConfig;
   @Mock private Cache cache;
@@ -76,13 +84,15 @@ public void setUp() {
     Mockito.when(buildContext.getBaseImageConfiguration()).thenReturn(imageConfiguration);
     Mockito.when(buildContext.getEventHandlers()).thenReturn(eventHandlers);
     Mockito.when(buildContext.getBaseImageLayersCache()).thenReturn(cache);
-    RegistryClient.Factory registryClientFactory = Mockito.mock(RegistryClient.Factory.class);
     Mockito.when(buildContext.newBaseImageRegistryClientFactory())
         .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);
 
     pullBaseImageStep = new PullBaseImageStep(buildContext, progressDispatcherFactory);
   }
@@ -409,4 +419,142 @@ public void testGetCachedBaseImages_v22ManifestListCached_onlyPlatforms()
     Assert.assertEquals(1, images.size());
     Assert.assertEquals("target-user", images.get(0).getUser());
   }
+
+  @Test
+  public void testTryMirrors_noMatchingMirrors()
+      throws LayerCountMismatchException, BadContainerConfigurationFormatException {
+    Mockito.when(imageConfiguration.getImageRegistry()).thenReturn("registry");
+    Mockito.when(buildContext.getRegistryMirrors())
+        .thenReturn(ImmutableListMultimap.of("unmatched1", "mirror1", "unmatched2", "mirror2"));
+
+    Optional<ImagesAndRegistryClient> result =
+        pullBaseImageStep.tryMirrors(buildContext, progressDispatcherFactory);
+    Assert.assertEquals(Optional.empty(), result);
+
+    InOrder inOrder = Mockito.inOrder(eventHandlers);
+    inOrder.verify(eventHandlers).dispatch(LogEvent.debug("mirror config: unmatched1 --> mirror1"));
+    inOrder.verify(eventHandlers).dispatch(LogEvent.debug("mirror config: unmatched2 --> mirror2"));
+    Mockito.verify(buildContext, Mockito.never()).newBaseImageRegistryClientFactory(Mockito.any());
+  }
+
+  @Test
+  public void testTryMirrors_mirrorIoError()
+      throws LayerCountMismatchException, BadContainerConfigurationFormatException, IOException,
+          RegistryException {
+    Mockito.when(imageConfiguration.getImageRegistry()).thenReturn("registry");
+    Mockito.when(buildContext.getRegistryMirrors())
+        .thenReturn(ImmutableListMultimap.of("registry", "gcr.io"));
+    Mockito.when(buildContext.newBaseImageRegistryClientFactory("gcr.io"))
+        .thenReturn(registryClientFactory);
+    Mockito.when(registryClient.pullManifest(Mockito.any()))
+        .thenThrow(new IOException("test exception"));
+
+    Optional<ImagesAndRegistryClient> result =
+        pullBaseImageStep.tryMirrors(buildContext, progressDispatcherFactory);
+    Assert.assertEquals(Optional.empty(), result);
+
+    InOrder inOrder = Mockito.inOrder(eventHandlers);
+    inOrder
+        .verify(eventHandlers)
+        .dispatch(LogEvent.info("trying mirror gcr.io for the base image"));
+    inOrder
+        .verify(eventHandlers)
+        .dispatch(LogEvent.debug("failed to get manifest from mirror gcr.io: test exception"));
+  }
+
+  @Test
+  public void testTryMirrors_multipleMirrors()
+      throws LayerCountMismatchException, BadContainerConfigurationFormatException, IOException,
+          RegistryException, InvalidImageReferenceException {
+    Mockito.when(imageConfiguration.getImage()).thenReturn(ImageReference.parse("registry/repo"));
+    Mockito.when(imageConfiguration.getImageRegistry()).thenReturn("registry");
+    Mockito.when(buildContext.getRegistryMirrors())
+        .thenReturn(ImmutableListMultimap.of("registry", "quay.io", "registry", "gcr.io"));
+
+    Mockito.when(buildContext.newBaseImageRegistryClientFactory("quay.io"))
+        .thenReturn(registryClientFactory);
+    Mockito.when(registryClient.pullManifest(Mockito.any()))
+        .thenThrow(new RegistryException("not found"));
+
+    RegistryClient.Factory gcrRegistryClientFactory = setUpWorkingRegistryClientFactory();
+    Mockito.when(buildContext.newBaseImageRegistryClientFactory("gcr.io"))
+        .thenReturn(gcrRegistryClientFactory);
+
+    Optional<ImagesAndRegistryClient> result =
+        pullBaseImageStep.tryMirrors(buildContext, progressDispatcherFactory);
+    Assert.assertTrue(result.isPresent());
+    Assert.assertEquals(gcrRegistryClientFactory.newRegistryClient(), result.get().registryClient);
+
+    InOrder inOrder = Mockito.inOrder(eventHandlers);
+    inOrder
+        .verify(eventHandlers)
+        .dispatch(LogEvent.info("trying mirror quay.io for the base image"));
+    inOrder
+        .verify(eventHandlers)
+        .dispatch(LogEvent.debug("failed to get manifest from mirror quay.io: not found"));
+    inOrder
+        .verify(eventHandlers)
+        .dispatch(LogEvent.info("trying mirror gcr.io for the base image"));
+    inOrder.verify(eventHandlers).dispatch(LogEvent.info("pulled manifest from mirror gcr.io"));
+  }
+
+  @Test
+  public void testCall_allMirrorsFail()
+      throws InvalidImageReferenceException, IOException, RegistryException,
+          LayerPropertyNotFoundException, LayerCountMismatchException,
+          BadContainerConfigurationFormatException, CacheCorruptedException,
+          CredentialRetrievalException {
+    Mockito.when(imageConfiguration.getImage()).thenReturn(ImageReference.parse("registry/repo"));
+    Mockito.when(imageConfiguration.getImageRegistry()).thenReturn("registry");
+    Mockito.when(buildContext.getRegistryMirrors())
+        .thenReturn(ImmutableListMultimap.of("registry", "quay.io", "registry", "gcr.io"));
+
+    Mockito.when(buildContext.newBaseImageRegistryClientFactory(Mockito.any()))
+        .thenReturn(registryClientFactory);
+    Mockito.when(registryClient.pullManifest(Mockito.any()))
+        .thenThrow(new RegistryException("not found"));
+
+    RegistryClient.Factory dockerHubRegistryClientFactory = setUpWorkingRegistryClientFactory();
+    Mockito.when(buildContext.newBaseImageRegistryClientFactory())
+        .thenReturn(dockerHubRegistryClientFactory);
+
+    ImagesAndRegistryClient result = pullBaseImageStep.call();
+    Assert.assertEquals(dockerHubRegistryClientFactory.newRegistryClient(), result.registryClient);
+
+    InOrder inOrder = Mockito.inOrder(eventHandlers);
+    inOrder
+        .verify(eventHandlers)
+        .dispatch(LogEvent.info("trying mirror quay.io for the base image"));
+    inOrder
+        .verify(eventHandlers)
+        .dispatch(LogEvent.debug("failed to get manifest from mirror quay.io: not found"));
+    inOrder
+        .verify(eventHandlers)
+        .dispatch(LogEvent.info("trying mirror gcr.io for the base image"));
+    inOrder
+        .verify(eventHandlers)
+        .dispatch(LogEvent.debug("failed to get manifest from mirror gcr.io: not found"));
+  }
+
+  private static RegistryClient.Factory setUpWorkingRegistryClientFactory()
+      throws IOException, RegistryException {
+    DescriptorDigest digest = Mockito.mock(DescriptorDigest.class);
+    V22ManifestTemplate manifest = new V22ManifestTemplate();
+    manifest.setContainerConfiguration(1234, digest);
+
+    RegistryClient.Factory clientFactory = Mockito.mock(RegistryClient.Factory.class);
+    RegistryClient client = Mockito.mock(RegistryClient.class);
+    Mockito.when(clientFactory.newRegistryClient()).thenReturn(client);
+    Mockito.when(client.pullManifest(Mockito.any()))
+        .thenReturn(new ManifestAndDigest<>(manifest, digest));
+    // mocking pulling container config json
+    Mockito.when(client.pullBlob(Mockito.any(), Mockito.any(), Mockito.any()))
+        .then(
+            invocation -> {
+              Consumer<Long> blobSizeListener = invocation.getArgument(1);
+              blobSizeListener.accept(1L);
+              return Blobs.from("{}");
+            });
+    return clientFactory;
+  }
 }

jib-gradle-plugin/src/main/java/com/google/cloud/tools/jib/gradle/BuildDockerTask.java

@@ -116,6 +116,7 @@ public void buildDocker()
               new GradleRawConfiguration(jibExtension),
               ignored -> java.util.Optional.empty(),
               projectProperties,
+              globalConfig,
               new GradleHelpfulSuggestions(HELPFUL_SUGGESTIONS_PREFIX))
           .runBuild();
 

jib-gradle-plugin/src/main/java/com/google/cloud/tools/jib/gradle/BuildTarTask.java

@@ -136,6 +136,7 @@ public void buildTar()
               new GradleRawConfiguration(jibExtension),
               ignored -> Optional.empty(),
               projectProperties,
+              globalConfig,
               new GradleHelpfulSuggestions(HELPFUL_SUGGESTIONS_PREFIX))
           .runBuild();
 

jib-maven-plugin/src/main/java/com/google/cloud/tools/jib/maven/BuildDockerMojo.java

@@ -88,14 +88,15 @@ public void execute() throws MojoExecutionException, MojoFailureException {
 
     Future<Optional<String>> updateCheckFuture = Futures.immediateFuture(Optional.empty());
     try {
-      updateCheckFuture =
-          MojoCommon.newUpdateChecker(projectProperties, GlobalConfig.readConfig(), getLog());
+      GlobalConfig globalConfig = GlobalConfig.readConfig();
+      updateCheckFuture = MojoCommon.newUpdateChecker(projectProperties, globalConfig, getLog());
 
       PluginConfigurationProcessor.createJibBuildRunnerForDockerDaemonImage(
               new MavenRawConfiguration(this),
               new MavenSettingsServerCredentials(
                   getSession().getSettings(), getSettingsDecrypter()),
               projectProperties,
+              globalConfig,
               new MavenHelpfulSuggestions(HELPFUL_SUGGESTIONS_PREFIX))
           .runBuild();