From a3b49f371ef13d3bc8b37f508f30296fb821ae5a Mon Sep 17 00:00:00 2001
From: Joe Lanford <joe.lanford@gmail.com>
Date: Thu, 17 Oct 2024 16:40:35 -0400
Subject: [PATCH] pkg/storage: handle more possible release label k/v pairs
 (#399)

* pkg/storage: handle more possible release label k/v pairs

Release labels are stored only in Kubernetes object metadata. In order
to support a wider variety of release labels, check release label
validity. If a release label is invalid for a kubernetes metadata.label
store it as an annotation instead.

Signed-off-by: Joe Lanford <joe.lanford@gmail.com>

* use a release wrapper to serialize custom labels into the release blob

Signed-off-by: Joe Lanford <joe.lanford@gmail.com>

---------

Signed-off-by: Joe Lanford <joe.lanford@gmail.com>
---
 pkg/storage/chunked.go | 60 +++++++++++++++++++++++++++++++++---------
 pkg/storage/labels.go  |  2 --
 2 files changed, 48 insertions(+), 14 deletions(-)

diff --git a/pkg/storage/chunked.go b/pkg/storage/chunked.go
index 4ebecbec..be411e66 100644
--- a/pkg/storage/chunked.go
+++ b/pkg/storage/chunked.go
@@ -94,6 +94,18 @@ type chunk struct {
 	data []byte
 }
 
+type releaseWrapper struct {
+	release.Release
+	Labels map[string]string `json:"labels"`
+}
+
+func wrapRelease(rls *release.Release) *releaseWrapper {
+	return &releaseWrapper{
+		Release: *rls,
+		Labels:  rls.Labels,
+	}
+}
+
 // encodeRelease encodes a release returning a base64 encoded
 // gzipped string representation, or error.
 func (c *chunkedSecrets) encodeReleaseAsChunks(key string, rls *release.Release) ([]chunk, error) {
@@ -105,7 +117,7 @@ func (c *chunkedSecrets) encodeReleaseAsChunks(key string, rls *release.Release)
 			return err
 		}
 		defer gzw.Close()
-		return json.NewEncoder(gzw).Encode(rls)
+		return json.NewEncoder(gzw).Encode(wrapRelease(rls))
 	}(); err != nil {
 		return nil, err
 	}
@@ -318,29 +330,51 @@ func (c *chunkedSecrets) Query(queryLabels map[string]string) ([]*release.Releas
 	c.Log("query: labels=%v", queryLabels)
 	defer c.Log("queried: labels=%v", queryLabels)
 
-	selector := newListIndicesLabelSelector(c.owner)
-	if queryRequirements, selectable := labels.Set(queryLabels).AsSelector().Requirements(); selectable {
-		selector = selector.Add(queryRequirements...)
+	// The only labels that get stored on the index secret are system labels, so we'll do a two-pass
+	// query. First, we'll request index secrets from the API server that match the query labels that
+	// are system labels. From there, we decode the releases that match, and then further filter those
+	// based on the rest of the query labels that are not system labels.
+	serverSelectorSet := labels.Set{}
+	clientSelectorSet := labels.Set{}
+	for k, v := range queryLabels {
+		if isSystemLabel(k) {
+			serverSelectorSet[k] = v
+		} else {
+			clientSelectorSet[k] = v
+		}
 	}
 
-	indexSecrets, err := c.client.List(context.Background(), metav1.ListOptions{LabelSelector: selector.String()})
-	if err != nil {
-		return nil, fmt.Errorf("query: %w", err)
+	// Pass 1: build the server selector and query for index secrets
+	serverSelector := newListIndicesLabelSelector(c.owner)
+	if queryRequirements, selectable := serverSelectorSet.AsSelector().Requirements(); selectable {
+		serverSelector = serverSelector.Add(queryRequirements...)
 	}
 
-	if len(indexSecrets.Items) == 0 {
-		return nil, driver.ErrReleaseNotFound
+	indexSecrets, err := c.client.List(context.Background(), metav1.ListOptions{LabelSelector: serverSelector.String()})
+	if err != nil {
+		return nil, fmt.Errorf("query: %w", err)
 	}
 
+	// Pass 2: decode the releases that matched the server selector and filter based on the client selector
 	results := make([]*release.Release, 0, len(indexSecrets.Items))
+	clientSelector := clientSelectorSet.AsSelector()
 	for _, indexSecret := range indexSecrets.Items {
 		indexSecret := indexSecret
 		rls, err := c.decodeRelease(context.Background(), &indexSecret)
 		if err != nil {
 			return nil, fmt.Errorf("query: failed to decode release: %w", err)
 		}
+
+		if !clientSelector.Matches(labels.Set(rls.Labels)) {
+			continue
+		}
 		results = append(results, rls)
 	}
+
+	if len(results) == 0 {
+		return nil, driver.ErrReleaseNotFound
+	}
+
 	return results, nil
 }
 
@@ -398,11 +432,13 @@ func (c *chunkedSecrets) decodeRelease(ctx context.Context, indexSecret *corev1.
 		return nil, fmt.Errorf("failed to create gzip reader: %w", err)
 	}
 	releaseDecoder := json.NewDecoder(gzr)
-	var r release.Release
-	if err := releaseDecoder.Decode(&r); err != nil {
+	var wrappedRelease releaseWrapper
+	if err := releaseDecoder.Decode(&wrappedRelease); err != nil {
 		return nil, fmt.Errorf("failed to decode release: %w", err)
 	}
-	r.Labels = filterSystemLabels(indexSecret.Labels)
+
+	r := wrappedRelease.Release
+	r.Labels = filterSystemLabels(wrappedRelease.Labels)
 	return &r, nil
 }
 
diff --git a/pkg/storage/labels.go b/pkg/storage/labels.go
index 1a54c3f7..65efad82 100644
--- a/pkg/storage/labels.go
+++ b/pkg/storage/labels.go
@@ -1,7 +1,6 @@
 package storage
 
 import (
-	"maps"
 	"strconv"
 
 	"helm.sh/helm/v3/pkg/release"
@@ -11,7 +10,6 @@ import (
 
 func newIndexLabels(owner, key string, rls *release.Release) map[string]string {
 	labels := map[string]string{}
-	maps.Copy(labels, rls.Labels)
 	labels["name"] = rls.Name
 	labels["owner"] = owner
 	labels["status"] = rls.Info.Status.String()