[Backport 6.2] Add regexp.replace support in role templates (#7250)

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## 6.2.4
+
+This release of Teleport contains multiple improvements.
+
+* Added support for `regexp.replace(variable, expression, replacement)` in role templates.
+
 ## 6.2.3
 
 This release of Teleport contains multiple improvements.

docs/pages/access-controls/guides/role-templates.mdx

@@ -281,8 +281,10 @@ spec:
 
     # Functions transform variables.
     database_users: ['{{email.local(external.email)}}']
+    database_labels:
+      'env': '{{regexp.replace(external.access["env"], "^(staging)$", "$1")}}'
     
-    # Labels can mix template and hard-coded values
+    # Labels can mix template and hard-coded values.
     node_labels:
       'env': '{{external.access["env"]}}'
       'region': 'us-west-2'
@@ -311,14 +313,25 @@ spec:
     # The variable external.groups gets replaced with a list.
     kubernetes_groups: ['devs', 'admins']
 
-    # The variable email.local will take a local part of the external.email attribute.
+    # The function email.local will take a local part of the external.email attribute.
     database_users: ['alice']
 
-    # Node labels have 'env' replaced from a variable
+    # The function regexp.replace will transform and filter only matching values.
+    database_labels:
+      'env': 'staging'
+
+    # Node labels have 'env' replaced from a variable and 'region' hard-coded.
     node_labels:
       'env': ['prod', 'staging']
       'region': 'us-west-2'
 
     kubernetes_labels:
       '*': '*'
 ```
+
+Available interpolation functions include:
+
+Function | Description
+--- | ---
+`email.local(variable)` | Extracts the local part of an email field, like `Alice <[email protected]>` or `[email protected]`.
+`regexp.replace(variable, expression, replacement)` | Finds all matches of `expression` and replaces them with `replacement`. This supports expansion, e.g. `regexp.replace(external.email, "^(.*)@example.com$", "$1")`. Values which do not match the expression will be filtered out.

lib/services/role_test.go

@@ -1548,6 +1548,16 @@ func TestApplyTraits(t *testing.T) {
 				outLogins: []string{"bar", "root"},
 			},
 		},
+		{
+			comment: "logins substitute in allow rule with regexp",
+			inTraits: map[string][]string{
+				"foo": {"bar-baz"},
+			},
+			allow: rule{
+				inLogins:  []string{`{{regexp.replace(external.foo, "^bar-(.*)$", "$1")}}`, "root"},
+				outLogins: []string{"baz", "root"},
+			},
+		},
 		{
 			comment: "logins substitute in deny rule",
 			inTraits: map[string][]string{
@@ -1588,6 +1598,16 @@ func TestApplyTraits(t *testing.T) {
 				outKubeUsers: []string{"IAM#bar;"},
 			},
 		},
+		{
+			comment: "kube user regexp interpolation in allow rule",
+			inTraits: map[string][]string{
+				"foo": {"bar-baz"},
+			},
+			allow: rule{
+				inKubeUsers:  []string{`IAM#{{regexp.replace(external.foo, "^bar-(.*)$", "$1")}};`},
+				outKubeUsers: []string{"IAM#baz;"},
+			},
+		},
 		{
 			comment: "kube users interpolation in deny rule",
 			inTraits: map[string][]string{
@@ -1694,6 +1714,25 @@ func TestApplyTraits(t *testing.T) {
 				inLogins: []string{`{{email.local(email.local)}}`, `{{email.local(email.local())}}`},
 			},
 		},
+		{
+			comment: "invalid regexp in logins does not get passed along",
+			inTraits: map[string][]string{
+				"foo": {"bar"},
+			},
+			allow: rule{
+				inLogins: []string{`{{regexp.replace(external.foo, "(()", "baz")}}`},
+			},
+		},
+		{
+			comment: "logins which to not match regexp get filtered out",
+			inTraits: map[string][]string{
+				"foo": {"dev-alice", "dev-bob", "prod-charlie"},
+			},
+			allow: rule{
+				inLogins:  []string{`{{regexp.replace(external.foo, "^dev-([a-zA-Z]+)$", "$1-admin")}}`},
+				outLogins: []string{"alice-admin", "bob-admin"},
+			},
+		},
 		{
 			comment: "variable in logins, none in traits",
 			inTraits: map[string][]string{

lib/utils/parse/parse.go

@@ -69,6 +69,34 @@ func (emailLocalTransformer) transform(in string) (string, error) {
 	return parts[0], nil
 }
 
+// regexpReplaceTransformer replaces all matches of re with replacement
+type regexpReplaceTransformer struct {
+	re          *regexp.Regexp
+	replacement string
+}
+
+// newRegexpReplaceTransformer attempts to create a regexpReplaceTransformer or
+// fails with error if the expression does not compile
+func newRegexpReplaceTransformer(expression, replacement string) (*regexpReplaceTransformer, error) {
+	re, err := regexp.Compile(expression)
+	if err != nil {
+		return nil, trace.BadParameter("failed parsing regexp %q: %v", expression, err)
+	}
+	return &regexpReplaceTransformer{
+		re:          re,
+		replacement: replacement,
+	}, nil
+}
+
+// transform applies the regexp replacement (with expansion)
+func (r regexpReplaceTransformer) transform(in string) (string, error) {
+	// filter out inputs which do not match the regexp at all
+	if !r.re.MatchString(in) {
+		return "", nil
+	}
+	return r.re.ReplaceAllString(in, r.replacement), nil
+}
+
 // Namespace returns a variable namespace, e.g. external or internal
 func (p *Expression) Namespace() string {
 	return p.namespace
@@ -90,7 +118,7 @@ func (p *Expression) Interpolate(traits map[string][]string) ([]string, error) {
 	if !ok {
 		return nil, trace.NotFound("variable is not found")
 	}
-	out := make([]string, len(values))
+	var out []string
 	for i := range values {
 		val := values[i]
 		var err error
@@ -100,7 +128,9 @@ func (p *Expression) Interpolate(traits map[string][]string) ([]string, error) {
 				return nil, trace.Wrap(err)
 			}
 		}
-		out[i] = p.prefix + val + p.suffix
+		if len(val) > 0 {
+			out = append(out, p.prefix+val+p.suffix)
+		}
 	}
 	return out, nil
 }
@@ -310,6 +340,8 @@ const (
 	RegexpMatchFnName = "match"
 	// RegexpNotMatchFnName is a name for regexp.not_match function.
 	RegexpNotMatchFnName = "not_match"
+	// RegexpReplaceFnName is a name for regexp.replace function.
+	RegexpReplaceFnName = "replace"
 )
 
 // transformer is an optional value transformer function that can take in
@@ -318,6 +350,24 @@ type transformer interface {
 	transform(in string) (string, error)
 }
 
+// getBasicString checks that arg is a properly quoted basic string and returns
+// it. If arg is not a properly quoted basic string, the second return value
+// will be false.
+func getBasicString(arg ast.Expr) (string, bool) {
+	basicLit, ok := arg.(*ast.BasicLit)
+	if !ok {
+		return "", false
+	}
+	if basicLit.Kind != token.STRING {
+		return "", false
+	}
+	str, err := strconv.Unquote(basicLit.Value)
+	if err != nil {
+		return "", false
+	}
+	return str, true
+}
+
 // maxASTDepth is the maximum depth of the AST that func walk will traverse.
 // The limit exists to protect against DoS via malicious inputs.
 const maxASTDepth = 1000
@@ -374,18 +424,12 @@ func walk(node ast.Node, depth int) (*walkResult, error) {
 					if len(n.Args) != 1 {
 						return nil, trace.BadParameter("expected 1 argument for %v.%v got %v", namespace, fn, len(n.Args))
 					}
-					re, ok := n.Args[0].(*ast.BasicLit)
+					re, ok := getBasicString(n.Args[0])
 					if !ok {
-						return nil, trace.BadParameter("argument to %v.%v must be a string literal", namespace, fn)
-					}
-					if re.Kind != token.STRING {
-						return nil, trace.BadParameter("argument to %v.%v must be a string literal", namespace, fn)
-					}
-					val, err := strconv.Unquote(re.Value)
-					if err != nil {
-						return nil, trace.BadParameter("regexp %q is not a properly quoted string: %v", re.Value, err)
+						return nil, trace.BadParameter("argument to %v.%v must be a properly quoted string literal", namespace, fn)
 					}
-					result.match, err = newRegexpMatcher(val, false)
+					var err error
+					result.match, err = newRegexpMatcher(re, false)
 					if err != nil {
 						return nil, trace.Wrap(err)
 					}
@@ -394,6 +438,28 @@ func walk(node ast.Node, depth int) (*walkResult, error) {
 						result.match = notMatcher{result.match}
 					}
 					return &result, nil
+				case RegexpReplaceFnName:
+					if len(n.Args) != 3 {
+						return nil, trace.BadParameter("expected 3 arguments for %v.%v got %v", namespace, fn, len(n.Args))
+					}
+					ret, err := walk(n.Args[0], depth+1)
+					if err != nil {
+						return nil, trace.Wrap(err)
+					}
+					result.parts = ret.parts
+					expression, ok := getBasicString(n.Args[1])
+					if !ok {
+						return nil, trace.BadParameter("second argument to %v.%v must be a properly quoted string literal", namespace, fn)
+					}
+					replacement, ok := getBasicString(n.Args[2])
+					if !ok {
+						return nil, trace.BadParameter("third argument to %v.%v must be a properly quoted string literal", namespace, fn)
+					}
+					result.transform, err = newRegexpReplaceTransformer(expression, replacement)
+					if err != nil {
+						return nil, trace.Wrap(err)
+					}
+					return &result, nil
 				default:
 					return nil, trace.BadParameter("unsupported function %v.%v, supported functions are: regexp.match, regexp.not_match", namespace, fn)
 				}

lib/utils/parse/parse_test.go

@@ -109,6 +109,28 @@ func TestVariable(t *testing.T) {
 			in:    "{{email.local(internal.bar)}}",
 			out:   Expression{namespace: "internal", variable: "bar", transform: emailLocalTransformer{}},
 		},
+		{
+			title: "regexp replace",
+			in:    `{{regexp.replace(internal.foo, "bar-(.*)", "$1")}}`,
+			out: Expression{
+				namespace: "internal",
+				variable:  "foo",
+				transform: &regexpReplaceTransformer{
+					re:          regexp.MustCompile("bar-(.*)"),
+					replacement: "$1",
+				},
+			},
+		},
+		{
+			title: "regexp replace with variable expression",
+			in:    `{{regexp.replace(internal.foo, internal.bar, "baz")}}`,
+			err:   trace.BadParameter(""),
+		},
+		{
+			title: "regexp replace with variable replacement",
+			in:    `{{regexp.replace(internal.foo, "bar", internal.baz)}}`,
+			err:   trace.BadParameter(""),
+		},
 	}
 
 	for _, tt := range tests {
@@ -119,7 +141,7 @@ func TestVariable(t *testing.T) {
 				return
 			}
 			require.NoError(t, err)
-			require.Empty(t, cmp.Diff(tt.out, *variable, cmp.AllowUnexported(Expression{})))
+			require.Equal(t, tt.out, *variable)
 		})
 	}
 }
