Only rewrite container image registry for requested images

1 files changed, 11 insertions, 18 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);
     }
 
 }