From 53ba1156ca98ec633e67f9af4c3d87a246c6916d Mon Sep 17 00:00:00 2001
From: Jiankun Lu <jiankun@google.com>
Date: Wed, 1 Feb 2023 18:20:42 -0800
Subject: [PATCH] [launcher] Fix image pulling in launcher

Now will only supply access token when pulling from GCP gcr.io or artifact
registries.
Also can pull images without a token from the public docker.io registry.

Signed-off-by: Jiankun Lu <jiankun@google.com>
---
 launcher/auth.go             |  7 ++++++-
 launcher/container_runner.go | 16 +++++++++-------
 launcher/launcher/main.go    |  2 +-
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/launcher/auth.go b/launcher/auth.go
index 31570014b..5d2130e3a 100644
--- a/launcher/auth.go
+++ b/launcher/auth.go
@@ -2,6 +2,7 @@ package launcher
 
 import (
 	"encoding/json"
+	"strings"
 
 	"cloud.google.com/go/compute/metadata"
 	"github.com/containerd/containerd/remotes"
@@ -31,7 +32,11 @@ func Resolver(token string) remotes.Resolver {
 	options := docker.ResolverOptions{}
 
 	credentials := func(host string) (string, string, error) {
-		return "_token", token, nil
+		// append the token if is talking to Artifact Registry or GCR Registry
+		if strings.HasSuffix(host, "docker.pkg.dev") || strings.HasSuffix(host, "gcr.io") {
+			return "_token", token, nil
+		}
+		return "", "", nil
 	}
 	authOpts := []docker.AuthorizerOpt{docker.WithAuthCreds(credentials)}
 	options.Authorizer = docker.NewDockerAuthorizer(authOpts...)
diff --git a/launcher/container_runner.go b/launcher/container_runner.go
index c95acb7c2..ad584d7b4 100644
--- a/launcher/container_runner.go
+++ b/launcher/container_runner.go
@@ -516,16 +516,18 @@ func (r *ContainerRunner) Run(ctx context.Context) error {
 }
 
 func initImage(ctx context.Context, cdClient *containerd.Client, launchSpec spec.LaunchSpec, token oauth2.Token, logger *log.Logger) (containerd.Image, error) {
-	var remoteOpt containerd.RemoteOpt
 	if token.Valid() {
-		remoteOpt = containerd.WithResolver(Resolver(token.AccessToken))
-	} else {
-		logger.Println("invalid auth token, will use empty auth")
-	}
+		remoteOpt := containerd.WithResolver(Resolver(token.AccessToken))
 
-	image, err := cdClient.Pull(ctx, launchSpec.ImageRef, containerd.WithPullUnpack, remoteOpt)
+		image, err := cdClient.Pull(ctx, launchSpec.ImageRef, containerd.WithPullUnpack, remoteOpt)
+		if err != nil {
+			return nil, fmt.Errorf("cannot pull the image: %w", err)
+		}
+		return image, nil
+	}
+	image, err := cdClient.Pull(ctx, launchSpec.ImageRef, containerd.WithPullUnpack)
 	if err != nil {
-		return nil, fmt.Errorf("cannot pull image: %w", err)
+		return nil, fmt.Errorf("cannot pull the image (no token, only works for a public image): %w", err)
 	}
 	return image, nil
 }
diff --git a/launcher/launcher/main.go b/launcher/launcher/main.go
index 8e4481e38..a912b541d 100644
--- a/launcher/launcher/main.go
+++ b/launcher/launcher/main.go
@@ -168,7 +168,7 @@ func startLauncher() error {
 
 	token, err := launcher.RetrieveAuthToken(mdsClient)
 	if err != nil {
-		logger.Printf("failed to retrieve auth token: %v, using empty auth", err)
+		logger.Printf("failed to retrieve auth token: %v, using empty auth for image pulling\n", err)
 	}
 
 	ctx := namespaces.WithNamespace(context.Background(), namespaces.Default)