From 4b07a2657a9e6d9de5d4c3fb15639cb950b502a8 Mon Sep 17 00:00:00 2001
From: Binbin Li <libinbin@microsoft.com>
Date: Wed, 3 Apr 2024 12:14:02 +0800
Subject: [PATCH] feat: add GetNamespace utils method for context
 [multi-tenancy PR 1] (#1356)

---
 httpserver/handlers.go         |  8 +++++++-
 internal/context/utils.go      |  9 +++++++++
 internal/context/utils_test.go | 37 ++++++++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/httpserver/handlers.go b/httpserver/handlers.go
index 92796ef2c..57c999bad 100644
--- a/httpserver/handlers.go
+++ b/httpserver/handlers.go
@@ -197,13 +197,19 @@ func (server *Server) mutate(ctx context.Context, w http.ResponseWriter, r *http
 				returnItem.Error = err.Error()
 				return
 			}
-			parsedReference, err := pkgUtils.ParseSubjectReference(image)
+			requestKey, err := pkgUtils.ParseRequestKey(image)
+			if err != nil {
+				returnItem.Error = err.Error()
+				return
+			}
+			parsedReference, err := pkgUtils.ParseSubjectReference(requestKey.Subject)
 			if err != nil {
 				err = errors.ErrorCodeReferenceInvalid.WithError(err).WithDetail(fmt.Sprintf("failed to parse image reference %s", image))
 				logger.GetLogger(ctx, server.LogOption).Error(err)
 				returnItem.Error = err.Error()
 				return
 			}
+			ctx = ctxUtils.SetContextWithNamespace(ctx, requestKey.Namespace)
 
 			if parsedReference.Digest == "" {
 				var selectedStore referrerstore.ReferrerStore
diff --git a/internal/context/utils.go b/internal/context/utils.go
index deb764d27..39c8ae5b7 100644
--- a/internal/context/utils.go
+++ b/internal/context/utils.go
@@ -29,6 +29,15 @@ func SetContextWithNamespace(ctx context.Context, namespace string) context.Cont
 	return context.WithValue(ctx, contextKeyNamespace, namespace)
 }
 
+// GetNamespace returns the embedded namespace from the context.
+func GetNamespace(ctx context.Context) string {
+	namespace := ctx.Value(contextKeyNamespace)
+	if namespace == nil {
+		return ""
+	}
+	return namespace.(string)
+}
+
 // CreateCacheKey creates a new cache key prefixed with embedded namespace.
 func CreateCacheKey(ctx context.Context, key string) string {
 	namespace := ctx.Value(contextKeyNamespace)
diff --git a/internal/context/utils_test.go b/internal/context/utils_test.go
index a446a04b7..1b2d94a0a 100644
--- a/internal/context/utils_test.go
+++ b/internal/context/utils_test.go
@@ -78,3 +78,40 @@ func TestCreateCacheKey(t *testing.T) {
 		})
 	}
 }
+
+func TestGetNamespace(t *testing.T) {
+	testCases := []struct {
+		name         string
+		namespaceSet bool
+		namespace    string
+	}{
+		{
+			name:         "no namespace",
+			namespaceSet: false,
+		},
+		{
+			name:         "empty namespace",
+			namespaceSet: true,
+			namespace:    "",
+		},
+		{
+			name:         "with namespace",
+			namespaceSet: true,
+			namespace:    testNamespace,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := context.Background()
+			if tc.namespaceSet {
+				ctx = SetContextWithNamespace(ctx, tc.namespace)
+			}
+
+			namespace := GetNamespace(ctx)
+			if namespace != tc.namespace {
+				t.Fatalf("expected namespace %s, got %s", tc.namespace, namespace)
+			}
+		})
+	}
+}