apikey: cache policy info lookup & last usage updates (#3291)

* cache policy and rate-limit last used update

* comments

* fix comments

* remove error return (no possible err)

apikey/lastusedcache.go

@@ -0,0 +1,36 @@
+package apikey
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"github.com/golang/groupcache/lru"
+	"github.com/google/uuid"
+)
+
+type lastUsedCache struct {
+	lru *lru.Cache
+
+	mx         sync.Mutex
+	updateFunc func(ctx context.Context, id uuid.UUID, ua, ip string) error
+}
+
+func newLastUsedCache(max int, updateFunc func(ctx context.Context, id uuid.UUID, ua, ip string) error) *lastUsedCache {
+	return &lastUsedCache{
+		lru:        lru.New(max),
+		updateFunc: updateFunc,
+	}
+}
+func (c *lastUsedCache) RecordUsage(ctx context.Context, id uuid.UUID, ua, ip string) error {
+	c.mx.Lock()
+	defer c.mx.Unlock()
+
+	// check if we've seen this key recently, and if it's been less than a minute
+	if t, ok := c.lru.Get(id); ok && time.Since(t.(time.Time)) < time.Minute {
+		return nil
+	}
+
+	c.lru.Add(id, time.Now())
+	return c.updateFunc(ctx, id, ua, ip)
+}

apikey/polcache.go

@@ -0,0 +1,109 @@
+package apikey
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"sync"
+
+	"github.com/golang/groupcache/lru"
+	"github.com/google/uuid"
+	"github.com/target/goalert/gadb"
+)
+
+// polCache handles caching of policyInfo objects, as well as negative caching
+// of invalid keys.
+type polCache struct {
+	lru *lru.Cache
+	neg *lru.Cache
+	mx  sync.Mutex
+
+	cfg polCacheConfig
+}
+
+type polCacheConfig struct {
+	FillFunc func(context.Context, uuid.UUID) (*policyInfo, bool, error)
+	Verify   func(context.Context, uuid.UUID) (bool, error)
+	MaxSize  int
+}
+
+func newPolCache(cfg polCacheConfig) *polCache {
+	return &polCache{
+		lru: lru.New(cfg.MaxSize),
+		neg: lru.New(cfg.MaxSize),
+		cfg: cfg,
+	}
+}
+
+// Revoke will add the key to the negative cache.
+func (c *polCache) Revoke(ctx context.Context, key uuid.UUID) {
+	c.mx.Lock()
+	defer c.mx.Unlock()
+
+	c.neg.Add(key, nil)
+	c.lru.Remove(key)
+}
+
+// Get will return the policyInfo for the given key.
+//
+// If the key is in the cache, it will be verified before returning.
+//
+// If it is not in the cache, it will be fetched and added to the cache.
+//
+// If either the key is invalid or the policy is invalid, the key will be
+// added to the negative cache.
+func (c *polCache) Get(ctx context.Context, key uuid.UUID) (value *policyInfo, ok bool, err error) {
+	c.mx.Lock()
+	defer c.mx.Unlock()
+
+	if _, ok := c.neg.Get(key); ok {
+		return value, false, nil
+	}
+
+	if v, ok := c.lru.Get(key); ok {
+		// Check if the key is still valid before returning it,
+		// if it is not valid, we can remove it from the cache.
+		isValid, err := c.cfg.Verify(ctx, key)
+		if err != nil {
+			return value, false, err
+		}
+
+		// Since each key has a unique ID and is signed, we can
+		// safely assume that an invalid key will always be invalid
+		// and can be negatively cached.
+		if !isValid {
+			c.neg.Add(key, nil)
+			c.lru.Remove(key)
+			return value, false, nil
+		}
+
+		return v.(*policyInfo), true, nil
+	}
+
+	// If the key is not in the cache, we need to fetch it,
+	// and add it to the cache. We can safely assume that
+	// the key is valid when returned from the FillFunc.
+	value, isValid, err := c.cfg.FillFunc(ctx, key)
+	if err != nil {
+		return value, false, err
+	}
+	if !isValid {
+		c.neg.Add(key, nil)
+		return value, false, nil
+	}
+
+	c.lru.Add(key, value)
+	return value, true, nil
+}
+
+func (s *Store) _verifyPolicyID(ctx context.Context, id uuid.UUID) (bool, error) {
+	valid, err := gadb.New(s.db).APIKeyAuthCheck(ctx, id)
+	if errors.Is(err, sql.ErrNoRows) {
+		return false, nil
+	}
+	if err != nil {
+		return false, err
+	}
+
+	return valid, nil
+}

apikey/store.go

@@ -26,6 +26,9 @@ import (
 type Store struct {
 	db  *sql.DB
 	key keyring.Keyring
+
+	polCache      *polCache
+	lastUsedCache *lastUsedCache
 }
 
 // NewStore will create a new Store.
@@ -35,6 +38,14 @@ func NewStore(ctx context.Context, db *sql.DB, key keyring.Keyring) (*Store, err
 		key: key,
 	}
 
+	s.polCache = newPolCache(polCacheConfig{
+		FillFunc: s._fetchPolicyInfo,
+		Verify:   s._verifyPolicyID,
+		MaxSize:  1000,
+	})
+
+	s.lastUsedCache = newLastUsedCache(1000, s._updateLastUsed)
+
 	return s, nil
 }
 
@@ -199,7 +210,7 @@ func (s *Store) AuthorizeGraphQL(ctx context.Context, tok, ua, ip string) (conte
 		return ctx, permission.Unauthorized()
 	}
 
-	info, valid, err := s._fetchPolicyInfo(ctx, id)
+	info, valid, err := s.polCache.Get(ctx, id)
 	if err != nil {
 		return nil, err
 	}
@@ -208,12 +219,15 @@ func (s *Store) AuthorizeGraphQL(ctx context.Context, tok, ua, ip string) (conte
 		return ctx, permission.Unauthorized()
 	}
 	if !bytes.Equal(info.Hash, claims.PolicyHash) {
+		// Successful cache lookup, but the policy has changed since the token was issued and so the token is no longer valid.
+		s.polCache.Revoke(ctx, id)
+
 		// We want to log this as a warning, because it is a potential security issue.
 		log.Log(ctx, fmt.Errorf("apikey: policy hash mismatch for key %s", id))
 		return ctx, permission.Unauthorized()
 	}
 
-	err = s._updateLastUsed(ctx, id, ua, ip)
+	err = s.lastUsedCache.RecordUsage(ctx, id, ua, ip)
 	if err != nil {
 		// Recording usage is not critical, so we log the error and continue.
 		log.Log(ctx, err)