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 provider != "anthropic" && provider != "openai" && provider != "kimi" {
82 return "", fmt.Errorf("provider must be anthropic, openai, or kimi")
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).
127func (m *Manager) StartOAuth(provider string) (string, error) {
128 if provider == "kimi" {
129 return m.StartKimiDeviceOAuth()
130 }
131 var op OAuthProvider
132 switch provider {
133 case "anthropic":
134 op = AnthropicOAuth
135 case "openai":
136 op = OpenAIOAuth
137 default:
138 return "", fmt.Errorf("provider must be anthropic, openai, or kimi")
139 }
140
141 m.mu.Lock()
142 if m.oauthServer != nil {
143 m.oauthServer.Shutdown()
144 m.oauthServer = nil
145 }
146 if m.oauthCancel != nil {
147 m.oauthCancel()
148 }
149 m.mu.Unlock()
150
151 pkce, err := NewPKCE()
152 if err != nil {
153 return "", err
154 }
155 authURL, state, err := op.AuthorizeURL(pkce)
156 if err != nil {
157 return "", err
158 }
159
160 cs, err := NewCallbackServer(op, state)
161 if err != nil {
162 return "", err
163 }
164
165 ctx, cancel := context.WithCancel(context.Background())
166 m.mu.Lock()
167 m.oauthServer = cs
168 m.oauthCtx = ctx
169 m.oauthCancel = cancel
170 m.mu.Unlock()
171
172 go m.awaitOAuth(ctx, op, cs, pkce, state)
173 go m.maybeOpen(authURL)
174 m.emit(Event{Kind: "started", Provider: provider, Method: "oauth", URL: authURL})
175 return authURL, nil
176}
177
178func (m *Manager) awaitOAuth(ctx context.Context, op OAuthProvider, cs *CallbackServer, pkce PKCE, state string) {
179 defer cs.Shutdown()
180
181 waitCtx, waitCancel := context.WithTimeout(ctx, 10*time.Minute)
182 defer waitCancel()
183 res, err := cs.Result(waitCtx)
184 if err != nil {
185 if ctx.Err() == nil {
186 m.emit(Event{Kind: "error", Provider: op.Name, Method: "oauth", Message: "timeout waiting for callback"})
187 }
188 return
189 }
190 if res.Err != nil {
191 m.emit(Event{Kind: "error", Provider: op.Name, Method: "oauth", Message: res.Err.Error()})
192 return
193 }
194
195 exCtx, exCancel := context.WithTimeout(context.Background(), 30*time.Second)
196 defer exCancel()
197 tok, err := op.Exchange(exCtx, res.Code, res.State, pkce)
198 if err != nil {
199 m.emit(Event{Kind: "error", Provider: op.Name, Method: "oauth", Message: err.Error()})
200 return
201 }
202 if err := m.store.SetOAuth(op.Name, *tok); err != nil {
203 m.emit(Event{Kind: "error", Provider: op.Name, Method: "oauth", Message: err.Error()})
204 return
205 }
206 m.emit(Event{Kind: "success", Provider: op.Name, Method: "oauth"})
207}
208
209// StartKimiDeviceOAuth starts Kimi Code's device-code subscription login.
210func (m *Manager) StartKimiDeviceOAuth() (string, error) {
211 ctx, cancel := context.WithCancel(context.Background())
212 m.mu.Lock()
213 if m.oauthCancel != nil {
214 m.oauthCancel()
215 }
216 m.oauthCtx = ctx
217 m.oauthCancel = cancel
218 m.mu.Unlock()
219
220 dev, err := RequestKimiDeviceAuthorization(ctx)
221 if err != nil {
222 return "", err
223 }
224 url := dev.VerificationURIComplete
225 go m.maybeOpen(url)
226 m.emit(Event{Kind: "started", Provider: "kimi", Method: "oauth", URL: url})
227 go func() {
228 tok, err := PollKimiDeviceToken(ctx, dev)
229 if err != nil {
230 if ctx.Err() == nil {
231 m.emit(Event{Kind: "error", Provider: "kimi", Method: "oauth", Message: err.Error()})
232 }
233 return
234 }
235 if err := m.store.SetOAuth("kimi", *tok); err != nil {
236 m.emit(Event{Kind: "error", Provider: "kimi", Method: "oauth", Message: err.Error()})
237 return
238 }
239 m.emit(Event{Kind: "success", Provider: "kimi", Method: "oauth"})
240 }()
241 return url, nil
242}
243
244// StartManualOAuth begins an OAuth flow but does NOT start a local
245// callback server or open a browser. The returned URL is shown to the
246// user so they can complete the authorization on another device; the
247// resulting code is pasted back via CompleteManualOAuth.
248func (m *Manager) StartManualOAuth(provider string) (string, error) {
249 if provider == "kimi" {
250 return m.StartKimiDeviceOAuth()
251 }
252 var op OAuthProvider
253 switch provider {
254 case "anthropic":
255 op = AnthropicManualOAuth
256 case "openai":
257 op = OpenAIOAuth
258 default:
259 return "", fmt.Errorf("provider must be anthropic, openai, or kimi")
260 }
261
262 pkce, err := NewPKCE()
263 if err != nil {
264 return "", err
265 }
266 authURL, state, err := op.AuthorizeURL(pkce)
267 if err != nil {
268 return "", err
269 }
270
271 m.mu.Lock()
272 m.manualOp = &op
273 m.manualPKCE = pkce
274 m.manualState = state
275 m.mu.Unlock()
276
277 m.emit(Event{Kind: "started", Provider: provider, Method: "oauth", URL: authURL})
278 return authURL, nil
279}
280
281// CompleteManualOAuth exchanges the user-pasted authorization code for
282// a token and stores it. Accepts either a raw code or a "code#state"
283// token shown by providers like Anthropic when code=true is set.
284func (m *Manager) CompleteManualOAuth(ctx context.Context, input string) error {
285 m.mu.Lock()
286 op := m.manualOp
287 pkce := m.manualPKCE
288 state := m.manualState
289 m.mu.Unlock()
290 if op == nil {
291 return fmt.Errorf("no manual oauth flow in progress")
292 }
293 code, pastedState := parseManualCodeInput(strings.TrimSpace(input))
294 if pastedState != "" {
295 state = pastedState
296 }
297 if code == "" {
298 return fmt.Errorf("empty code")
299 }
300 exCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
301 defer cancel()
302 tok, err := op.Exchange(exCtx, code, state, pkce)
303 if err != nil {
304 m.emit(Event{Kind: "error", Provider: op.Name, Method: "oauth", Message: err.Error()})
305 return err
306 }
307 if err := m.store.SetOAuth(op.Name, *tok); err != nil {
308 m.emit(Event{Kind: "error", Provider: op.Name, Method: "oauth", Message: err.Error()})
309 return err
310 }
311 m.mu.Lock()
312 m.manualOp = nil
313 m.manualPKCE = PKCE{}
314 m.manualState = ""
315 m.mu.Unlock()
316 m.emit(Event{Kind: "success", Provider: op.Name, Method: "oauth"})
317 return nil
318}
319
320// parseManualCodeInput accepts any of:
321// - a bare authorization code
322// - a "code#state" pair
323// - a full redirect URL like http(s)://host:port/callback?code=X&state=Y
324//
325// and returns the extracted code and (if any) state.
326func parseManualCodeInput(s string) (code, state string) {
327 if strings.HasPrefix(s, "http://") || strings.HasPrefix(s, "https://") {
328 if u, err := url.Parse(s); err == nil {
329 q := u.Query()
330 return q.Get("code"), q.Get("state")
331 }
332 }
333 if idx := strings.IndexByte(s, '#'); idx >= 0 {
334 return s[:idx], s[idx+1:]
335 }
336 return s, ""
337}
338
339// HasBrowser reports whether the current environment probably has a
340// working interactive browser reachable from localhost. Used by the
341// login flow to auto-switch to paste-code mode on headless boxes
342// (containers, SSH without display forwarding, etc.) instead of
343// trying to bind a callback port the user can never reach.
344func HasBrowser() bool {
345 if os.Getenv("ZOT_NO_BROWSER") != "" {
346 return false
347 }
348 if os.Getenv("ZOT_FORCE_BROWSER") != "" {
349 return true
350 }
351 if _, err := os.Stat("/.dockerenv"); err == nil {
352 return false
353 }
354 if b, err := os.ReadFile("/proc/1/cgroup"); err == nil {
355 txt := string(b)
356 if strings.Contains(txt, "docker") || strings.Contains(txt, "kubepods") || strings.Contains(txt, "containerd") {
357 return false
358 }
359 }
360 switch runtime.GOOS {
361 case "darwin", "windows":
362 return true
363 default:
364 if os.Getenv("DISPLAY") == "" && os.Getenv("WAYLAND_DISPLAY") == "" {
365 return false
366 }
367 return true
368 }
369}
370
371// CancelOAuth aborts any in-flight OAuth flow.
372func (m *Manager) CancelOAuth() {
373 m.mu.Lock()
374 defer m.mu.Unlock()
375 if m.oauthCancel != nil {
376 m.oauthCancel()
377 m.oauthCancel = nil
378 }
379 if m.oauthServer != nil {
380 m.oauthServer.Shutdown()
381 m.oauthServer = nil
382 }
383}
384
385// ---- shared ----
386
387func (m *Manager) emit(e Event) {
388 select {
389 case m.events <- e:
390 default:
391 }
392}
393
394// maybeOpen tries to open u in the system browser.
395func (m *Manager) maybeOpen(u string) {
396 if !m.openBrowser {
397 return
398 }
399 var cmd *exec.Cmd
400 switch runtime.GOOS {
401 case "darwin":
402 cmd = exec.Command("open", u)
403 case "windows":
404 cmd = exec.Command("rundll32", "url.dll,FileProtocolHandler", u)
405 default:
406 cmd = exec.Command("xdg-open", u)
407 }
408 _ = cmd.Start()
409 m.emit(Event{Kind: "browser_open", URL: u})
410}
411
412// RefreshIfNeeded returns the currently-usable credential for provider,
413// refreshing an expired OAuth access token if necessary.
414func (m *Manager) RefreshIfNeeded(ctx context.Context, provider string) (string, string, error) {
415 creds, err := m.store.Load()
416 if err != nil {
417 return "", "", err
418 }
419 p := creds.get(provider)
420 if p == nil {
421 return "", "", fmt.Errorf("unknown provider %q", provider)
422 }
423 if p.APIKey != "" {
424 return p.APIKey, "apikey", nil
425 }
426 if p.OAuth == nil {
427 return "", "", fmt.Errorf("no credentials for %s", provider)
428 }
429 if !p.OAuth.Expired() {
430 return p.OAuth.AccessToken, "oauth", nil
431 }
432 if p.OAuth.RefreshToken == "" {
433 return "", "", fmt.Errorf("%s access token expired and no refresh token is available; please /login again", provider)
434 }
435 var op OAuthProvider
436 switch provider {
437 case "anthropic":
438 op = AnthropicOAuth
439 case "openai":
440 op = OpenAIOAuth
441 }
442 tok, err := op.Refresh(ctx, p.OAuth.RefreshToken)
443 if err != nil {
444 return "", "", err
445 }
446 if tok.RefreshToken == "" {
447 tok.RefreshToken = p.OAuth.RefreshToken
448 }
449 if err := m.store.SetOAuth(provider, *tok); err != nil {
450 return "", "", err
451 }
452 return tok.AccessToken, "oauth", nil
453}