ppl: add data type, implement string and list matchers (#2228)

* ppl: add data type, implement string and list matchers

* update policy converter

config/policy_ppl.go

@@ -40,14 +40,18 @@ func (p *Policy) ToPPL() *parser.Policy {
 		allowRule.Or = append(allowRule.Or,
 			parser.Criterion{
 				Name: "domain",
-				Data: parser.String(ad),
+				Data: parser.Object{
+					"is": parser.String(ad),
+				},
 			})
 	}
 	for _, ag := range p.AllAllowedGroups() {
 		allowRule.Or = append(allowRule.Or,
 			parser.Criterion{
 				Name: "group",
-				Data: parser.String(ag),
+				Data: parser.Object{
+					"has": parser.String(ag),
+				},
 			})
 	}
 	for _, aic := range p.AllAllowedIDPClaims() {
@@ -64,11 +68,15 @@ func (p *Policy) ToPPL() *parser.Policy {
 		allowRule.Or = append(allowRule.Or,
 			parser.Criterion{
 				Name: "user",
-				Data: parser.String(au),
+				Data: parser.Object{
+					"is": parser.String(au),
+				},
 			},
 			parser.Criterion{
 				Name: "email",
-				Data: parser.String(au),
+				Data: parser.Object{
+					"is": parser.String(au),
+				},
 			})
 	}
 	ppl.Rules = append(ppl.Rules, allowRule)

config/policy_ppl_test.go

@@ -67,47 +67,41 @@ authenticated_user_0 {
 }
 
 domains_0 {
-	rule_data := "a.example.com"
 	session := get_session(input.session.id)
 	user := get_user(session)
 	domain := split(get_user_email(session, user), "@")[1]
-	domain == rule_data
+	domain == "a.example.com"
 }
 
 domains_1 {
-	rule_data := "b.example.com"
 	session := get_session(input.session.id)
 	user := get_user(session)
 	domain := split(get_user_email(session, user), "@")[1]
-	domain == rule_data
+	domain == "b.example.com"
 }
 
 domains_2 {
-	rule_data := "c.example.com"
 	session := get_session(input.session.id)
 	user := get_user(session)
 	domain := split(get_user_email(session, user), "@")[1]
-	domain == rule_data
+	domain == "c.example.com"
 }
 
 domains_3 {
-	rule_data := "d.example.com"
 	session := get_session(input.session.id)
 	user := get_user(session)
 	domain := split(get_user_email(session, user), "@")[1]
-	domain == rule_data
+	domain == "d.example.com"
 }
 
 domains_4 {
-	rule_data := "e.example.com"
 	session := get_session(input.session.id)
 	user := get_user(session)
 	domain := split(get_user_email(session, user), "@")[1]
-	domain == rule_data
+	domain == "e.example.com"
 }
 
 groups_0 {
-	rule_data := "group1"
 	session := get_session(input.session.id)
 	directory_user := get_directory_user(session)
 	group_ids := get_group_ids(session, directory_user)
@@ -126,13 +120,10 @@ groups_0 {
 		directory_group.email != null
 	]
 	groups = array.concat(group_ids, array.concat(group_names, group_emails))
-	some group
-	group = groups[_0]
-	group == rule_data
+	count([true | some v; v = groups[_0]; v == "group1"]) > 0
 }
 
 groups_1 {
-	rule_data := "group2"
 	session := get_session(input.session.id)
 	directory_user := get_directory_user(session)
 	group_ids := get_group_ids(session, directory_user)
@@ -151,13 +142,10 @@ groups_1 {
 		directory_group.email != null
 	]
 	groups = array.concat(group_ids, array.concat(group_names, group_emails))
-	some group
-	group = groups[_0]
-	group == rule_data
+	count([true | some v; v = groups[_0]; v == "group2"]) > 0
 }
 
 groups_2 {
-	rule_data := "group3"
 	session := get_session(input.session.id)
 	directory_user := get_directory_user(session)
 	group_ids := get_group_ids(session, directory_user)
@@ -176,13 +164,10 @@ groups_2 {
 		directory_group.email != null
 	]
 	groups = array.concat(group_ids, array.concat(group_names, group_emails))
-	some group
-	group = groups[_0]
-	group == rule_data
+	count([true | some v; v = groups[_0]; v == "group3"]) > 0
 }
 
 groups_3 {
-	rule_data := "group4"
 	session := get_session(input.session.id)
 	directory_user := get_directory_user(session)
 	group_ids := get_group_ids(session, directory_user)
@@ -201,13 +186,10 @@ groups_3 {
 		directory_group.email != null
 	]
 	groups = array.concat(group_ids, array.concat(group_names, group_emails))
-	some group
-	group = groups[_0]
-	group == rule_data
+	count([true | some v; v = groups[_0]; v == "group4"]) > 0
 }
 
 groups_4 {
-	rule_data := "group5"
 	session := get_session(input.session.id)
 	directory_user := get_directory_user(session)
 	group_ids := get_group_ids(session, directory_user)
@@ -226,9 +208,7 @@ groups_4 {
 		directory_group.email != null
 	]
 	groups = array.concat(group_ids, array.concat(group_names, group_emails))
-	some group
-	group = groups[_0]
-	group == rule_data
+	count([true | some v; v = groups[_0]; v == "group5"]) > 0
 }
 
 claims_0 {
@@ -268,78 +248,73 @@ claims_2 {
 }
 
 users_0 {
-	rule_data := "user1"
 	session := get_session(input.session.id)
 	user := get_user(session)
-	user.id == rule_data
+	user_id := user.id
+	user_id == "user1"
 }
 
 emails_0 {
-	rule_data := "user1"
 	session := get_session(input.session.id)
 	user := get_user(session)
 	email := get_user_email(session, user)
-	email == rule_data
+	email == "user1"
 }
 
 users_1 {
-	rule_data := "user2"
 	session := get_session(input.session.id)
 	user := get_user(session)
-	user.id == rule_data
+	user_id := user.id
+	user_id == "user2"
 }
 
 emails_1 {
-	rule_data := "user2"
 	session := get_session(input.session.id)
 	user := get_user(session)
 	email := get_user_email(session, user)
-	email == rule_data
+	email == "user2"
 }
 
 users_2 {
-	rule_data := "user3"
 	session := get_session(input.session.id)
 	user := get_user(session)
-	user.id == rule_data
+	user_id := user.id
+	user_id == "user3"
 }
 
 emails_2 {
-	rule_data := "user3"
 	session := get_session(input.session.id)
 	user := get_user(session)
 	email := get_user_email(session, user)
-	email == rule_data
+	email == "user3"
 }
 
 users_3 {
-	rule_data := "user4"
 	session := get_session(input.session.id)
 	user := get_user(session)
-	user.id == rule_data
+	user_id := user.id
+	user_id == "user4"
 }
 
 emails_3 {
-	rule_data := "user4"
 	session := get_session(input.session.id)
 	user := get_user(session)
 	email := get_user_email(session, user)
-	email == rule_data
+	email == "user4"
 }
 
 users_4 {
-	rule_data := "user5"
 	session := get_session(input.session.id)
 	user := get_user(session)
-	user.id == rule_data
+	user_id := user.id
+	user_id == "user5"
 }
 
 emails_4 {
-	rule_data := "user5"
 	session := get_session(input.session.id)
 	user := get_user(session)
 	email := get_user_email(session, user)
-	email == rule_data
+	email == "user5"
 }
 
 or_0 = v1 {

pkg/policy/criteria/accept.go

@@ -3,6 +3,7 @@ package criteria
 import (
 	"github.com/open-policy-agent/opa/ast"
 
+	"github.com/pomerium/pomerium/pkg/policy/generator"
 	"github.com/pomerium/pomerium/pkg/policy/parser"
 )
 
@@ -14,6 +15,10 @@ type acceptCriterion struct {
 	g *Generator
 }
 
+func (acceptCriterion) DataType() CriterionDataType {
+	return generator.CriterionDataTypeUnused
+}
+
 func (acceptCriterion) Names() []string {
 	return []string{"accept"}
 }

pkg/policy/criteria/authenticated_user.go

@@ -3,6 +3,7 @@ package criteria
 import (
 	"github.com/open-policy-agent/opa/ast"
 
+	"github.com/pomerium/pomerium/pkg/policy/generator"
 	"github.com/pomerium/pomerium/pkg/policy/parser"
 	"github.com/pomerium/pomerium/pkg/policy/rules"
 )
@@ -17,6 +18,10 @@ type authenticatedUserCriterion struct {
 	g *Generator
 }
 
+func (authenticatedUserCriterion) DataType() CriterionDataType {
+	return generator.CriterionDataTypeUnused
+}
+
 func (authenticatedUserCriterion) Names() []string {
 	return []string{"authenticated_user"}
 }

pkg/policy/criteria/claims.go

@@ -3,6 +3,7 @@ package criteria
 import (
 	"github.com/open-policy-agent/opa/ast"
 
+	"github.com/pomerium/pomerium/pkg/policy/generator"
 	"github.com/pomerium/pomerium/pkg/policy/parser"
 	"github.com/pomerium/pomerium/pkg/policy/rules"
 )
@@ -35,6 +36,10 @@ type claimsCriterion struct {
 	g *Generator
 }
 
+func (claimsCriterion) DataType() CriterionDataType {
+	return generator.CriterionDataTypeUnknown
+}
+
 func (claimsCriterion) Names() []string {
 	return []string{"claim", "claims"}
 }

pkg/policy/criteria/cors_preflight.go

@@ -3,6 +3,7 @@ package criteria
 import (
 	"github.com/open-policy-agent/opa/ast"
 
+	"github.com/pomerium/pomerium/pkg/policy/generator"
 	"github.com/pomerium/pomerium/pkg/policy/parser"
 )
 
@@ -16,6 +17,10 @@ type corsPreflightCriterion struct {
 	g *Generator
 }
 
+func (corsPreflightCriterion) DataType() CriterionDataType {
+	return generator.CriterionDataTypeUnused
+}
+
 func (corsPreflightCriterion) Names() []string {
 	return []string{"cors_preflight"}
 }

pkg/policy/criteria/criteria.go

@@ -15,6 +15,8 @@ type (
 	Criterion = generator.Criterion
 	// A CriterionConstructor is a function which returns a Criterion for a Generator.
 	CriterionConstructor = generator.CriterionConstructor
+	// The CriterionDataType indicates the expected type of data for the criterion.
+	CriterionDataType = generator.CriterionDataType
 )
 
 var allCriteria struct {
@@ -39,3 +41,10 @@ func Register(criterionConstructor CriterionConstructor) {
 	allCriteria.a = a
 	allCriteria.Unlock()
 }
+
+const (
+	// CriterionDataTypeStringListMatcher indicates the expected data type is a string list matcher.
+	CriterionDataTypeStringListMatcher CriterionDataType = "string_list_matcher"
+	// CriterionDataTypeStringMatcher indicates the expected data type is a string matcher.
+	CriterionDataTypeStringMatcher CriterionDataType = "string_matcher"
+)

pkg/policy/criteria/domains.go

@@ -1,8 +1,6 @@
 package criteria
 
 import (
-	"fmt"
-
 	"github.com/open-policy-agent/opa/ast"
 
 	"github.com/pomerium/pomerium/pkg/policy/parser"
@@ -25,20 +23,21 @@ type domainsCriterion struct {
 	g *Generator
 }
 
+func (domainsCriterion) DataType() CriterionDataType {
+	return CriterionDataTypeStringMatcher
+}
+
 func (domainsCriterion) Names() []string {
 	return []string{"domain", "domains"}
 }
 
 func (c domainsCriterion) GenerateRule(_ string, data parser.Value) (*ast.Rule, []*ast.Rule, error) {
 	r := c.g.NewRule("domains")
-	r.Body = append(r.Body, ast.Assign.Expr(ast.VarTerm("rule_data"), ast.NewTerm(data.RegoValue())))
 	r.Body = append(r.Body, domainsBody...)
 
-	switch data.(type) {
-	case parser.String:
-		r.Body = append(r.Body, ast.MustParseExpr(`domain == rule_data`))
-	default:
-		return nil, nil, fmt.Errorf("unsupported value type: %T", data)
+	err := matchString(&r.Body, ast.VarTerm("domain"), data)
+	if err != nil {
+		return nil, nil, err
 	}
 
 	return r, []*ast.Rule{

pkg/policy/criteria/domains_test.go

@@ -15,7 +15,8 @@ func TestDomains(t *testing.T) {
 		res, err := evaluate(t, `
 allow:
   and:
-    - domain: example.com
+    - domain:
+        is: example.com
 `, []dataBrokerRecord{}, Input{Session: InputSession{ID: "SESSION_ID"}})
 		require.NoError(t, err)
 		require.Equal(t, false, res["allow"])
@@ -25,7 +26,8 @@ allow:
 		res, err := evaluate(t, `
 allow:
   and:
-    - domain: example.com
+    - domain:
+        is: example.com
 `,
 			[]dataBrokerRecord{
 				&session.Session{
@@ -46,7 +48,8 @@ allow:
 		res, err := evaluate(t, `
 allow:
   and:
-    - domain: example.com
+    - domain:
+        is: example.com
 `,
 			[]dataBrokerRecord{
 				&session.Session{

pkg/policy/criteria/emails.go

@@ -1,10 +1,9 @@
 package criteria
 
 import (
-	"fmt"
-
 	"github.com/open-policy-agent/opa/ast"
 
+	"github.com/pomerium/pomerium/pkg/policy/generator"
 	"github.com/pomerium/pomerium/pkg/policy/parser"
 	"github.com/pomerium/pomerium/pkg/policy/rules"
 )
@@ -25,20 +24,21 @@ type emailsCriterion struct {
 	g *Generator
 }
 
+func (emailsCriterion) DataType() generator.CriterionDataType {
+	return CriterionDataTypeStringMatcher
+}
+
 func (emailsCriterion) Names() []string {
 	return []string{"email", "emails"}
 }
 
 func (c emailsCriterion) GenerateRule(_ string, data parser.Value) (*ast.Rule, []*ast.Rule, error) {
 	r := c.g.NewRule("emails")
-	r.Body = append(r.Body, ast.Assign.Expr(ast.VarTerm("rule_data"), ast.NewTerm(data.RegoValue())))
 	r.Body = append(r.Body, emailsBody...)
 
-	switch data.(type) {
-	case parser.String:
-		r.Body = append(r.Body, ast.MustParseExpr(`email == rule_data`))
-	default:
-		return nil, nil, fmt.Errorf("unsupported value type: %T", data)
+	err := matchString(&r.Body, ast.VarTerm("email"), data)
+	if err != nil {
+		return nil, nil, err
 	}
 
 	return r, []*ast.Rule{

pkg/policy/criteria/emails_test.go

@@ -15,7 +15,8 @@ func TestEmails(t *testing.T) {
 		res, err := evaluate(t, `
 allow:
   and:
-    - email: test@example.com
+    - email:
+        is: test@example.com
 `, []dataBrokerRecord{}, Input{Session: InputSession{ID: "SESSION_ID"}})
 		require.NoError(t, err)
 		require.Equal(t, false, res["allow"])
@@ -25,7 +26,8 @@ allow:
 		res, err := evaluate(t, `
 allow:
   and:
-    - email: test@example.com
+    - email:
+        is: test@example.com
 `,
 			[]dataBrokerRecord{
 				&session.Session{
@@ -46,7 +48,8 @@ allow:
 		res, err := evaluate(t, `
 allow:
   and:
-    - email: test2@example.com
+    - email:
+        is: test2@example.com
 `,
 			[]dataBrokerRecord{
 				&session.Session{

pkg/policy/criteria/groups.go

@@ -1,10 +1,9 @@
 package criteria
 
 import (
-	"fmt"
-
 	"github.com/open-policy-agent/opa/ast"
 
+	"github.com/pomerium/pomerium/pkg/policy/generator"
 	"github.com/pomerium/pomerium/pkg/policy/parser"
 	"github.com/pomerium/pomerium/pkg/policy/rules"
 )
@@ -38,32 +37,27 @@ var groupsBody = ast.Body{
 	ast.MustParseExpr(`
 		groups = array.concat(group_ids, array.concat(group_names, group_emails))
 	`),
-	ast.MustParseExpr(`
-		some group
-	`),
-	ast.MustParseExpr(`
-		group = groups[_]
-	`),
 }
 
 type groupsCriterion struct {
 	g *Generator
 }
 
+func (groupsCriterion) DataType() generator.CriterionDataType {
+	return CriterionDataTypeStringListMatcher
+}
+
 func (groupsCriterion) Names() []string {
 	return []string{"group", "groups"}
 }
 
 func (c groupsCriterion) GenerateRule(_ string, data parser.Value) (*ast.Rule, []*ast.Rule, error) {
 	r := c.g.NewRule("groups")
-	r.Body = append(r.Body, ast.Assign.Expr(ast.VarTerm("rule_data"), ast.NewTerm(data.RegoValue())))
 	r.Body = append(r.Body, groupsBody...)
 
-	switch data.(type) {
-	case parser.String:
-		r.Body = append(r.Body, ast.MustParseExpr(`group == rule_data`))
-	default:
-		return nil, nil, fmt.Errorf("unsupported value type: %T", data)
+	err := matchStringList(&r.Body, ast.VarTerm("groups"), data)
+	if err != nil {
+		return nil, nil, err
 	}
 
 	return r, []*ast.Rule{

pkg/policy/criteria/groups_test.go

@@ -14,8 +14,10 @@ func TestGroups(t *testing.T) {
 		res, err := evaluate(t, `
 allow:
   and:
-    - groups: group1
-    - groups: group2
+    - groups:
+        has: group1
+    - groups:
+        has: group2
 `, []dataBrokerRecord{}, Input{Session: InputSession{ID: "SESSION_ID"}})
 		require.NoError(t, err)
 		require.Equal(t, false, res["allow"])
@@ -25,7 +27,8 @@ allow:
 		res, err := evaluate(t, `
 allow:
   and:
-    - groups: group1
+    - groups:
+        has: group1
 `,
 			[]dataBrokerRecord{
 				&session.Session{
@@ -46,7 +49,8 @@ allow:
 		res, err := evaluate(t, `
 allow:
   and:
-    - groups: "group1@example.com"
+    - groups:
+        has: "group1@example.com"
 `,
 			[]dataBrokerRecord{
 				&session.Session{
@@ -71,7 +75,8 @@ allow:
 		res, err := evaluate(t, `
 allow:
   and:
-    - groups: "Group 1"
+    - groups:
+        has: "Group 1"
 `,
 			[]dataBrokerRecord{
 				&session.Session{

pkg/policy/criteria/invalid_client_certificate.go

@@ -3,6 +3,7 @@ package criteria
 import (
 	"github.com/open-policy-agent/opa/ast"
 
+	"github.com/pomerium/pomerium/pkg/policy/generator"
 	"github.com/pomerium/pomerium/pkg/policy/parser"
 )
 
@@ -16,6 +17,10 @@ type invalidClientCertificateCriterion struct {
 	g *Generator
 }
 
+func (invalidClientCertificateCriterion) DataType() CriterionDataType {
+	return generator.CriterionDataTypeUnused
+}
+
 func (invalidClientCertificateCriterion) Names() []string {
 	return []string{"invalid_client_certificate"}
 }

pkg/policy/criteria/matchers.go

@@ -0,0 +1,99 @@
+package criteria
+
+import (
+	"fmt"
+
+	"github.com/open-policy-agent/opa/ast"
+
+	"github.com/pomerium/pomerium/pkg/policy/parser"
+)
+
+type matcher func(*ast.Body, *ast.Term, parser.Value) error
+
+func matchString(dst *ast.Body, left *ast.Term, right parser.Value) error {
+	obj, ok := right.(parser.Object)
+	if !ok {
+		return fmt.Errorf("expected object for string matcher, got: %T", right)
+	}
+
+	lookup := map[string]matcher{
+		"contains":    matchStringContains,
+		"ends_with":   matchStringEndsWith,
+		"is":          matchStringIs,
+		"starts_with": matchStringStartsWith,
+	}
+	for k, v := range obj {
+		f, ok := lookup[k]
+		if !ok {
+			return fmt.Errorf("unknown string matcher operator: %s", k)
+		}
+		err := f(dst, left, v)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func matchStringContains(dst *ast.Body, left *ast.Term, right parser.Value) error {
+	*dst = append(*dst, ast.Contains.Expr(left, ast.NewTerm(right.RegoValue())))
+	return nil
+}
+
+func matchStringEndsWith(dst *ast.Body, left *ast.Term, right parser.Value) error {
+	*dst = append(*dst, ast.EndsWith.Expr(left, ast.NewTerm(right.RegoValue())))
+	return nil
+}
+
+func matchStringIs(dst *ast.Body, left *ast.Term, right parser.Value) error {
+	*dst = append(*dst, ast.Equal.Expr(left, ast.NewTerm(right.RegoValue())))
+	return nil
+}
+
+func matchStringStartsWith(dst *ast.Body, left *ast.Term, right parser.Value) error {
+	*dst = append(*dst, ast.StartsWith.Expr(left, ast.NewTerm(right.RegoValue())))
+	return nil
+}
+
+func matchStringList(dst *ast.Body, left *ast.Term, right parser.Value) error {
+	obj, ok := right.(parser.Object)
+	if !ok {
+		return fmt.Errorf("expected object for string list matcher, got: %T", right)
+	}
+
+	lookup := map[string]matcher{
+		"has": matchStringListHas,
+	}
+	for k, v := range obj {
+		f, ok := lookup[k]
+		if !ok {
+			return fmt.Errorf("unknown string list matcher operator: %s", k)
+		}
+		err := f(dst, left, v)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func matchStringListHas(dst *ast.Body, left *ast.Term, right parser.Value) error {
+	body := ast.Body{
+		ast.MustParseExpr("some v"),
+		ast.Equality.Expr(ast.VarTerm("v"), ast.RefTerm(left, ast.VarTerm("$0"))),
+	}
+	err := matchStringIs(&body, ast.VarTerm("v"), right)
+	if err != nil {
+		return err
+	}
+	*dst = append(*dst, ast.GreaterThan.Expr(
+		ast.Count.Call(
+			ast.ArrayComprehensionTerm(
+				ast.BooleanTerm(true),
+				body,
+			),
+		),
+		ast.IntNumberTerm(0),
+	))
+	return nil
+}

pkg/policy/criteria/matchers_test.go

@@ -0,0 +1,69 @@
+package criteria
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/format"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/pomerium/pomerium/pkg/policy/parser"
+)
+
+func TestStringMatcher(t *testing.T) {
+	str := func(x interface{}) string {
+		bs := format.MustAst(x)
+		return strings.TrimSpace(string(bs))
+	}
+
+	t.Run("contains", func(t *testing.T) {
+		var body ast.Body
+		err := matchString(&body, ast.VarTerm("example"), parser.Object{
+			"contains": parser.String("test"),
+		})
+		require.NoError(t, err)
+		assert.Equal(t, `contains(example, "test")`, str(body))
+	})
+	t.Run("ends_with", func(t *testing.T) {
+		var body ast.Body
+		err := matchString(&body, ast.VarTerm("example"), parser.Object{
+			"ends_with": parser.String("test"),
+		})
+		require.NoError(t, err)
+		assert.Equal(t, `endswith(example, "test")`, str(body))
+	})
+	t.Run("is", func(t *testing.T) {
+		var body ast.Body
+		err := matchString(&body, ast.VarTerm("example"), parser.Object{
+			"is": parser.String("test"),
+		})
+		require.NoError(t, err)
+		assert.Equal(t, `example == "test"`, str(body))
+	})
+	t.Run("starts_with", func(t *testing.T) {
+		var body ast.Body
+		err := matchString(&body, ast.VarTerm("example"), parser.Object{
+			"starts_with": parser.String("test"),
+		})
+		require.NoError(t, err)
+		assert.Equal(t, `startswith(example, "test")`, str(body))
+	})
+}
+
+func TestStringListMatcher(t *testing.T) {
+	str := func(x interface{}) string {
+		bs := format.MustAst(x)
+		return strings.TrimSpace(string(bs))
+	}
+
+	t.Run("has", func(t *testing.T) {
+		var body ast.Body
+		err := matchStringList(&body, ast.VarTerm("example"), parser.Object{
+			"has": parser.String("test"),
+		})
+		require.NoError(t, err)
+		assert.Equal(t, `count([true | some v; v = example[_]; v == "test"]) > 0`, str(body))
+	})
+}

pkg/policy/criteria/pomerium_routes.go

@@ -3,6 +3,7 @@ package criteria
 import (
 	"github.com/open-policy-agent/opa/ast"
 
+	"github.com/pomerium/pomerium/pkg/policy/generator"
 	"github.com/pomerium/pomerium/pkg/policy/parser"
 )
 
@@ -16,6 +17,10 @@ type pomeriumRoutesCriterion struct {
 	g *Generator
 }
 
+func (pomeriumRoutesCriterion) DataType() generator.CriterionDataType {
+	return generator.CriterionDataTypeUnused
+}
+
 func (pomeriumRoutesCriterion) Names() []string {
 	return []string{"pomerium_routes"}
 }

pkg/policy/criteria/reject.go

@@ -3,6 +3,7 @@ package criteria
 import (
 	"github.com/open-policy-agent/opa/ast"
 
+	"github.com/pomerium/pomerium/pkg/policy/generator"
 	"github.com/pomerium/pomerium/pkg/policy/parser"
 )
 
@@ -14,6 +15,10 @@ type rejectMatcher struct {
 	g *Generator
 }
 
+func (rejectMatcher) DataType() CriterionDataType {
+	return generator.CriterionDataTypeUnused
+}
+
 func (rejectMatcher) Names() []string {
 	return []string{"reject"}
 }

pkg/policy/criteria/users.go

@@ -1,10 +1,9 @@
 package criteria
 
 import (
-	"fmt"
-
 	"github.com/open-policy-agent/opa/ast"
 
+	"github.com/pomerium/pomerium/pkg/policy/generator"
 	"github.com/pomerium/pomerium/pkg/policy/parser"
 	"github.com/pomerium/pomerium/pkg/policy/rules"
 )
@@ -16,26 +15,30 @@ var usersBody = ast.Body{
 	ast.MustParseExpr(`
 		user := get_user(session)
 	`),
+	ast.MustParseExpr(`
+		user_id := user.id
+	`),
 }
 
 type usersCriterion struct {
 	g *Generator
 }
 
+func (usersCriterion) DataType() generator.CriterionDataType {
+	return CriterionDataTypeStringMatcher
+}
+
 func (usersCriterion) Names() []string {
 	return []string{"user", "users"}
 }
 
 func (c usersCriterion) GenerateRule(_ string, data parser.Value) (*ast.Rule, []*ast.Rule, error) {
 	r := c.g.NewRule("users")
-	r.Body = append(r.Body, ast.Assign.Expr(ast.VarTerm("rule_data"), ast.NewTerm(data.RegoValue())))
 	r.Body = append(r.Body, usersBody...)
 
-	switch data.(type) {
-	case parser.String:
-		r.Body = append(r.Body, ast.MustParseExpr(`user.id == rule_data`))
-	default:
-		return nil, nil, fmt.Errorf("unsupported value type: %T", data)
+	err := matchString(&r.Body, ast.VarTerm("user_id"), data)
+	if err != nil {
+		return nil, nil, err
 	}
 
 	return r, []*ast.Rule{

pkg/policy/generator/criterion.go

@@ -8,6 +8,7 @@ import (
 
 // A Criterion generates rego rules based on data.
 type Criterion interface {
+	DataType() CriterionDataType
 	Names() []string
 	GenerateRule(subPath string, data parser.Value) (rule *ast.Rule, additionalRules []*ast.Rule, err error)
 }
@@ -17,10 +18,16 @@ type CriterionConstructor func(*Generator) Criterion
 
 // A criterionFunc is a criterion implemented as a function and a list of names.
 type criterionFunc struct {
+	dataType     CriterionDataType
 	names        []string
 	generateRule func(subPath string, data parser.Value) (rule *ast.Rule, additionalRules []*ast.Rule, err error)
 }
 
+// DataType returns the criterion data type.
+func (c criterionFunc) DataType() CriterionDataType {
+	return c.dataType
+}
+
 // Names returns the names of the criterion.
 func (c criterionFunc) Names() []string {
 	return c.names
@@ -33,6 +40,7 @@ func (c criterionFunc) GenerateRule(subPath string, data parser.Value) (rule *as
 
 // NewCriterionFunc creates a new Criterion from a function.
 func NewCriterionFunc(
+	dataType CriterionDataType,
 	names []string,
 	f func(subPath string, data parser.Value) (rule *ast.Rule, additionalRules []*ast.Rule, err error),
 ) Criterion {
@@ -41,3 +49,14 @@ func NewCriterionFunc(
 		generateRule: f,
 	}
 }
+
+// A CriterionDataType describes the expected format of the data to be sent to the criterion.
+type CriterionDataType string
+
+const (
+	// CriterionDataTypeUnknown indicates that the type of data is unknown.
+	CriterionDataTypeUnknown CriterionDataType = ""
+
+	// CriterionDataTypeUnused indicates that the data is unused.
+	CriterionDataTypeUnused CriterionDataType = "unused"
+)

pkg/policy/generator/generator_test.go

@@ -12,7 +12,7 @@ import (
 
 func Test(t *testing.T) {
 	g := New(WithCriterion(func(g *Generator) Criterion {
-		return NewCriterionFunc([]string{"accept"}, func(subPath string, data parser.Value) (rule *ast.Rule, additionalRules []*ast.Rule, err error) {
+		return NewCriterionFunc(CriterionDataTypeUnused, []string{"accept"}, func(subPath string, data parser.Value) (rule *ast.Rule, additionalRules []*ast.Rule, err error) {
 			rule = g.NewRule("accept")
 			rule.Body = append(rule.Body, ast.MustParseExpr("1 == 1"))
 			return rule, nil, nil