Only rewrite container image registry for requested images

2 files changed, 19 insertions, 19 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImages.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImages.java
index 59dbb0b3241..583045083a9 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImages.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImages.java
@@ -37,24 +37,17 @@ public class ContainerImages {
         Optional<DockerImage> requestedImage = node.allocation()
                                                    .flatMap(allocation -> allocation.membership().cluster().dockerImageRepo());
         NodeType nodeType = node.type().isHost() ? node.type().childNodeType() : node.type();
-        final DockerImage image;
-        if (requestedImage.isPresent()) {
-            image = requestedImage.get();
-        } else if (nodeType == NodeType.tenant) {
-            if (!node.resources().gpuResources().isZero()) {
-                image = tenantGpuImage.orElseThrow(() -> new IllegalArgumentException(node + " has GPU resources, but there is no GPU container image available"));
-            } else {
-                image = tenantImage.orElse(defaultImage);
-            }
-        } else {
-            image = defaultImage;
-        }
-        return rewriteRegistry(image);
-    }
-
-    /** Rewrite the registry part of given image, using this zone's default image */
-    private DockerImage rewriteRegistry(DockerImage image) {
-        return image.withRegistry(defaultImage.registry());
+        DockerImage wantedImage =
+                nodeType != NodeType.tenant ?
+                        defaultImage :
+                node.resources().gpuResources().isZero() ?
+                        tenantImage.orElse(defaultImage) :
+                        tenantGpuImage.orElseThrow(() -> new IllegalArgumentException(node + " has GPU resources, but there is no GPU container image available"));
+
+        return requestedImage
+                // Rewrite requested images to make sure they come from a trusted registry
+                .map(image -> image.withRegistry(wantedImage.registry()))
+                .orElse(wantedImage);
     }
 
 }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java
index bb7ea52ca0e..4537aaef45b 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java
@@ -24,7 +24,7 @@ public class ContainerImagesTest {
 
     @Test
     public void image_selection() {
-        DockerImage defaultImage = DockerImage.fromString("registry.example.com/vespa/default");
+        DockerImage defaultImage = DockerImage.fromString("different.example.com/vespa/default");
         DockerImage tenantImage = DockerImage.fromString("registry.example.com/vespa/tenant");
         DockerImage gpuImage = DockerImage.fromString("registry.example.com/vespa/tenant-gpu");
         ContainerImages images = new ContainerImages(defaultImage, Optional.of(tenantImage), Optional.of(gpuImage));
@@ -45,6 +45,13 @@ public class ContainerImagesTest {
         DockerImage requested = DockerImage.fromString("registry.example.com/vespa/special");
         assertEquals(requested, images.get(node(NodeType.tenant, requested)));
 
+        // Malicious registry is rewritten to the trusted one
+        DockerImage malicious = DockerImage.fromString("malicious.example.com/vespa/special");
+        assertEquals(requested, images.get(node(NodeType.tenant, malicious)));
+
+        // Requested image registry for config is rewritten to the defaultImage registry
+        assertEquals(DockerImage.fromString("different.example.com/vespa/special"), images.get(node(NodeType.config, requested)));
+
         // When there is no custom tenant image, the default one is used
         images = new ContainerImages(defaultImage, Optional.empty(), Optional.of(gpuImage));
         assertEquals(defaultImage, images.get(node(NodeType.host)));