mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-13 08:15:29 +00:00
298 lines
9.1 KiB
Go
298 lines
9.1 KiB
Go
package mcpoauth
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/rand"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/nextlevelbuilder/goclaw/internal/security"
|
|
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
const (
|
|
flowTTL = 10 * time.Minute
|
|
flowCleanupEvery = time.Minute
|
|
)
|
|
|
|
// OAuthTokens is the token response from the Authorization Server.
|
|
type OAuthTokens struct {
|
|
AccessToken string `json:"access_token"`
|
|
RefreshToken string `json:"refresh_token,omitempty"`
|
|
TokenType string `json:"token_type"`
|
|
ExpiresIn int `json:"expires_in,omitempty"`
|
|
Scope string `json:"scope,omitempty"`
|
|
}
|
|
|
|
// PendingFlow holds state for an in-flight authorization code flow.
|
|
type PendingFlow struct {
|
|
ServerID uuid.UUID
|
|
TenantID uuid.UUID
|
|
UserID string // empty = global/tenant-level, non-empty = per-user
|
|
InitiatingUserID string // admin who called /start — used for WS event routing
|
|
CodeVerifier string
|
|
UsePKCE bool // true → send code_verifier on token exchange
|
|
RedirectURI string
|
|
TokenEndpoint string
|
|
ClientID string
|
|
ClientSecret string // empty for public clients (PKCE)
|
|
ResourceURI string // RFC 8707
|
|
Issuer string // AS issuer — persisted with the token for status display
|
|
CreatedAt time.Time
|
|
}
|
|
|
|
// StartFlowParams holds the parameters required to begin an OAuth flow.
|
|
type StartFlowParams struct {
|
|
ServerID uuid.UUID
|
|
TenantID uuid.UUID
|
|
UserID string
|
|
InitiatingUserID string // admin who called /start — stored for WS event routing on callback
|
|
DiscoveryResult *DiscoveryResult
|
|
ClientID string
|
|
ClientSecret string
|
|
RedirectURI string
|
|
Scopes string
|
|
// GrantType controls the authorization flow variant.
|
|
// "pkce" (default/empty) — Authorization Code + PKCE S256, public client (no ClientSecret).
|
|
// "authorization_code" — Authorization Code + PKCE S256, confidential client (sends ClientSecret).
|
|
// "client_credentials" — not handled by StartFlow; use ClientCredentials() instead.
|
|
GrantType string
|
|
|
|
// ExtraAuthParams carries provider-specific authorization-request parameters
|
|
// that goclaw's generic flow does not otherwise set (e.g. Dropbox's
|
|
// token_access_type=offline, which the AS requires to issue a refresh token;
|
|
// without it Dropbox returns an online-only token that cannot be refreshed).
|
|
// Standard OAuth keys already set by StartFlow are never overridden.
|
|
ExtraAuthParams url.Values
|
|
}
|
|
|
|
// FlowManager manages in-progress OAuth authorization code flows.
|
|
type FlowManager struct {
|
|
mu sync.Mutex
|
|
pending map[string]*PendingFlow // state → flow
|
|
client *http.Client
|
|
}
|
|
|
|
func NewFlowManager(client *http.Client) *FlowManager {
|
|
fm := &FlowManager{
|
|
pending: make(map[string]*PendingFlow),
|
|
client: client,
|
|
}
|
|
go fm.cleanupLoop()
|
|
return fm
|
|
}
|
|
|
|
// StartFlow initiates an authorization code flow.
|
|
// GrantType "pkce" — PKCE S256, public client (no client_secret).
|
|
// GrantType "authorization_code" — PKCE S256, confidential client (sends client_secret too).
|
|
// OAuth 2.1 mandates PKCE for ALL authorization code flows; the grant type only controls
|
|
// whether a client_secret is included, not whether PKCE is used.
|
|
// Returns the authorization URL and CSRF state token.
|
|
func (fm *FlowManager) StartFlow(ctx context.Context, p StartFlowParams) (authURL, state string, err error) {
|
|
_ = ctx
|
|
|
|
// OAuth 2.1 §4.1.1: PKCE is mandatory for all authorization code flows.
|
|
// "authorization_code" differs from "pkce" only in that it sends client_secret too.
|
|
usePKCE := p.GrantType != "client_credentials"
|
|
|
|
var verifier, challenge string
|
|
if usePKCE {
|
|
verifier, challenge, err = generatePKCE()
|
|
if err != nil {
|
|
return "", "", fmt.Errorf("mcpoauth: generate PKCE: %w", err)
|
|
}
|
|
}
|
|
|
|
state, err = generateState()
|
|
if err != nil {
|
|
return "", "", fmt.Errorf("mcpoauth: generate state: %w", err)
|
|
}
|
|
|
|
fm.mu.Lock()
|
|
fm.pending[state] = &PendingFlow{
|
|
ServerID: p.ServerID,
|
|
TenantID: p.TenantID,
|
|
UserID: p.UserID,
|
|
InitiatingUserID: p.InitiatingUserID,
|
|
CodeVerifier: verifier,
|
|
UsePKCE: usePKCE,
|
|
RedirectURI: p.RedirectURI,
|
|
TokenEndpoint: p.DiscoveryResult.TokenEndpoint,
|
|
ClientID: p.ClientID,
|
|
ClientSecret: p.ClientSecret,
|
|
ResourceURI: p.DiscoveryResult.ResourceURI,
|
|
Issuer: p.DiscoveryResult.Issuer,
|
|
CreatedAt: time.Now(),
|
|
}
|
|
fm.mu.Unlock()
|
|
|
|
q := url.Values{
|
|
"response_type": {"code"},
|
|
"client_id": {p.ClientID},
|
|
"redirect_uri": {p.RedirectURI},
|
|
"state": {state},
|
|
}
|
|
if usePKCE {
|
|
q.Set("code_challenge", challenge)
|
|
q.Set("code_challenge_method", "S256")
|
|
}
|
|
if p.Scopes != "" {
|
|
q.Set("scope", p.Scopes)
|
|
}
|
|
if p.DiscoveryResult.ResourceURI != "" {
|
|
q.Set("resource", p.DiscoveryResult.ResourceURI) // RFC 8707
|
|
}
|
|
for k, vs := range p.ExtraAuthParams {
|
|
if q.Has(k) {
|
|
continue // standard OAuth params win — never overridden
|
|
}
|
|
for _, v := range vs {
|
|
q.Add(k, v)
|
|
}
|
|
}
|
|
|
|
authURL = p.DiscoveryResult.AuthorizationEndpoint + "?" + q.Encode()
|
|
return authURL, state, nil
|
|
}
|
|
|
|
// ExchangeCode validates the CSRF state, exchanges the authorization code for tokens,
|
|
// and returns the tokens along with the original flow parameters.
|
|
func (fm *FlowManager) ExchangeCode(ctx context.Context, state, code string) (*OAuthTokens, *PendingFlow, error) {
|
|
fm.mu.Lock()
|
|
flow, ok := fm.pending[state]
|
|
if ok {
|
|
delete(fm.pending, state)
|
|
}
|
|
fm.mu.Unlock()
|
|
|
|
if !ok || time.Since(flow.CreatedAt) > flowTTL {
|
|
return nil, nil, fmt.Errorf("mcpoauth: invalid or expired OAuth state")
|
|
}
|
|
|
|
tokens, err := fm.exchangeCodeHTTP(ctx, flow, code)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
return tokens, flow, nil
|
|
}
|
|
|
|
func (fm *FlowManager) exchangeCodeHTTP(ctx context.Context, flow *PendingFlow, code string) (*OAuthTokens, error) {
|
|
params := url.Values{
|
|
"grant_type": {"authorization_code"},
|
|
"code": {code},
|
|
"redirect_uri": {flow.RedirectURI},
|
|
"client_id": {flow.ClientID},
|
|
}
|
|
if flow.UsePKCE {
|
|
params.Set("code_verifier", flow.CodeVerifier)
|
|
}
|
|
if flow.ClientSecret != "" {
|
|
params.Set("client_secret", flow.ClientSecret)
|
|
}
|
|
if flow.ResourceURI != "" {
|
|
params.Set("resource", flow.ResourceURI) // RFC 8707
|
|
}
|
|
|
|
return postTokenRequest(ctx, fm.client, flow.TokenEndpoint, params)
|
|
}
|
|
|
|
// ClientCredentials performs a client credentials grant (no user interaction).
|
|
func (fm *FlowManager) ClientCredentials(ctx context.Context, tokenEndpoint, clientID, clientSecret, scopes, resourceURI string) (*OAuthTokens, error) {
|
|
params := url.Values{
|
|
"grant_type": {"client_credentials"},
|
|
"client_id": {clientID},
|
|
}
|
|
if clientSecret != "" {
|
|
params.Set("client_secret", clientSecret)
|
|
}
|
|
if scopes != "" {
|
|
params.Set("scope", scopes)
|
|
}
|
|
if resourceURI != "" {
|
|
params.Set("resource", resourceURI)
|
|
}
|
|
return postTokenRequest(ctx, fm.client, tokenEndpoint, params)
|
|
}
|
|
|
|
func (fm *FlowManager) cleanupLoop() {
|
|
ticker := time.NewTicker(flowCleanupEvery)
|
|
defer ticker.Stop()
|
|
for range ticker.C {
|
|
cutoff := time.Now().Add(-flowTTL)
|
|
fm.mu.Lock()
|
|
for state, flow := range fm.pending {
|
|
if flow.CreatedAt.Before(cutoff) {
|
|
delete(fm.pending, state)
|
|
}
|
|
}
|
|
fm.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
// postTokenRequest sends a POST to the token endpoint with form-encoded params.
|
|
func postTokenRequest(ctx context.Context, client *http.Client, tokenEndpoint string, params url.Values) (*OAuthTokens, error) {
|
|
_, ip, err := security.Validate(tokenEndpoint)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("mcpoauth: SSRF validation failed for token endpoint: %w", err)
|
|
}
|
|
ctx = security.WithPinnedIP(ctx, ip)
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenEndpoint, bytes.NewBufferString(params.Encode()))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
|
req.Header.Set("Accept", "application/json")
|
|
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("mcpoauth: token request failed: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
body, _ := io.ReadAll(io.LimitReader(resp.Body, 64*1024))
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("mcpoauth: token endpoint returned HTTP %d: %s", resp.StatusCode, string(body))
|
|
}
|
|
var tokens OAuthTokens
|
|
if err := json.Unmarshal(body, &tokens); err != nil {
|
|
return nil, fmt.Errorf("mcpoauth: token response parse error: %w", err)
|
|
}
|
|
if tokens.AccessToken == "" {
|
|
return nil, fmt.Errorf("mcpoauth: token response missing access_token")
|
|
}
|
|
if tokens.TokenType == "" {
|
|
tokens.TokenType = "Bearer"
|
|
}
|
|
return &tokens, nil
|
|
}
|
|
|
|
// generatePKCE returns (code_verifier, code_challenge_S256).
|
|
func generatePKCE() (verifier, challenge string, err error) {
|
|
buf := make([]byte, 32)
|
|
if _, err = rand.Read(buf); err != nil {
|
|
return "", "", err
|
|
}
|
|
verifier = base64.RawURLEncoding.EncodeToString(buf)
|
|
h := sha256.Sum256([]byte(verifier))
|
|
challenge = base64.RawURLEncoding.EncodeToString(h[:])
|
|
return verifier, challenge, nil
|
|
}
|
|
|
|
// generateState returns a cryptographically random CSRF state token.
|
|
func generateState() (string, error) {
|
|
buf := make([]byte, 16)
|
|
if _, err := rand.Read(buf); err != nil {
|
|
return "", err
|
|
}
|
|
return base64.RawURLEncoding.EncodeToString(buf), nil
|
|
}
|