Copy expression package out of /e

lib/expression/dict.go

@@ -0,0 +1,88 @@
+/*
+ * Teleport
+ * Copyright (C) 2024  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package expression
+
+import "github.com/gravitational/trace"
+
+// Dict is a map of type string key and Set values.
+type Dict map[string]Set
+
+// NewDict returns a dict initialized with the key-value pairs as specified in
+// [pairs].
+func NewDict(pairs ...pair) (Dict, error) {
+	d := make(Dict, len(pairs))
+	for _, p := range pairs {
+		k, ok := p.first.(string)
+		if !ok {
+			return nil, trace.BadParameter("dict keys must have type string, got %T", p.first)
+		}
+		v, ok := p.second.(Set)
+		if !ok {
+			return nil, trace.BadParameter("dict values must have type set, got %T", p.second)
+		}
+		d[k] = v
+	}
+	return d, nil
+}
+
+func (d Dict) addValues(key string, values ...string) Dict {
+	out := d.clone()
+	s := out[key]
+	if s == nil {
+		out[key] = NewSet(values...)
+		return out
+	}
+	// Calling set.add would do an unnecessary extra copy, add the values
+	// "manually".
+	for _, value := range values {
+		s[value] = struct{}{}
+	}
+	return out
+}
+
+func (d Dict) put(key string, value Set) Dict {
+	out := d.clone()
+	out[key] = value
+	return out
+}
+
+func (d Dict) remove(keys ...string) any {
+	out := d.clone()
+	for _, key := range keys {
+		delete(out, key)
+	}
+	return out
+}
+
+func (d Dict) clone() Dict {
+	out := make(Dict, len(d))
+	for key, set := range d {
+		out[key] = set.clone()
+	}
+	return out
+}
+
+// Get implements typical.Getter[set]
+func (d Dict) Get(key string) (Set, error) {
+	return d[key], nil
+}
+
+type pair struct {
+	first, second any
+}

lib/expression/evaluator.go

@@ -0,0 +1,70 @@
+/*
+ * Teleport
+ * Copyright (C) 2024  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package expression
+
+import (
+	"errors"
+
+	"github.com/gravitational/trace"
+
+	"github.com/gravitational/teleport/lib/utils/typical"
+)
+
+// EvaluateTraitsMap evaluates expression that must evaluate to either string or Set.
+// traitsMap: key is name of the trait and values are list of predicate expressions.
+func EvaluateTraitsMap[TEnv any](env TEnv, traitsMap map[string][]string, parseExpression func(input string) (typical.Expression[TEnv, any], error)) (Dict, error) {
+	d, err := NewDict()
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+	for key, values := range traitsMap {
+		for _, expr := range values {
+			e, err := parseExpression(expr)
+			if err != nil {
+				var u typical.UnknownIdentifierError
+				if errors.As(err, &u) {
+					id := u.Identifier()
+					if id == expr {
+						// If the entire expression evaluates to a single unknown
+						// identifier, treat it as a string. This is to support rules like
+						//   groups: [devs]
+						// instead of requiring extra quotes like
+						//   groups: ['"devs"']
+						d[key] = union(d[key], NewSet(id))
+						continue
+					}
+				}
+				return nil, trace.Wrap(err, "error parsing expression: %q", expr)
+			}
+
+			result, err := e.Evaluate(env)
+			if err != nil {
+				return nil, trace.Wrap(err, "error evaluating expression: %q", expr)
+			}
+
+			s, err := traitsMapResultToSet(result, expr)
+			if err != nil {
+				return nil, trace.Wrap(err)
+			}
+
+			d[key] = union(d[key], s)
+		}
+	}
+	return d, nil
+}