Commilitia-Drop/desktop/platform/oauth.go

// Package platform implements cdrop desktop's native capabilities — the Go
// shell that the shared web UI drives over Wails bindings.
//
// This file owns the login flow: the Auth Broker device-authorization flow (RFC
// 8252 loopback redirect with PKCE, as a public client — no secret). The desktop
// opens the system browser at the broker's consent page; on approval the broker
// redirects to the loopback /callback with a one-time code, which Go exchanges
// (with the PKCE verifier) for a broker machine token. The verifier and the
// resulting tokens never enter the WebView / JS context — refresh is Go-driven.
package platform

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// OAuthConfig carries the Auth Broker coordinates. BrokerURL is the broker's PUBLIC
// origin (e.g. https://sso.commilitia.net); App is this app's key in the broker apps
// registry ("cdrop"). The app's loopback redirect must be registered in the broker's
// apps.json (redirect_uris) for the device flow to accept it.
type OAuthConfig struct {
	BrokerURL string
	App       string
}

// TokenResult is the subset of the broker token response the shell keeps. ExpiresIn
// is derived from the broker's access_expires (unix) at parse time.
type TokenResult struct {
	AccessToken  string
	RefreshToken string
	ExpiresIn    int
}

// Flow runs one interactive login. openURL is injected so production uses Wails'
// runtime.BrowserOpenURL while tests substitute a fake user-agent.
type Flow struct {
	cfg     OAuthConfig
	openURL func(string)
	client  *http.Client
	timeout time.Duration
}

// NewFlow builds a Flow with sane defaults: a 30s HTTP client for the token
// exchange and a 3-minute ceiling on the whole interactive login.
func NewFlow(cfg OAuthConfig, openURL func(string)) *Flow {
	return &Flow{
		cfg:     cfg,
		openURL: openURL,
		client:  &http.Client{Timeout: 30 * time.Second},
		timeout: 3 * time.Minute,
	}
}

// Login performs the broker device-authorization loopback PKCE flow and returns the
// minted machine token.
//
// Steps: generate verifier/challenge/state → bind an ephemeral 127.0.0.1 listener →
// open the system browser at the broker's /device/authorize (which sends the user
// through broker SSO if needed, then shows a consent page) → wait for the browser to
// hit the loopback /callback with the code → validate state → exchange code+verifier
// at /device/token. The verifier never leaves Go.
func (f *Flow) Login(ctx context.Context) (*TokenResult, error) {
	if f.cfg.BrokerURL == "" || f.cfg.App == "" {
		return nil, errors.New("oauth: incomplete config (broker_url / app)")
	}

	verifier, challenge, err := newPKCE()
	if err != nil {
		return nil, fmt.Errorf("oauth: pkce: %w", err)
	}
	state, err := randString(24)
	if err != nil {
		return nil, fmt.Errorf("oauth: state: %w", err)
	}

	// RFC 8252 §7.3: bind the IPv4 loopback literal (not "localhost") on an
	// OS-assigned ephemeral port.
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, fmt.Errorf("oauth: loopback listen: %w", err)
	}
	defer ln.Close()
	redirectURI := fmt.Sprintf("http://127.0.0.1:%d/callback", ln.Addr().(*net.TCPAddr).Port)

	type cbResult struct {
		code string
		err  error
	}
	resCh := make(chan cbResult, 1)
	send := func(r cbResult) {
		// non-blocking: only the first callback wins, later ones are dropped.
		select {
		case resCh <- r:
		default:
		}
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/callback", func(w http.ResponseWriter, r *http.Request) {
		q := r.URL.Query()
		if e := q.Get("error"); e != "" {
			writeClosePage(w, false)
			send(cbResult{err: fmt.Errorf("oauth: authorization denied: %s", e)})
			return
		}
		if q.Get("state") != state {
			writeClosePage(w, false)
			send(cbResult{err: errors.New("oauth: state mismatch (possible CSRF)")})
			return
		}
		code := q.Get("code")
		if code == "" {
			writeClosePage(w, false)
			send(cbResult{err: errors.New("oauth: callback missing code")})
			return
		}
		writeClosePage(w, true)
		send(cbResult{code: code})
	})

	srv := &http.Server{Handler: mux}
	go func() { _ = srv.Serve(ln) }()
	defer srv.Close()

	authURL := strings.TrimRight(f.cfg.BrokerURL, "/") + "/device/authorize?" + url.Values{
		"app":                   {f.cfg.App},
		"redirect_uri":          {redirectURI},
		"state":                 {state},
		"code_challenge":        {challenge},
		"code_challenge_method": {"S256"},
		"description":           {"cdrop 桌面客户端"},
	}.Encode()
	f.openURL(authURL)

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-time.After(f.timeout):
		return nil, errors.New("oauth: login timed out")
	case res := <-resCh:
		if res.err != nil {
			return nil, res.err
		}
		return f.exchange(ctx, res.code, verifier, redirectURI)
	}
}

