authorize: custom rego policies (#1123)

* add support for custom rego policies

* add support for passing custom policies

authorize/evaluator/custom.go

@@ -0,0 +1,116 @@
+package evaluator
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+
+	"github.com/open-policy-agent/opa/rego"
+	"github.com/open-policy-agent/opa/storage"
+)
+
+// A CustomEvaluatorRequest is the data needed to evaluate a custom rego policy.
+type CustomEvaluatorRequest struct {
+	RegoPolicy string
+	HTTP       RequestHTTP    `json:"http"`
+	Session    RequestSession `json:"session"`
+}
+
+// A CustomEvaluatorResponse is the response from the evaluation of a custom rego policy.
+type CustomEvaluatorResponse struct {
+	Allowed bool
+	Denied  bool
+	Reason  string
+}
+
+// A CustomEvaluator evaluates custom rego policies.
+type CustomEvaluator struct {
+	store   storage.Store
+	mu      sync.Mutex
+	queries map[string]rego.PreparedEvalQuery
+}
+
+// NewCustomEvaluator creates a new CustomEvaluator.
+func NewCustomEvaluator(store storage.Store) *CustomEvaluator {
+	ce := &CustomEvaluator{
+		store:   store,
+		queries: map[string]rego.PreparedEvalQuery{},
+	}
+	return ce
+}
+
+// Evaluate evaluates the custom rego policy.
+func (ce *CustomEvaluator) Evaluate(ctx context.Context, req *CustomEvaluatorRequest) (*CustomEvaluatorResponse, error) {
+	q, err := ce.getPreparedEvalQuery(ctx, req.RegoPolicy)
+	if err != nil {
+		return nil, err
+	}
+
+	resultSet, err := q.Eval(ctx, rego.EvalInput(struct {
+		HTTP    RequestHTTP    `json:"http"`
+		Session RequestSession `json:"session"`
+	}{HTTP: req.HTTP, Session: req.Session}))
+	if err != nil {
+		return nil, err
+	}
+
+	vars, ok := resultSet[0].Bindings.WithoutWildcards()["result"].(map[string]interface{})
+	if !ok {
+		vars = make(map[string]interface{})
+	}
+
+	res := &CustomEvaluatorResponse{}
+	res.Allowed, _ = vars["allow"].(bool)
+	if v, ok := vars["deny"]; ok {
+		// support `deny = true`
+		if b, ok := v.(bool); ok {
+			res.Denied = b
+		}
+
+		// support `deny[reason] = true`
+		if m, ok := v.(map[string]interface{}); ok {
+			for mk, mv := range m {
+				if b, ok := mv.(bool); ok {
+					res.Denied = b
+					res.Reason = mk
+				}
+			}
+		}
+	}
+	return res, nil
+}
+
+func (ce *CustomEvaluator) getPreparedEvalQuery(ctx context.Context, src string) (rego.PreparedEvalQuery, error) {
+	ce.mu.Lock()
+	defer ce.mu.Unlock()
+
+	q, ok := ce.queries[src]
+	if ok {
+		return q, nil
+	}
+
+	r := rego.New(
+		rego.Store(ce.store),
+		rego.Module("pomerium.custom_policy", src),
+		rego.Query("result = data.pomerium.custom_policy"),
+	)
+	q, err := r.PrepareForEval(ctx)
+	if err != nil {
+		// if no package is in the src, add it
+		if strings.Contains(err.Error(), "package expected") {
+			r = rego.New(
+				rego.Store(ce.store),
+				rego.Module("pomerium.custom_policy", "package pomerium.custom_policy\n\n"+src),
+				rego.Query("result = data.pomerium.custom_policy"),
+			)
+			q, err = r.PrepareForEval(ctx)
+		}
+	}
+	if err != nil {
+		return q, fmt.Errorf("invalid rego policy: %w", err)
+	}
+
+	ce.queries[src] = q
+	return q, nil
+}

authorize/evaluator/custom_test.go

