ProofedMesh/services/admin-api/internal/attest/pocket_id_token_source.go

package attest

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"
)

// PocketIDTokenSource issues Bearer tokens using the Pocket ID OIDC
// `client_credentials` grant and caches each token until shortly before its
// expiry. It is safe for concurrent use.
//
// Flow:
//
//  1. On first use, GET {Issuer}/.well-known/openid-configuration and read
//     `token_endpoint`.
//  2. POST to the token endpoint with grant_type=client_credentials, using
//     HTTP Basic auth (`client_secret_basic`) for ClientID/ClientSecret.
//  3. Cache the access token until RefreshLeeway before `expires_in` elapses.
//
// The admin-api uses the same instance as the TokenSource for both the
// Signing Service and the Transparency Log client, so a single Pocket ID
// client credentials grant covers both backends.
type PocketIDTokenSource struct {
	// Issuer is the Pocket ID base URL, e.g. "https://auth.basicstack.de".
	// The trailing slash is optional.
	Issuer string
	// ClientID / ClientSecret identify the Pocket ID OIDC client that the
	// admin-api is registered as. They are typically sourced from env vars
	// (POCKET_ID_CLIENT_ID / POCKET_ID_CLIENT_SECRET) rather than flags so
	// they do not end up in process listings.
	ClientID     string
	ClientSecret string
	// Scope and Audience are optional. Pocket ID accepts both as standard
	// OIDC parameters; leave empty when not required by the server policy.
	Scope    string
	Audience string
	// Client overrides the default HTTP client. The default has a 10s
	// timeout and no custom transport.
	Client *http.Client
	// RefreshLeeway is the window before token expiry at which the cache
	// is considered invalid and a new token must be fetched. Zero defaults
	// to 30 seconds, matching the task spec ("verwirft bei exp-30s").
	RefreshLeeway time.Duration
	// Now is injected for deterministic tests. Zero defaults to time.Now.
	Now func() time.Time

	mu            sync.Mutex
	tokenEndpoint string
	cachedToken   string
	cachedExpiry  time.Time
}

type pocketIDTokenResponse struct {
	AccessToken string `json:"access_token"`
	TokenType   string `json:"token_type"`
	ExpiresIn   int    `json:"expires_in"`
}

type pocketIDDiscoveryDoc struct {
	TokenEndpoint string `json:"token_endpoint"`
}

// Token returns a cached access token when still valid, or fetches a fresh
// one from Pocket ID. It satisfies the attest.TokenSource interface.
func (p *PocketIDTokenSource) Token(ctx context.Context) (string, error) {
	if p == nil {
		return "", errors.New("attest.PocketIDTokenSource: nil receiver")
	}
	if p.Issuer == "" {
		return "", errors.New("attest.PocketIDTokenSource: Issuer is required")
	}
	if p.ClientID == "" || p.ClientSecret == "" {
		return "", errors.New("attest.PocketIDTokenSource: ClientID and ClientSecret are required")
	}

	p.mu.Lock()
	defer p.mu.Unlock()

	leeway := p.RefreshLeeway
	if leeway <= 0 {
		leeway = 30 * time.Second
	}
	now := p.nowLocked()
	if p.cachedToken != "" && now.Add(leeway).Before(p.cachedExpiry) {
		return p.cachedToken, nil
	}

	if p.tokenEndpoint == "" {
		if err := p.discoverLocked(ctx); err != nil {
			return "", err
		}
	}

	tok, expiry, err := p.fetchLocked(ctx, now)
	if err != nil {
		return "", err
	}
	p.cachedToken = tok
	p.cachedExpiry = expiry
	return tok, nil
}

func (p *PocketIDTokenSource) nowLocked() time.Time {
	if p.Now != nil {
		return p.Now()
	}
	return time.Now()
}

func (p *PocketIDTokenSource) httpClient() *http.Client {
	if p.Client != nil {
		return p.Client
	}
	return &http.Client{Timeout: 10 * time.Second}
}

func (p *PocketIDTokenSource) discoverLocked(ctx context.Context) error {
	base := strings.TrimRight(p.Issuer, "/")
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+"/.well-known/openid-configuration", nil)
	if err != nil {
		return fmt.Errorf("pocket id discovery: %w", err)
	}
	req.Header.Set("Accept", "application/json")
	resp, err := p.httpClient().Do(req)
	if err != nil {
		return fmt.Errorf("pocket id discovery: %w", err)
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("pocket id discovery: status %d: %s", resp.StatusCode, string(raw))
	}
	var doc pocketIDDiscoveryDoc
	if err := json.Unmarshal(raw, &doc); err != nil {
		return fmt.Errorf("pocket id discovery: decode: %w", err)
	}
	if doc.TokenEndpoint == "" {
		return errors.New("pocket id discovery: token_endpoint missing from OIDC configuration")
	}
	p.tokenEndpoint = doc.TokenEndpoint
	return nil
}

func (p *PocketIDTokenSource) fetchLocked(ctx context.Context, now time.Time) (string, time.Time, error) {
	form := url.Values{}
	form.Set("grant_type", "client_credentials")
	if p.Scope != "" {
		form.Set("scope", p.Scope)
	}
	if p.Audience != "" {
		form.Set("audience", p.Audience)
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.tokenEndpoint, strings.NewReader(form.Encode()))
	if err != nil {
		return "", time.Time{}, fmt.Errorf("pocket id token: %w", err)
	}
	req.SetBasicAuth(p.ClientID, p.ClientSecret)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Set("Accept", "application/json")
	resp, err := p.httpClient().Do(req)
	if err != nil {
		return "", time.Time{}, fmt.Errorf("pocket id token: %w", err)
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode != http.StatusOK {
		return "", time.Time{}, fmt.Errorf("pocket id token: status %d: %s", resp.StatusCode, string(raw))
	}
	var out pocketIDTokenResponse
	if err := json.Unmarshal(raw, &out); err != nil {
		return "", time.Time{}, fmt.Errorf("pocket id token: decode: %w", err)
	}
	if out.AccessToken == "" {
		return "", time.Time{}, errors.New("pocket id token: response missing access_token")
	}
	if out.ExpiresIn <= 0 {
		return "", time.Time{}, errors.New("pocket id token: response has non-positive expires_in")
	}
	return out.AccessToken, now.Add(time.Duration(out.ExpiresIn) * time.Second), nil
}