Oct 16, 2024 · Oct 15, 2024 · Oct 15, 2024 · Oct 16, 2024 · Oct 16, 2024 · Oct 16, 2024
diff --git a/coderd/cryptokeys/cache.go b/coderd/cryptokeys/cache.go
 package cryptokeys

 import (
 "context"
 "encoding/hex"
 "io"
 "strconv"
 "sync"
 "time"

 "golang.org/x/xerrors"

 "cdr.dev/slog"
 "github.com/coder/coder/v2/coderd/database"
 "github.com/coder/coder/v2/coderd/database/db2sdk"
 "github.com/coder/coder/v2/codersdk"
 "github.com/coder/quartz"
 )

 var (
 ErrKeyNotFound    = xerrors.New("key not found")
 ErrKeyInvalid     = xerrors.New("key is invalid for use")
 ErrClosed         = xerrors.New("closed")
 ErrInvalidFeature = xerrors.New("invalid feature for this operation")
 )

 type Fetcher interface {
 Fetch(ctx context.Context) ([]codersdk.CryptoKey, error)
 }

 type EncryptionKeycache interface {
 // EncryptingKey returns the latest valid key for encrypting payloads. A valid
 // key is one that is both past its start time and before its deletion time.
 EncryptingKey(ctx context.Context) (id string, key interface{}, err error)
 // DecryptingKey returns the key with the provided id which maps to its sequence
 // number. The key is valid for decryption as long as it is not deleted or past
 // its deletion date. We must allow for keys prior to their start time to
 // account for clock skew between peers (one key may be past its start time on
 // one machine while another is not).
 DecryptingKey(ctx context.Context, id string) (key interface{}, err error)
 io.Closer
 }

 type SigningKeycache interface {
 // SigningKey returns the latest valid key for signing. A valid key is one
 // that is both past its start time and before its deletion time.
 SigningKey(ctx context.Context) (id string, key interface{}, err error)
 // VerifyingKey returns the key with the provided id which should map to its
 // sequence number. The key is valid for verifying as long as it is not deleted
 // or past its deletion date. We must allow for keys prior to their start time
 // to account for clock skew between peers (one key may be past its start time
 // on one machine while another is not).
 VerifyingKey(ctx context.Context, id string) (key interface{}, err error)
 io.Closer
 }

 const (
 // latestSequence is a special sequence number that represents the latest key.
 latestSequence = -1
 // refreshInterval is the interval at which the key cache will refresh.
 refreshInterval = time.Minute * 10
 )

 type DBFetcher struct {
 DB      database.Store
 Feature database.CryptoKeyFeature
 }

 func (d *DBFetcher) Fetch(ctx context.Context) ([]codersdk.CryptoKey, error) {
 keys, err := d.DB.GetCryptoKeysByFeature(ctx, d.Feature)
 if err != nil {
 return nil, xerrors.Errorf("get crypto keys by feature: %w", err)
 }

 return db2sdk.CryptoKeys(keys), nil
 }

 // cache implements the caching functionality for both signing and encryption keys.
 type cache struct {
 clock         quartz.Clock
 refreshCtx    context.Context
 refreshCancel context.CancelFunc
 fetcher       Fetcher
 logger        slog.Logger
 feature       codersdk.CryptoKeyFeature

 mu        sync.Mutex
 keys      map[int32]codersdk.CryptoKey
 lastFetch time.Time
 refresher *quartz.Timer
 fetching  bool
 closed    bool
 cond      *sync.Cond
 }

 type CacheOption func(*cache)

 func WithCacheClock(clock quartz.Clock) CacheOption {
 return func(d *cache) {
 d.clock = clock
 }
 }

 // NewSigningCache instantiates a cache. Close should be called to release resources
 // associated with its internal timer.
 func NewSigningCache(ctx context.Context, logger slog.Logger, fetcher Fetcher,
 feature codersdk.CryptoKeyFeature, opts ...func(*cache),
 ) (SigningKeycache, error) {
 if !isSigningKeyFeature(feature) {
 return nil, xerrors.Errorf("invalid feature: %s", feature)
 }
 return newCache(ctx, logger, fetcher, feature, opts...)
 }

 func NewEncryptionCache(ctx context.Context, logger slog.Logger, fetcher Fetcher,
 feature codersdk.CryptoKeyFeature, opts ...func(*cache),
 ) (EncryptionKeycache, error) {
 if !isEncryptionKeyFeature(feature) {
 return nil, xerrors.Errorf("invalid feature: %s", feature)
 }
 return newCache(ctx, logger, fetcher, feature, opts...)
 }

 func newCache(ctx context.Context, logger slog.Logger, fetcher Fetcher, feature codersdk.CryptoKeyFeature, opts ...func(*cache)) (*cache, error) {
 cache := &cache{
 clock:   quartz.NewReal(),
 logger:  logger,
 fetcher: fetcher,
 feature: feature,
 }

 for _, opt := range opts {
 opt(cache)
 }

 cache.cond = sync.NewCond(&cache.mu)
 cache.refreshCtx, cache.refreshCancel = context.WithCancel(ctx)
 cache.refresher = cache.clock.AfterFunc(refreshInterval, cache.refresh)

 keys, err := cache.cryptoKeys(ctx)
 if err != nil {
 cache.refreshCancel()
 return nil, xerrors.Errorf("initial fetch: %w", err)
 }
 cache.keys = keys
 return cache, nil
 }

 func (c *cache) EncryptingKey(ctx context.Context) (string, interface{}, error) {
 if !isEncryptionKeyFeature(c.feature) {
 return "", nil, ErrInvalidFeature
 }

 return c.cryptoKey(ctx, latestSequence)
 }

 func (c *cache) DecryptingKey(ctx context.Context, id string) (interface{}, error) {
 if !isEncryptionKeyFeature(c.feature) {
 return nil, ErrInvalidFeature
 }

 seq, err := strconv.ParseInt(id, 10, 64)
 if err != nil {
 return nil, xerrors.Errorf("parse id: %w", err)
 }

 _, secret, err := c.cryptoKey(ctx, int32(seq))
 if err != nil {
 return nil, xerrors.Errorf("crypto key: %w", err)
 }
 return secret, nil
 }

 func (c *cache) SigningKey(ctx context.Context) (string, interface{}, error) {
 if !isSigningKeyFeature(c.feature) {
 return "", nil, ErrInvalidFeature
 }

 return c.cryptoKey(ctx, latestSequence)
 }

 func (c *cache) VerifyingKey(ctx context.Context, id string) (interface{}, error) {
 if !isSigningKeyFeature(c.feature) {
 return nil, ErrInvalidFeature
 }

 seq, err := strconv.ParseInt(id, 10, 64)
 if err != nil {
 return nil, xerrors.Errorf("parse id: %w", err)
 }

 _, secret, err := c.cryptoKey(ctx, int32(seq))
 if err != nil {
 return nil, xerrors.Errorf("crypto key: %w", err)
 }

 return secret, nil
 }

 func isEncryptionKeyFeature(feature codersdk.CryptoKeyFeature) bool {
 return feature == codersdk.CryptoKeyFeatureWorkspaceApp
 }

 func isSigningKeyFeature(feature codersdk.CryptoKeyFeature) bool {
 switch feature {
 case codersdk.CryptoKeyFeatureTailnetResume, codersdk.CryptoKeyFeatureOIDCConvert:
 return true
 default:
 return false
 }
 }

 func idSecret(k codersdk.CryptoKey) (string, []byte, error) {
 key, err := hex.DecodeString(k.Secret)
 if err != nil {
 return "", nil, xerrors.Errorf("decode key: %w", err)
 }

 return strconv.FormatInt(int64(k.Sequence), 10), key, nil
 }

 func (c *cache) cryptoKey(ctx context.Context, sequence int32) (string, []byte, error) {
 c.mu.Lock()
 defer c.mu.Unlock()

 if c.closed {
 return "", nil, ErrClosed
 }

 var key codersdk.CryptoKey
 var ok bool
 for key, ok = c.key(sequence); !ok && c.fetching && !c.closed; {
 c.cond.Wait()
 }

 if c.closed {
 return "", nil, ErrClosed
 }

 if ok {
 return checkKey(key, sequence, c.clock.Now())
 }

 c.fetching = true
 c.mu.Unlock()

 keys, err := c.cryptoKeys(ctx)
 if err != nil {
 return "", nil, xerrors.Errorf("get keys: %w", err)
 }

 c.mu.Lock()
 c.lastFetch = c.clock.Now()
 c.refresher.Reset(refreshInterval)
 c.keys = keys
 c.fetching = false
 c.cond.Broadcast()

 key, ok = c.key(sequence)
 if !ok {
 return "", nil, ErrKeyNotFound
 }

 return checkKey(key, sequence, c.clock.Now())
 }

 func (c *cache) key(sequence int32) (codersdk.CryptoKey, bool) {
 if sequence == latestSequence {
 return c.keys[latestSequence], c.keys[latestSequence].CanSign(c.clock.Now())
 }

 key, ok := c.keys[sequence]
 return key, ok
 }

 func checkKey(key codersdk.CryptoKey, sequence int32, now time.Time) (string, []byte, error) {
 if sequence == latestSequence {
 if !key.CanSign(now) {
 return "", nil, ErrKeyInvalid
 }
 return idSecret(key)
 }

 if !key.CanVerify(now) {
 return "", nil, ErrKeyInvalid
 }

 return idSecret(key)
 }

 // refresh fetches the keys and updates the cache.
 func (c *cache) refresh() {
 now := c.clock.Now("CryptoKeyCache", "refresh")
 c.mu.Lock()
 defer c.mu.Unlock()

 if c.closed {
 return
 }

 // If something's already fetching, we don't need to do anything.
 if c.fetching {
 return
 }

 // There's a window we must account for where the timer fires while a fetch
 // is ongoing but prior to the timer getting reset. In this case we want to
 // avoid double fetching.
 if now.Sub(c.lastFetch) < refreshInterval {
 return
 }

 c.fetching = true

 c.mu.Unlock()
 keys, err := c.cryptoKeys(c.refreshCtx)
 if err != nil {
 c.logger.Error(c.refreshCtx, "fetch crypto keys", slog.Error(err))
 return
 }

 // We don't defer an unlock here due to the deferred unlock at the top of the function.
 c.mu.Lock()

 c.lastFetch = c.clock.Now()
 c.refresher.Reset(refreshInterval)
 c.keys = keys
 c.fetching = false
 c.cond.Broadcast()
 }

 // cryptoKeys queries the control plane for the crypto keys.
 // Outside of initialization, this should only be called by fetch.
 func (c *cache) cryptoKeys(ctx context.Context) (map[int32]codersdk.CryptoKey, error) {
 keys, err := c.fetcher.Fetch(ctx)
 if err != nil {
 return nil, xerrors.Errorf("crypto keys: %w", err)
 }
 cache := toKeyMap(keys, c.clock.Now())
 return cache, nil
 }

 func toKeyMap(keys []codersdk.CryptoKey, now time.Time) map[int32]codersdk.CryptoKey {
 m := make(map[int32]codersdk.CryptoKey)
 var latest codersdk.CryptoKey
 for _, key := range keys {
 m[key.Sequence] = key
 if key.Sequence > latest.Sequence && key.CanSign(now) {
 m[latestSequence] = key
 }
 }
 return m
 }

 func (c *cache) Close() error {
 c.mu.Lock()
 defer c.mu.Unlock()

 if c.closed {
 return nil
 }

 c.closed = true
 c.refreshCancel()
 c.refresher.Stop()
 c.cond.Broadcast()

 return nil
 }
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,369 @@
		package cryptokeys

		import (
		"context"
		"encoding/hex"
		"io"
		"strconv"
		"sync"
		"time"

		"golang.org/x/xerrors"

		"cdr.dev/slog"
		"github.com/coder/coder/v2/coderd/database"
		"github.com/coder/coder/v2/coderd/database/db2sdk"
		"github.com/coder/coder/v2/codersdk"
		"github.com/coder/quartz"
		)

		var (
		ErrKeyNotFound = xerrors.New("key not found")
		ErrKeyInvalid = xerrors.New("key is invalid for use")
		ErrClosed = xerrors.New("closed")
		ErrInvalidFeature = xerrors.New("invalid feature for this operation")
		)

		type Fetcher interface {
		Fetch(ctx context.Context) ([]codersdk.CryptoKey, error)
		}

		type EncryptionKeycache interface {
		// EncryptingKey returns the latest valid key for encrypting payloads. A valid
		// key is one that is both past its start time and before its deletion time.
		EncryptingKey(ctx context.Context) (id string, key interface{}, err error)
		// DecryptingKey returns the key with the provided id which maps to its sequence
		// number. The key is valid for decryption as long as it is not deleted or past
		// its deletion date. We must allow for keys prior to their start time to
		// account for clock skew between peers (one key may be past its start time on
		// one machine while another is not).
		DecryptingKey(ctx context.Context, id string) (key interface{}, err error)
		io.Closer
		}

		type SigningKeycache interface {
		// SigningKey returns the latest valid key for signing. A valid key is one
		// that is both past its start time and before its deletion time.
		SigningKey(ctx context.Context) (id string, key interface{}, err error)
		// VerifyingKey returns the key with the provided id which should map to its
		// sequence number. The key is valid for verifying as long as it is not deleted
		// or past its deletion date. We must allow for keys prior to their start time
		// to account for clock skew between peers (one key may be past its start time
		// on one machine while another is not).
		VerifyingKey(ctx context.Context, id string) (key interface{}, err error)
		io.Closer
		}

		const (
		// latestSequence is a special sequence number that represents the latest key.
		latestSequence = -1
		// refreshInterval is the interval at which the key cache will refresh.
		refreshInterval = time.Minute * 10
		)

		type DBFetcher struct {
		DB database.Store
		Feature database.CryptoKeyFeature
		}

		func (d *DBFetcher) Fetch(ctx context.Context) ([]codersdk.CryptoKey, error) {
		keys, err := d.DB.GetCryptoKeysByFeature(ctx, d.Feature)
		if err != nil {
		return nil, xerrors.Errorf("get crypto keys by feature: %w", err)
		}

		return db2sdk.CryptoKeys(keys), nil
		}

		// cache implements the caching functionality for both signing and encryption keys.
		type cache struct {
		clock quartz.Clock
		refreshCtx context.Context
		refreshCancel context.CancelFunc
		fetcher Fetcher
		logger slog.Logger
		feature codersdk.CryptoKeyFeature

		mu sync.Mutex
		keys map[int32]codersdk.CryptoKey
		lastFetch time.Time
		refresher *quartz.Timer
		fetching bool
		closed bool
		cond *sync.Cond
		}

		type CacheOption func(*cache)

		func WithCacheClock(clock quartz.Clock) CacheOption {
		return func(d *cache) {
		d.clock = clock
		}
		}

		// NewSigningCache instantiates a cache. Close should be called to release resources
		// associated with its internal timer.
		func NewSigningCache(ctx context.Context, logger slog.Logger, fetcher Fetcher,
		feature codersdk.CryptoKeyFeature, opts ...func(*cache),
		) (SigningKeycache, error) {
		if !isSigningKeyFeature(feature) {
		return nil, xerrors.Errorf("invalid feature: %s", feature)
		}
		return newCache(ctx, logger, fetcher, feature, opts...)
		}

		func NewEncryptionCache(ctx context.Context, logger slog.Logger, fetcher Fetcher,
		feature codersdk.CryptoKeyFeature, opts ...func(*cache),
		) (EncryptionKeycache, error) {
		if !isEncryptionKeyFeature(feature) {
		return nil, xerrors.Errorf("invalid feature: %s", feature)
		}
		return newCache(ctx, logger, fetcher, feature, opts...)
		}

		func newCache(ctx context.Context, logger slog.Logger, fetcher Fetcher, feature codersdk.CryptoKeyFeature, opts ...func(cache)) (cache, error) {
		cache := &cache{
		clock: quartz.NewReal(),
		logger: logger,
		fetcher: fetcher,
		feature: feature,
		}

		for _, opt := range opts {
		opt(cache)
		}

		cache.cond = sync.NewCond(&cache.mu)
		cache.refreshCtx, cache.refreshCancel = context.WithCancel(ctx)
		cache.refresher = cache.clock.AfterFunc(refreshInterval, cache.refresh)

		keys, err := cache.cryptoKeys(ctx)
		if err != nil {
		cache.refreshCancel()
		return nil, xerrors.Errorf("initial fetch: %w", err)
		}
		cache.keys = keys
		return cache, nil
		}

		func (c *cache) EncryptingKey(ctx context.Context) (string, interface{}, error) {
		if !isEncryptionKeyFeature(c.feature) {
		return "", nil, ErrInvalidFeature
		}

		return c.cryptoKey(ctx, latestSequence)
		}

		func (c *cache) DecryptingKey(ctx context.Context, id string) (interface{}, error) {
		if !isEncryptionKeyFeature(c.feature) {
		return nil, ErrInvalidFeature
		}

		seq, err := strconv.ParseInt(id, 10, 64)
		if err != nil {
		return nil, xerrors.Errorf("parse id: %w", err)
		}

		_, secret, err := c.cryptoKey(ctx, int32(seq))
		if err != nil {
		return nil, xerrors.Errorf("crypto key: %w", err)
		}
		return secret, nil
		}

		func (c *cache) SigningKey(ctx context.Context) (string, interface{}, error) {
		if !isSigningKeyFeature(c.feature) {
		return "", nil, ErrInvalidFeature
		}

		return c.cryptoKey(ctx, latestSequence)
		}

		func (c *cache) VerifyingKey(ctx context.Context, id string) (interface{}, error) {
		if !isSigningKeyFeature(c.feature) {
		return nil, ErrInvalidFeature
		}

		seq, err := strconv.ParseInt(id, 10, 64)
		if err != nil {
		return nil, xerrors.Errorf("parse id: %w", err)
		}

		_, secret, err := c.cryptoKey(ctx, int32(seq))
		if err != nil {
		return nil, xerrors.Errorf("crypto key: %w", err)
		}

		return secret, nil
		}

		func isEncryptionKeyFeature(feature codersdk.CryptoKeyFeature) bool {
		return feature == codersdk.CryptoKeyFeatureWorkspaceApp
sreya marked this conversation as resolved. Show resolvedHide resolved
		}

		func isSigningKeyFeature(feature codersdk.CryptoKeyFeature) bool {
		switch feature {
		case codersdk.CryptoKeyFeatureTailnetResume, codersdk.CryptoKeyFeatureOIDCConvert:
		return true
		default:
		return false
		}
		}

		func idSecret(k codersdk.CryptoKey) (string, []byte, error) {
		key, err := hex.DecodeString(k.Secret)
		if err != nil {
		return "", nil, xerrors.Errorf("decode key: %w", err)
		}

		return strconv.FormatInt(int64(k.Sequence), 10), key, nil
		}

		func (c *cache) cryptoKey(ctx context.Context, sequence int32) (string, []byte, error) {
		c.mu.Lock()
		defer c.mu.Unlock()

		if c.closed {
		return "", nil, ErrClosed
		}

		var key codersdk.CryptoKey
		var ok bool
		for key, ok = c.key(sequence); !ok && c.fetching && !c.closed; {
		c.cond.Wait()
		}

		if c.closed {
		return "", nil, ErrClosed
		}

		if ok {
		return checkKey(key, sequence, c.clock.Now())
		}

		c.fetching = true
		c.mu.Unlock()

		keys, err := c.cryptoKeys(ctx)
		if err != nil {
		return "", nil, xerrors.Errorf("get keys: %w", err)
		}

		c.mu.Lock()
		c.lastFetch = c.clock.Now()
		c.refresher.Reset(refreshInterval)
		c.keys = keys
		c.fetching = false
		c.cond.Broadcast()

		key, ok = c.key(sequence)
		if !ok {
		return "", nil, ErrKeyNotFound
		}

		return checkKey(key, sequence, c.clock.Now())
		}

		func (c *cache) key(sequence int32) (codersdk.CryptoKey, bool) {
		if sequence == latestSequence {
		return c.keys[latestSequence], c.keys[latestSequence].CanSign(c.clock.Now())
		}

		key, ok := c.keys[sequence]
		return key, ok
		}

		func checkKey(key codersdk.CryptoKey, sequence int32, now time.Time) (string, []byte, error) {
		if sequence == latestSequence {
		if !key.CanSign(now) {
		return "", nil, ErrKeyInvalid
		}
		return idSecret(key)
		}

		if !key.CanVerify(now) {
		return "", nil, ErrKeyInvalid
		}

		return idSecret(key)
		}

		// refresh fetches the keys and updates the cache.
		func (c *cache) refresh() {
		now := c.clock.Now("CryptoKeyCache", "refresh")
		c.mu.Lock()
		defer c.mu.Unlock()

		if c.closed {
		return
		}

		// If something's already fetching, we don't need to do anything.
		if c.fetching {
		return
		}

		// There's a window we must account for where the timer fires while a fetch
		// is ongoing but prior to the timer getting reset. In this case we want to
		// avoid double fetching.
		if now.Sub(c.lastFetch) < refreshInterval {
		return
		}

		c.fetching = true

		c.mu.Unlock()
		keys, err := c.cryptoKeys(c.refreshCtx)
		if err != nil {
		c.logger.Error(c.refreshCtx, "fetch crypto keys", slog.Error(err))
		return
		}

		// We don't defer an unlock here due to the deferred unlock at the top of the function.
		c.mu.Lock()

		c.lastFetch = c.clock.Now()
		c.refresher.Reset(refreshInterval)
		c.keys = keys
		c.fetching = false
		c.cond.Broadcast()
		}

		// cryptoKeys queries the control plane for the crypto keys.
		// Outside of initialization, this should only be called by fetch.
		func (c *cache) cryptoKeys(ctx context.Context) (map[int32]codersdk.CryptoKey, error) {
		keys, err := c.fetcher.Fetch(ctx)
		if err != nil {
		return nil, xerrors.Errorf("crypto keys: %w", err)
		}
		cache := toKeyMap(keys, c.clock.Now())
		return cache, nil
		}

		func toKeyMap(keys []codersdk.CryptoKey, now time.Time) map[int32]codersdk.CryptoKey {
		m := make(map[int32]codersdk.CryptoKey)
		var latest codersdk.CryptoKey
		for _, key := range keys {
		m[key.Sequence] = key
		if key.Sequence > latest.Sequence && key.CanSign(now) {
		m[latestSequence] = key
		}
		}
		return m
		}

		func (c *cache) Close() error {
		c.mu.Lock()
		defer c.mu.Unlock()

		if c.closed {
		return nil
		}

		c.closed = true
		c.refreshCancel()
		c.refresher.Stop()
		c.cond.Broadcast()

		return nil
		}