authorize: avoid serializing databroker data map to improve performance (#995)

authorize/evaluator/evaluator.go

@@ -30,6 +30,12 @@ import (
 	"github.com/pomerium/pomerium/internal/log"
 )
 
+const (
+	sessionTypeURL       = "type.googleapis.com/session.Session"
+	userTypeURL          = "type.googleapis.com/user.User"
+	directoryUserTypeURL = "type.googleapis.com/directory.User"
+)
+
 // Evaluator specifies the interface for a policy engine.
 type Evaluator struct {
 	rego  *rego.Rego
@@ -202,15 +208,25 @@ func (e *Evaluator) SignedJWT(req *Request) (string, error) {
 }
 
 type input struct {
-	DataBrokerData           DataBrokerData `json:"databroker_data"`
+	DataBrokerData           dataBrokerDataInput `json:"databroker_data"`
 	HTTP                     RequestHTTP         `json:"http"`
 	Session                  RequestSession      `json:"session"`
 	IsValidClientCertificate bool                `json:"is_valid_client_certificate"`
 }
 
+type dataBrokerDataInput struct {
+	Session       interface{} `json:"session,omitempty"`
+	User          interface{} `json:"user,omitempty"`
+	DirectoryUser interface{} `json:"directory_user,omitempty"`
+}
+
 func (e *Evaluator) newInput(req *Request, isValidClientCertificate bool) *input {
 	i := new(input)
-	i.DataBrokerData = req.DataBrokerData
+	i.DataBrokerData.Session = req.DataBrokerData.Get(sessionTypeURL, req.Session.ID)
+	if obj, ok := i.DataBrokerData.Session.(interface{ GetUserId() string }); ok {
+		i.DataBrokerData.User = req.DataBrokerData.Get(userTypeURL, obj.GetUserId())
+		i.DataBrokerData.DirectoryUser = req.DataBrokerData.Get(directoryUserTypeURL, obj.GetUserId())
+	}
 	i.HTTP = req.HTTP
 	i.Session = req.Session
 	i.IsValidClientCertificate = isValidClientCertificate

authorize/evaluator/evaluator_test.go

@@ -1,16 +1,23 @@
 package evaluator
 
 import (
+	"context"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"net/url"
 	"testing"
 
+	"github.com/golang/protobuf/ptypes"
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/pomerium/pomerium/config"
+	"github.com/pomerium/pomerium/internal/grpc/databroker"
 	"github.com/pomerium/pomerium/internal/grpc/directory"
+	"github.com/pomerium/pomerium/internal/grpc/session"
+	"github.com/pomerium/pomerium/internal/grpc/user"
 )
 
 func TestJSONMarshal(t *testing.T) {
@@ -25,7 +32,7 @@ func TestJSONMarshal(t *testing.T) {
 		"type.googleapis.com/user.User":       map[string]interface{}{},
 	}
 
-	bs, _ := json.Marshal(input{
+	bs, _ := json.Marshal(new(Evaluator).newInput(&Request{
 		DataBrokerData: dbd,
 		HTTP: RequestHTTP{
 			Method: "GET",
@@ -40,18 +47,9 @@ func TestJSONMarshal(t *testing.T) {
 			ImpersonateEmail:  "y@example.com",
 			ImpersonateGroups: []string{"group1"},
 		},
-		IsValidClientCertificate: true,
+	}, true))
-	})
 	assert.JSONEq(t, `{
 		"databroker_data": {
-			"type.googleapis.com/directory.User": {
-				"user1": {
-					"id": "user1",
-					"groups": ["group1", "group2"]
-				}
-			},
-			"type.googleapis.com/session.Session": {},
-			"type.googleapis.com/user.User": {}
 		},
 		"http": {
 			"client_certificate": "CLIENT_CERTIFICATE",
@@ -97,3 +95,68 @@ func mustParseURL(str string) *url.URL {
 	}
 	return u
 }
+
+func BenchmarkEvaluator_Evaluate(b *testing.B) {
+	e, err := New(&config.Options{
+		AuthenticateURL: mustParseURL("https://authn.example.com"),
+	})
+	if !assert.NoError(b, err) {
+		return
+	}
+
+	lastSessionID := ""
+
+	dbd := make(DataBrokerData)
+	for i := 0; i < 100; i++ {
+		sessionID := uuid.New().String()
+		lastSessionID = sessionID
+		userID := uuid.New().String()
+		data, _ := ptypes.MarshalAny(&session.Session{
+			Version: fmt.Sprint(i),
+			Id:      sessionID,
+			UserId:  userID,
+			IdToken: &session.IDToken{
+				Issuer:   "benchmark",
+				Subject:  userID,
+				IssuedAt: ptypes.TimestampNow(),
+			},
+			OauthToken: &session.OAuthToken{
+				AccessToken:  "ACCESS TOKEN",
+				TokenType:    "Bearer",
+				RefreshToken: "REFRESH TOKEN",
+			},
+		})
+		dbd.Update(&databroker.Record{
+			Version: fmt.Sprint(i),
+			Type:    "type.googleapis.com/session.Session",
+			Id:      sessionID,
+			Data:    data,
+		})
+		data, _ = ptypes.MarshalAny(&user.User{
+			Version: fmt.Sprint(i),
+			Id:      userID,
+		})
+		dbd.Update(&databroker.Record{
+			Version: fmt.Sprint(i),
+			Type:    "type.googleapis.com/user.User",
+			Id:      userID,
+			Data:    data,
+		})
+	}
+
+	b.ResetTimer()
+	ctx := context.Background()
+	for i := 0; i < b.N; i++ {
+		e.Evaluate(ctx, &Request{
+			DataBrokerData: dbd,
+			HTTP: RequestHTTP{
+				Method:  "GET",
+				URL:     "https://example.com/path",
+				Headers: map[string]string{},
+			},
+			Session: RequestSession{
+				ID: lastSessionID,
+			},
+		})
+	}
+}

authorize/evaluator/opa/policy/authz.rego

@@ -4,9 +4,9 @@ default allow = false
 
 
 route := first_allowed_route(input.http.url)
-session := input.databroker_data["type.googleapis.com/session.Session"][input.session.id]
+session := input.databroker_data.session
-user := input.databroker_data["type.googleapis.com/user.User"][session.user_id]
+user := input.databroker_data.user
-directory_user := input.databroker_data["type.googleapis.com/directory.User"][session.user_id]
+directory_user := input.databroker_data.directory_user
 
 
 # allow public

authorize/evaluator/opa/policy/authz_test.rego

@@ -7,16 +7,12 @@ test_email_allowed {
 			"allowed_users": ["x@example.com"]
 		}] with
 		input.databroker_data as {
-			"type.googleapis.com/session.Session": {
+			"session": {
-				"session1": {
                 "user_id": "user1"
-				}
             },
-			"type.googleapis.com/user.User": {
+			"user": {
-				"user1": {
 				"email": "x@example.com"
 			}
-			}
 		} with
 		input.http as { "url": "http://example.com" } with
 		input.session as { "id": "session1" }