sentryagent-idp/sdk-go/token_manager.go

package agentidp

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

const refreshBufferSeconds = 60

// cachedToken holds an access token and its expiry time.
type cachedToken struct {
	accessToken string
	expiresAt   time.Time
}

// isValid returns true if the token will not expire within the refresh buffer.
func (c *cachedToken) isValid() bool {
	return time.Now().Add(refreshBufferSeconds * time.Second).Before(c.expiresAt)
}

// TokenManager obtains and caches OAuth 2.0 client credentials tokens.
// It is safe for concurrent use by multiple goroutines.
type TokenManager struct {
	baseURL      string
	clientID     string
	clientSecret string
	scope        string
	httpClient   *http.Client
	mu           sync.Mutex
	cached       *cachedToken
}

// NewTokenManager creates a TokenManager that fetches tokens from baseURL
// using the given client credentials and scope.
func NewTokenManager(baseURL, clientID, clientSecret, scope string) *TokenManager {
	return &TokenManager{
		baseURL:      strings.TrimRight(baseURL, "/"),
		clientID:     clientID,
		clientSecret: clientSecret,
		scope:        scope,
		httpClient:   &http.Client{Timeout: 10 * time.Second},
	}
}

// GetToken returns a valid access token, fetching a new one if the cache is
// empty or the cached token is within the refresh buffer window.
// It is goroutine-safe.
func (tm *TokenManager) GetToken(ctx context.Context) (string, error) {
	tm.mu.Lock()
	defer tm.mu.Unlock()

	if tm.cached != nil && tm.cached.isValid() {
		return tm.cached.accessToken, nil
	}

	token, err := tm.fetchToken(ctx)
	if err != nil {
		return "", err
	}

	tm.cached = token
	return token.accessToken, nil
}

// ClearCache invalidates the cached token. The next call to GetToken will
// fetch a fresh token from the server.
func (tm *TokenManager) ClearCache() {
	tm.mu.Lock()
	defer tm.mu.Unlock()
	tm.cached = nil
}

// fetchToken performs the OAuth 2.0 client credentials grant.
// Must be called with mu held.
func (tm *TokenManager) fetchToken(ctx context.Context) (*cachedToken, error) {
	form := url.Values{}
	form.Set("grant_type", "client_credentials")
	form.Set("client_id", tm.clientID)
	form.Set("client_secret", tm.clientSecret)
	form.Set("scope", tm.scope)

	tokenURL := tm.baseURL + "/api/v1/token"
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, bytes.NewBufferString(form.Encode()))
	if err != nil {
		return nil, &AgentIdPError{
			Code:       "REQUEST_BUILD_ERROR",
			Message:    fmt.Sprintf("failed to build token request: %s", err.Error()),
			HTTPStatus: 0,
		}
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := tm.httpClient.Do(req)
	if err != nil {
		return nil, newNetworkError(err)
	}
	defer resp.Body.Close() //nolint:errcheck

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, newNetworkError(err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, parseOAuth2Error(respBody, resp.StatusCode)
	}

	var tr TokenResponse
	if err := json.Unmarshal(respBody, &tr); err != nil {
		return nil, &AgentIdPError{
			Code:       "PARSE_ERROR",
			Message:    fmt.Sprintf("failed to parse token response: %s", err.Error()),
			HTTPStatus: resp.StatusCode,
		}
	}

	return &cachedToken{
		accessToken: tr.AccessToken,
		expiresAt:   time.Now().Add(time.Duration(tr.ExpiresIn) * time.Second),
	}, nil
}