sherlock/internal/authn/source.go

//go:build unix

package authn

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"time"

	"gitea.alexandru.macocian.me/amacocian/sherlock/internal/keyring"
)

// minRenewInterval bounds how often Run attempts a renewal, even when
// the access-token lifetime is very short (the sherlock-cli provider
// issues ~5-minute tokens). Prevents a hot loop on sub-minute TTLs.
const minRenewInterval = 5 * time.Second

// renewBackoff is how long Run waits after a failed refresh before
// re-evaluating. Refresh failures are usually transient (a network
// blip); the wallet still holds the soon-to-expire token in the
// meantime.
const renewBackoff = 15 * time.Second

// OnToken is the callback every MCP must supply to a TokenSource. The
// source invokes it once with the initial token (from Start) and again
// on every renewal (from Run), each call on the source's goroutine.
//
// Implementations MUST be safe to call concurrently with the MCP's
// request path — the token rotates underneath in-flight work. The
// canonical implementation is TokenHolder.Set (see below), which the
// request path then reads lock-free via TokenHolder.AccessToken.
type OnToken func(keyring.TokenSet)

// TokenSource keeps one service's access token fresh and pushes every
// new token to the MCP via a mandatory callback.
//
// The problem it solves: Authentik access tokens are short-lived
// (~5 min for the sherlock-cli provider). An MCP that captured a token
// once at startup would start getting 401s mid-session. A TokenSource
// renews the token before it expires and hands each new one to the
// MCP's callback, so the MCP never has to poll or reason about expiry.
//
// Lifecycle — eager (e.g. --probe):
//
//	src := authn.NewTokenSource(store, "grafana", cfg, opts, holder.Set)
//	ts, err := src.Start(ctx)   // initial OAuth (browser on first use)
//	go src.Run(ctx)             // renews + calls the callback forever
//
// Lifecycle — lazy (an MCP serving an agent): do NOT call Start/Run at
// startup. Bind a lifetime context, start the MCP server immediately so
// its handshake completes, and let the first token read trigger auth:
//
//	src := authn.NewTokenSource(...).Lifetime(ctx)
//	// in the request path, before using a token:
//	if err := src.EnsureStarted(reqCtx); err != nil { ... }
//
// EnsureStarted runs the (possibly browser-opening) initial auth on its
// first successful call and launches the renewer in the background. This
// keeps the interactive flow off the MCP startup path — the browser pops
// during the first tool call, not during the connect handshake (which
// the agent would time out).
type TokenSource struct {
	store   keyring.Store
	service string
	cfg     Config
	opts    EnsureOptions
	onToken OnToken

	// skew/minWait/clock are tunable for tests; production uses the
	// package defaults.
	skew    time.Duration
	minWait time.Duration
	clock   func() time.Time

	mu      sync.Mutex
	current keyring.TokenSet

	// lifeCtx bounds the background renewer launched by EnsureStarted.
	// Set via Lifetime; defaults to context.Background() if unset.
	lifeCtx context.Context

	// startMu serialises EnsureStarted; started flips true once the
	// initial auth has succeeded and the renewer is running.
	startMu sync.Mutex
	started bool
}

// NewTokenSource builds a TokenSource for service. cfg/opts mirror what
// Ensure takes. Both opts.LockPath and onToken are required — a missing
// lock path or a nil callback is always a wiring bug, so we panic
// rather than silently degrade.
func NewTokenSource(store keyring.Store, service string, cfg Config, opts EnsureOptions, onToken OnToken) *TokenSource {
	if opts.LockPath == "" {
		panic("authn: NewTokenSource requires opts.LockPath")
	}
	if onToken == nil {
		panic("authn: NewTokenSource requires a non-nil onToken callback")
	}
	return &TokenSource{
		store:   store,
		service: service,
		cfg:     cfg,
		opts:    opts,
		onToken: onToken,
		skew:    RefreshSkew,
		minWait: minRenewInterval,
		clock:   time.Now,
	}
}

// Lifetime sets the context that bounds the background renewer launched
// by EnsureStarted. MCPs pass their long-lived (signal-cancelled) main
// context so the renewer outlives any single request. Returns the
// receiver for chaining. Without it, EnsureStarted falls back to
// context.Background() (fine for a short-lived process that exits when
// the agent kills it).
func (s *TokenSource) Lifetime(ctx context.Context) *TokenSource {
	s.lifeCtx = ctx
	return s
}

// EnsureStarted lazily performs the initial token acquisition on its
// first successful call — which may open a browser — and then launches
// the background renewer. It is idempotent and safe to call from every
// request handler: once started, it returns nil immediately.
//
// ctx bounds only the initial auth flow; the renewer uses the Lifetime
// context. If the initial auth fails (e.g. the operator cancels the
// browser, or ctx is cancelled), started stays false so the next call
// retries.
func (s *TokenSource) EnsureStarted(ctx context.Context) error {
	s.startMu.Lock()
	defer s.startMu.Unlock()
	if s.started {
		return nil
	}
	if _, err := s.Start(ctx); err != nil {
		return err
	}
	s.started = true
	life := s.lifeCtx
	if life == nil {
		life = context.Background()
	}
	go func() { _ = s.Run(life) }()
	return nil
}