@@ -173,6 +195,54 @@ func TestInterpolate(t *testing.T) {
 			traits: map[string][]string{"foo": []string{"a", "b"}, "bar": []string{"c"}},
 			res:    result{values: []string{"foo"}},
 		},
+		{
+			title: "regexp replacement with numeric match",
+			in: Expression{
+				variable: "foo",
+				transform: regexpReplaceTransformer{
+					re:          regexp.MustCompile("bar-(.*)"),
+					replacement: "$1",
+				},
+			},
+			traits: map[string][]string{"foo": []string{"bar-baz"}},
+			res:    result{values: []string{"baz"}},
+		},
+		{
+			title: "regexp replacement with named match",
+			in: Expression{
+				variable: "foo",
+				transform: regexpReplaceTransformer{
+					re:          regexp.MustCompile("bar-(?P<suffix>.*)"),
+					replacement: "${suffix}",
+				},
+			},
+			traits: map[string][]string{"foo": []string{"bar-baz"}},
+			res:    result{values: []string{"baz"}},
+		},
+		{
+			title: "regexp replacement with multiple matches",
+			in: Expression{
+				variable: "foo",
+				transform: regexpReplaceTransformer{
+					re:          regexp.MustCompile("foo-(.*)-(.*)"),
+					replacement: "$1.$2",
+				},
+			},
+			traits: map[string][]string{"foo": []string{"foo-bar-baz"}},
+			res:    result{values: []string{"bar.baz"}},
+		},
+		{
+			title: "regexp replacement with no match",
+			in: Expression{
+				variable: "foo",
+				transform: regexpReplaceTransformer{
+					re:          regexp.MustCompile("^bar-(.*)$"),
+					replacement: "$1-matched",
+				},
+			},
+			traits: map[string][]string{"foo": []string{"foo-test1", "bar-test2"}},
+			res:    result{values: []string{"test2-matched"}},
+		},
 	}
 
 	for _, tt := range tests {
@@ -184,7 +254,7 @@ func TestInterpolate(t *testing.T) {
 				return
 			}
 			require.NoError(t, err)
-			require.Empty(t, cmp.Diff(tt.res.values, values))
+			require.Equal(t, tt.res.values, values)
 		})
 	}
 }