From 2309fbbd2e8a29f6008c7ae1c9b045a2ebf764e3 Mon Sep 17 00:00:00 2001
From: Daishan <daishan@rancher.com>
Date: Fri, 2 Apr 2021 15:51:57 -0700
Subject: [PATCH] Add ability to support registryOverride

---
 pkg/podimpersonation/podimpersonation.go | 12 +++++++++---
 pkg/server/server.go                     |  3 +++
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/pkg/podimpersonation/podimpersonation.go b/pkg/podimpersonation/podimpersonation.go
index 360b03b..5bc8736 100644
--- a/pkg/podimpersonation/podimpersonation.go
+++ b/pkg/podimpersonation/podimpersonation.go
@@ -122,6 +122,7 @@ type PodOptions struct {
 	ConfigMapsToCreate []*v1.ConfigMap
 	SecretsToCreate    []*v1.Secret
 	Wait               bool
+	ImageOverride      string
 }
 
 // CreatePod will create a pod with a service account that impersonates as user. Corresponding
@@ -338,7 +339,7 @@ func (s *PodImpersonation) createPod(ctx context.Context, user user.Info, role *
 		return nil, err
 	}
 
-	pod = s.augmentPod(pod, sa)
+	pod = s.augmentPod(pod, sa, podOptions.ImageOverride)
 
 	if err := s.createConfigMaps(ctx, user, role, pod, podOptions, client); err != nil {
 		return nil, err
@@ -488,7 +489,7 @@ func (s *PodImpersonation) adminKubeConfig(user user.Info, role *rbacv1.ClusterR
 	}, nil
 }
 
-func (s *PodImpersonation) augmentPod(pod *v1.Pod, sa *v1.ServiceAccount) *v1.Pod {
+func (s *PodImpersonation) augmentPod(pod *v1.Pod, sa *v1.ServiceAccount, imageOverride string) *v1.Pod {
 	var (
 		zero = int64(0)
 		t    = true
@@ -547,9 +548,14 @@ func (s *PodImpersonation) augmentPod(pod *v1.Pod, sa *v1.ServiceAccount) *v1.Po
 		}
 	}
 
+	image := imageOverride
+	if image == "" {
+		image = s.imageName()
+	}
+
 	pod.Spec.Containers = append(pod.Spec.Containers, v1.Container{
 		Name:            "proxy",
-		Image:           s.imageName(),
+		Image:           image,
 		ImagePullPolicy: v1.PullIfNotPresent,
 		Env: []v1.EnvVar{
 			{
diff --git a/pkg/server/server.go b/pkg/server/server.go
index db491be..7e63963 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -36,6 +36,7 @@ type Server struct {
 	BaseSchemas     *types.APISchemas
 	AccessSetLookup accesscontrol.AccessSetLookup
 	APIServer       *apiserver.Server
+	ClusterRegistry string
 
 	authMiddleware      auth.Middleware
 	controllers         *Controllers
@@ -57,6 +58,7 @@ type Options struct {
 	Router                     router.RouterFunc
 	AggregationSecretNamespace string
 	AggregationSecretName      string
+	ClusterRegistry            string
 }
 
 func New(ctx context.Context, restConfig *rest.Config, opts *Options) (*Server, error) {
@@ -74,6 +76,7 @@ func New(ctx context.Context, restConfig *rest.Config, opts *Options) (*Server,
 		router:                     opts.Router,
 		aggregationSecretNamespace: opts.AggregationSecretNamespace,
 		aggregationSecretName:      opts.AggregationSecretName,
+		ClusterRegistry:            opts.ClusterRegistry,
 	}
 
 	if err := setup(ctx, server); err != nil {