// Start performs the initial token acquisition, running a full OAuth
// flow if the wallet is empty (this is the call that may open a
// browser). It caches the token and invokes the callback with it, then
// returns it so callers can do a one-shot check (e.g. --probe). Call
// Start once, before Run.
func (s *TokenSource) Start(ctx context.Context) (keyring.TokenSet, error) {
	ts, err := Ensure(ctx, s.store, s.service, s.cfg, s.opts)
	if err != nil {
		return keyring.TokenSet{}, err
	}
	s.mu.Lock()
	s.current = ts
	s.mu.Unlock()
	s.onToken(ts)
	return ts, nil
}

// Run renews the token ahead of expiry until ctx is cancelled,
// invoking the callback on every rotation. It performs refresh-only
// renewals (it never opens a browser); if the refresh token has itself
// expired the renewal fails and Run backs off, leaving re-login to the
// operator (next sherlock invocation). Returns ctx.Err() on exit.
func (s *TokenSource) Run(ctx context.Context) error {
	for {
		timer := time.NewTimer(s.nextWait())
		select {
		case <-ctx.Done():
			timer.Stop()
			return ctx.Err()
		case <-timer.C:
		}

		if err := s.renew(ctx); err != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			// Transient failure: bounded backoff, then re-evaluate.
			b := time.NewTimer(renewBackoff)
			select {
			case <-ctx.Done():
				b.Stop()
				return ctx.Err()
			case <-b.C:
			}
		}
	}
}

// nextWait computes how long to sleep before the next proactive
// renewal: skew before the bearer token expires, clamped to minWait so
// a very short TTL doesn't spin.
func (s *TokenSource) nextWait() time.Duration {
	s.mu.Lock()
	cur := s.current
	s.mu.Unlock()
	if cur.Empty() {
		return s.minWait
	}
	wait := expiryForRefresh(cur).Sub(s.clock()) - s.skew
	if wait < s.minWait {
		return s.minWait
	}
	return wait
}

// renew refreshes the token under the cross-process flock and, if the
// access token actually rotated, pushes it to the callback. Never opens
// a browser.
func (s *TokenSource) renew(ctx context.Context) error {
	fresh, err := EnsureFresh(ctx, s.store, s.service, EnsureFreshOptions{
		LockPath:   s.opts.LockPath,
		Discoverer: s.opts.Discoverer,
		Clock:      s.clock,
	})
	if err != nil {
		return err
	}
	s.mu.Lock()
	changed := fresh.AccessToken != s.current.AccessToken
	s.current = fresh
	s.mu.Unlock()
	if changed {
		s.onToken(fresh)
	}
	return nil
}

// Current returns a snapshot of the cached TokenSet without any I/O.
func (s *TokenSource) Current() keyring.TokenSet {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.current
}

// errNoToken is returned by a TokenHolder accessor when no token has
// been delivered yet (should not happen after Start).
var errNoToken = errors.New("authn: no token delivered yet")

// TokenHolder is the canonical MCP-side glue between a TokenSource's
// push callback and a request path that needs the current bearer. Its
// Set method is the OnToken callback; AccessToken / Bearer are the
// lock-free reads the request path uses.
//
//	holder := authn.NewTokenHolder()
//	src := authn.NewTokenSource(store, svc, cfg, opts, holder.Set)
//	... // request path: req.Header.Set("Authorization", holder.Bearer())
type TokenHolder struct {
	v atomic.Pointer[keyring.TokenSet]
}

// NewTokenHolder returns an empty holder.
func NewTokenHolder() *TokenHolder { return &TokenHolder{} }

// Set stores ts as the current token. It satisfies OnToken and is safe
// for concurrent use.
func (h *TokenHolder) Set(ts keyring.TokenSet) {
	cp := ts
	h.v.Store(&cp)
}

// Current returns the last token delivered, or a zero TokenSet.
func (h *TokenHolder) Current() keyring.TokenSet {
	if p := h.v.Load(); p != nil {
		return *p
	}
	return keyring.TokenSet{}
}

// AccessToken returns the current access token, or "" if none yet.
func (h *TokenHolder) AccessToken() string {
	if p := h.v.Load(); p != nil {
		return p.AccessToken
	}
	return ""
}

// Bearer returns a ready-to-use "Bearer <token>" header value, or an
// error if no token has been delivered yet.
func (h *TokenHolder) Bearer() (string, error) {
	tok := h.AccessToken()
	if tok == "" {
		return "", errNoToken
	}
	return "Bearer " + tok, nil
}