@@ -0,0 +1,57 @@
+package evaluator
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCustomEvaluator(t *testing.T) {
+	ctx, clearTimeout := context.WithTimeout(context.Background(), time.Second*10)
+	defer clearTimeout()
+
+	store := NewStore()
+	t.Run("bool deny", func(t *testing.T) {
+		ce := NewCustomEvaluator(store.opaStore)
+		res, err := ce.Evaluate(ctx, &CustomEvaluatorRequest{
+			RegoPolicy: `
+				package pomerium.custom_policy
+
+				deny = true
+			`,
+		})
+		if !assert.NoError(t, err) {
+			return
+		}
+		assert.Equal(t, true, res.Denied)
+		assert.Empty(t, res.Reason)
+	})
+	t.Run("set deny", func(t *testing.T) {
+		ce := NewCustomEvaluator(store.opaStore)
+		res, err := ce.Evaluate(ctx, &CustomEvaluatorRequest{
+			RegoPolicy: `
+				package pomerium.custom_policy
+
+				deny["test"] = true
+			`,
+		})
+		if !assert.NoError(t, err) {
+			return
+		}
+		assert.Equal(t, true, res.Denied)
+		assert.Equal(t, "test", res.Reason)
+	})
+	t.Run("missing package", func(t *testing.T) {
+		ce := NewCustomEvaluator(store.opaStore)
+		res, err := ce.Evaluate(ctx, &CustomEvaluatorRequest{
+			RegoPolicy: `allow = true`,
+		})
+		if !assert.NoError(t, err) {
+			return
+		}
+		assert.NotNil(t, res)
+	})
+
+}

authorize/evaluator/evaluator.go

