Add support for image registry proxy in Kibishii installation (#9063)

* Add support for image registry proxy in Kibishii installation
Signed-off-by: Priyansh Choudhary <im1706@gmail.com>

test/util/kibishii/kibishii_utils.go

@@ -323,22 +323,36 @@ func installKibishii(
 			fmt.Printf("targetKustomizeDir for windows %s\n", targetKustomizeDir)
 		}
 		fmt.Printf("The installed Kibishii Kustomize package directory is %s.\n", targetKustomizeDir)
+	}
 
-		kibishiiImage := readBaseKibishiiImage(path.Join(kibishiiDirectory, "base", "kibishii.yaml"))
-		if err := generateKibishiiImagePatch(
-			path.Join(imageRegistryProxy, kibishiiImage),
-			path.Join(targetKustomizeDir, "worker-image-patch.yaml"),
-		); err != nil {
-			return nil
-		}
+	// update kibishi images with image registry proxy if it is set
+	baseDir := resolveBasePath(kibishiiDirectory)
+	fmt.Printf("Using image registry proxy %s to patch Kibishii images. Base Dir: %s\n", imageRegistryProxy, baseDir)
 
-		jumpPadImage := readBaseJumpPadImage(path.Join(kibishiiDirectory, "base", "jump-pad.yaml"))
-		if err := generateJumpPadPatch(
-			path.Join(imageRegistryProxy, jumpPadImage),
-			path.Join(targetKustomizeDir, "jump-pad-image-patch.yaml"),
-		); err != nil {
-			return nil
-		}
+	sanitizedTargetKustomizeDir := strings.ReplaceAll(targetKustomizeDir, "overlays/sc-reclaim-policy", "")
+
+	kibishiiImage := readBaseKibishiiImage(path.Join(baseDir, "base", "kibishii.yaml"))
+	if err := generateKibishiiImagePatch(
+		path.Join(imageRegistryProxy, kibishiiImage),
+		path.Join(sanitizedTargetKustomizeDir, "worker-image-patch.yaml"),
+	); err != nil {
+		return nil
+	}
+
+	jumpPadImage := readBaseJumpPadImage(path.Join(baseDir, "base", "jump-pad.yaml"))
+	if err := generateJumpPadPatch(
+		path.Join(imageRegistryProxy, jumpPadImage),
+		path.Join(sanitizedTargetKustomizeDir, "jump-pad-image-patch.yaml"),
+	); err != nil {
+		return nil
+	}
+
+	etcdImage := readBaseEtcdImage(path.Join(baseDir, "base", "etcd.yaml"))
+	if err := generateEtcdImagePatch(
+		path.Join(imageRegistryProxy, etcdImage),
+		path.Join(sanitizedTargetKustomizeDir, "etcd-image-patch.yaml"),
+	); err != nil {
+		return nil
 	}
 
 	// We use kustomize to generate YAML for Kibishii from the checked-in yaml directories
@@ -397,6 +411,24 @@ func installKibishii(
 	return err
 }
 
