1package auth
2
3import (
4 "context"
5 "io"
6 "net"
7 "net/http"
8 "net/url"
9 "path/filepath"
10 "strings"
11 "testing"
12 "time"
13)
14
15func TestProbeAPIKeyAcceptsExtraProviderWithoutProbe(t *testing.T) {
16 SetExtraAPIKeyProviders([]string{"my-company"})
17 t.Cleanup(func() { SetExtraAPIKeyProviders(nil) })
18
19 if err := ProbeAPIKey(context.Background(), "my-company", "test-key"); err != nil {
20 t.Fatalf("ProbeAPIKey custom provider: %v", err)
21 }
22}
23
24func TestProbeGondolaAPIKeyUsesAuthenticatedBalanceEndpoint(t *testing.T) {
25 client := &http.Client{Transport: authRoundTripFunc(func(req *http.Request) (*http.Response, error) {
26 if req.URL.String() != "https://api.gondola-ai.com/v1/balance" {
27 t.Fatalf("probe URL = %q", req.URL.String())
28 }
29 if got := req.Header.Get("authorization"); got != "Bearer gnd_test" {
30 t.Fatalf("authorization = %q", got)
31 }
32 return &http.Response{
33 StatusCode: http.StatusUnauthorized,
34 Body: io.NopCloser(strings.NewReader(`{"error":"invalid key"}`)),
35 Header: make(http.Header),
36 Request: req,
37 }, nil
38 })}
39
40 err := probeAPIKey(context.Background(), "gondola", "gnd_test", client)
41 if err == nil || !strings.Contains(err.Error(), "rejected the key") {
42 t.Fatalf("probeAPIKey error = %v", err)
43 }
44}
45
46type authRoundTripFunc func(*http.Request) (*http.Response, error)
47
48func (f authRoundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
49 return f(req)
50}
51
52func TestStoreRoundTrip(t *testing.T) {
53 dir := t.TempDir()
54 s := NewStore(filepath.Join(dir, "auth.json"))
55
56 if err := s.SetAPIKey("anthropic", "sk-ant-xyz"); err != nil {
57 t.Fatal(err)
58 }
59 c, err := s.Load()
60 if err != nil {
61 t.Fatal(err)
62 }
63 if c.Anthropic.APIKey != "sk-ant-xyz" {
64 t.Fatalf("got %+v", c)
65 }
66 if !c.Has("anthropic") || c.Method("anthropic") != "apikey" {
67 t.Fatalf("method mismatch: %q", c.Method("anthropic"))
68 }
69
70 tok := OAuthToken{AccessToken: "tok", RefreshToken: "ref", Expiry: time.Now().Add(1 * time.Hour)}
71 if err := s.SetOAuth("openai", tok); err != nil {
72 t.Fatal(err)
73 }
74 c, _ = s.Load()
75 if c.OpenAI.OAuth == nil || c.OpenAI.OAuth.AccessToken != "tok" {
76 t.Fatalf("got %+v", c.OpenAI)
77 }
78 if c.Method("openai") != "oauth" {
79 t.Fatalf("method=%q", c.Method("openai"))
80 }
81
82 if err := s.Clear("anthropic"); err != nil {
83 t.Fatal(err)
84 }
85 c, _ = s.Load()
86 if c.Has("anthropic") {
87 t.Fatalf("expected cleared, got %+v", c.Anthropic)
88 }
89}
90
91func TestPKCE(t *testing.T) {
92 p, err := NewPKCE()
93 if err != nil {
94 t.Fatal(err)
95 }
96 if len(p.Verifier) < 40 || len(p.Challenge) < 40 {
97 t.Fatalf("short pkce: %+v", p)
98 }
99 if p.Verifier == p.Challenge {
100 t.Fatal("verifier == challenge")
101 }
102}
103
104func TestAuthorizeURL(t *testing.T) {
105 p, _ := NewPKCE()
106 u, state, err := AnthropicOAuth.AuthorizeURL(p)
107 if err != nil {
108 t.Fatal(err)
109 }
110 parsed, err := url.Parse(u)
111 if err != nil {
112 t.Fatal(err)
113 }
114 q := parsed.Query()
115 if q.Get("client_id") != AnthropicOAuth.ClientID {
116 t.Fatal("missing client_id")
117 }
118 if q.Get("code_challenge_method") != "S256" {
119 t.Fatal("wrong challenge method")
120 }
121 // anthropic uses state == verifier
122 if q.Get("state") != p.Verifier {
123 t.Fatalf("state mismatch: %q vs verifier %q", q.Get("state"), p.Verifier)
124 }
125 if state != p.Verifier {
126 t.Fatalf("returned state != verifier")
127 }
128 if q.Get("redirect_uri") != AnthropicOAuth.RedirectURI() {
129 t.Fatalf("redirect_uri=%q", q.Get("redirect_uri"))
130 }
131 if q.Get("code") != "true" {
132 t.Fatalf("missing code=true extra arg")
133 }
134}
135
136func TestOpenAIAuthorizeURL(t *testing.T) {
137 p, _ := NewPKCE()
138 u, state, err := OpenAIOAuth.AuthorizeURL(p)
139 if err != nil {
140 t.Fatal(err)
141 }
142 parsed, _ := url.Parse(u)
143 q := parsed.Query()
144 if q.Get("redirect_uri") != "http://localhost:1455/auth/callback" {
145 t.Fatalf("redirect_uri=%q", q.Get("redirect_uri"))
146 }
147 if state == p.Verifier {
148 t.Fatal("openai state should be random, not verifier")
149 }
150 if q.Get("originator") != "zot" {
151 t.Fatalf("originator=%q", q.Get("originator"))
152 }
153}
154
155func TestServerIndexAndAPIKey(t *testing.T) {
156 s, err := NewServer()
157 if err != nil {
158 t.Fatal(err)
159 }
160 defer s.Shutdown(context.Background())
161
162 // Swap the probe out for a fake that always succeeds.
163 s.probeFn = func(ctx context.Context, provider, key string) error { return nil }
164
165 resp, err := http.Get(s.URL() + "/")
166 if err != nil {
167 t.Fatal(err)
168 }
169 defer resp.Body.Close()
170 if resp.StatusCode != 200 {
171 t.Fatalf("index status %d", resp.StatusCode)
172 }
173
174 resp, err = http.Get(s.URL() + "/apikey?provider=anthropic")
175 if err != nil {
176 t.Fatal(err)
177 }
178 resp.Body.Close()
179 if resp.StatusCode != 200 {
180 t.Fatalf("apikey form status %d", resp.StatusCode)
181 }
182
183 // POST a fake key and expect a success result on the channel.
184 _, err = http.PostForm(s.URL()+"/apikey", url.Values{
185 "provider": {"anthropic"},
186 "api_key": {"sk-ant-test"},
187 })
188 if err != nil {
189 t.Fatal(err)
190 }
191 select {
192 case r := <-s.Result():
193 if r.Err != nil {
194 t.Fatalf("unexpected err: %v", r.Err)
195 }
196 if r.Provider != "anthropic" || r.APIKey != "sk-ant-test" || r.Method != "apikey" {
197 t.Fatalf("bad result: %+v", r)
198 }
199 case <-time.After(2 * time.Second):
200 t.Fatal("no result")
201 }
202}
203
204func TestServerCallback(t *testing.T) {
205 s, err := NewServer()
206 if err != nil {
207 t.Fatal(err)
208 }
209 defer s.Shutdown(context.Background())
210
211 u := s.URL() + "/callback?code=abc123&state=anthropic:xyz"
212 resp, err := http.Get(u)
213 if err != nil {
214 t.Fatal(err)
215 }
216 resp.Body.Close()
217
218 select {
219 case r := <-s.Result():
220 if r.Provider != "anthropic" || r.Code != "abc123" || r.State != "anthropic:xyz" || r.Method != "oauth" {
221 t.Fatalf("bad result: %+v", r)
222 }
223 case <-time.After(2 * time.Second):
224 t.Fatal("no result")
225 }
226}
227
228func TestCallbackServerRoundTrip(t *testing.T) {
229 // Use a provider config with an ephemeral port so tests stay hermetic.
230 op := OAuthProvider{
231 Name: "anthropic",
232 RedirectHost: "127.0.0.1",
233 RedirectPort: 0, // let kernel pick
234 RedirectPath: "/callback",
235 }
236 // We need a real fixed port for CallbackServer.NewCallbackServer; find one.
237 l, err := net.Listen("tcp", "127.0.0.1:0")
238 if err != nil {
239 t.Fatal(err)
240 }
241 port := l.Addr().(*net.TCPAddr).Port
242 l.Close()
243 op.RedirectPort = port
244
245 cs, err := NewCallbackServer(op, "thestate")
246 if err != nil {
247 t.Fatal(err)
248 }
249 defer cs.Shutdown()
250
251 go func() {
252 time.Sleep(50 * time.Millisecond)
253 _, _ = http.Get(cs.URL() + "?code=abc&state=thestate")
254 }()
255 ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
256 defer cancel()
257 r, err := cs.Result(ctx)
258 if err != nil {
259 t.Fatal(err)
260 }
261 if r.Code != "abc" || r.State != "thestate" {
262 t.Fatalf("bad result: %+v", r)
263 }
264}
265
266func TestCallbackServerStateMismatch(t *testing.T) {
267 l, _ := net.Listen("tcp", "127.0.0.1:0")
268 port := l.Addr().(*net.TCPAddr).Port
269 l.Close()
270 op := OAuthProvider{Name: "anthropic", RedirectHost: "127.0.0.1", RedirectPort: port, RedirectPath: "/callback"}
271 cs, err := NewCallbackServer(op, "expected")
272 if err != nil {
273 t.Fatal(err)
274 }
275 defer cs.Shutdown()
276
277 go func() {
278 time.Sleep(50 * time.Millisecond)
279 _, _ = http.Get(cs.URL() + "?code=abc&state=wrong")
280 }()
281 ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
282 defer cancel()
283 r, err := cs.Result(ctx)
284 if err != nil {
285 t.Fatal(err)
286 }
287 if r.Err == nil || !strings.Contains(r.Err.Error(), "state") {
288 t.Fatalf("want state error, got %+v", r)
289 }
290}
291
292func TestServerCallbackError(t *testing.T) {
293 s, err := NewServer()
294 if err != nil {
295 t.Fatal(err)
296 }
297 defer s.Shutdown(context.Background())
298
299 _, _ = http.Get(s.URL() + "/callback?error=access_denied&state=openai:xyz")
300 select {
301 case r := <-s.Result():
302 if r.Err == nil || r.Provider != "openai" {
303 t.Fatalf("bad result: %+v", r)
304 }
305 if !strings.Contains(r.Err.Error(), "access_denied") {
306 t.Fatalf("missing reason: %v", r.Err)
307 }
308 case <-time.After(2 * time.Second):
309 t.Fatal("no result")
310 }
311}