1package auth
2
3import (
4 "context"
5 "errors"
6 "net"
7 "net/http"
8 "net/http/httptest"
9 "net/url"
10 "path/filepath"
11 "strings"
12 "testing"
13 "time"
14)
15
16func TestAPIKeyProvidersIncludesOllama(t *testing.T) {
17 found := false
18 for _, p := range APIKeyProviders() {
19 if p == "ollama" {
20 found = true
21 break
22 }
23 }
24 if !found {
25 t.Fatalf("APIKeyProviders() = %v, want ollama listed", APIKeyProviders())
26 }
27}
28
29func TestProbeAPIKeyAcceptsExtraProviderWithoutProbe(t *testing.T) {
30 SetExtraAPIKeyProviders([]string{"my-company"})
31 t.Cleanup(func() { SetExtraAPIKeyProviders(nil) })
32
33 if err := ProbeAPIKey(context.Background(), "my-company", "test-key"); err != nil {
34 t.Fatalf("ProbeAPIKey custom provider: %v", err)
35 }
36}
37
38func TestStoreRoundTrip(t *testing.T) {
39 dir := t.TempDir()
40 s := NewStore(filepath.Join(dir, "auth.json"))
41
42 if err := s.SetAPIKey("anthropic", "sk-ant-xyz"); err != nil {
43 t.Fatal(err)
44 }
45 c, err := s.Load()
46 if err != nil {
47 t.Fatal(err)
48 }
49 if c.Anthropic.APIKey != "sk-ant-xyz" {
50 t.Fatalf("got %+v", c)
51 }
52 if !c.Has("anthropic") || c.Method("anthropic") != "apikey" {
53 t.Fatalf("method mismatch: %q", c.Method("anthropic"))
54 }
55
56 tok := OAuthToken{AccessToken: "tok", RefreshToken: "ref", Expiry: time.Now().Add(1 * time.Hour)}
57 if err := s.SetOAuth("openai", tok); err != nil {
58 t.Fatal(err)
59 }
60 c, _ = s.Load()
61 if c.OpenAI.OAuth == nil || c.OpenAI.OAuth.AccessToken != "tok" {
62 t.Fatalf("got %+v", c.OpenAI)
63 }
64 if c.Method("openai") != "oauth" {
65 t.Fatalf("method=%q", c.Method("openai"))
66 }
67
68 if err := s.Clear("anthropic"); err != nil {
69 t.Fatal(err)
70 }
71 c, _ = s.Load()
72 if c.Has("anthropic") {
73 t.Fatalf("expected cleared, got %+v", c.Anthropic)
74 }
75}
76
77func TestPKCE(t *testing.T) {
78 p, err := NewPKCE()
79 if err != nil {
80 t.Fatal(err)
81 }
82 if len(p.Verifier) < 40 || len(p.Challenge) < 40 {
83 t.Fatalf("short pkce: %+v", p)
84 }
85 if p.Verifier == p.Challenge {
86 t.Fatal("verifier == challenge")
87 }
88}
89
90func TestAuthorizeURL(t *testing.T) {
91 p, _ := NewPKCE()
92 u, state, err := AnthropicOAuth.AuthorizeURL(p)
93 if err != nil {
94 t.Fatal(err)
95 }
96 parsed, err := url.Parse(u)
97 if err != nil {
98 t.Fatal(err)
99 }
100 q := parsed.Query()
101 if q.Get("client_id") != AnthropicOAuth.ClientID {
102 t.Fatal("missing client_id")
103 }
104 if q.Get("code_challenge_method") != "S256" {
105 t.Fatal("wrong challenge method")
106 }
107 // anthropic uses state == verifier
108 if q.Get("state") != p.Verifier {
109 t.Fatalf("state mismatch: %q vs verifier %q", q.Get("state"), p.Verifier)
110 }
111 if state != p.Verifier {
112 t.Fatalf("returned state != verifier")
113 }
114 if q.Get("redirect_uri") != AnthropicOAuth.RedirectURI() {
115 t.Fatalf("redirect_uri=%q", q.Get("redirect_uri"))
116 }
117 if q.Get("code") != "true" {
118 t.Fatalf("missing code=true extra arg")
119 }
120}
121
122func TestOpenAIAuthorizeURL(t *testing.T) {
123 p, _ := NewPKCE()
124 u, state, err := OpenAIOAuth.AuthorizeURL(p)
125 if err != nil {
126 t.Fatal(err)
127 }
128 parsed, _ := url.Parse(u)
129 q := parsed.Query()
130 if q.Get("redirect_uri") != "http://localhost:1455/auth/callback" {
131 t.Fatalf("redirect_uri=%q", q.Get("redirect_uri"))
132 }
133 if state == p.Verifier {
134 t.Fatal("openai state should be random, not verifier")
135 }
136 if q.Get("originator") != "zot" {
137 t.Fatalf("originator=%q", q.Get("originator"))
138 }
139}
140
141func TestServerIndexAndAPIKey(t *testing.T) {
142 s, err := NewServer()
143 if err != nil {
144 t.Fatal(err)
145 }
146 defer s.Shutdown(context.Background())
147
148 // Swap the probe out for a fake that always succeeds.
149 s.probeFn = func(ctx context.Context, provider, key string) error { return nil }
150
151 resp, err := http.Get(s.URL() + "/")
152 if err != nil {
153 t.Fatal(err)
154 }
155 defer resp.Body.Close()
156 if resp.StatusCode != 200 {
157 t.Fatalf("index status %d", resp.StatusCode)
158 }
159
160 resp, err = http.Get(s.URL() + "/apikey?provider=anthropic")
161 if err != nil {
162 t.Fatal(err)
163 }
164 resp.Body.Close()
165 if resp.StatusCode != 200 {
166 t.Fatalf("apikey form status %d", resp.StatusCode)
167 }
168
169 // POST a fake key and expect a success result on the channel.
170 _, err = http.PostForm(s.URL()+"/apikey", url.Values{
171 "provider": {"anthropic"},
172 "api_key": {"sk-ant-test"},
173 })
174 if err != nil {
175 t.Fatal(err)
176 }
177 select {
178 case r := <-s.Result():
179 if r.Err != nil {
180 t.Fatalf("unexpected err: %v", r.Err)
181 }
182 if r.Provider != "anthropic" || r.APIKey != "sk-ant-test" || r.Method != "apikey" {
183 t.Fatalf("bad result: %+v", r)
184 }
185 case <-time.After(2 * time.Second):
186 t.Fatal("no result")
187 }
188}
189
190func TestServerCallback(t *testing.T) {
191 s, err := NewServer()
192 if err != nil {
193 t.Fatal(err)
194 }
195 defer s.Shutdown(context.Background())
196
197 u := s.URL() + "/callback?code=abc123&state=anthropic:xyz"
198 resp, err := http.Get(u)
199 if err != nil {
200 t.Fatal(err)
201 }
202 resp.Body.Close()
203
204 select {
205 case r := <-s.Result():
206 if r.Provider != "anthropic" || r.Code != "abc123" || r.State != "anthropic:xyz" || r.Method != "oauth" {
207 t.Fatalf("bad result: %+v", r)
208 }
209 case <-time.After(2 * time.Second):
210 t.Fatal("no result")
211 }
212}
213
214func TestCallbackServerRoundTrip(t *testing.T) {
215 // Use a provider config with an ephemeral port so tests stay hermetic.
216 op := OAuthProvider{
217 Name: "anthropic",
218 RedirectHost: "127.0.0.1",
219 RedirectPort: 0, // let kernel pick
220 RedirectPath: "/callback",
221 }
222 // We need a real fixed port for CallbackServer.NewCallbackServer; find one.
223 l, err := net.Listen("tcp", "127.0.0.1:0")
224 if err != nil {
225 t.Fatal(err)
226 }
227 port := l.Addr().(*net.TCPAddr).Port
228 l.Close()
229 op.RedirectPort = port
230
231 cs, err := NewCallbackServer(op, "thestate")
232 if err != nil {
233 t.Fatal(err)
234 }
235 defer cs.Shutdown()
236
237 go func() {
238 time.Sleep(50 * time.Millisecond)
239 _, _ = http.Get(cs.URL() + "?code=abc&state=thestate")
240 }()
241 ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
242 defer cancel()
243 r, err := cs.Result(ctx)
244 if err != nil {
245 t.Fatal(err)
246 }
247 if r.Code != "abc" || r.State != "thestate" {
248 t.Fatalf("bad result: %+v", r)
249 }
250}
251
252func TestCallbackServerStateMismatch(t *testing.T) {
253 l, _ := net.Listen("tcp", "127.0.0.1:0")
254 port := l.Addr().(*net.TCPAddr).Port
255 l.Close()
256 op := OAuthProvider{Name: "anthropic", RedirectHost: "127.0.0.1", RedirectPort: port, RedirectPath: "/callback"}
257 cs, err := NewCallbackServer(op, "expected")
258 if err != nil {
259 t.Fatal(err)
260 }
261 defer cs.Shutdown()
262
263 go func() {
264 time.Sleep(50 * time.Millisecond)
265 _, _ = http.Get(cs.URL() + "?code=abc&state=wrong")
266 }()
267 ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
268 defer cancel()
269 r, err := cs.Result(ctx)
270 if err != nil {
271 t.Fatal(err)
272 }
273 if r.Err == nil || !strings.Contains(r.Err.Error(), "state") {
274 t.Fatalf("want state error, got %+v", r)
275 }
276}
277
278func TestServerCallbackError(t *testing.T) {
279 s, err := NewServer()
280 if err != nil {
281 t.Fatal(err)
282 }
283 defer s.Shutdown(context.Background())
284
285 _, _ = http.Get(s.URL() + "/callback?error=access_denied&state=openai:xyz")
286 select {
287 case r := <-s.Result():
288 if r.Err == nil || r.Provider != "openai" {
289 t.Fatalf("bad result: %+v", r)
290 }
291 if !strings.Contains(r.Err.Error(), "access_denied") {
292 t.Fatalf("missing reason: %v", r.Err)
293 }
294 case <-time.After(2 * time.Second):
295 t.Fatal("no result")
296 }
297}
298
299func TestStartManualOAuthRejectsOpenAIWithoutCallback(t *testing.T) {
300 m := NewManager(NewStore(filepath.Join(t.TempDir(), "auth.json")))
301 if _, err := m.StartManualOAuth("openai-codex"); err == nil {
302 t.Fatal("expected manual OpenAI OAuth to be rejected")
303 } else if !strings.Contains(err.Error(), "manual code login is not supported") {
304 t.Fatalf("unexpected error: %v", err)
305 }
306}
307
308func TestBrowserOAuthSharesTransactionWithPastedCode(t *testing.T) {
309 for _, provider := range []string{"anthropic", "openai-codex"} {
310 t.Run(provider, func(t *testing.T) {
311 m := NewManager(NewStore(filepath.Join(t.TempDir(), "auth.json")))
312 m.openBrowser = false
313 t.Cleanup(m.Close)
314 loginURL, err := m.StartOAuth(provider)
315 if err != nil {
316 if strings.Contains(err.Error(), "bind ") {
317 t.Skipf("callback port unavailable: %v", err)
318 }
319 t.Fatal(err)
320 }
321 u, err := url.Parse(loginURL)
322 if err != nil {
323 t.Fatal(err)
324 }
325 state := u.Query().Get("state")
326 if m.manualOp == nil || m.manualState != state || m.manualPKCE.Challenge != u.Query().Get("code_challenge") {
327 t.Fatal("pasted-code path does not share the browser transaction")
328 }
329 if err := m.CompleteManualOAuth(context.Background(), "code#wrong-state"); err == nil || !strings.Contains(err.Error(), "state mismatch") {
330 t.Fatalf("wrong state accepted: %v", err)
331 }
332 if err := m.CompleteManualOAuth(context.Background(), ""); err == nil || err.Error() != "empty code" {
333 t.Fatalf("shared transaction unavailable after bad code: %v", err)
334 }
335 m.CancelOAuth()
336 if err := m.CompleteManualOAuth(context.Background(), "code#"+state); err == nil || !strings.Contains(err.Error(), "no manual oauth flow") {
337 t.Fatalf("canceled flow still accepts code: %v", err)
338 }
339 })
340 }
341}
342
343func TestPastedOAuthCodeUsesBrowserPKCE(t *testing.T) {
344 requests := make(chan url.Values, 2)
345 tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
346 _ = r.ParseForm()
347 requests <- r.PostForm
348 w.Header().Set("content-type", "application/json")
349 _, _ = w.Write([]byte(`{"access_token":"test-token","refresh_token":"test-refresh","expires_in":3600}`))
350 }))
351 defer tokenServer.Close()
352 old := OpenAIOAuth.TokenURL
353 OpenAIOAuth.TokenURL = tokenServer.URL
354 t.Cleanup(func() { OpenAIOAuth.TokenURL = old })
355
356 for _, method := range []string{"paste", "callback"} {
357 t.Run(method, func(t *testing.T) {
358 store := NewStore(filepath.Join(t.TempDir(), "auth.json"))
359 m := NewManager(store)
360 m.openBrowser = false
361 t.Cleanup(m.Close)
362 loginURL, err := m.StartOAuth("openai-codex")
363 if err != nil {
364 if strings.Contains(err.Error(), "bind ") {
365 t.Skipf("callback port unavailable: %v", err)
366 }
367 t.Fatal(err)
368 }
369 u, err := url.Parse(loginURL)
370 if err != nil {
371 t.Fatal(err)
372 }
373 verifier := m.manualPKCE.Verifier
374 redirect := u.Query().Get("redirect_uri")
375 codeURL := redirect + "?code=test-code&state=" + url.QueryEscape(u.Query().Get("state"))
376 if method == "paste" {
377 if err := m.CompleteManualOAuth(context.Background(), codeURL); err != nil {
378 t.Fatal(err)
379 }
380 } else {
381 resp, err := http.Get(codeURL)
382 if err != nil {
383 t.Fatal(err)
384 }
385 resp.Body.Close()
386 if resp.StatusCode != http.StatusOK {
387 t.Fatalf("callback status = %d", resp.StatusCode)
388 }
389 }
390 select {
391 case ev := <-m.Events():
392 if ev.Kind != "started" {
393 t.Fatalf("event = %+v", ev)
394 }
395 case <-time.After(2 * time.Second):
396 t.Fatal("no started event")
397 }
398 select {
399 case ev := <-m.Events():
400 if ev.Kind != "success" || ev.Provider != "openai-codex" {
401 t.Fatalf("event = %+v", ev)
402 }
403 case <-time.After(2 * time.Second):
404 t.Fatal("no success event")
405 }
406 select {
407 case form := <-requests:
408 if form.Get("code_verifier") != verifier || form.Get("redirect_uri") != redirect || form.Get("code") != "test-code" {
409 t.Fatalf("exchange did not use browser transaction: %v", form)
410 }
411 case <-time.After(2 * time.Second):
412 t.Fatal("no token exchange")
413 }
414 creds, err := store.Load()
415 if err != nil {
416 t.Fatal(err)
417 }
418 if creds.OpenAI.OAuth == nil || creds.OpenAI.OAuth.AccessToken != "test-token" {
419 t.Fatalf("stored credential: %v", creds.OpenAI.OAuth)
420 }
421 })
422 }
423}
424
425func TestCanceledBrowserOAuthDoesNotExchangeOrEmit(t *testing.T) {
426 exchanged := make(chan struct{}, 1)
427 tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
428 exchanged <- struct{}{}
429 w.Header().Set("content-type", "application/json")
430 _, _ = w.Write([]byte(`{"access_token":"test-token","refresh_token":"test-refresh","expires_in":3600}`))
431 }))
432 defer tokenServer.Close()
433 old := OpenAIOAuth.TokenURL
434 OpenAIOAuth.TokenURL = tokenServer.URL
435 t.Cleanup(func() { OpenAIOAuth.TokenURL = old })
436
437 store := NewStore(filepath.Join(t.TempDir(), "auth.json"))
438 m := NewManager(store)
439 m.openBrowser = false
440 t.Cleanup(m.Close)
441 loginURL, err := m.StartOAuth("openai-codex")
442 if err != nil {
443 if strings.Contains(err.Error(), "bind ") {
444 t.Skipf("callback port unavailable: %v", err)
445 }
446 t.Fatal(err)
447 }
448 u, _ := url.Parse(loginURL)
449 state := u.Query().Get("state")
450 <-m.Events() // started
451 m.CancelOAuth()
452 // A pasted code after cancellation must not resurrect the transaction.
453 if err := m.CompleteManualOAuth(context.Background(), "code#"+state); err == nil {
454 t.Fatal("canceled transaction accepted a pasted code")
455 }
456 select {
457 case <-exchanged:
458 t.Fatal("token exchange ran after cancellation")
459 case ev := <-m.Events():
460 t.Fatalf("unexpected event after cancellation: %+v", ev)
461 case <-time.After(100 * time.Millisecond):
462 }
463 // The callback port must be released so a new login can bind it.
464 if _, err := m.StartOAuth("openai-codex"); err != nil {
465 t.Fatalf("callback port not released after cancel: %v", err)
466 }
467}
468
469func TestPastedOAuthCodeConsumesTransactionOnce(t *testing.T) {
470 tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
471 w.Header().Set("content-type", "application/json")
472 _, _ = w.Write([]byte(`{"access_token":"test-token","refresh_token":"test-refresh","expires_in":3600}`))
473 }))
474 defer tokenServer.Close()
475 old := OpenAIOAuth.TokenURL
476 OpenAIOAuth.TokenURL = tokenServer.URL
477 t.Cleanup(func() { OpenAIOAuth.TokenURL = old })
478
479 m := NewManager(NewStore(filepath.Join(t.TempDir(), "auth.json")))
480 m.openBrowser = false
481 t.Cleanup(m.Close)
482 loginURL, err := m.StartOAuth("openai-codex")
483 if err != nil {
484 if strings.Contains(err.Error(), "bind ") {
485 t.Skipf("callback port unavailable: %v", err)
486 }
487 t.Fatal(err)
488 }
489 u, _ := url.Parse(loginURL)
490 state := u.Query().Get("state")
491 if err := m.CompleteManualOAuth(context.Background(), "code#"+state); err != nil {
492 t.Fatal(err)
493 }
494 if err := m.CompleteManualOAuth(context.Background(), "code#"+state); err == nil {
495 t.Fatal("second pasted code accepted")
496 }
497 var success int
498 deadline := time.After(300 * time.Millisecond)
499loop:
500 for {
501 select {
502 case ev := <-m.Events():
503 if ev.Kind == "success" {
504 success++
505 }
506 case <-deadline:
507 break loop
508 }
509 }
510 if success != 1 {
511 t.Fatalf("success events = %d, want 1", success)
512 }
513}
514
515// A pasted-code exchange that is blocked at the token endpoint must be
516// canceled by CancelOAuth, Close, or a newer login: no credential, no success
517// event.
518func TestCancelAbortsInFlightPastedCodeExchange(t *testing.T) {
519 for _, name := range []string{"cancel", "new login", "close"} {
520 t.Run(name, func(t *testing.T) {
521 arrived := make(chan struct{}, 4)
522 release := make(chan struct{})
523 tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
524 arrived <- struct{}{}
525 select {
526 case <-release:
527 case <-r.Context().Done():
528 return
529 }
530 w.Header().Set("content-type", "application/json")
531 _, _ = w.Write([]byte(`{"access_token":"late-token","refresh_token":"late-refresh","expires_in":3600}`))
532 }))
533 defer tokenServer.Close()
534 defer close(release)
535 old := AnthropicManualOAuth.TokenURL
536 AnthropicManualOAuth.TokenURL = tokenServer.URL
537 t.Cleanup(func() { AnthropicManualOAuth.TokenURL = old })
538
539 store := NewStore(filepath.Join(t.TempDir(), "auth.json"))
540 m := NewManager(store)
541 m.openBrowser = false
542 t.Cleanup(m.Close)
543 loginURL, err := m.StartManualOAuth("anthropic")
544 if err != nil {
545 t.Fatal(err)
546 }
547 u, _ := url.Parse(loginURL)
548 state := u.Query().Get("state")
549 <-m.Events() // started
550
551 // The caller's context is never canceled: only the manager may
552 // abort the exchange.
553 done := make(chan error, 1)
554 go func() { done <- m.CompleteManualOAuth(context.Background(), "code#"+state) }()
555 select {
556 case <-arrived:
557 case <-time.After(2 * time.Second):
558 t.Fatal("token exchange never started")
559 }
560 switch name {
561 case "new login":
562 if _, err := m.StartManualOAuth("anthropic"); err != nil {
563 t.Fatal(err)
564 }
565 case "close":
566 m.Close()
567 default:
568 m.CancelOAuth()
569 }
570 select {
571 case err := <-done:
572 if !errors.Is(err, ErrOAuthCanceled) {
573 t.Fatalf("CompleteManualOAuth = %v, want ErrOAuthCanceled", err)
574 }
575 case <-time.After(2 * time.Second):
576 t.Fatal("canceled exchange did not return")
577 }
578 creds, err := store.Load()
579 if err == nil && creds.Anthropic.OAuth != nil {
580 t.Fatalf("canceled exchange stored a credential")
581 }
582 for {
583 select {
584 case ev := <-m.Events():
585 if ev.Kind == "success" || ev.Kind == "error" {
586 t.Fatalf("stale event after cancel: %+v", ev)
587 }
588 continue
589 default:
590 }
591 break
592 }
593 })
594 }
595}