1package auth
2
3import (
4 "context"
5 "fmt"
6 "net/url"
7 "os"
8 "os/exec"
9 "runtime"
10 "strings"
11 "sync"
12 "time"
13)
14
15// Event is delivered on Manager.Events().
16type Event struct {
17 Kind string // "started" | "browser_open" | "success" | "error" | "canceled"
18 Provider string
19 Method string
20 URL string // login URL (for "started"/"browser_open")
21 Message string // on error
22}
23
24// Manager drives a login flow end-to-end. It owns a local web server
25// (for api-key form) plus provider-specific OAuth callback servers.
26type Manager struct {
27 store *Store
28 keyServer *Server // random-port web form server (api-key flow)
29 oauthServer *CallbackServer // fixed-port callback server (oauth flow, only one at a time)
30 mu sync.Mutex
31 events chan Event
32 openBrowser bool
33
34 oauthCtx context.Context
35 oauthCancel context.CancelFunc
36
37 manualOp *OAuthProvider
38 manualPKCE PKCE
39 manualState string
40}
41
42// NewManager returns a Manager bound to store.
43func NewManager(store *Store) *Manager {
44 return &Manager{
45 store: store,
46 events: make(chan Event, 16),
47 openBrowser: true,
48 }
49}
50
51// Store returns the underlying credential store.
52func (m *Manager) Store() *Store { return m.store }
53
54// Events returns the read-only event channel.
55func (m *Manager) Events() <-chan Event { return m.events }
56
57// Close shuts down any running servers and cancels pending flows.
58func (m *Manager) Close() {
59 m.mu.Lock()
60 defer m.mu.Unlock()
61 if m.keyServer != nil {
62 ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
63 _ = m.keyServer.Shutdown(ctx)
64 cancel()
65 m.keyServer = nil
66 }
67 if m.oauthServer != nil {
68 m.oauthServer.Shutdown()
69 m.oauthServer = nil
70 }
71 if m.oauthCancel != nil {
72 m.oauthCancel()
73 m.oauthCancel = nil
74 }
75}
76
77// ---- API key flow ----
78
79// StartAPIKey launches the API-key login flow.
80func (m *Manager) StartAPIKey(provider string) (string, error) {
81 if !isKnownAPIKeyProvider(provider) {
82 return "", fmt.Errorf("provider must be anthropic, openai, kimi, or google")
83 }
84 if err := m.ensureKeyServer(); err != nil {
85 return "", err
86 }
87 u := m.keyServer.URL() + "/apikey?provider=" + provider
88 go m.maybeOpen(u)
89 m.emit(Event{Kind: "started", Provider: provider, Method: "apikey", URL: u})
90 return u, nil
91}
92
93func (m *Manager) ensureKeyServer() error {
94 m.mu.Lock()
95 defer m.mu.Unlock()
96 if m.keyServer != nil {
97 return nil
98 }
99 s, err := NewServer()
100 if err != nil {
101 return err
102 }
103 m.keyServer = s
104 go m.consumeKeyServerResults()
105 return nil
106}
107
108func (m *Manager) consumeKeyServerResults() {
109 for res := range m.keyServer.Result() {
110 if res.Err != nil {
111 m.emit(Event{Kind: "error", Provider: res.Provider, Method: res.Method, Message: res.Err.Error()})
112 continue
113 }
114 if err := m.store.SetAPIKey(res.Provider, res.APIKey); err != nil {
115 m.emit(Event{Kind: "error", Provider: res.Provider, Method: "apikey", Message: err.Error()})
116 continue
117 }
118 m.emit(Event{Kind: "success", Provider: res.Provider, Method: "apikey"})
119 }
120}
121
122// ---- OAuth flow ----
123
124// StartOAuth launches the subscription OAuth flow for provider.
125// Only one oauth flow may be in progress at a time (because the
126// callback port is fixed per provider and re-used by the official CLIs).
127//
128// Note: "google" is intentionally not supported here. Google does not
129// offer a subscription OAuth that exchanges a Google One AI / Gemini
130// Advanced login for usable Generative Language API credentials, so
131// the only supported google login path is the API-key flow.
132func (m *Manager) StartOAuth(provider string) (string, error) {
133 if provider == "kimi" {
134 return m.StartKimiDeviceOAuth()
135 }
136 var op OAuthProvider
137 switch provider {
138 case "anthropic":
139 op = AnthropicOAuth
140 case "openai":
141 op = OpenAIOAuth
142 case "google":
143 return "", fmt.Errorf("google login is api-key only; use api key login for gemini")
144 default:
145 return "", fmt.Errorf("provider must be anthropic, openai, kimi, or google")
146 }
147
148 m.mu.Lock()
149 if m.oauthServer != nil {
150 m.oauthServer.Shutdown()
151 m.oauthServer = nil
152 }
153 if m.oauthCancel != nil {
154 m.oauthCancel()
155 }
156 m.mu.Unlock()
157
158 pkce, err := NewPKCE()
159 if err != nil {
160 return "", err
161 }
162 authURL, state, err := op.AuthorizeURL(pkce)
163 if err != nil {
164 return "", err
165 }
166
167 cs, err := NewCallbackServer(op, state)
168 if err != nil {
169 return "", err
170 }
171
172 ctx, cancel := context.WithCancel(context.Background())
173 m.mu.Lock()
174 m.oauthServer = cs
175 m.oauthCtx = ctx
176 m.oauthCancel = cancel
177 m.mu.Unlock()
178
179 go m.awaitOAuth(ctx, op, cs, pkce, state)
180 go m.maybeOpen(authURL)
181 m.emit(Event{Kind: "started", Provider: provider, Method: "oauth", URL: authURL})
182 return authURL, nil
183}
184
185func (m *Manager) awaitOAuth(ctx context.Context, op OAuthProvider, cs *CallbackServer, pkce PKCE, state string) {
186 defer cs.Shutdown()
187
188 waitCtx, waitCancel := context.WithTimeout(ctx, 10*time.Minute)
189 defer waitCancel()
190 res, err := cs.Result(waitCtx)
191 if err != nil {
192 if ctx.Err() == nil {
193 m.emit(Event{Kind: "error", Provider: op.Name, Method: "oauth", Message: "timeout waiting for callback"})
194 }
195 return
196 }
197 if res.Err != nil {
198 m.emit(Event{Kind: "error", Provider: op.Name, Method: "oauth", Message: res.Err.Error()})
199 return
200 }
201
202 exCtx, exCancel := context.WithTimeout(context.Background(), 30*time.Second)
203 defer exCancel()
204 tok, err := op.Exchange(exCtx, res.Code, res.State, pkce)
205 if err != nil {
206 m.emit(Event{Kind: "error", Provider: op.Name, Method: "oauth", Message: err.Error()})
207 return
208 }
209 if err := m.store.SetOAuth(op.Name, *tok); err != nil {
210 m.emit(Event{Kind: "error", Provider: op.Name, Method: "oauth", Message: err.Error()})
211 return
212 }
213 m.emit(Event{Kind: "success", Provider: op.Name, Method: "oauth"})
214}
215
216// StartKimiDeviceOAuth starts Kimi Code's device-code subscription login.
217func (m *Manager) StartKimiDeviceOAuth() (string, error) {
218 ctx, cancel := context.WithCancel(context.Background())
219 m.mu.Lock()
220 if m.oauthCancel != nil {
221 m.oauthCancel()
222 }
223 m.oauthCtx = ctx
224 m.oauthCancel = cancel
225 m.mu.Unlock()
226
227 dev, err := RequestKimiDeviceAuthorization(ctx)
228 if err != nil {
229 return "", err
230 }
231 url := dev.VerificationURIComplete
232 go m.maybeOpen(url)
233 m.emit(Event{Kind: "started", Provider: "kimi", Method: "oauth", URL: url})
234 go func() {
235 tok, err := PollKimiDeviceToken(ctx, dev)
236 if err != nil {
237 if ctx.Err() == nil {
238 m.emit(Event{Kind: "error", Provider: "kimi", Method: "oauth", Message: err.Error()})
239 }
240 return
241 }
242 if err := m.store.SetOAuth("kimi", *tok); err != nil {
243 m.emit(Event{Kind: "error", Provider: "kimi", Method: "oauth", Message: err.Error()})
244 return
245 }
246 m.emit(Event{Kind: "success", Provider: "kimi", Method: "oauth"})
247 }()
248 return url, nil
249}
250
251// StartManualOAuth begins an OAuth flow but does NOT start a local
252// callback server or open a browser. The returned URL is shown to the
253// user so they can complete the authorization on another device; the
254// resulting code is pasted back via CompleteManualOAuth.
255func (m *Manager) StartManualOAuth(provider string) (string, error) {
256 if provider == "kimi" {
257 return m.StartKimiDeviceOAuth()
258 }
259 var op OAuthProvider
260 switch provider {
261 case "anthropic":
262 op = AnthropicManualOAuth
263 case "openai":
264 op = OpenAIOAuth
265 case "google":
266 return "", fmt.Errorf("google login is api-key only; use api key login for gemini")
267 default:
268 return "", fmt.Errorf("provider must be anthropic, openai, kimi, or google")
269 }
270
271 pkce, err := NewPKCE()
272 if err != nil {
273 return "", err
274 }
275 authURL, state, err := op.AuthorizeURL(pkce)
276 if err != nil {
277 return "", err
278 }
279
280 m.mu.Lock()
281 m.manualOp = &op
282 m.manualPKCE = pkce
283 m.manualState = state
284 m.mu.Unlock()
285
286 m.emit(Event{Kind: "started", Provider: provider, Method: "oauth", URL: authURL})
287 return authURL, nil
288}
289
290// CompleteManualOAuth exchanges the user-pasted authorization code for
291// a token and stores it. Accepts either a raw code or a "code#state"
292// token shown by providers like Anthropic when code=true is set.
293func (m *Manager) CompleteManualOAuth(ctx context.Context, input string) error {
294 m.mu.Lock()
295 op := m.manualOp
296 pkce := m.manualPKCE
297 state := m.manualState
298 m.mu.Unlock()
299 if op == nil {
300 return fmt.Errorf("no manual oauth flow in progress")
301 }
302 code, pastedState := parseManualCodeInput(strings.TrimSpace(input))
303 if pastedState != "" {
304 state = pastedState
305 }
306 if code == "" {
307 return fmt.Errorf("empty code")
308 }
309 exCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
310 defer cancel()
311 tok, err := op.Exchange(exCtx, code, state, pkce)
312 if err != nil {
313 m.emit(Event{Kind: "error", Provider: op.Name, Method: "oauth", Message: err.Error()})
314 return err
315 }
316 if err := m.store.SetOAuth(op.Name, *tok); err != nil {
317 m.emit(Event{Kind: "error", Provider: op.Name, Method: "oauth", Message: err.Error()})
318 return err
319 }
320 m.mu.Lock()
321 m.manualOp = nil
322 m.manualPKCE = PKCE{}
323 m.manualState = ""
324 m.mu.Unlock()
325 m.emit(Event{Kind: "success", Provider: op.Name, Method: "oauth"})
326 return nil
327}
328
329// parseManualCodeInput accepts any of:
330// - a bare authorization code
331// - a "code#state" pair
332// - a full redirect URL like http(s)://host:port/callback?code=X&state=Y
333//
334// and returns the extracted code and (if any) state.
335func parseManualCodeInput(s string) (code, state string) {
336 if strings.HasPrefix(s, "http://") || strings.HasPrefix(s, "https://") {
337 if u, err := url.Parse(s); err == nil {
338 q := u.Query()
339 return q.Get("code"), q.Get("state")
340 }
341 }
342 if idx := strings.IndexByte(s, '#'); idx >= 0 {
343 return s[:idx], s[idx+1:]
344 }
345 return s, ""
346}
347
348// HasBrowser reports whether the current environment probably has a
349// working interactive browser reachable from localhost. Used by the
350// login flow to auto-switch to paste-code mode on headless boxes
351// (containers, SSH without display forwarding, etc.) instead of
352// trying to bind a callback port the user can never reach.
353func HasBrowser() bool {
354 if os.Getenv("ZOT_NO_BROWSER") != "" {
355 return false
356 }
357 if os.Getenv("ZOT_FORCE_BROWSER") != "" {
358 return true
359 }
360 if _, err := os.Stat("/.dockerenv"); err == nil {
361 return false
362 }
363 if b, err := os.ReadFile("/proc/1/cgroup"); err == nil {
364 txt := string(b)
365 if strings.Contains(txt, "docker") || strings.Contains(txt, "kubepods") || strings.Contains(txt, "containerd") {
366 return false
367 }
368 }
369 switch runtime.GOOS {
370 case "darwin", "windows":
371 return true
372 default:
373 if os.Getenv("DISPLAY") == "" && os.Getenv("WAYLAND_DISPLAY") == "" {
374 return false
375 }
376 return true
377 }
378}
379
380// CancelOAuth aborts any in-flight OAuth flow.
381func (m *Manager) CancelOAuth() {
382 m.mu.Lock()
383 defer m.mu.Unlock()
384 if m.oauthCancel != nil {
385 m.oauthCancel()
386 m.oauthCancel = nil
387 }
388 if m.oauthServer != nil {
389 m.oauthServer.Shutdown()
390 m.oauthServer = nil
391 }
392}
393
394// ---- shared ----
395
396func (m *Manager) emit(e Event) {
397 select {
398 case m.events <- e:
399 default:
400 }
401}
402
403// maybeOpen tries to open u in the system browser.
404func (m *Manager) maybeOpen(u string) {
405 if !m.openBrowser {
406 return
407 }
408 var cmd *exec.Cmd
409 switch runtime.GOOS {
410 case "darwin":
411 cmd = exec.Command("open", u)
412 case "windows":
413 cmd = exec.Command("rundll32", "url.dll,FileProtocolHandler", u)
414 default:
415 cmd = exec.Command("xdg-open", u)
416 }
417 _ = cmd.Start()
418 m.emit(Event{Kind: "browser_open", URL: u})
419}
420
421// RefreshIfNeeded returns the currently-usable credential for provider,
422// refreshing an expired OAuth access token if necessary.
423func (m *Manager) RefreshIfNeeded(ctx context.Context, provider string) (string, string, error) {
424 creds, err := m.store.Load()
425 if err != nil {
426 return "", "", err
427 }
428 p := creds.get(provider)
429 if p == nil {
430 return "", "", fmt.Errorf("unknown provider %q", provider)
431 }
432 if p.APIKey != "" {
433 return p.APIKey, "apikey", nil
434 }
435 if p.OAuth == nil {
436 return "", "", fmt.Errorf("no credentials for %s", provider)
437 }
438 if !p.OAuth.Expired() {
439 return p.OAuth.AccessToken, "oauth", nil
440 }
441 if p.OAuth.RefreshToken == "" {
442 return "", "", fmt.Errorf("%s access token expired and no refresh token is available; please /login again", provider)
443 }
444 var op OAuthProvider
445 switch provider {
446 case "anthropic":
447 op = AnthropicOAuth
448 case "openai":
449 op = OpenAIOAuth
450 }
451 tok, err := op.Refresh(ctx, p.OAuth.RefreshToken)
452 if err != nil {
453 return "", "", err
454 }
455 if tok.RefreshToken == "" {
456 tok.RefreshToken = p.OAuth.RefreshToken
457 }
458 if err := m.store.SetOAuth(provider, *tok); err != nil {
459 return "", "", err
460 }
461 return tok.AccessToken, "oauth", nil
462}