From 828f3c12ba843ed72200e9b96db3669982d85e34 Mon Sep 17 00:00:00 2001
From: Joao Pereira <joaod@vmware.com>
Date: Tue, 17 Jan 2023 14:14:05 -0600
Subject: [PATCH] Enable imgpkg keychains when environment variable is provided

Signed-off-by: Joao Pereira <joaod@vmware.com>
---
 pkg/vendir/fetch/image/imgpkg.go      | 21 +++++++++++
 pkg/vendir/fetch/image/imgpkg_test.go | 54 +++++++++++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/pkg/vendir/fetch/image/imgpkg.go b/pkg/vendir/fetch/image/imgpkg.go
index 294d8f18..f6cf7074 100644
--- a/pkg/vendir/fetch/image/imgpkg.go
+++ b/pkg/vendir/fetch/image/imgpkg.go
@@ -7,10 +7,12 @@ import (
 	"bytes"
 	"fmt"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/google/go-containerregistry/pkg/name"
 	"github.com/vmware-tanzu/carvel-imgpkg/pkg/imgpkg/registry"
+	"github.com/vmware-tanzu/carvel-imgpkg/pkg/imgpkg/registry/auth"
 	"github.com/vmware-tanzu/carvel-imgpkg/pkg/imgpkg/v1"
 	ctlconf "github.com/vmware-tanzu/carvel-vendir/pkg/vendir/config"
 	ctlfetch "github.com/vmware-tanzu/carvel-vendir/pkg/vendir/fetch"
@@ -154,11 +156,30 @@ func (t *Imgpkg) RegistryOpts() (registry.Opts, error) {
 		return registry.Opts{}, err
 	}
 
+	var activeKeychains []auth.IAASKeychain
+	for _, envVar := range t.opts.EnvironFunc() {
+		if strings.HasPrefix(envVar, "IMGPKG_ACTIVE_KEYCHAINS") {
+			keychains := strings.SplitN(envVar, "=", 2)
+			if len(keychains) != 2 {
+				return registry.Opts{}, fmt.Errorf("Expected 'IMGPKG_ACTIVE_KEYCHAINS' environment variable to have a list of keychains but got '%s'", envVar)
+			}
+
+			if strings.Contains(keychains[1], ",") {
+				for _, keychainName := range strings.Split(keychains[1], ",") {
+					activeKeychains = append(activeKeychains, auth.IAASKeychain(strings.TrimSpace(keychainName)))
+				}
+			} else {
+				activeKeychains = append(activeKeychains, auth.IAASKeychain(strings.TrimSpace(keychains[1])))
+			}
+		}
+	}
+
 	return registry.Opts{
 		VerifyCerts:           !t.opts.DangerousSkipTLSVerify,
 		Insecure:              false,
 		ResponseHeaderTimeout: 30 * time.Second,
 		RetryCount:            5,
+		ActiveKeychains:       activeKeychains,
 		EnvironFunc: func() []string {
 			return append(envVariables, t.opts.EnvironFunc()...)
 		},
diff --git a/pkg/vendir/fetch/image/imgpkg_test.go b/pkg/vendir/fetch/image/imgpkg_test.go
index 9df10564..5b670fcb 100644
--- a/pkg/vendir/fetch/image/imgpkg_test.go
+++ b/pkg/vendir/fetch/image/imgpkg_test.go
@@ -20,6 +20,7 @@ import (
 	"github.com/phayes/freeport"
 	"github.com/stretchr/testify/require"
 	ctlregistry "github.com/vmware-tanzu/carvel-imgpkg/pkg/imgpkg/registry"
+	"github.com/vmware-tanzu/carvel-imgpkg/pkg/imgpkg/registry/auth"
 	ctlconf "github.com/vmware-tanzu/carvel-vendir/pkg/vendir/config"
 	ctlfetch "github.com/vmware-tanzu/carvel-vendir/pkg/vendir/fetch"
 	ctlcache "github.com/vmware-tanzu/carvel-vendir/pkg/vendir/fetch/cache"
@@ -139,6 +140,59 @@ func TestImgpkgAuth(t *testing.T) {
 
 		requireImgpkgEnv(t, nil, opts.EnvironFunc())
 	})
+
+	t.Run("enable keychain auth with list of keychains", func(t *testing.T) {
+		cache, err := ctlcache.NewCache("", "1Mi")
+		require.NoError(t, err)
+
+		imgpkg := ctlimg.NewImgpkg(
+			ctlimg.ImgpkgOpts{
+				EnvironFunc: func() []string {
+					return []string{"IMGPKG_ACTIVE_KEYCHAINS=gcr,ecr"}
+				},
+			},
+			ctlfetch.SingleSecretRefFetcher{},
+			cache,
+		)
+
+		opts, err := imgpkg.RegistryOpts()
+		require.NoError(t, err)
+		require.Equal(t, []auth.IAASKeychain{"gcr", "ecr"}, opts.ActiveKeychains)
+	})
+
+	t.Run("enable keychain auth with single keychain", func(t *testing.T) {
+		cache, err := ctlcache.NewCache("", "1Mi")
+		require.NoError(t, err)
+
+		imgpkg := ctlimg.NewImgpkg(
+			ctlimg.ImgpkgOpts{
+				EnvironFunc: func() []string {
+					return []string{"IMGPKG_ACTIVE_KEYCHAINS=single"}
+				},
+			},
+			ctlfetch.SingleSecretRefFetcher{},
+			cache,
+		)
+
+		opts, err := imgpkg.RegistryOpts()
+		require.NoError(t, err)
+		require.Equal(t, []auth.IAASKeychain{"single"}, opts.ActiveKeychains)
+	})
+
+	t.Run("no keychain enable when environment variable not set", func(t *testing.T) {
+		cache, err := ctlcache.NewCache("", "1Mi")
+		require.NoError(t, err)
+
+		imgpkg := ctlimg.NewImgpkg(
+			ctlimg.ImgpkgOpts{},
+			ctlfetch.SingleSecretRefFetcher{},
+			cache,
+		)
+
+		opts, err := imgpkg.RegistryOpts()
+		require.NoError(t, err)
+		require.Nil(t, opts.ActiveKeychains)
+	})
 }
 
 func TestImgpkgCache(t *testing.T) {