connectivity: Add annotations map helper type

This type will be used in a subsequent commit to specify the wanted
annotations for a specific deployment or namespace. Unit tests are added
to ensure the wildcard matching behavior.

Signed-off-by: Sebastian Wicki <[email protected]>

connectivity/check/check.go

@@ -4,9 +4,11 @@
 package check
 
 import (
+	"encoding/json"
 	"fmt"
 	"io"
 	"regexp"
+	"strings"
 	"time"
 
 	"github.com/cilium/cilium/api/v1/flow"
@@ -87,6 +89,102 @@ type nodesWithoutCiliumIP struct {
 	Mask int
 }
 
+type annotations map[string]string
+
+func marshalMap[M ~map[K]V, K comparable, V any](m *M) string {
+	if m == nil || len(*m) == 0 {
+		return "{}" // avoids printing "null" for nil map
+	}
+
+	b, err := json.Marshal(*m)
+	if err != nil {
+		return fmt.Sprintf("error: %s", err)
+	}
+	return string(b)
+}
+
+// String implements pflag.Value
+func (a *annotations) String() string {
+	return marshalMap(a)
+}
+
+// Set implements pflag.Value
+func (a *annotations) Set(s string) error {
+	return json.Unmarshal([]byte(s), a)
+}
+
+// Type implements pflag.Value
+func (a *annotations) Type() string {
+	return "json"
+}
+
+type annotationsMap map[string]annotations
+
+// String implements pflag.Value
+func (a *annotationsMap) String() string {
+	return marshalMap(a)
+}
+
+// Set implements pflag.Value
+func (a *annotationsMap) Set(s string) error {
+	var target annotationsMap
+	err := json.Unmarshal([]byte(s), &target)
+	if err != nil {
+		return err
+	} else if a == nil {
+		return nil
+	}
+
+	// Validate keys for Match function, `*` is only allowed at the end of the string
+	for key := range target {
+		_, suffix, ok := strings.Cut(key, "*")
+		if ok && len(suffix) > 0 {
+			return fmt.Errorf("invalid match key %q: wildcard only allowed at end of key", key)
+		}
+	}
+
+	*a = target
+	return nil
+}
+
+// Type implements pflag.Value
+func (a *annotationsMap) Type() string {
+	return "json"
+}
+
+// Match extracts the annotations for the matching component s. If the
+// annotation map contains s as a key, the corresponding value will be returned.
+// Otherwise, every map key containing a `*` character will be treated as
+// prefix pattern, i.e. a map key `foo*` will match the name `foobar`.
+func (a *annotationsMap) Match(name string) annotations {
+	// Invalid map or component name that contains a wildcard
+	if a == nil || strings.Contains(name, "*") {
+		return nil
+	}
+
+	// Direct match
+	if match, ok := (*a)[name]; ok {
+		return match
+	}
+
+	// Find the longest prefix match
+	var longestPrefix string
+	var longestMatch annotations
+	for pattern, match := range *a {
+		prefix, _, ok := strings.Cut(pattern, "*")
+		if !ok || !strings.HasPrefix(name, prefix) {
+			continue // not a matching pattern
+		}
+
+		if len(prefix) >= len(longestPrefix) {
+			longestPrefix = prefix
+			longestMatch = match
+		}
+	}
+
+	return longestMatch
+}
+
 func (p Parameters) ciliumEndpointTimeout() time.Duration {
 	return 5 * time.Minute
 }

connectivity/check/check_test.go

@@ -0,0 +1,71 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright Authors of Cilium
+
+package check
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAnnotationMap(t *testing.T) {
+	var a annotationsMap
+	err := a.Set(`not json`)
+	assert.IsType(t, err, &json.SyntaxError{})
+
+	err = a.Set(`{"foo*bar":{}}`)
+	assert.ErrorContains(t, err, "wildcard only allowed at end of key")
+
+	err = a.Set(`{"baz*qux*":{}}`)
+	assert.ErrorContains(t, err, "wildcard only allowed at end of key")
+
+	err = a.Set(`{"**":{}}`)
+	assert.ErrorContains(t, err, "wildcard only allowed at end of key")
+
+	var clientAnnotations = annotations{"baz": "qux"}
+	var echoSameNodeAnnotations = annotations{"quux": "corge"}
+	var echoWildcardAnnotations = annotations{"grault": "grault"}
+	var wildcardAnnotations = annotations{"waldo": "fred"}
+
+	err = a.Set(`{
+		"client": ` + clientAnnotations.String() + `,
+		"echo-same-node": ` + echoSameNodeAnnotations.String() + `,
+		"echo*": ` + echoWildcardAnnotations.String() + `,
+		"*": ` + wildcardAnnotations.String() + `
+	}`)
+	assert.NoError(t, err)
+	assert.Equal(t, annotationsMap{
+		"client":         clientAnnotations,
+		"echo-same-node": echoSameNodeAnnotations,
+		"echo*":          echoWildcardAnnotations,
+		"*":              wildcardAnnotations,
+	}, a)
+
+	// Test wildcard fallback
+	assert.Equal(t, a.Match("echo*"), annotations(nil)) // wildcard not allowed here
+	assert.Equal(t, a.Match("*"), annotations(nil))     // wildcard not allowed here
+
+	assert.Equal(t, a.Match("client"), clientAnnotations)
+	assert.Equal(t, a.Match("echo-same-node"), echoSameNodeAnnotations)
+	assert.Equal(t, a.Match("echo-other-node"), echoWildcardAnnotations)
+	assert.Equal(t, a.Match("other"), wildcardAnnotations)
+
+	err = a.Set(`{
+		"echo-same-*": ` + echoSameNodeAnnotations.String() + `,
+		"echo*": ` + echoWildcardAnnotations.String() + `
+	}`)
+	assert.NoError(t, err)
+	assert.Equal(t, annotationsMap{
+		"echo-same-*": echoSameNodeAnnotations,
+		"echo*":       echoWildcardAnnotations,
+	}, a)
+
+	// Tests longest prefix match
+	assert.Equal(t, a.Match("echo-same-node"), echoSameNodeAnnotations)
+	assert.Equal(t, a.Match("echo-other-node"), echoWildcardAnnotations)
+	assert.Equal(t, a.Match("echo"), echoWildcardAnnotations)
+	assert.Equal(t, a.Match("other"), annotations(nil))
+
+}