// exchange trades the authorization code + PKCE verifier (+ the same redirect_uri)
// for the broker machine token at /device/token.
func (f *Flow) exchange(ctx context.Context, code, verifier, redirectURI string) (*TokenResult, error) {
	return f.postJSON(ctx, strings.TrimRight(f.cfg.BrokerURL, "/")+"/device/token", map[string]string{
		"code":          code,
		"code_verifier": verifier,
		"redirect_uri":  redirectURI,
	})
}

// Refresh exchanges a refresh credential for a fresh access token (+ rotated refresh)
// at the broker's /refresh. Call it when the access token nears expiry; the refresh
// credential never leaves Go, so refresh is driven from the shell, not the WebView.
func (f *Flow) Refresh(ctx context.Context, refreshToken string) (*TokenResult, error) {
	if f.cfg.BrokerURL == "" {
		return nil, errors.New("oauth: incomplete config (broker_url)")
	}
	if refreshToken == "" {
		return nil, errors.New("oauth: empty refresh token")
	}
	return f.postJSON(ctx, strings.TrimRight(f.cfg.BrokerURL, "/")+"/refresh", map[string]string{
		"refresh": refreshToken,
	})
}

// brokerTokenResp is the broker's response shape for /device/token and /refresh.
type brokerTokenResp struct {
	Access        string `json:"access"`
	Refresh       string `json:"refresh"`
	AccessExpires int64  `json:"access_expires"`
}

// postJSON POSTs a JSON body to a broker token endpoint and decodes the response.
// Shared by the device-token exchange and refresh.
func (f *Flow) postJSON(ctx context.Context, endpoint string, body map[string]string) (*TokenResult, error) {
	buf, err := json.Marshal(body)
	if err != nil {
		return nil, fmt.Errorf("oauth: marshal token request: %w", err)
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(string(buf)))
	if err != nil {
		return nil, fmt.Errorf("oauth: build token request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	resp, err := f.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("oauth: token request: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("oauth: token endpoint status %d", resp.StatusCode)
	}
	var tr brokerTokenResp
	if err := json.NewDecoder(resp.Body).Decode(&tr); err != nil {
		return nil, fmt.Errorf("oauth: decode token: %w", err)
	}
	if tr.Access == "" {
		return nil, errors.New("oauth: token response missing access")
	}
	expiresIn := int(tr.AccessExpires - time.Now().Unix())
	if expiresIn < 0 {
		expiresIn = 0
	}
	return &TokenResult{AccessToken: tr.Access, RefreshToken: tr.Refresh, ExpiresIn: expiresIn}, nil
}

// --- PKCE (RFC 7636) ---

func newPKCE() (verifier, challenge string, err error) {
	verifier, err = randString(64) // 43–128 chars required; 64 is comfortable
	if err != nil {
		return "", "", err
	}
	sum := sha256.Sum256([]byte(verifier))
	challenge = base64.RawURLEncoding.EncodeToString(sum[:])
	return verifier, challenge, nil
}

// randString returns n characters of base64url-encoded crypto-random data. The
// base64url alphabet (A–Z a–z 0–9 - _) is a subset of the PKCE unreserved set,
// so the output is a valid code_verifier / state.
func randString(n int) (string, error) {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	s := base64.RawURLEncoding.EncodeToString(b)
	if len(s) > n {
		s = s[:n]
	}
	return s, nil
}

// writeClosePage is what the browser tab lands on after the redirect. The tab
// usually can't be closed by script (it wasn't script-opened), so the text is
// the real affordance; the close attempt is best-effort.
func writeClosePage(w http.ResponseWriter, ok bool) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	msg := "已登录，可关闭本页返回 cdrop。"
	if !ok {
		msg = "登录未完成，可关闭本页返回 cdrop 重试。"
	}
	fmt.Fprintf(w, `<!doctype html><html lang="zh-Hans"><head><meta charset="utf-8"><title>cdrop</title></head>`+
		`<body style="font-family:system-ui,sans-serif;text-align:center;margin-top:20vh">`+
		`<p>%s</p><script>setTimeout(function(){window.close();},800);</script></body></html>`, msg)
}