+func resolveBasePath(dir string) string {
+	// If the path includes "overlays", strip everything up to "overlays/"
+	parts := strings.Split(dir, "overlays")
+	if len(parts) > 1 {
+		// Assume root of repo is before "overlays", add "/base"
+		return parts[0]
+	}
+	return dir // no "overlays" found, return original path
+}
+
+func stripRegistry(image string) string {
+	// If the image includes a registry (quay.io, docker.io, etc.), strip it
+	if parts := strings.SplitN(image, "/", 2); len(parts) == 2 && strings.Contains(parts[0], ".") {
+		return parts[1] // remove the registry
+	}
+	return image // already no registry
+}
+
 func readBaseKibishiiImage(kibishiiFilePath string) string {
 	bytes, err := os.ReadFile(kibishiiFilePath)
 	if err != nil {
@@ -435,6 +467,51 @@ func readBaseJumpPadImage(jumpPadFilePath string) string {
 	return jumpPadImage
 }
 
+func readBaseEtcdImage(etcdFilePath string) string {
+	bytes, err := os.ReadFile(etcdFilePath)
+	if err != nil {
+		fmt.Printf("Failed to read etcd pod yaml file: %v\n", err)
+		return ""
+	}
+
+	// Split on document marker
+	docs := strings.Split(string(bytes), "---")
+
+	for _, doc := range docs {
+		doc = strings.TrimSpace(doc)
+		if doc == "" {
+			continue
+		}
+
+		var typeMeta corev1api.TypedLocalObjectReference
+		if err := yaml.Unmarshal([]byte(doc), &typeMeta); err != nil {
+			continue
+		}
+
+		if typeMeta.Kind != "Pod" {
+			continue
+		}
+
+		var pod corev1api.Pod
+		if err := yaml.Unmarshal([]byte(doc), &pod); err != nil {
+			fmt.Printf("Failed to unmarshal pod: %v\n", err)
+			continue
+		}
+
+		if len(pod.Spec.Containers) > 0 {
+			fullImage := pod.Spec.Containers[0].Image
+			fmt.Printf("Full etcd image: %s\n", fullImage)
+
+			imageWithoutRegistry := stripRegistry(fullImage)
+			fmt.Printf("Stripped etcd image: %s\n", imageWithoutRegistry)
+			return imageWithoutRegistry
+		}
+	}
+
+	fmt.Println("No etcd pod with container image found.")
+	return ""
+}
+
 type patchImageData struct {
 	Image string
 }
@@ -454,11 +531,10 @@ spec:
 `
 
 	file, err := os.OpenFile(patchDirectory, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
-	defer file.Close()
-
 	if err != nil {
 		return err
 	}
+	defer file.Close()
 
 	patchTemplate, err := template.New("imagePatch").Parse(patchString)
 	if err != nil {
@@ -484,11 +560,10 @@ spec:
       image: {{.Image}}
 `
 	file, err := os.OpenFile(patchDirectory, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
-	defer file.Close()
-
 	if err != nil {
 		return err
 	}
+	defer file.Close()
 
 	patchTemplate, err := template.New("imagePatch").Parse(patchString)
 	if err != nil {
@@ -502,6 +577,54 @@ spec:
 	return nil
 }
 
+func generateEtcdImagePatch(etcdImage string, patchPath string) error {
+	patchString := `
+apiVersion: v1
+kind: Pod
+metadata:
+  name: etcd0
+spec:
+  containers:
+  - name: etcd0
+    image: {{.Image}}
+---
+apiVersion: v1
+kind: Pod
+metadata:
+  name: etcd1
+spec:
+  containers:
+  - name: etcd1
+    image: {{.Image}}
+---
+apiVersion: v1
+kind: Pod
+metadata:
+  name: etcd2
+spec:
+  containers:
+  - name: etcd2
+    image: {{.Image}}
+`
+
+	file, err := os.OpenFile(patchPath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	patchTemplate, err := template.New("imagePatch").Parse(patchString)
+	if err != nil {
+		return err
+	}
+
+	if err := patchTemplate.Execute(file, patchImageData{Image: etcdImage}); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func generateData(ctx context.Context, namespace string, kibishiiData *KibishiiData) error {
 	timeout := 30 * time.Minute
 	interval := 1 * time.Second

test/util/kibishii/kibishii_utils_test.go

@@ -0,0 +1,89 @@
+/*
+Copyright the Velero contributors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package kibishii
+
+import "testing"
+
+func TestResolveBasePath(t *testing.T) {
+	tests := []struct {
+		input    string
+		expected string
+	}{
+		{
+			input:    "/home/user/project/overlays/sc-reclaim-policy/azure",
+			expected: "/home/user/project/",
+		},
+		{
+			input:    "/go/src/github.com/org/repo/base",
+			expected: "/go/src/github.com/org/repo/base",
+		},
+		{
+			input:    "/some/dir/overlays",
+			expected: "/some/dir/",
+		},
+		{
+			input:    "overlays/sc-reclaim-policy/azure",
+			expected: "",
+		},
+	}
+
+	for _, tt := range tests {
+		actual := resolveBasePath(tt.input)
+		if actual != tt.expected {
+			t.Errorf("resolveBasePath(%q) = %q; want %q", tt.input, actual, tt.expected)
+		}
+	}
+}
+
+func TestStripRegistry(t *testing.T) {
+	tests := []struct {
+		input    string
+		expected string
+	}{
+		{
+			input:    "quay.io/example/image:tag",
+			expected: "example/image:tag",
+		},
+		{
+			input:    "docker.io/library/nginx:latest",
+			expected: "library/nginx:latest",
+		},
+		{
+			input:    "gcr.io/project/app",
+			expected: "project/app",
+		},
+		{
+			input:    "my-custom-reg.io/myapp",
+			expected: "myapp",
+		},
+		{
+			input:    "ubuntu:20.04",
+			expected: "ubuntu:20.04",
+		},
+		{
+			input:    "library/nginx",
+			expected: "library/nginx",
+		},
+	}
+
+	for _, tt := range tests {
+		actual := stripRegistry(tt.input)
+		if actual != tt.expected {
+			t.Errorf("stripRegistry(%q) = %q; want %q", tt.input, actual, tt.expected)
+		}
+	}
+}