authorize: reuse policy evaluators where possible

authorize/authorize.go

@@ -8,6 +8,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/rs/zerolog"
 	"golang.org/x/sync/errgroup"
 
 	"github.com/pomerium/pomerium/authorize/evaluator"
@@ -46,7 +47,7 @@ func New(cfg *config.Config) (*Authorize, error) {
 	}
 	a.accessTracker = NewAccessTracker(a, accessTrackerMaxSize, accessTrackerDebouncePeriod)
 
-	state, err := newAuthorizeStateFromConfig(cfg, a.store)
+	state, err := newAuthorizeStateFromConfig(cfg, a.store, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -86,11 +87,15 @@ func validateOptions(o *config.Options) error {
 }
 
 // newPolicyEvaluator returns an policy evaluator.
-func newPolicyEvaluator(opts *config.Options, store *store.Store) (*evaluator.Evaluator, error) {
+func newPolicyEvaluator(
+	opts *config.Options, store *store.Store, previous *evaluator.Evaluator,
+) (*evaluator.Evaluator, error) {
 	metrics.AddPolicyCountCallback("pomerium-authorize", func() int64 {
 		return int64(len(opts.GetAllPolicies()))
 	})
-	ctx := context.Background()
+	ctx := log.WithContext(context.Background(), func(c zerolog.Context) zerolog.Context {
+		return c.Str("service", "authorize")
+	})
 	ctx, span := trace.StartSpan(ctx, "authorize.newPolicyEvaluator")
 	defer span.End()
 
@@ -126,7 +131,7 @@ func newPolicyEvaluator(opts *config.Options, store *store.Store) (*evaluator.Ev
 			"authorize: internal error: couldn't build client cert constraints: %w", err)
 	}
 
-	return evaluator.New(ctx, store,
+	return evaluator.New(ctx, store, previous,
 		evaluator.WithPolicies(opts.GetAllPolicies()),
 		evaluator.WithClientCA(clientCA),
 		evaluator.WithAddDefaultClientCertificateRule(addDefaultClientCertificateRule),
@@ -141,8 +146,9 @@ func newPolicyEvaluator(opts *config.Options, store *store.Store) (*evaluator.Ev
 
 // OnConfigChange updates internal structures based on config.Options
 func (a *Authorize) OnConfigChange(ctx context.Context, cfg *config.Config) {
+	currentState := a.state.Load()
 	a.currentOptions.Store(cfg.Options)
-	if state, err := newAuthorizeStateFromConfig(cfg, a.store); err != nil {
+	if state, err := newAuthorizeStateFromConfig(cfg, a.store, currentState.evaluator); err != nil {
 		log.Error(ctx).Err(err).Msg("authorize: error updating state")
 	} else {
 		a.state.Store(state)

authorize/authorize_test.go

@@ -179,7 +179,7 @@ func TestNewPolicyEvaluator_addDefaultClientCertificateRule(t *testing.T) {
 			c.opts.Policies = []config.Policy{{
 				To: mustParseWeightedURLs(t, "http://example.com"),
 			}}
-			e, err := newPolicyEvaluator(c.opts, store)
+			e, err := newPolicyEvaluator(c.opts, store, nil)
 			require.NoError(t, err)
 
 			r, err := e.Evaluate(context.Background(), &evaluator.Request{

authorize/check_response_test.go

@@ -131,7 +131,7 @@ func TestAuthorize_okResponse(t *testing.T) {
 	a := &Authorize{currentOptions: config.NewAtomicOptions(), state: atomicutil.NewValue(new(authorizeState))}
 	a.currentOptions.Store(opt)
 	a.store = store.New()
-	pe, err := newPolicyEvaluator(opt, a.store)
+	pe, err := newPolicyEvaluator(opt, a.store, nil)
 	require.NoError(t, err)
 	a.state.Load().evaluator = pe
 

authorize/evaluator/config.go

@@ -2,18 +2,24 @@ package evaluator
 
 import (
 	"github.com/pomerium/pomerium/config"
+	"github.com/pomerium/pomerium/internal/hashutil"
 )
 
 type evaluatorConfig struct {
-	policies                                          []config.Policy
-	clientCA                                          []byte
-	clientCRL                                         []byte
-	addDefaultClientCertificateRule                   bool
-	clientCertConstraints                             ClientCertConstraints
-	signingKey                                        []byte
-	authenticateURL                                   string
-	googleCloudServerlessAuthenticationServiceAccount string
-	jwtClaimsHeaders                                  config.JWTClaimHeaders
+	Policies                                          []config.Policy `hash:"-"`
+	ClientCA                                          []byte
+	ClientCRL                                         []byte
+	AddDefaultClientCertificateRule                   bool
+	ClientCertConstraints                             ClientCertConstraints
+	SigningKey                                        []byte
+	AuthenticateURL                                   string
+	GoogleCloudServerlessAuthenticationServiceAccount string
+	JwtClaimsHeaders                                  config.JWTClaimHeaders
+}
+
+// cacheKey() returns a hash over the configuration, except for the policies.
+func (e *evaluatorConfig) cacheKey() uint64 {
+	return hashutil.MustHash(e)
 }
 
 // An Option customizes the evaluator config.
@@ -30,64 +36,64 @@ func getConfig(options ...Option) *evaluatorConfig {
 // WithPolicies sets the policies in the config.
 func WithPolicies(policies []config.Policy) Option {
 	return func(cfg *evaluatorConfig) {
-		cfg.policies = policies
+		cfg.Policies = policies
 	}
 }
 
 // WithClientCA sets the client CA in the config.
 func WithClientCA(clientCA []byte) Option {
 	return func(cfg *evaluatorConfig) {
-		cfg.clientCA = clientCA
+		cfg.ClientCA = clientCA
 	}
 }
 
 // WithClientCRL sets the client CRL in the config.
 func WithClientCRL(clientCRL []byte) Option {
 	return func(cfg *evaluatorConfig) {
-		cfg.clientCRL = clientCRL
+		cfg.ClientCRL = clientCRL
 	}
 }
 
 // WithAddDefaultClientCertificateRule sets whether to add a default
 // invalid_client_certificate deny rule to all policies.
 func WithAddDefaultClientCertificateRule(addDefaultClientCertificateRule bool) Option {
 	return func(cfg *evaluatorConfig) {
-		cfg.addDefaultClientCertificateRule = addDefaultClientCertificateRule
+		cfg.AddDefaultClientCertificateRule = addDefaultClientCertificateRule
 	}
 }
 
 // WithClientCertConstraints sets addition client certificate constraints.
 func WithClientCertConstraints(constraints *ClientCertConstraints) Option {
 	return func(cfg *evaluatorConfig) {
-		cfg.clientCertConstraints = *constraints
+		cfg.ClientCertConstraints = *constraints
 	}
 }
 
 // WithSigningKey sets the signing key and algorithm in the config.
 func WithSigningKey(signingKey []byte) Option {
 	return func(cfg *evaluatorConfig) {
-		cfg.signingKey = signingKey
+		cfg.SigningKey = signingKey
 	}
 }
 
 // WithAuthenticateURL sets the authenticate URL in the config.
 func WithAuthenticateURL(authenticateURL string) Option {
 	return func(cfg *evaluatorConfig) {
-		cfg.authenticateURL = authenticateURL
+		cfg.AuthenticateURL = authenticateURL
 	}
 }
 
 // WithGoogleCloudServerlessAuthenticationServiceAccount sets the google cloud serverless authentication service
 // account in the config.
 func WithGoogleCloudServerlessAuthenticationServiceAccount(serviceAccount string) Option {
 	return func(cfg *evaluatorConfig) {
-		cfg.googleCloudServerlessAuthenticationServiceAccount = serviceAccount
+		cfg.GoogleCloudServerlessAuthenticationServiceAccount = serviceAccount
 	}
 }
 
 // WithJWTClaimsHeaders sets the JWT claims headers in the config.
 func WithJWTClaimsHeaders(headers config.JWTClaimHeaders) Option {
 	return func(cfg *evaluatorConfig) {
-		cfg.jwtClaimsHeaders = headers
+		cfg.JwtClaimsHeaders = headers
 	}
 }

authorize/evaluator/evaluator.go

@@ -95,44 +95,78 @@ type Evaluator struct {
 	clientCA              []byte
 	clientCRL             []byte
 	clientCertConstraints ClientCertConstraints
+
+	cfgCacheKey uint64
 }
 
 // New creates a new Evaluator.
-func New(ctx context.Context, store *store.Store, options ...Option) (*Evaluator, error) {
-	e := &Evaluator{store: store}
-
+func New(
+	ctx context.Context, store *store.Store, previous *Evaluator, options ...Option,
+) (*Evaluator, error) {
 	cfg := getConfig(options...)
 
-	err := e.updateStore(cfg)
+	err := updateStore(store, cfg)
 	if err != nil {
 		return nil, err
 	}
 
-	e.headersEvaluators, err = NewHeadersEvaluator(ctx, store)
+	e := &Evaluator{
+		store:                 store,
+		clientCA:              cfg.ClientCA,
+		clientCRL:             cfg.ClientCRL,
+		clientCertConstraints: cfg.ClientCertConstraints,
+		cfgCacheKey:           cfg.cacheKey(),
+	}
+
+	// If there is a previous Evaluator constructed from the same settings, we
+	// can reuse the HeadersEvaluator along with any PolicyEvaluators for
+	// unchanged policies.
+	var cachedPolicyEvaluators map[uint64]*PolicyEvaluator
+	if previous != nil && previous.cfgCacheKey == e.cfgCacheKey {
+		e.headersEvaluators = previous.headersEvaluators
+		cachedPolicyEvaluators = previous.policyEvaluators
+	} else {
+		e.headersEvaluators, err = NewHeadersEvaluator(ctx, store)
+		if err != nil {
+			return nil, err
+		}
+	}
+	e.policyEvaluators, err = getOrCreatePolicyEvaluators(ctx, cfg, store, cachedPolicyEvaluators)
 	if err != nil {
 		return nil, err
 	}
 
-	e.clientCA = cfg.clientCA
-	e.clientCRL = cfg.clientCRL
-	e.clientCertConstraints = cfg.clientCertConstraints
+	return e, nil
+}
 
-	e.policyEvaluators = make(map[uint64]*PolicyEvaluator)
-	for i := range cfg.policies {
-		configPolicy := cfg.policies[i]
+func getOrCreatePolicyEvaluators(
+	ctx context.Context, cfg *evaluatorConfig, store *store.Store,
+	cachedPolicyEvaluators map[uint64]*PolicyEvaluator,
+) (map[uint64]*PolicyEvaluator, error) {
+	var newCount, reusedCount int
+	m := make(map[uint64]*PolicyEvaluator)
+	for i := range cfg.Policies {
+		configPolicy := cfg.Policies[i]
 		id, err := configPolicy.RouteID()
 		if err != nil {
 			return nil, fmt.Errorf("authorize: error computing policy route id: %w", err)
 		}
+		p := cachedPolicyEvaluators[id]
+		if p != nil && p.policyChecksum == configPolicy.Checksum() {
+			m[id] = p
+			reusedCount++
+			continue
+		}
 		policyEvaluator, err :=
-			NewPolicyEvaluator(ctx, store, &configPolicy, cfg.addDefaultClientCertificateRule)
+			NewPolicyEvaluator(ctx, store, &configPolicy, cfg.AddDefaultClientCertificateRule)
 		if err != nil {
 			return nil, err
 		}
-		e.policyEvaluators[id] = policyEvaluator
+		m[id] = policyEvaluator
+		newCount++
 	}
-
-	return e, nil
+	log.Info(ctx).Msgf("updated policy evaluators: %d created, %d reused", newCount, reusedCount)
+	return m, nil
 }
 
 // Evaluate evaluates the rego for the given policy and generates the identity headers.
@@ -251,26 +285,26 @@ func (e *Evaluator) getClientCA(policy *config.Policy) (string, error) {
 	return string(e.clientCA), nil
 }
 
-func (e *Evaluator) updateStore(cfg *evaluatorConfig) error {
+func updateStore(store *store.Store, cfg *evaluatorConfig) error {
 	jwk, err := getJWK(cfg)
 	if err != nil {
 		return fmt.Errorf("authorize: couldn't create signer: %w", err)
 	}
 
-	e.store.UpdateGoogleCloudServerlessAuthenticationServiceAccount(
-		cfg.googleCloudServerlessAuthenticationServiceAccount,
+	store.UpdateGoogleCloudServerlessAuthenticationServiceAccount(
+		cfg.GoogleCloudServerlessAuthenticationServiceAccount,
 	)
-	e.store.UpdateJWTClaimHeaders(cfg.jwtClaimsHeaders)
-	e.store.UpdateRoutePolicies(cfg.policies)
-	e.store.UpdateSigningKey(jwk)
+	store.UpdateJWTClaimHeaders(cfg.JwtClaimsHeaders)
+	store.UpdateRoutePolicies(cfg.Policies)
+	store.UpdateSigningKey(jwk)
 
 	return nil
 }
 
 func getJWK(cfg *evaluatorConfig) (*jose.JSONWebKey, error) {
 	var decodedCert []byte
 	// if we don't have a signing key, generate one
-	if len(cfg.signingKey) == 0 {
+	if len(cfg.SigningKey) == 0 {
 		key, err := cryptutil.NewSigningKey()
 		if err != nil {
 			return nil, fmt.Errorf("couldn't generate signing key: %w", err)
@@ -280,7 +314,7 @@ func getJWK(cfg *evaluatorConfig) (*jose.JSONWebKey, error) {
 			return nil, fmt.Errorf("bad signing key: %w", err)
 		}
 	} else {
-		decodedCert = cfg.signingKey
+		decodedCert = cfg.SigningKey
 	}
 
 	jwk, err := cryptutil.PrivateJWKFromBytes(decodedCert)

authorize/evaluator/evaluator_test.go

@@ -36,7 +36,7 @@ func TestEvaluator(t *testing.T) {
 		store := store.New()
 		store.UpdateJWTClaimHeaders(config.NewJWTClaimHeaders("email", "groups", "user", "CUSTOM_KEY"))
 		store.UpdateSigningKey(privateJWK)
-		e, err := New(ctx, store, options...)
+		e, err := New(ctx, store, nil, options...)
 		require.NoError(t, err)
 		return e.Evaluate(ctx, req)
 	}

authorize/evaluator/policy_evaluator.go

@@ -103,7 +103,8 @@ func (q policyQuery) checksum() string {
 
 // A PolicyEvaluator evaluates policies.
 type PolicyEvaluator struct {
-	queries []policyQuery
+	queries        []policyQuery
+	policyChecksum uint64
 }
 
 // NewPolicyEvaluator creates a new PolicyEvaluator.
@@ -112,6 +113,7 @@ func NewPolicyEvaluator(
 	addDefaultClientCertificateRule bool,
 ) (*PolicyEvaluator, error) {
 	e := new(PolicyEvaluator)
+	e.policyChecksum = configPolicy.Checksum()
 
 	// generate the base rego script for the policy
 	ppl := configPolicy.ToPPL()

authorize/state.go

@@ -28,7 +28,9 @@ type authorizeState struct {
 	authenticateKeyFetcher     hpke.KeyFetcher
 }
 
-func newAuthorizeStateFromConfig(cfg *config.Config, store *store.Store) (*authorizeState, error) {
+func newAuthorizeStateFromConfig(
+	cfg *config.Config, store *store.Store, previousPolicyEvaluator *evaluator.Evaluator,
+) (*authorizeState, error) {
 	if err := validateOptions(cfg.Options); err != nil {
 		return nil, fmt.Errorf("authorize: bad options: %w", err)
 	}
@@ -37,7 +39,7 @@ func newAuthorizeStateFromConfig(cfg *config.Config, store *store.Store) (*autho
 
 	var err error
 
-	state.evaluator, err = newPolicyEvaluator(cfg.Options, store)
+	state.evaluator, err = newPolicyEvaluator(cfg.Options, store, previousPolicyEvaluator)
 	if err != nil {
 		return nil, fmt.Errorf("authorize: failed to update policy with options: %w", err)
 	}