rego: Fragment specialization

internal/regopolicyinterpreter/regopolicyinterpreter.go

@@ -8,6 +8,8 @@
 	"fmt"
 	"log"
 	"os"
+	"reflect"
+	"slices"
 	"sync"
 
 	"github.com/open-policy-agent/opa/ast"
@@ -61,12 +63,28 @@
 
 /* See README for more details on Metadata */
 
-type regoMetadata map[string]map[string]interface{}
+// This is conceptually a map[string](regoMetadataMap|regoMetadataSet)
+type regoMetadata map[string]interface{}
 
 const metadataRootKey = "metadata"
 const metadataOperationsKey = "metadata"
 
 type regoMetadataAction string
+type regoMetadataType string
+
+type regoMetadataMap map[string]interface{}
+
+// This uses a pointer to a slice so that we can update it after getting a
+// reference
+type regoMetadataSet []interface{}
+
+const (
+	// A string to anything map
+	metadataTypeMap regoMetadataType = "map"
+
+	// A list of anything that we treat as a set
+	metadataTypeSet regoMetadataType = "set"
+)
 
 const (
 	metadataAdd    regoMetadataAction = "add"
@@ -76,9 +94,13 @@
 
 type regoMetadataOperation struct {
 	Action regoMetadataAction `json:"action"`
-	Name   string             `json:"name"`
-	Key    string             `json:"key"`
-	Value  interface{}        `json:"value"`
+
+	// Defaults to map
+	Type regoMetadataType `json:"type"`
+
+	Name  string      `json:"name"`
+	Key   string      `json:"key"`
+	Value interface{} `json:"value"`
 }
 
 // The result from a policy query
@@ -121,22 +143,44 @@
 	return valueCopy, nil
 }
 
-// deep copy for regoMetadata.
-// We cannot use copyObject for this due to the fact that map[string]interface{}
-// is a concrete type and a map of it cannot be used as a map of interface{}.
-func copyRegoMetadata(value regoMetadata) (regoMetadata, error) {
-	valueJSON, err := json.Marshal(value)
-	if err != nil {
-		return nil, err
-	}
-
-	var valueCopy regoMetadata
-	err = json.Unmarshal(valueJSON, &valueCopy)
-	if err != nil {
-		return nil, err
+// deep copy for regoMetadata, preserves inner regoMetadataMap/Set types
+func copyRegoMetadata(metadata regoMetadata) (regoMetadata, error) {
+	metadataCopy := make(regoMetadata)
+	for key, val := range metadata {
+		switch v := val.(type) {
+		case regoMetadataMap:
+			var copyVal regoMetadataMap
+			valueJSON, err := json.Marshal(v)
+			if err != nil {
+				return nil, err
+			}
+			err = json.Unmarshal(valueJSON, &copyVal)
+			if err != nil {
+				return nil, err
+			}
+			metadataCopy[key] = copyVal
+		case regoMetadataSet:
+			var copyVal regoMetadataSet
+			valueJSON, err := json.Marshal(v)
+			if err != nil {
+				return nil, err
+			}
+			err = json.Unmarshal(valueJSON, &copyVal)
+			if err != nil {
+				return nil, err
+			}
+			metadataCopy[key] = copyVal
+		default:
+			// We technically shouldn't reach here
+			copyVal, err := copyValue(v)
+			if err != nil {
+				return nil, err
+			}
+			metadataCopy[key] = copyVal
+		}
 	}
 
-	return valueCopy, nil
+	return metadataCopy, nil
 }
 
 // NewRegoPolicyInterpreter creates a new RegoPolicyInterpreter, using the code provided.
@@ -228,29 +272,56 @@
 	}
 }
 
