package provider import ( "context" "crypto/aes" "crypto/cipher" "encoding/base64" "encoding/json" "net/http" "net/http/httptest" "net/url" "os" "path/filepath" "strings" "sync/atomic" "testing" "time" ) // zcodeEncrypt stores a value as ZCode's credential store does. func zcodeEncrypt(t *testing.T, v string) string { t.Helper() block, _ := aes.NewCipher(zcodeSecret()) gcm, _ := cipher.NewGCM(block) iv := []byte("0123456789ab") sealed := gcm.Seal(nil, iv, []byte(v), nil) ct, tag := sealed[:len(sealed)-16], sealed[len(sealed)-16:] b := base64.RawURLEncoding.EncodeToString return "enc:v1:" + b(iv) + "." + b(tag) + "." + b(ct) } // ZCode's own account is read from its credential store; a second, signed // in from magpie with ZCode's sign-in, is kept beside it with its own key. func TestZCodeAccounts(t *testing.T) { home := signIn(t) t.Setenv("ZCODE_CREDENTIAL_SECRET", "test-secret") writeFile(t, filepath.Join(home, ".zcode", "v2", "credentials.json"), map[string]any{ "oauth:zai:user_info": zcodeEncrypt(t, `{"user_id":"u1","email":"13800000000@phone.local","name":"旅行者0000"}`), "account-provider:coding-plan:account:zai-individual-coding-plan:account:abc:api-key": zcodeEncrypt(t, "own.secret"), "oauth:zai:access_token": zcodeEncrypt(t, "jwt"), }) var polls atomic.Int32 var srv *httptest.Server srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { auth := r.Header.Get("Authorization") ok := func(data any) { json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": data}) } switch { case r.URL.Path == "/api/v1/oauth/cli/init": if !strings.HasPrefix(auth, "Bearer ") { w.WriteHeader(401) return } ok(map[string]any{"flow_id": "f1", "authorize_url": "https://chat.z.ai/oauth?client_id=c&redirect_uri=x", "expires_at": time.Now().Add(time.Minute).Unix(), "poll_interval_sec": 0}) case r.URL.Path == "/api/v1/oauth/cli/poll/f1": if polls.Add(1) < 2 { ok(map[string]any{"status": "pending"}) return } ok(map[string]any{"status": "ready", "token": "zc", "zai": map[string]any{"access_token": "zai-tok"}, "user": map[string]any{"user_id": "u2", "email": "two@example.com", "name": "Two"}}) case r.URL.Path == "/api/auth/z/login": var body map[string]string json.NewDecoder(r.Body).Decode(&body) if body["token"] != "zai-tok" { w.WriteHeader(401) return } ok(map[string]any{"access_token": "biz"}) case auth == "Bearer biz" && r.URL.Path == "/api/biz/customer/getCustomerInfo": ok(map[string]any{"organizations": []any{ map[string]any{"organizationId": "o0", "organizationName": "other", "projects": []any{map[string]any{"projectId": "p0", "projectName": "x", "projectType": "1"}}}, map[string]any{"organizationId": "o1", "organizationName": "默认机构", "projects": []any{ map[string]any{"projectId": "p9", "projectName": "api", "projectType": "2"}, map[string]any{"projectId": "p1", "projectName": "默认项目", "projectType": "1"}, }}, }}) case auth == "Bearer biz" && r.URL.Path == "/api/biz/v1/organization/o1/projects/p1/api_keys": if r.Method == http.MethodPost { ok(map[string]any{"name": "zcode-api-key", "apiKey": "two"}) return } ok([]any{map[string]any{"name": "other", "apiKey": "nope"}}) case auth == "Bearer biz" && r.URL.Path == "/api/biz/v1/organization/o1/projects/p1/api_keys/copy/two": ok(map[string]any{"secretKey": "secret2"}) case r.URL.Path == "/api/biz/subscription/list": if auth != "two.secret2" && auth != "own.secret" { w.WriteHeader(401) return } ok([]any{map[string]any{"productName": "GLM Coding Pro", "status": "VALID"}}) case r.URL.Path == "/api/monitor/usage/quota/limit": if auth != "two.secret2" && auth != "own.secret" { w.WriteHeader(401) return } reset := time.Now().Add(time.Hour).UnixMilli() ok(map[string]any{"level": "pro", "limits": []any{ map[string]any{"type": "CREDIT_LIMIT", "unit": 3, "number": 5, "usage": 2000, "remaining": 1500, "percentage": 25, "nextResetTime": reset}, map[string]any{"type": "CREDIT_LIMIT", "unit": 6, "number": 1, "usage": 10000, "remaining": 9000, "percentage": 10, "nextResetTime": reset}, }}) default: w.WriteHeader(404) } })) defer srv.Close() oldAPI, oldZai, oldBase := zcodeAPI, zcodeZaiAPI, ZCodeZaiBase zcodeAPI, zcodeZaiAPI, ZCodeZaiBase = srv.URL, srv.URL, srv.URL+"/api/anthropic" defer func() { zcodeAPI, zcodeZaiAPI, ZCodeZaiBase = oldAPI, oldZai, oldBase }() // ZCode's own, named by its phone number who, k, ok := zcodeOwn() if !ok || who != "13800000000" || k.Key != "own.secret" { t.Fatalf("own: %v %q %+v", ok, who, k) } st, err := StartSignIn("zcode") if err != nil { t.Fatal(err) } u, _ := url.Parse(st.URL) if u.Host != "chat.z.ai" || !strings.Contains(u.Query().Get("redirect_uri"), "zcode%3A%2F%2Foauth%2Fcallback") && !strings.Contains(u.Query().Get("redirect_uri"), "zcode://oauth/callback") { t.Fatalf("sign-in page: %s", st.URL) } ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() st, _ = WaitSignIn(ctx, st.ID) if st.State != "done" || st.User != "two@example.com" || st.Plan != "GLM Coding Pro" || st.Using { t.Fatalf("done: %+v", st) } var users []string for _, l := range Logins("zcode") { users = append(users, l.User+map[bool]string{true: "*", false: ""}[l.Active]+map[bool]string{true: "+", false: ""}[l.On]) } if got := strings.Join(users, " "); got != "13800000000*+ two@example.com+" { t.Fatalf("logins: %s", got) } p, ok := find(All(), "zcode") if !ok || p.Account.User != "13800000000" || p.Anthropic != srv.URL+"/api/anthropic" { t.Fatalf("first: %+v", p) } also := p.AlsoOn() if len(also) != 1 || also[0].Account.User != "two@example.com" { t.Fatalf("also on: %+v", also) } req, _ := http.NewRequest("POST", also[0].Anthropic+"/v1/messages", nil) req.Header.Set("Authorization", "Bearer magpie") if err := also[0].Sign(context.Background(), req, Anthropic, []byte(`{}`)); err != nil || req.Header.Get("x-api-key") != "two.secret2" || req.Header.Get("Authorization") != "Bearer two.secret2" { t.Fatalf("signs with its own key: %v %v", err, req.Header) } if ms := also[0].Account.models(); len(ms) != len(zcodeModels) || ms[0].Context != 1_000_000 { t.Fatalf("models: %+v", ms) } use := LoginUsage(context.Background(), "zcode") q := use["two@example.com"] if q.Error != "" || q.Plan != "GLM Coding Pro" || len(q.Windows) != 2 || q.Windows[0].Name != "5 hours" || q.Windows[0].Used != 25 || q.Windows[1].Name != "Weekly" || q.Windows[1].Used != 10 || q.Windows[0].ResetsAt == nil { t.Fatalf("usage: %+v", q) } if err := SwitchLogin("zcode", "two@example.com"); err != nil { t.Fatal(err) } if p, _ := find(All(), "zcode"); p.Account.User != "two@example.com" { t.Fatalf("after switch: %+v", p.Account) } if err := SwitchLogin("zcode", "13800000000"); err != nil { t.Fatal(err) } if err := ForgetLogin("zcode", "two@example.com"); err != nil { t.Fatal(err) } if ls := Logins("zcode"); len(ls) != 1 { t.Fatalf("after forget: %+v", ls) } // signed out of ZCode, with none of magpie's: no ZCode subscription os.Remove(filepath.Join(home, ".zcode", "v2", "credentials.json")) if _, ok := find(All(), "zcode"); ok { t.Fatal("zcode without an account") } }