depot/go/tokend/tokencache.go

package main

import (
	"context"
	"os/user"
	"sync"
	"time"

	log "github.com/golang/glog"
	vapi "github.com/hashicorp/vault/api"
)

type TokenSecret = vapi.Secret

type ttledSecret struct {
	*TokenSecret

	expiration time.Time

	isRenewable    bool
	renewThreshold time.Time
}

func wrapToken(s *TokenSecret) ttledSecret {
	now := time.Now()

	ttl := s.Auth.LeaseDuration
	renewable := s.Auth.Renewable

	var ttlBuffer time.Duration
	if ttl < 300 {
		// Give ourselves an extra two minute buffer for renewal.
		ttlBuffer = 2 * time.Minute
	}

	return ttledSecret{
		TokenSecret: s,

		expiration: now.Add(time.Duration(ttl) * time.Second),

		isRenewable:    renewable,
		renewThreshold: now.Add(time.Duration(ttl/2) * time.Second).Add(-ttlBuffer),
	}
}

func (s ttledSecret) Expired() bool {
	// We use !After rather than Before so that if it's _exactly_ now then we still return true.
	return !s.expiration.After(time.Now())
}

func (s ttledSecret) ShouldRenew() bool {
	if !s.isRenewable {
		return false
	}

	return time.Now().After(s.renewThreshold)
}

type tokenInteractor interface {
	Revoke(ctx context.Context, tokenSecret *TokenSecret) error
	Issue(ctx context.Context, username string, isPlainUser bool) (*TokenSecret, error)
	Renew(ctx context.Context, tokenSecret *TokenSecret) (*TokenSecret, error)
}

type tokenUserCache struct {
	l sync.RWMutex
	m map[string]ttledSecret
	i tokenInteractor
}

func newCache(i tokenInteractor) *tokenUserCache {
	return &tokenUserCache{
		m: make(map[string]ttledSecret),
		i: i,
	}
}

func (c *tokenUserCache) expire() (revoke map[string]*TokenSecret) {
	c.l.Lock()
	defer c.l.Unlock()

	var remove []string
	revoke = make(map[string]*TokenSecret)
	for username, s := range c.m {
		// If the token has expired, remove it.
		if s.Expired() {
			remove = append(remove, username)
			continue
		}

		// If the user no longer exists on the system, revoke it.
		if _, err := user.Lookup(username); err != nil {
			log.Infof("token for %v will be revoked because user lookup returned error: %v", username, err)
			remove = append(remove, username)
			revoke[username] = s.TokenSecret
			continue
		}

		// Otherwise, leave it alone.
	}

	for _, username := range remove {
		delete(c.m, username)
	}

	return revoke
}

func (c *tokenUserCache) renew() {
	ctx := context.Background()

	toRenew := make(map[string]*TokenSecret)
	c.l.Lock()
	for username, s := range c.m {
		if s.ShouldRenew() {
			toRenew[username] = s.TokenSecret
		}
	}
	c.l.Unlock()

	for username, s := range toRenew {
		log.Infof("renewing token for %v", username)
		newS, err := c.i.Renew(ctx, s)
		if err != nil {
			log.Errorf("renewing token for %v: %w (discarding cached token)", username, err)

			// Discard the token in cache.
			c.l.Lock()
			delete(c.m, username)
			c.l.Unlock()
			continue
		}

		c.l.Lock()
		if oldS, ok := c.m[username]; ok && oldS.TokenSecret == s {
			c.m[username] = wrapToken(newS)
		} else if ok {
			log.Warningf("after renewing token for %v discovered that the token in cache had changed in the meantime; dropping refreshed token")
		}
		c.l.Unlock()
	}
}

func (c *tokenUserCache) tick(ctx context.Context) error {
	log.Info("token cache is ticking")
	revoke := c.expire()
	for username, tokenSecret := range revoke {
		if err := c.i.Revoke(ctx, tokenSecret); err != nil {
			log.Errorf("unable to revoke token for %v: %w", username, err)
		}
	}

	c.renew()

	return nil
}

func (c *tokenUserCache) Get(ctx context.Context, username string, isPlainUser bool) (*TokenSecret, error) {
	c.l.RLock()
	token, ok := c.m[username]
	c.l.RUnlock()

	// If we got a token, but it's expired, delete it and pretend it didn't exist.
	if ok && token.Expired() {
		c.l.Lock()
		delete(c.m, username)
		c.l.Unlock()

		ok = false
	}

	if ok {
		// We have a non-expired pre-existing token!
		return token.TokenSecret, nil
	}

	// Issue a new token.
	issuedToken, err := c.i.Issue(ctx, username, isPlainUser)
	if err != nil {
		return nil, err
	}

	// OK, now we check if someone has issued a token for us in the
	// meantime...  We could (but don't) do any more complex coordination
	// than this, so we might actually issue two tokens for the same user
	// concurrently. Even if we do this though, one of the tokens will
	// "win" and we'll revoke the other one.
	c.l.Lock()
	defer c.l.Unlock()

	if bgToken, ok := c.m[username]; ok && !bgToken.Expired() {
		// Just use this one, I guess.
		go c.i.Revoke(context.Background(), issuedToken)
		return bgToken.TokenSecret, nil
	}

	token = wrapToken(issuedToken)
	c.m[username] = token

	return token.TokenSecret, nil
}

func (c *tokenUserCache) Purge(ctx context.Context, username string) {
	c.l.Lock()
	token, ok := c.m[username]
	delete(c.m, username)
	c.l.Unlock()

	if !ok {
		return
	}

	// Revoke the token as well, for good measure, in case the parent hasn't already done that.
	c.i.Revoke(ctx, token.TokenSecret)
}