Add logic on nbc-webhook to pickup internal or external registry

components/odh-notebook-controller/controllers/notebook_webhook.go

@@ -20,6 +20,10 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"os"
+	"path/filepath"
+	"sort"
+	"strings"
 
 	"github.com/go-logr/logr"
 	nbv1 "github.com/kubeflow/kubeflow/components/notebook-controller/api/v1"
@@ -28,7 +32,11 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/util/intstr"
+	"k8s.io/client-go/dynamic"
+	"k8s.io/client-go/rest"
+	"k8s.io/client-go/tools/clientcmd"
 	"k8s.io/utils/pointer"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
@@ -246,6 +254,12 @@ func (w *NotebookWebhook) Handle(ctx context.Context, req admission.Request) adm
 		if err != nil {
 			return admission.Errored(http.StatusInternalServerError, err)
 		}
+
+		// Check Imagestream Info
+		err = SetContainerImageFromRegistry(ctx, w.Client, notebook, log)
+		if err != nil {
+			return admission.Errored(http.StatusInternalServerError, err)
+		}
 	}
 
 	// Inject the OAuth proxy if the annotation is present but only if Service Mesh is disabled
@@ -438,3 +452,112 @@ func InjectCertConfig(notebook *nbv1.Notebook, configMapName string) error {
 	}
 	return nil
 }
+
+// This function checks if there is an internal registry and takes the corresponding actions to set the container.image value.
+// If an internal registry is detected, it uses the default values specified in the Notebook Custom Resource (CR).
+// Otherwise, it checks the last-image-selection annotation to find the image stream and fetches the image from status.dockerImageReference,
+// assigning it to the container.image value.
+func SetContainerImageFromRegistry(ctx context.Context, cli client.Client, notebook *nbv1.Notebook, log logr.Logger) error {
+
+	// Load kubeconfig
+	config, err := rest.InClusterConfig()
+	if err != nil {
+		kubeconfig := os.Getenv("KUBECONFIG")
+		if kubeconfig == "" {
+			kubeconfig = filepath.Join(os.Getenv("HOME"), ".kube", "config")
+		}
+		config, err = clientcmd.BuildConfigFromFlags("", kubeconfig)
+		if err != nil {
+			log.Error(err, "Error creating config")
+			return err
+		}
+	}
+
+	// Create a dynamic client
+	dynamicClient, err := dynamic.NewForConfig(config)
+	if err != nil {
+		log.Error(err, "Error creating dynamic client")
+		return err
+	}
+	// Specify the GroupVersionResource for imagestreams
+	ims := schema.GroupVersionResource{
+		Group:    "image.openshift.io",
+		Version:  "v1",
+		Resource: "imagestreams",
+	}
+
+	annotations := notebook.GetAnnotations()
+	if annotations != nil {
+		if imageSelection, exists := annotations["notebooks.opendatahub.io/last-image-selection"]; exists {
+
+			// Check if the image selection has an internal registry, if so  will pickup this. This value constructed on the initialization of the Notebook CR.
+			if strings.Contains(notebook.Spec.Template.Spec.Containers[0].Image, "image-registry.openshift-image-registry.svc:5000") {
+				log.Info("Internal registry found. Will pickup the default value from image field.")
+				return nil
+			} else {
+				// Split the imageSelection to imagestream and tag
+				parts := strings.Split(imageSelection, ":")
+				if len(parts) != 2 {
+					log.Error(nil, "Invalid image selection format")
+					return fmt.Errorf("invalid image selection format")
+				}
+
+				imagestreamName := parts[0]
+				tag := parts[1]
+
+				// Specify the namespaces to search in
+				namespaces := []string{"opendatahub", "redhat-ods-applications"}
+
+				for _, namespace := range namespaces {
+					// List imagestreams in the specified namespace
+					imagestreams, err := dynamicClient.Resource(ims).Namespace(namespace).List(ctx, metav1.ListOptions{})
+					if err != nil {
+						log.Error(err, "Cannot list imagestreams", "namespace", namespace)
+						continue
+					}
+
+					// Iterate through the imagestreams to find matches
+					for _, item := range imagestreams.Items {
+						metadata := item.Object["metadata"].(map[string]interface{})
+						name := metadata["name"].(string)
+
+						if name == imagestreamName {
+							status := item.Object["status"].(map[string]interface{})
+
+							log.Info("No Internal registry found, pick up imageHash from status.tag.dockerImageReference")
+
+							tags := status["tags"].([]interface{})
+							for _, t := range tags {
+								tagMap := t.(map[string]interface{})
+								tagName := tagMap["tag"].(string)
+								if tagName == tag {
+									items := tagMap["items"].([]interface{})
+									if len(items) > 0 {
+										// Sort items by creationTimestamp to get the most recent one
+										sort.Slice(items, func(i, j int) bool {
+											iTime := items[i].(map[string]interface{})["created"].(string)
+											jTime := items[j].(map[string]interface{})["created"].(string)
+											return iTime > jTime
+										})
+										imageHash := items[0].(map[string]interface{})["dockerImageReference"].(string)
+										notebook.Spec.Template.Spec.Containers[0].Image = imageHash
+										// Update the JUPYTER_IMAGE environment variable
+										for i, envVar := range notebook.Spec.Template.Spec.Containers[0].Env {
+											if envVar.Name == "JUPYTER_IMAGE" {
+												notebook.Spec.Template.Spec.Containers[0].Env[i].Value = imageHash
+												break
+											}
+										}
+										return nil
+									}
+								}
+							}
+						}
+					}
+				}
+			}
+		}
+	}
+
+	return nil
+}

components/odh-notebook-controller/controllers/suite_test.go

@@ -1,5 +1,4 @@
 /*
-
 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
@@ -19,7 +18,9 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"io/ioutil"
 	"net"
+	"os"
 	"path/filepath"
 	"testing"
 	"time"
@@ -83,6 +84,36 @@ var _ = BeforeSuite(func() {
 	}
 	logf.SetLogger(zap.New(zap.WriteTo(GinkgoWriter), zap.UseFlagOptions(&opts)))
 
+	// Create a temporary kubeconfig file
+	kubeconfigContent := []byte(`
+apiVersion: v1
+clusters:
+- cluster:
+    server: https://localhost:6443
+  name: local
+contexts:
+- context:
+    cluster: local
+    user: user
+  name: local
+current-context: local
+kind: Config
+preferences: {}
+users:
+- name: user
+  user:
+    token: fake-token
+`)
+	tmpDir, err := ioutil.TempDir("", "kubeconfig")
+	Expect(err).NotTo(HaveOccurred())
+
+	kubeconfigPath := filepath.Join(tmpDir, "config")
+	err = ioutil.WriteFile(kubeconfigPath, kubeconfigContent, 0644)
+	Expect(err).NotTo(HaveOccurred())
+
+	// Set the KUBECONFIG environment variable
+	os.Setenv("KUBECONFIG", kubeconfigPath)
+
 	// Initiliaze test environment:
 	// https://pkg.go.dev/sigs.k8s.io/controller-runtime/pkg/envtest#Environment.Start
 	By("Bootstrapping test environment")
@@ -98,7 +129,7 @@ var _ = BeforeSuite(func() {
 		},
 	}
 
-	cfg, err := envTest.Start()
+	cfg, err = envTest.Start()
 	Expect(err).NotTo(HaveOccurred())
 	Expect(cfg).NotTo(BeNil())