ProofedMesh/services/dummy-sink/internal/auth/pocket_id.go

package auth

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. Safe for concurrent use.
//
// Shape matches services/testagent-transform/internal/auth so the
// signing-service's Pocket-ID validator sees the same token format
// admin-api, approval-service, and testagent-transform already send.
type PocketIDTokenSource struct {
	Issuer        string
	ClientID      string
	ClientSecret  string
	Scope         string
	Audience      string
	Client        *http.Client
	RefreshLeeway time.Duration
	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.
func (p *PocketIDTokenSource) Token(ctx context.Context) (string, error) {
	if p == nil {
		return "", errors.New("auth.PocketIDTokenSource: nil receiver")
	}
	if p.Issuer == "" {
		return "", errors.New("auth.PocketIDTokenSource: Issuer is required")
	}
	if p.ClientID == "" || p.ClientSecret == "" {
		return "", errors.New("auth.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
}