verae-mcp-timestamping-server-go/internal/tools/oauth.go

package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

// Default portal base for OAuth (authentik-brokered via edge).
const defaultOauthPortalURL = "https://customer.verae.ai"

// OauthStartArgs begins a browser OAuth window flow.
type OauthStartArgs struct {
	Provider  string `json:"provider,omitempty" jsonschema:"google|github|apple|x (default google); authentik Source slug"`
	AcceptTos bool   `json:"accept_tos,omitempty" jsonschema:"required true when creating a new account via social login"`
}

// OauthStartOut is returned by oauth_start.
type OauthStartOut struct {
	AuthorizeURL string `json:"authorize_url"`
	State        string `json:"state"`
	ExpiresAt    string `json:"expires_at,omitempty"`
	Provider     string `json:"provider"`
	HowToUse     string `json:"how_to_use"`
}

// OauthCompleteArgs finishes OAuth after the browser callback (poll edge result).
type OauthCompleteArgs struct {
	State   string `json:"state" jsonschema:"state from oauth_start"`
	WaitSec int    `json:"wait_sec,omitempty" jsonschema:"seconds to poll (default 45, max 120)"`
}

// OauthCompleteOut is returned when social login succeeds.
type OauthCompleteOut struct {
	Username    string `json:"username,omitempty"`
	CustomerID  string `json:"customer_id,omitempty"`
	Created     bool   `json:"created,omitempty"`
	AccessToken string `json:"access_token,omitempty"`
	HowToUse    string `json:"how_to_use,omitempty"`
	Status      string `json:"status,omitempty"` // pending when still waiting (short poll)
}

func oauthPortalBase() string {
	u := strings.TrimRight(strings.TrimSpace(os.Getenv("VERAE_OAUTH_PORTAL_URL")), "/")
	if u == "" {
		u = defaultOauthPortalURL
	}
	return u
}

func (d *Deps) oauthHTTPClient() *http.Client {
	return &http.Client{Timeout: 30 * time.Second}
}

func (d *Deps) oauthStart(ctx context.Context, _ *mcp.CallToolRequest, in OauthStartArgs) (*mcp.CallToolResult, OauthStartOut, error) {
	provider := strings.ToLower(strings.TrimSpace(in.Provider))
	if provider == "" {
		provider = "google"
	}
	switch provider {
	case "google", "github", "apple", "x":
	default:
		return nil, OauthStartOut{}, fmt.Errorf("provider must be google, github, apple, or x")
	}
	q := url.Values{
		"provider": {provider},
		"popup":    {"1"},
	}
	if in.AcceptTos {
		q.Set("accept_tos", "1")
	}
	endpoint := oauthPortalBase() + "/portal/api/oauth/start?" + q.Encode()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
	if err != nil {
		return nil, OauthStartOut{}, err
	}
	req.Header.Set("Accept", "application/json")
	res, err := d.oauthHTTPClient().Do(req)
	if err != nil {
		return nil, OauthStartOut{}, fmt.Errorf("oauth start: %w", err)
	}
	defer res.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(res.Body, 1<<20))
	if res.StatusCode == http.StatusServiceUnavailable {
		return nil, OauthStartOut{}, fmt.Errorf("oauth not configured on portal (%s)", strings.TrimSpace(string(body)))
	}
	if res.StatusCode >= 400 {
		return nil, OauthStartOut{}, fmt.Errorf("oauth start HTTP %d: %s", res.StatusCode, strings.TrimSpace(string(body)))
	}
	var raw struct {
		AuthorizeURL string `json:"authorize_url"`
		State        string `json:"state"`
		ExpiresAt    string `json:"expires_at"`
		Provider     string `json:"provider"`
		Error        string `json:"error"`
		Detail       string `json:"detail"`
	}
	if err := json.Unmarshal(body, &raw); err != nil {
		return nil, OauthStartOut{}, fmt.Errorf("oauth start parse: %w", err)
	}
	if raw.AuthorizeURL == "" || raw.State == "" {
		msg := raw.Error
		if raw.Detail != "" {
			msg = raw.Detail
		}
		if msg == "" {
			msg = "missing authorize_url/state"
		}
		return nil, OauthStartOut{}, fmt.Errorf("oauth start: %s", msg)
	}
	how := "Open authorize_url in a browser window (or system browser). " +
		"Complete sign-in at auth.verae.ai. Then call tool oauth_complete with the same state. " +
		"For a new account, call oauth_start again with accept_tos=true after accepting Terms on the portal if needed."
	return nil, OauthStartOut{
		AuthorizeURL: raw.AuthorizeURL,
		State:        raw.State,
		ExpiresAt:    raw.ExpiresAt,
		Provider:     provider,
		HowToUse:     how,
	}, nil
}