-// GetMetadata retrieves a copy of a single metadata item from the policy.
-func (r *RegoPolicyInterpreter) GetMetadata(name string, key string) (interface{}, error) {
-	r.dataAndModulesMutex.Lock()
-	defer r.dataAndModulesMutex.Unlock()
-
+// Does not make any copies, caller must hold dataAndModulesMutex.
+// Returns the zero value of T if the metadata item does not exist.
+func _getMetadata[T any](r *RegoPolicyInterpreter, name string) (T, error) {
 	metadataRoot, ok := r.data[metadataRootKey].(regoMetadata)
 	if !ok {
-		return nil, errors.New("illegal interpreter state: invalid metadata object type")
+		return *new(T), errors.New("illegal interpreter state: invalid metadata object type")
 	}
 
 	if metadata, ok := metadataRoot[name]; ok {
-		if value, ok := metadata[key]; ok {
-			value, err := copyValue(value) //nolint:govet // shadow
-			if err != nil {
-				return nil, fmt.Errorf("unable to copy value: %w", err)
-			}
+		value, ok := metadata.(T)
+		if !ok {
+			return *new(T), fmt.Errorf("metadata %s has the wrong type (wanted %T, but saved type was %T)", name, value, metadata)
+		}
+		return value, nil
+	}
 
-			return value, nil
-		} else {
-			return nil, fmt.Errorf("value not found in %s for key %s", name, key)
+	return *new(T), nil
+}
+
+func _setMetadata[T any](r *RegoPolicyInterpreter, name string, value T) error {
+	metadataRoot, ok := r.data[metadataRootKey].(regoMetadata)
+	if !ok {
+		return errors.New("illegal interpreter state: invalid metadata object type")
+	}
+	metadataRoot[name] = value
+	return nil
+}
+
+// GetMetadata retrieves a copy of a single metadata item from the policy.
+func (r *RegoPolicyInterpreter) GetMetadataMapValue(name string, key string) (interface{}, error) {
+	r.dataAndModulesMutex.Lock()
+	defer r.dataAndModulesMutex.Unlock()
+
+	metadata, err := _getMetadata[regoMetadataMap](r, name)
+	if err != nil {
+		return nil, err
+	}
+	if metadata == nil {
+		return nil, fmt.Errorf("value not found in %s for key %s (map has not been initialized)", name, key)
+	}
+
+	if value, ok := metadata[key]; ok {
+		value, err := copyValue(value) //nolint:govet // shadow
+		if err != nil {
+			return nil, fmt.Errorf("unable to copy value: %w", err)
 		}
+
+		return value, nil
 	} else {
-		return nil, fmt.Errorf("metadata not found for name %s", name)
+		return nil, fmt.Errorf("value not found in %s for key %s", name, key)
 	}
 }
 
@@ -287,6 +358,16 @@
 	if !ok {
 		return nil, errors.New("unable to load metadata object")
 	}
+	dataType, ok := data["type"]
+	if !ok {
+		metadataOp.Type = metadataTypeMap
+	} else {
+		dataTypeStr, ok := dataType.(string)
+		if !ok {
+			return nil, errors.New("unable to load metadata type")
+		}
+		metadataOp.Type = regoMetadataType(dataTypeStr)
+	}
 	metadataOp.Name, ok = data["name"].(string)
 	if !ok {
 		return nil, errors.New("unable to load metadata name")
@@ -296,29 +377,27 @@
 		return nil, errors.New("unable to load metadata action")
 	}
 	metadataOp.Action = regoMetadataAction(action)
-	metadataOp.Key, ok = data["key"].(string)
-	if !ok {
-		return nil, errors.New("unable to load metadata key")
-	}
 
-	if metadataOp.Action != metadataRemove {
-		metadataOp.Value, ok = data["value"]
+	var hasKey, hasValue bool
+	key, hasKey := data["key"]
+	if hasKey {
+		metadataOp.Key, ok = key.(string)
 		if !ok {
-			return nil, errors.New("unable to load metadata value")
+			return nil, errors.New("unable to load metadata key")
 		}
 	}
 
-	return &metadataOp, nil
-}
+	metadataOp.Value, hasValue = data["value"]
+
+	if (metadataOp.Action != metadataRemove || metadataOp.Type == metadataTypeSet) && !hasValue {
+		return nil, errors.New("missing metadata value")
+	}
 
-func (m regoMetadata) getOrCreate(name string) map[string]interface{} {
-	if metadata, ok := m[name]; ok {
-		return metadata
+	if metadataOp.Type == metadataTypeMap && !hasKey {
+		return nil, errors.New("missing metadata key")
 	}
 
-	metadata := make(map[string]interface{})
-	m[name] = metadata
-	return metadata
+	return &metadataOp, nil
 }
 
 func (r *RegoPolicyInterpreter) UpdateOSType(os string) error {
@@ -327,6 +406,7 @@
 	ops := []*regoMetadataOperation{
 		{
 			Action: metadataAdd,
+			Type:   metadataTypeMap,
 			Name:   "operatingsystem",
 			Key:    "ostype",
 			Value:  os,
@@ -335,32 +415,66 @@
 	return r.updateMetadata(ops)
 }
 
+// dataAndModulesMutex must be held before calling this
 func (r *RegoPolicyInterpreter) updateMetadata(ops []*regoMetadataOperation) error {
-	// dataAndModulesMutex must be held before calling this
-
-	metadataRoot, ok := r.data[metadataRootKey].(regoMetadata)
-	if !ok {
-		return errors.New("illegal interpreter state: invalid metadata object type")
-	}
-
 	for _, op := range ops {
-		metadata := metadataRoot.getOrCreate(op.Name)
-		switch op.Action {
-		case metadataAdd:
-			if _, ok := metadata[op.Key]; ok {
-				return fmt.Errorf("cannot add metadata value, key %s[%s] already exists", op.Name, op.Key)
-			} else {
-				metadata[op.Key] = op.Value
+		switch op.Type {
+		case metadataTypeMap:
+			metadata, err := _getMetadata[regoMetadataMap](r, op.Name)
+			if err != nil {
+				return fmt.Errorf("unable to get metadata: %w", err)
+			}
+			if metadata == nil {
+				metadata = make(regoMetadataMap)
 			}
+			switch op.Action {
+			case metadataAdd:
+				if _, ok := metadata[op.Key]; ok {
+					return fmt.Errorf("cannot add metadata value, key %s[%s] already exists", op.Name, op.Key)
+				} else {
+					metadata[op.Key] = op.Value
+				}
+
+			case metadataUpdate:
+				metadata[op.Key] = op.Value
 
-		case metadataUpdate:
-			metadata[op.Key] = op.Value
+			case metadataRemove:
+				delete(metadata, op.Key)
 
-		case metadataRemove:
-			delete(metadata, op.Key)
+			default:
+				return fmt.Errorf("unrecognized metadata action: %s for map", op.Action)
+			}
 
+			if err := _setMetadata[regoMetadataMap](r, op.Name, metadata); err != nil {
+				return fmt.Errorf("unable to set metadata: %w", err)
+			}
+		case metadataTypeSet:
+			metadata, err := _getMetadata[regoMetadataSet](r, op.Name)
+			if err != nil {
+				return fmt.Errorf("unable to get metadata: %w", err)
+			}
+			if metadata == nil {
+				metadata = make(regoMetadataSet, 0)
+			}
+			switch op.Action {
+			case metadataAdd:
+				if !slices.ContainsFunc(metadata, func(e interface{}) bool {
+					return reflect.DeepEqual(e, op.Value)
+				}) {
+					metadata = append(metadata, op.Value)
+				}
+			case metadataRemove:
+				metadata = slices.DeleteFunc(metadata, func(e interface{}) bool {
+					return reflect.DeepEqual(e, op.Value)
+				})
+			default:
+				return fmt.Errorf("unrecognized metadata action: %s for set", op.Action)
+			}
+			if err := _setMetadata[regoMetadataSet](r, op.Name, metadata); err != nil {
+				return fmt.Errorf("unable to set metadata: %w", err)
+			}
 		default:
-			return fmt.Errorf("unrecognized metadata action: %s", op.Action)
+			return fmt.Errorf("unrecognized metadata type: %s", op.Type)
 		}
 	}
 
@@ -513,6 +627,19 @@
 	}
 }
 
+// Array attempts to interpret the result value as an array.
+func (r RegoQueryResult) Array(key string) ([]interface{}, error) {
+	if value, ok := r[key]; ok {
+		if arr, ok := value.([]interface{}); ok {
+			return arr, nil
+		} else {
+			return nil, fmt.Errorf("value for '%s' is not an array", key)
+		}
+	} else {
+		return nil, fmt.Errorf("unable to find value for key '%s'", key)
+	}
+}
+
 // Bool attempts to interpret a result value as a boolean.
 func (r RegoQueryResult) Bool(key string) (bool, error) {
 	if value, ok := r[key]; ok {