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.