feat: make proxy and proxy-init image configurable (#1443)

Signed-off-by: Anish Ramasekar <anish.ramasekar@gmail.com>

pkg/config/config.go

@@ -7,8 +7,10 @@ import (
 
 // Config holds configuration from the env variables
 type Config struct {
-	Cloud    string `envconfig:"AZURE_ENVIRONMENT"`
-	TenantID string `envconfig:"AZURE_TENANT_ID"`
+	Cloud          string `envconfig:"AZURE_ENVIRONMENT" default:"AzurePublicCloud"`
+	TenantID       string `envconfig:"AZURE_TENANT_ID" required:"true"`
+	ProxyImage     string `envconfig:"PROXY_IMAGE"`
+	ProxyInitImage string `envconfig:"PROXY_INIT_IMAGE"`
 }
 
 // ParseConfig parses the configuration from env variables
@@ -28,7 +30,7 @@ func ParseConfig() (*Config, error) {
 // validateConfig validates the configuration
 func validateConfig(c *Config) error {
 	if c.TenantID == "" {
-		return errors.New("tenant ID is required")
+		return errors.New("AZURE_TENANT_ID is required")
 	}
 	return nil
 }

pkg/config/config_test.go

@@ -0,0 +1,62 @@
+package config
+
+import (
+	"os"
+	"testing"
+)
+
+func TestParseConfig(t *testing.T) {
+	tests := []struct {
+		name      string
+		cloud     string
+		tenantID  string
+		wantErr   bool
+		wantCloud string
+	}{
+		{
+			name:      "cloud name defaulting to AzurePublicCloud",
+			cloud:     "",
+			tenantID:  "tenant-id",
+			wantCloud: "AzurePublicCloud",
+			wantErr:   false,
+		},
+		{
+			name:      "cloud name set to AzureChinaCloud",
+			cloud:     "AzureChinaCloud",
+			tenantID:  "tenant-id",
+			wantCloud: "AzureChinaCloud",
+			wantErr:   false,
+		},
+		{
+			name:      "missing tenant id should return error",
+			cloud:     "AzureChinaCloud",
+			tenantID:  "",
+			wantCloud: "",
+			wantErr:   true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			os.Setenv("AZURE_TENANT_ID", tt.tenantID)
+			os.Setenv("AZURE_ENVIRONMENT", tt.cloud)
+			defer func() {
+				os.Unsetenv("AZURE_TENANT_ID")
+				os.Unsetenv("AZURE_ENVIRONMENT")
+			}()
+
+			c, err := ParseConfig()
+			if (err != nil) != tt.wantErr {
+				t.Fatalf("ParseConfig() error = %v, wantErr %v", err, tt.wantErr)
+			}
+			if !tt.wantErr {
+				if c.Cloud != tt.cloud {
+					t.Errorf("ParseConfig() got = %v, want %v", c.Cloud, tt.cloud)
+				}
+				if c.TenantID != tt.tenantID {
+					t.Errorf("ParseConfig() got = %v, want %v", c.TenantID, tt.tenantID)
+				}
+			}
+		})
+	}
+}

pkg/webhook/webhook.go

@@ -50,6 +50,8 @@ type podMutator struct {
 	decoder            *admission.Decoder
 	audience           string
 	azureAuthorityHost string
+	proxyImage         string
+	proxyInitImage     string
 }
 
 // NewPodMutator returns a pod mutation handler
@@ -67,6 +69,14 @@ func NewPodMutator(client client.Client, reader client.Reader, audience string,
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to get AAD endpoint")
 	}
+	proxyImage := c.ProxyImage
+	if len(proxyImage) == 0 {
+		proxyImage = fmt.Sprintf("%s/%s:%s", ProxyImageRegistry, ProxySidecarImageName, ProxyImageVersion)
+	}
+	proxyInitImage := c.ProxyInitImage
+	if len(proxyInitImage) == 0 {
+		proxyInitImage = fmt.Sprintf("%s/%s:%s", ProxyImageRegistry, ProxyInitImageName, ProxyImageVersion)
+	}
 
 	if err := registerMetrics(); err != nil {
 		return nil, errors.Wrap(err, "failed to register metrics")
@@ -79,6 +89,8 @@ func NewPodMutator(client client.Client, reader client.Reader, audience string,
 		decoder:            admission.NewDecoder(scheme),
 		audience:           audience,
 		azureAuthorityHost: azureAuthorityHost,
+		proxyImage:         proxyImage,
+		proxyInitImage:     proxyInitImage,
 	}, nil
 }
 
@@ -188,16 +200,14 @@ func (m *podMutator) mutateContainers(containers []corev1.Container, clientID st
 }
 
 func (m *podMutator) injectProxyInitContainer(containers []corev1.Container, proxyPort int32) []corev1.Container {
-	imageRepository := strings.Join([]string{ProxyImageRegistry, ProxyInitImageName}, "/")
 	for _, container := range containers {
-		if strings.HasPrefix(container.Image, imageRepository) || container.Name == ProxyInitContainerName {
+		if container.Name == ProxyInitContainerName {
 			return containers
 		}
 	}
-
 	containers = append(containers, corev1.Container{
 		Name:            ProxyInitContainerName,
-		Image:           strings.Join([]string{imageRepository, ProxyImageVersion}, ":"),
+		Image:           m.proxyInitImage,
 		ImagePullPolicy: corev1.PullIfNotPresent,
 		SecurityContext: &corev1.SecurityContext{
 			Capabilities: &corev1.Capabilities{
@@ -218,17 +228,15 @@ func (m *podMutator) injectProxyInitContainer(containers []corev1.Container, pro
 }
 
 func (m *podMutator) injectProxySidecarContainer(containers []corev1.Container, proxyPort int32) []corev1.Container {
-	imageRepository := strings.Join([]string{ProxyImageRegistry, ProxySidecarImageName}, "/")
 	for _, container := range containers {
-		if strings.HasPrefix(container.Image, imageRepository) || container.Name == ProxySidecarContainerName {
+		if container.Name == ProxySidecarContainerName {
 			return containers
 		}
 	}
-
 	logLevel := currentLogLevel() // run the proxy at the same log level as the webhook
 	containers = append([]corev1.Container{{
 		Name:            ProxySidecarContainerName,
-		Image:           strings.Join([]string{imageRepository, ProxyImageVersion}, ":"),
+		Image:           m.proxyImage,
 		ImagePullPolicy: corev1.PullIfNotPresent,
 		Args: []string{
 			fmt.Sprintf("--proxy-port=%d", proxyPort),

pkg/webhook/webhook_test.go

@@ -1044,7 +1044,7 @@ func TestInjectProxyInitContainer(t *testing.T) {
 		},
 	}
 
-	m := &podMutator{}
+	m := &podMutator{proxyInitImage: imageURL}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			containers := m.injectProxyInitContainer(test.containers, proxyPort)
@@ -1147,7 +1147,7 @@ func TestInjectProxySidecarContainer(t *testing.T) {
 		},
 	}
 
-	m := &podMutator{}
+	m := &podMutator{proxyImage: imageURL}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			containers := m.injectProxySidecarContainer(test.containers, proxyPort)