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

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

@@ -20,6 +20,9 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"os"
+	"sort"
+	"strings"
 
 	"github.com/go-logr/logr"
 	nbv1 "github.com/kubeflow/kubeflow/components/notebook-controller/api/v1"
@@ -28,7 +31,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 +253,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 +451,110 @@ 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")
+		config, err = clientcmd.BuildConfigFromFlags("", kubeconfig)
+		if err != nil {
+			fmt.Printf("Error creating config: %v\n", err)
+			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 strings.Contains(notebook.Spec.Template.Spec.Containers[0].Image, "image-registry.openshift-image-registry.svc:5000") {
+				log.Info("Internal registry found in image, no need to look further.")
+				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
+}