member-console/internal/workflows/oauth_token_provider.go

package workflows

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

const (
	defaultTokenSafetyBuffer   = 30 * time.Second
	defaultTokenRequestTimeout = 10 * time.Second
	authHeaderKey              = "authorization"
	bearerPrefix               = "Bearer "
	maxTokenErrorBodyBytes     = 2048
)

// OAuthTokenProviderConfig configures OAuth2 client credentials token fetching.
type OAuthTokenProviderConfig struct {
	TokenURL     string
	ClientID     string
	ClientSecret string
	Scopes       []string
}

// OAuthTokenProvider implements Temporal's HeadersProvider using OAuth2 client credentials.
type OAuthTokenProvider struct {
	TokenURL     string
	ClientID     string
	ClientSecret string
	Scopes       []string

	mu           sync.Mutex
	accessToken  string
	expiresAt    time.Time
	httpClient   *http.Client
	safetyBuffer time.Duration
}

// NewOAuthTokenProvider validates configuration and returns a ready provider.
func NewOAuthTokenProvider(cfg OAuthTokenProviderConfig) (*OAuthTokenProvider, error) {
	tokenURL := strings.TrimSpace(cfg.TokenURL)
	if tokenURL == "" {
		return nil, errors.New("oauth token url is required")
	}
	if _, err := url.ParseRequestURI(tokenURL); err != nil {
		return nil, fmt.Errorf("invalid oauth token url: %w", err)
	}
	if strings.TrimSpace(cfg.ClientID) == "" {
		return nil, errors.New("oauth client id is required")
	}
	if strings.TrimSpace(cfg.ClientSecret) == "" {
		return nil, errors.New("oauth client secret is required")
	}

	scopes := append([]string(nil), cfg.Scopes...)

	return &OAuthTokenProvider{
		TokenURL:     tokenURL,
		ClientID:     cfg.ClientID,
		ClientSecret: cfg.ClientSecret,
		Scopes:       scopes,
		httpClient:   &http.Client{Timeout: defaultTokenRequestTimeout},
		safetyBuffer: defaultTokenSafetyBuffer,
	}, nil
}

// GetHeaders returns Authorization headers for each outgoing Temporal request.
func (p *OAuthTokenProvider) GetHeaders(ctx context.Context) (map[string]string, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if err := p.ensureDefaultsLocked(); err != nil {
		return nil, err
	}

	if p.isTokenValidLocked() {
		return map[string]string{
			authHeaderKey: bearerPrefix + p.accessToken,
		}, nil
	}

	token, expiresAt, err := p.fetchTokenLocked(ctx)
	if err != nil {
		return nil, err
	}

	p.accessToken = token
	p.expiresAt = expiresAt

	return map[string]string{
		authHeaderKey: bearerPrefix + token,
	}, nil
}

func (p *OAuthTokenProvider) ensureDefaultsLocked() error {
	if strings.TrimSpace(p.TokenURL) == "" {
		return errors.New("oauth token url is required")
	}
	if strings.TrimSpace(p.ClientID) == "" {
		return errors.New("oauth client id is required")
	}
	if strings.TrimSpace(p.ClientSecret) == "" {
		return errors.New("oauth client secret is required")
	}
	if p.httpClient == nil {
		p.httpClient = &http.Client{Timeout: defaultTokenRequestTimeout}
	}
	if p.safetyBuffer <= 0 {
		p.safetyBuffer = defaultTokenSafetyBuffer
	}
	return nil
}

func (p *OAuthTokenProvider) isTokenValidLocked() bool {
	if p.accessToken == "" || p.expiresAt.IsZero() {
		return false
	}
	return time.Until(p.expiresAt) > p.safetyBuffer
}

type oauthTokenResponse struct {
	AccessToken string `json:"access_token"`
	ExpiresIn   int64  `json:"expires_in"`
	TokenType   string `json:"token_type"`
}

func (p *OAuthTokenProvider) fetchTokenLocked(ctx context.Context) (string, time.Time, error) {
	form := url.Values{}
	form.Set("grant_type", "client_credentials")
	form.Set("client_id", p.ClientID)
	form.Set("client_secret", p.ClientSecret)
	if len(p.Scopes) > 0 {
		form.Set("scope", strings.Join(p.Scopes, " "))
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.TokenURL, strings.NewReader(form.Encode()))
	if err != nil {
		return "", time.Time{}, fmt.Errorf("failed to build oauth token request: %w", err)
	}
	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("failed to request oauth token: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", time.Time{}, fmt.Errorf("failed to read oauth token response: %w", err)
	}

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return "", time.Time{}, fmt.Errorf("oauth token request failed: %s: %s", resp.Status, truncateTokenBody(body))
	}

	var tokenResp oauthTokenResponse
	if err := json.Unmarshal(body, &tokenResp); err != nil {
		return "", time.Time{}, fmt.Errorf("failed to decode oauth token response: %w", err)
	}
	if tokenResp.AccessToken == "" {
		return "", time.Time{}, errors.New("oauth token response missing access_token")
	}

	expiresAt := time.Now()
	if tokenResp.ExpiresIn > 0 {
		expiresAt = expiresAt.Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
	}

	return tokenResp.AccessToken, expiresAt, nil
}

func truncateTokenBody(body []byte) string {
	trimmed := strings.TrimSpace(string(body))
	if len(trimmed) <= maxTokenErrorBodyBytes {
		return trimmed
	}
	return trimmed[:maxTokenErrorBodyBytes] + "..."
}