func (d *Deps) oauthComplete(ctx context.Context, _ *mcp.CallToolRequest, in OauthCompleteArgs) (*mcp.CallToolResult, OauthCompleteOut, error) {
	state := strings.TrimSpace(in.State)
	if state == "" {
		return nil, OauthCompleteOut{}, fmt.Errorf("state is required (from oauth_start)")
	}
	wait := in.WaitSec
	if wait <= 0 {
		wait = 45
	}
	if wait > 120 {
		wait = 120
	}
	deadline := time.Now().Add(time.Duration(wait) * time.Second)
	var lastPending bool
	for {
		out, pending, err := d.pollOauthResult(ctx, state)
		if err != nil {
			return nil, OauthCompleteOut{}, err
		}
		if !pending {
			if !d.Public && out.AccessToken != "" {
				_ = d.persistOauthBearer(out.Username, out.AccessToken)
			}
			out.HowToUse = tokenHowToUse
			return nil, out, nil
		}
		lastPending = true
		if time.Now().After(deadline) {
			break
		}
		select {
		case <-ctx.Done():
			return nil, OauthCompleteOut{}, ctx.Err()
		case <-time.After(2 * time.Second):
		}
	}
	if lastPending {
		return nil, OauthCompleteOut{
			Status:   "pending",
			HowToUse: "Browser sign-in not finished yet. Open authorize_url from oauth_start, then call oauth_complete again with the same state.",
		}, nil
	}
	return nil, OauthCompleteOut{}, fmt.Errorf("oauth complete: unexpected empty result")
}

func (d *Deps) pollOauthResult(ctx context.Context, state string) (OauthCompleteOut, bool, error) {
	endpoint := oauthPortalBase() + "/portal/api/oauth/result?state=" + url.QueryEscape(state)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
	if err != nil {
		return OauthCompleteOut{}, false, err
	}
	req.Header.Set("Accept", "application/json")
	res, err := d.oauthHTTPClient().Do(req)
	if err != nil {
		return OauthCompleteOut{}, false, fmt.Errorf("oauth result: %w", err)
	}
	defer res.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(res.Body, 1<<20))
	switch res.StatusCode {
	case http.StatusAccepted:
		return OauthCompleteOut{Status: "pending"}, true, nil
	case http.StatusNotFound:
		return OauthCompleteOut{}, false, fmt.Errorf("unknown_or_expired_state")
	}
	if res.StatusCode >= 400 {
		return OauthCompleteOut{}, false, fmt.Errorf("oauth result HTTP %d: %s", res.StatusCode, strings.TrimSpace(string(body)))
	}
	var raw struct {
		OK         bool   `json:"ok"`
		Token      string `json:"token"`
		APIKey     string `json:"apiKey"`
		Created    bool   `json:"created"`
		Email      string `json:"email"`
		CustomerID string `json:"customerId"`
		Error      string `json:"error"`
		Detail     string `json:"detail"`
		Status     string `json:"status"`
	}
	if err := json.Unmarshal(body, &raw); err != nil {
		return OauthCompleteOut{}, false, fmt.Errorf("oauth result parse: %w", err)
	}
	if raw.Status == "pending" {
		return OauthCompleteOut{Status: "pending"}, true, nil
	}
	if !raw.OK {
		msg := raw.Error
		if raw.Detail != "" {
			msg = raw.Detail
		}
		if msg == "" {
			msg = "oauth failed"
		}
		return OauthCompleteOut{}, false, fmt.Errorf("%s", msg)
	}
	// Prefer long-lived xVerae- portal key for commercial MCP stamps when EDGE_URL is set.
	tok := strings.TrimSpace(raw.APIKey)
	if tok == "" {
		tok = strings.TrimSpace(raw.Token)
	}
	return OauthCompleteOut{
		Username:    raw.Email,
		CustomerID:  raw.CustomerID,
		Created:     raw.Created,
		AccessToken: tok,
	}, false, nil
}

func (d *Deps) persistOauthBearer(username, token string) error {
	if d.Global == nil || token == "" {
		return nil
	}
	return d.persistCallerToken(username, "", token, "")
}