chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,310 @@
|
||||
package oauth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"golang.org/x/oauth2"
|
||||
)
|
||||
|
||||
// DefaultAuthTimeout bounds how long a single authorization attempt waits for
|
||||
// the user to complete the browser or device flow.
|
||||
const DefaultAuthTimeout = 5 * time.Minute
|
||||
|
||||
// tokenRefreshTimeout bounds each background refresh of an expiring token so a
|
||||
// stalled GitHub token endpoint cannot block a tool call indefinitely.
|
||||
const tokenRefreshTimeout = 30 * time.Second
|
||||
|
||||
// flowStatus tracks the manager's single-flight authorization state.
|
||||
type flowStatus int
|
||||
|
||||
const (
|
||||
statusIdle flowStatus = iota // no flow running
|
||||
statusStarting // a flow is being prepared (brief)
|
||||
statusInProgress // a flow is running on a secure channel; callers may join
|
||||
statusAwaitingUser // a flow is running but the user must act out-of-band
|
||||
)
|
||||
|
||||
// Outcome reports the result of an authorization attempt that did not
|
||||
// immediately yield a token.
|
||||
type Outcome struct {
|
||||
// UserAction, when non-nil, must be surfaced to the user. The authorization
|
||||
// flow continues in the background; the user should retry once they have
|
||||
// completed it.
|
||||
UserAction *UserAction
|
||||
}
|
||||
|
||||
// UserAction is an instruction for the user to complete authorization out of
|
||||
// band (the last-resort channel, used when neither a browser nor URL
|
||||
// elicitation is available).
|
||||
type UserAction struct {
|
||||
// Message is ready to display to the user.
|
||||
Message string
|
||||
// URL is the authorization URL or device verification URI.
|
||||
URL string
|
||||
// UserCode is the device-flow code to enter, if any.
|
||||
UserCode string
|
||||
}
|
||||
|
||||
// Manager owns the OAuth login flows and the resulting (refreshing) token for a
|
||||
// single stdio session. It is safe for concurrent use; only one authorization
|
||||
// flow runs at a time.
|
||||
type Manager struct {
|
||||
config Config
|
||||
refreshConfig *oauth2.Config
|
||||
logger *slog.Logger
|
||||
|
||||
// Test seams, set by NewManager to real implementations.
|
||||
openURL func(string) error
|
||||
inDocker func() bool
|
||||
|
||||
mu sync.Mutex
|
||||
source oauth2.TokenSource // refreshing source, set once authorized
|
||||
status flowStatus
|
||||
pending *UserAction
|
||||
done chan struct{}
|
||||
lastErr error
|
||||
refreshErrLogged bool // true once a refresh failure has been logged, reset on re-auth
|
||||
}
|
||||
|
||||
// NewManager builds a Manager for the given configuration. A nil logger logs to
|
||||
// stderr.
|
||||
func NewManager(cfg Config, logger *slog.Logger) *Manager {
|
||||
if logger == nil {
|
||||
logger = slog.New(slog.NewTextHandler(os.Stderr, nil))
|
||||
}
|
||||
m := &Manager{
|
||||
config: cfg,
|
||||
logger: logger,
|
||||
openURL: openBrowser,
|
||||
inDocker: isRunningInDocker,
|
||||
}
|
||||
m.refreshConfig = m.oauth2Config("")
|
||||
return m
|
||||
}
|
||||
|
||||
// AccessToken returns a currently valid access token, refreshing it if needed,
|
||||
// or "" if the session is not authorized (or a refresh has failed and
|
||||
// re-authorization is required). It is cheap to call repeatedly: the underlying
|
||||
// token source caches and only refreshes when the token has expired.
|
||||
func (m *Manager) AccessToken() string {
|
||||
m.mu.Lock()
|
||||
src := m.source
|
||||
m.mu.Unlock()
|
||||
if src == nil {
|
||||
return ""
|
||||
}
|
||||
// Refresh (if needed) happens here, off the lock, because ReuseTokenSource may
|
||||
// make a blocking network call and holding m.mu would serialize every tool call.
|
||||
tok, err := src.Token()
|
||||
if err != nil {
|
||||
// A refresh failure (expired GitHub App refresh token, revoked grant, or a
|
||||
// network blip) leaves the session unauthorized and forces a re-login.
|
||||
// Surface it once, otherwise it only manifests as a surprise re-authorization
|
||||
// prompt. The oauth2 error carries the token endpoint's response, not the
|
||||
// access or refresh token.
|
||||
m.mu.Lock()
|
||||
if !m.refreshErrLogged {
|
||||
m.refreshErrLogged = true
|
||||
m.logger.Warn("OAuth token refresh failed; re-authorization required", "error", err)
|
||||
}
|
||||
m.mu.Unlock()
|
||||
return ""
|
||||
}
|
||||
if !tok.Valid() {
|
||||
return ""
|
||||
}
|
||||
return tok.AccessToken
|
||||
}
|
||||
|
||||
// HasToken reports whether a valid token is currently available.
|
||||
func (m *Manager) HasToken() bool {
|
||||
return m.AccessToken() != ""
|
||||
}
|
||||
|
||||
// Authenticate ensures the session is authorized.
|
||||
//
|
||||
// It returns (nil, nil) once a token is available, so the caller may proceed.
|
||||
// It returns (&Outcome{UserAction}, nil) when the user must complete the flow
|
||||
// out of band; the flow continues in the background and the caller should show
|
||||
// the action and have the user retry. It returns (nil, err) on failure.
|
||||
//
|
||||
// Only one flow runs at a time. Concurrent callers either join a running secure
|
||||
// flow, receive the pending user action, or are told to retry shortly.
|
||||
func (m *Manager) Authenticate(ctx context.Context, prompter Prompter) (*Outcome, error) {
|
||||
if m.AccessToken() != "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
switch m.status {
|
||||
case statusAwaitingUser:
|
||||
ua := m.pending
|
||||
m.mu.Unlock()
|
||||
return &Outcome{UserAction: ua}, nil
|
||||
case statusStarting:
|
||||
m.mu.Unlock()
|
||||
return &Outcome{UserAction: &UserAction{
|
||||
Message: "GitHub authorization is already in progress. Please retry your request in a few seconds.",
|
||||
}}, nil
|
||||
case statusInProgress:
|
||||
done := m.done
|
||||
m.mu.Unlock()
|
||||
return m.joinWait(ctx, done)
|
||||
}
|
||||
|
||||
// Idle: this call owns the new flow.
|
||||
m.status = statusStarting
|
||||
m.lastErr = nil
|
||||
m.done = make(chan struct{})
|
||||
done := m.done
|
||||
m.mu.Unlock()
|
||||
|
||||
plan, err := m.begin(prompter)
|
||||
if err != nil {
|
||||
m.complete(nil, err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
if plan.userAction != nil {
|
||||
m.status = statusAwaitingUser
|
||||
m.pending = plan.userAction
|
||||
} else {
|
||||
m.status = statusInProgress
|
||||
}
|
||||
m.mu.Unlock()
|
||||
|
||||
bgCtx, cancel := context.WithTimeout(context.Background(), DefaultAuthTimeout)
|
||||
go m.runFlow(bgCtx, cancel, plan)
|
||||
|
||||
if plan.userAction != nil {
|
||||
return &Outcome{UserAction: plan.userAction}, nil
|
||||
}
|
||||
return m.joinWait(ctx, done)
|
||||
}
|
||||
|
||||
// runFlow executes a prepared flow in the background and records the result. The
|
||||
// optional display prompt runs concurrently: a decline (or other failure) aborts
|
||||
// the flow, while an undeliverable prompt degrades to the manual fallback without
|
||||
// tearing the flow down, so the user can still authorize out of band.
|
||||
func (m *Manager) runFlow(ctx context.Context, cancel context.CancelFunc, plan *flowPlan) {
|
||||
defer cancel()
|
||||
|
||||
if plan.display != nil {
|
||||
go func() {
|
||||
err := plan.display(ctx)
|
||||
switch {
|
||||
case err == nil:
|
||||
// Prompt shown; the flow completes when the token arrives.
|
||||
case ctx.Err() != nil:
|
||||
// The flow is already ending (timed out or cancelled elsewhere),
|
||||
// so there is nothing to fall back to. Checking this before the
|
||||
// fallback also prevents misreading a context-cancelled prompt as
|
||||
// a transport failure.
|
||||
case errors.Is(err, ErrPromptUnavailable) && plan.fallback != nil:
|
||||
// The client advertised the capability but could not deliver the
|
||||
// prompt. Surface the manual instructions instead of failing, and
|
||||
// keep the background flow alive so the user can still authorize.
|
||||
m.logger.Debug("authorization prompt undeliverable; falling back to manual instructions", "reason", err)
|
||||
m.fallBackToUserAction(plan.fallback)
|
||||
default:
|
||||
// A user decline (ErrPromptDeclined) or any other prompt failure
|
||||
// ends the flow.
|
||||
m.logger.Debug("authorization prompt closed", "reason", err)
|
||||
cancel()
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
tok, err := plan.run(ctx)
|
||||
m.complete(tok, err)
|
||||
}
|
||||
|
||||
// fallBackToUserAction promotes a running secure flow to the manual user-action
|
||||
// channel after its prompt could not be delivered. The background flow keeps
|
||||
// running, so the user can complete authorization out of band and retry. It is a
|
||||
// no-op if the flow has already resolved.
|
||||
func (m *Manager) fallBackToUserAction(ua *UserAction) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if m.status != statusInProgress {
|
||||
return
|
||||
}
|
||||
m.status = statusAwaitingUser
|
||||
m.pending = ua
|
||||
// Wake any callers joined on this flow so they receive the action, and clear
|
||||
// done so complete() does not double-close it when run() later finishes.
|
||||
if m.done != nil {
|
||||
close(m.done)
|
||||
m.done = nil
|
||||
}
|
||||
}
|
||||
|
||||
// complete records the flow result, installing a refreshing token source on
|
||||
// success, and wakes any joined callers.
|
||||
func (m *Manager) complete(tok *oauth2.Token, err error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
m.status = statusIdle
|
||||
m.pending = nil
|
||||
if err != nil {
|
||||
m.lastErr = err
|
||||
m.logger.Debug("oauth flow failed", "error", err)
|
||||
} else {
|
||||
m.lastErr = nil
|
||||
// Config.TokenSource returns a ReuseTokenSource that refreshes expired
|
||||
// tokens using the refresh token — this is what makes GitHub App
|
||||
// (expiring) tokens work transparently. The refresh uses a bounded HTTP
|
||||
// client so a stalled token endpoint can't block a tool call forever.
|
||||
refreshCtx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{Timeout: tokenRefreshTimeout})
|
||||
m.source = m.refreshConfig.TokenSource(refreshCtx, tok)
|
||||
m.refreshErrLogged = false
|
||||
m.logger.Info("github authorization complete")
|
||||
}
|
||||
if m.done != nil {
|
||||
close(m.done)
|
||||
m.done = nil
|
||||
}
|
||||
}
|
||||
|
||||
// joinWait blocks until the running flow finishes or ctx is cancelled. If the
|
||||
// flow was promoted to the manual channel while waiting (its prompt could not be
|
||||
// delivered), it returns that user action rather than an error.
|
||||
func (m *Manager) joinWait(ctx context.Context, done chan struct{}) (*Outcome, error) {
|
||||
select {
|
||||
case <-done:
|
||||
if m.AccessToken() != "" {
|
||||
return nil, nil
|
||||
}
|
||||
m.mu.Lock()
|
||||
pending := m.pending
|
||||
err := m.lastErr
|
||||
m.mu.Unlock()
|
||||
if pending != nil {
|
||||
return &Outcome{UserAction: pending}, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, errors.New("authorization did not complete")
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) oauth2Config(redirectURL string) *oauth2.Config {
|
||||
return &oauth2.Config{
|
||||
ClientID: m.config.ClientID,
|
||||
ClientSecret: m.config.ClientSecret,
|
||||
RedirectURL: redirectURL,
|
||||
Scopes: m.config.Scopes,
|
||||
Endpoint: m.config.Endpoint,
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user