@@ -37,6 +37,7 @@ const (
 
 // Evaluator specifies the interface for a policy engine.
 type Evaluator struct {
+	custom   *CustomEvaluator
 	rego     *rego.Rego
 	query    rego.PreparedEvalQuery
 	policies []config.Policy
@@ -50,6 +51,7 @@ type Evaluator struct {
 // New creates a new Evaluator.
 func New(options *config.Options, store *Store) (*Evaluator, error) {
 	e := &Evaluator{
+		custom:           NewCustomEvaluator(store.opaStore),
 		authenticateHost: options.AuthenticateURL.Host,
 		policies:         options.Policies,
 	}
@@ -149,6 +151,23 @@ func (e *Evaluator) Evaluate(ctx context.Context, req *Request) (*Result, error)
 	}
 
 	allow := allowed(res[0].Bindings.WithoutWildcards())
+	// evaluate any custom policies
+	if allow {
+		for _, src := range req.CustomPolicies {
+			cres, err := e.custom.Evaluate(ctx, &CustomEvaluatorRequest{
+				RegoPolicy: src,
+				HTTP:       req.HTTP,
+				Session:    req.Session,
+			})
+			if err != nil {
+				return nil, err
+			}
+			allow = allow && (!cres.Allowed || cres.Denied)
+			if cres.Reason != "" {
+				evalResult.Message = cres.Reason
+			}
+		}
+	}
 	if allow {
 		evalResult.Status = http.StatusOK
 		evalResult.Message = "OK"
@@ -162,7 +181,9 @@ func (e *Evaluator) Evaluate(ctx context.Context, req *Request) (*Result, error)
 	}
 
 	evalResult.Status = http.StatusForbidden
-	evalResult.Message = "forbidden"
+	if evalResult.Message == "" {
+		evalResult.Message = "forbidden"
+	}
 	return evalResult, nil
 }
 
@@ -258,6 +279,7 @@ type (
 		DataBrokerData DataBrokerData `json:"databroker_data"`
 		HTTP           RequestHTTP    `json:"http"`
 		Session        RequestSession `json:"session"`
+		CustomPolicies []string
 	}
 
 	// RequestHTTP is the HTTP field in the request.

go.sum

@@ -385,8 +385,6 @@ github.com/onsi/gocleanup v0.0.0-20140331211545-c1a5478700b5/go.mod h1:tHaogb+iP
 github.com/onsi/gomega v1.4.3/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY=
 github.com/onsi/gomega v1.8.1 h1:C5Dqfs/LeauYDX0jJXIe2SWmwCbGzx9yF8C8xy3Lh34=
 github.com/onsi/gomega v1.8.1/go.mod h1:Ho0h+IUsWyvy1OpqCwxlQ/21gkhVunqlU8fDGcoTdcA=
-github.com/open-policy-agent/opa v0.21.1 h1:c4lUnB0mO2KssiUnyh6Y9IGhggvXI3EgObkmhVTvEqQ=
-github.com/open-policy-agent/opa v0.21.1/go.mod h1:cZaTfhxsj7QdIiUI0U9aBtOLLTqVNe+XE60+9kZKLHw=
 github.com/open-policy-agent/opa v0.22.0 h1:KZvn0uMQIorBIwYk8Vc89dp8No9FIEF8eFl0sc1r/1U=
 github.com/open-policy-agent/opa v0.22.0/go.mod h1:rrwxoT/b011T0cyj+gg2VvxqTtn6N3gp/jzmr3fjW44=
 github.com/openzipkin/zipkin-go v0.1.6/go.mod h1:QgAqvLzwWbR/WpD4A3cGpPtJrZXNIiJc5AZX7/PBEpw=
@@ -496,6 +494,7 @@ github.com/spf13/cast v1.3.1 h1:nFm6S0SMdyzrzcmThSipiEubIDy8WEXKNZ0UOgiRpng=
 github.com/spf13/cast v1.3.1/go.mod h1:Qx5cxh0v+4UWYiBimWS+eyWzqEqokIECu5etghLkUJE=
 github.com/spf13/cobra v0.0.0-20181021141114-fe5e611709b0 h1:BgSbPgT2Zu8hDen1jJDGLWO8voaSRVrwsk18Q/uSh5M=
 github.com/spf13/cobra v0.0.0-20181021141114-fe5e611709b0/go.mod h1:1l0Ry5zgKvJasoi3XT1TypsSe7PqH0Sj9dhYf7v3XqQ=
+github.com/spf13/cobra v1.0.0 h1:6m/oheQuQ13N9ks4hubMG6BnvwOeaJrqSPLahSnczz8=
 github.com/spf13/cobra v1.0.0/go.mod h1:/6GTrnGXV9HjY+aR4k0oJ5tcvakLuG6EuKReYlHNrgE=
 github.com/spf13/jwalterweatherman v1.0.0/go.mod h1:cQK4TGJAtQXfYWX+Ddv3mKDzgVb68N+wFjFa4jdeBTo=
 github.com/spf13/jwalterweatherman v1.1.0 h1:ue6voC5bR5F8YxI5S67j9i582FU4Qvo2bmqnqMYADFk=
@@ -784,8 +783,6 @@ google.golang.org/genproto v0.0.0-20200224152610-e50cd9704f63/go.mod h1:55QSHmfG
 google.golang.org/genproto v0.0.0-20200305110556-506484158171/go.mod h1:55QSHmfGQM9UVYDPBsyGGes0y52j32PQ3BqQfXhyH3c=
 google.golang.org/genproto v0.0.0-20200331122359-1ee6d9798940/go.mod h1:55QSHmfGQM9UVYDPBsyGGes0y52j32PQ3BqQfXhyH3c=
 google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEYHJ7i3ixzK3sjbqSGDJWnxyFXZblF3eUsNvo=
-google.golang.org/genproto v0.0.0-20200711021454-869866162049 h1:YFTFpQhgvrLrmxtiIncJxFXeCyq84ixuKWVCaCAi9Oc=
-google.golang.org/genproto v0.0.0-20200711021454-869866162049/go.mod h1:FWY/as6DDZQgahTzZj3fqbO1CbirC29ZNUFHwi0/+no=
 google.golang.org/genproto v0.0.0-20200715011427-11fb19a81f2c h1:6DWnZZ6EY/59QRRQttZKiktVL23UuQYs7uy75MhhLRM=
 google.golang.org/genproto v0.0.0-20200715011427-11fb19a81f2c/go.mod h1:FWY/as6DDZQgahTzZj3fqbO1CbirC29ZNUFHwi0/+no=
 google.golang.org/grpc v1.17.0/go.mod h1:6QZJwpn2B+Zp71q/5VxRsJ6NXXVCE5NRUHRo+f3cWCs=