From 62e55fd58a6414355f37836adb39d8472b1bc448 Mon Sep 17 00:00:00 2001
From: Derek McGowan <derek@mcg.dev>
Date: Wed, 13 Sep 2023 16:14:24 -0700
Subject: [PATCH] Add support for registry token in containerd pull logic

When registry token is provided, the authorization header can be
directly applied to the registry request. No other type of
authorization will be attempted when the registry token is provided.

Signed-off-by: Derek McGowan <derek@mcg.dev>
---
 daemon/containerd/resolver.go | 40 +++++++++++++++++++++++++++++++----
 1 file changed, 36 insertions(+), 4 deletions(-)

diff --git a/daemon/containerd/resolver.go b/daemon/containerd/resolver.go
index 7f7ee134560b0..306fafa06068c 100644
--- a/daemon/containerd/resolver.go
+++ b/daemon/containerd/resolver.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"net/http"
 
+	cerrdefs "github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/log"
 	"github.com/containerd/containerd/remotes"
 	"github.com/containerd/containerd/remotes/docker"
@@ -33,7 +34,7 @@ func (i *ImageService) newResolverFromAuthConfig(ctx context.Context, authConfig
 func hostsWrapper(hostsFn docker.RegistryHosts, optAuthConfig *registrytypes.AuthConfig, regService RegistryConfigProvider) docker.RegistryHosts {
 	var authorizer docker.Authorizer
 	if optAuthConfig != nil {
-		authorizer = docker.NewDockerAuthorizer(authorizationCredsFromAuthConfig(*optAuthConfig))
+		authorizer = authorizerFromAuthConfig(*optAuthConfig)
 	}
 
 	return func(n string) ([]docker.RegistryHost, error) {
@@ -55,13 +56,20 @@ func hostsWrapper(hostsFn docker.RegistryHosts, optAuthConfig *registrytypes.Aut
 	}
 }
 
-func authorizationCredsFromAuthConfig(authConfig registrytypes.AuthConfig) docker.AuthorizerOpt {
+func authorizerFromAuthConfig(authConfig registrytypes.AuthConfig) docker.Authorizer {
 	cfgHost := registry.ConvertToHostname(authConfig.ServerAddress)
 	if cfgHost == "" || cfgHost == registry.IndexHostname {
 		cfgHost = registry.DefaultRegistryHost
 	}
 
-	return docker.WithAuthCreds(func(host string) (string, string, error) {
+	if authConfig.RegistryToken != "" {
+		return &bearerAuthorizer{
+			host:   cfgHost,
+			bearer: authConfig.RegistryToken,
+		}
+	}
+
+	return docker.NewDockerAuthorizer(docker.WithAuthCreds(func(host string) (string, string, error) {
 		if cfgHost != host {
 			log.G(context.TODO()).WithFields(log.Fields{
 				"host":    host,
@@ -73,7 +81,31 @@ func authorizationCredsFromAuthConfig(authConfig registrytypes.AuthConfig) docke
 			return "", authConfig.IdentityToken, nil
 		}
 		return authConfig.Username, authConfig.Password, nil
-	})
+	}))
+}
+
+type bearerAuthorizer struct {
+	host   string
+	bearer string
+}
+
+func (a *bearerAuthorizer) Authorize(ctx context.Context, req *http.Request) error {
+	if req.Host != a.host {
+		log.G(ctx).WithFields(log.Fields{
+			"host":    req.Host,
+			"cfgHost": a.host,
+		}).Warn("Host doesn't match for bearer token")
+		return nil
+	}
+
+	req.Header.Set("Authorization", "Bearer "+a.bearer)
+
+	return nil
+}
+
+func (a *bearerAuthorizer) AddResponses(context.Context, []*http.Response) error {
+	// Return not implemented to prevent retry of the request when bearer did not succeed
+	return cerrdefs.ErrNotImplemented
 }
 
 type httpFallback struct {