-
-
Notifications
You must be signed in to change notification settings - Fork 1.1k
fix(management): refresh antigravity token for api-call $TOKEN$ #888
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -33,6 +33,13 @@ var geminiOAuthScopes = []string{ | |
| "https://www.googleapis.com/auth/userinfo.profile", | ||
| } | ||
|
|
||
| const ( | ||
| antigravityOAuthClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com" | ||
| antigravityOAuthClientSecret = "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" | ||
| ) | ||
|
|
||
| var antigravityOAuthTokenURL = "https://oauth2.googleapis.com/token" | ||
|
|
||
| type apiCallRequest struct { | ||
| AuthIndexSnake *string `json:"auth_index"` | ||
| AuthIndexCamel *string `json:"authIndex"` | ||
|
|
@@ -251,6 +258,10 @@ func (h *Handler) resolveTokenForAuth(ctx context.Context, auth *coreauth.Auth) | |
| token, errToken := h.refreshGeminiOAuthAccessToken(ctx, auth) | ||
| return token, errToken | ||
| } | ||
| if provider == "antigravity" { | ||
| token, errToken := h.refreshAntigravityOAuthAccessToken(ctx, auth) | ||
| return token, errToken | ||
| } | ||
|
|
||
| return tokenValueForAuth(auth), nil | ||
| } | ||
|
|
@@ -325,6 +336,161 @@ func (h *Handler) refreshGeminiOAuthAccessToken(ctx context.Context, auth *corea | |
| return strings.TrimSpace(currentToken.AccessToken), nil | ||
| } | ||
|
|
||
| func (h *Handler) refreshAntigravityOAuthAccessToken(ctx context.Context, auth *coreauth.Auth) (string, error) { | ||
| if ctx == nil { | ||
| ctx = context.Background() | ||
| } | ||
| if auth == nil { | ||
| return "", nil | ||
| } | ||
|
|
||
| metadata := auth.Metadata | ||
| if len(metadata) == 0 { | ||
| return "", fmt.Errorf("antigravity oauth metadata missing") | ||
| } | ||
|
|
||
| current := strings.TrimSpace(tokenValueFromMetadata(metadata)) | ||
| if current != "" && !antigravityTokenNeedsRefresh(metadata) { | ||
| return current, nil | ||
| } | ||
|
|
||
| refreshToken := stringValue(metadata, "refresh_token") | ||
| if refreshToken == "" { | ||
| return "", fmt.Errorf("antigravity refresh token missing") | ||
| } | ||
|
|
||
| tokenURL := strings.TrimSpace(antigravityOAuthTokenURL) | ||
| if tokenURL == "" { | ||
| tokenURL = "https://oauth2.googleapis.com/token" | ||
| } | ||
| form := url.Values{} | ||
| form.Set("client_id", antigravityOAuthClientID) | ||
| form.Set("client_secret", antigravityOAuthClientSecret) | ||
| form.Set("grant_type", "refresh_token") | ||
| form.Set("refresh_token", refreshToken) | ||
|
|
||
| req, errReq := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL, strings.NewReader(form.Encode())) | ||
| if errReq != nil { | ||
| return "", errReq | ||
| } | ||
| req.Header.Set("Content-Type", "application/x-www-form-urlencoded") | ||
|
|
||
| httpClient := &http.Client{ | ||
| Timeout: defaultAPICallTimeout, | ||
| Transport: h.apiCallTransport(auth), | ||
| } | ||
| resp, errDo := httpClient.Do(req) | ||
| if errDo != nil { | ||
| return "", errDo | ||
| } | ||
| defer func() { | ||
| if errClose := resp.Body.Close(); errClose != nil { | ||
| log.Errorf("response body close error: %v", errClose) | ||
| } | ||
| }() | ||
|
|
||
| bodyBytes, errRead := io.ReadAll(resp.Body) | ||
| if errRead != nil { | ||
| return "", errRead | ||
| } | ||
| if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { | ||
| return "", fmt.Errorf("antigravity oauth token refresh failed: status %d: %s", resp.StatusCode, strings.TrimSpace(string(bodyBytes))) | ||
| } | ||
|
|
||
| var tokenResp struct { | ||
| AccessToken string `json:"access_token"` | ||
| RefreshToken string `json:"refresh_token"` | ||
| ExpiresIn int64 `json:"expires_in"` | ||
| TokenType string `json:"token_type"` | ||
| } | ||
| if errUnmarshal := json.Unmarshal(bodyBytes, &tokenResp); errUnmarshal != nil { | ||
| return "", errUnmarshal | ||
| } | ||
|
|
||
| if strings.TrimSpace(tokenResp.AccessToken) == "" { | ||
| return "", fmt.Errorf("antigravity oauth token refresh returned empty access_token") | ||
| } | ||
|
|
||
| if auth.Metadata == nil { | ||
| auth.Metadata = make(map[string]any) | ||
| } | ||
| now := time.Now() | ||
| auth.Metadata["access_token"] = strings.TrimSpace(tokenResp.AccessToken) | ||
| if strings.TrimSpace(tokenResp.RefreshToken) != "" { | ||
| auth.Metadata["refresh_token"] = strings.TrimSpace(tokenResp.RefreshToken) | ||
| } | ||
| if tokenResp.ExpiresIn > 0 { | ||
| auth.Metadata["expires_in"] = tokenResp.ExpiresIn | ||
| auth.Metadata["timestamp"] = now.UnixMilli() | ||
| auth.Metadata["expired"] = now.Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339) | ||
| } | ||
| auth.Metadata["type"] = "antigravity" | ||
|
|
||
| if h != nil && h.authManager != nil { | ||
| auth.LastRefreshedAt = now | ||
| auth.UpdatedAt = now | ||
| _, _ = h.authManager.Update(ctx, auth) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The error returned by if _, err := h.authManager.Update(ctx, auth); err != nil {
log.Warnf("failed to persist refreshed antigravity token: %v", err)
} |
||
| } | ||
|
|
||
| return strings.TrimSpace(tokenResp.AccessToken), nil | ||
| } | ||
|
|
||
| func antigravityTokenNeedsRefresh(metadata map[string]any) bool { | ||
| // Refresh a bit early to avoid requests racing token expiry. | ||
| const skew = 30 * time.Second | ||
|
|
||
| if metadata == nil { | ||
| return true | ||
| } | ||
| if expStr, ok := metadata["expired"].(string); ok { | ||
| if ts, errParse := time.Parse(time.RFC3339, strings.TrimSpace(expStr)); errParse == nil { | ||
| return !ts.After(time.Now().Add(skew)) | ||
| } | ||
| } | ||
| expiresIn := int64Value(metadata["expires_in"]) | ||
| timestampMs := int64Value(metadata["timestamp"]) | ||
| if expiresIn > 0 && timestampMs > 0 { | ||
| exp := time.UnixMilli(timestampMs).Add(time.Duration(expiresIn) * time.Second) | ||
| return !exp.After(time.Now().Add(skew)) | ||
| } | ||
| return true | ||
| } | ||
|
|
||
| func int64Value(raw any) int64 { | ||
| switch typed := raw.(type) { | ||
| case int: | ||
| return int64(typed) | ||
| case int32: | ||
| return int64(typed) | ||
| case int64: | ||
| return typed | ||
| case uint: | ||
| return int64(typed) | ||
| case uint32: | ||
| return int64(typed) | ||
| case uint64: | ||
| if typed > uint64(^uint64(0)>>1) { | ||
| return 0 | ||
| } | ||
| return int64(typed) | ||
| case float32: | ||
| return int64(typed) | ||
| case float64: | ||
| return int64(typed) | ||
| case json.Number: | ||
| if i, errParse := typed.Int64(); errParse == nil { | ||
| return i | ||
| } | ||
| case string: | ||
| if s := strings.TrimSpace(typed); s != "" { | ||
| if i, errParse := json.Number(s).Int64(); errParse == nil { | ||
| return i | ||
| } | ||
| } | ||
| } | ||
| return 0 | ||
| } | ||
|
Comment on lines
+459
to
+492
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This |
||
|
|
||
| func geminiOAuthMetadata(auth *coreauth.Auth) (map[string]any, func(map[string]any)) { | ||
| if auth == nil { | ||
| return nil, nil | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,173 @@ | ||
| package management | ||
|
|
||
| import ( | ||
| "context" | ||
| "encoding/json" | ||
| "io" | ||
| "net/http" | ||
| "net/http/httptest" | ||
| "net/url" | ||
| "strings" | ||
| "sync" | ||
| "testing" | ||
| "time" | ||
|
|
||
| coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" | ||
| ) | ||
|
|
||
| type memoryAuthStore struct { | ||
| mu sync.Mutex | ||
| items map[string]*coreauth.Auth | ||
| } | ||
|
|
||
| func (s *memoryAuthStore) List(ctx context.Context) ([]*coreauth.Auth, error) { | ||
| _ = ctx | ||
| s.mu.Lock() | ||
| defer s.mu.Unlock() | ||
| out := make([]*coreauth.Auth, 0, len(s.items)) | ||
| for _, a := range s.items { | ||
| out = append(out, a.Clone()) | ||
| } | ||
| return out, nil | ||
| } | ||
|
|
||
| func (s *memoryAuthStore) Save(ctx context.Context, auth *coreauth.Auth) (string, error) { | ||
| _ = ctx | ||
| if auth == nil { | ||
| return "", nil | ||
| } | ||
| s.mu.Lock() | ||
| if s.items == nil { | ||
| s.items = make(map[string]*coreauth.Auth) | ||
| } | ||
| s.items[auth.ID] = auth.Clone() | ||
| s.mu.Unlock() | ||
| return auth.ID, nil | ||
| } | ||
|
|
||
| func (s *memoryAuthStore) Delete(ctx context.Context, id string) error { | ||
| _ = ctx | ||
| s.mu.Lock() | ||
| delete(s.items, id) | ||
| s.mu.Unlock() | ||
| return nil | ||
| } | ||
|
|
||
| func TestResolveTokenForAuth_Antigravity_RefreshesExpiredToken(t *testing.T) { | ||
| var callCount int | ||
| srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | ||
| callCount++ | ||
| if r.Method != http.MethodPost { | ||
| t.Fatalf("expected POST, got %s", r.Method) | ||
| } | ||
| if ct := r.Header.Get("Content-Type"); !strings.HasPrefix(ct, "application/x-www-form-urlencoded") { | ||
| t.Fatalf("unexpected content-type: %s", ct) | ||
| } | ||
| bodyBytes, _ := io.ReadAll(r.Body) | ||
| _ = r.Body.Close() | ||
| values, err := url.ParseQuery(string(bodyBytes)) | ||
| if err != nil { | ||
| t.Fatalf("parse form: %v", err) | ||
| } | ||
| if values.Get("grant_type") != "refresh_token" { | ||
| t.Fatalf("unexpected grant_type: %s", values.Get("grant_type")) | ||
| } | ||
| if values.Get("refresh_token") != "rt" { | ||
| t.Fatalf("unexpected refresh_token: %s", values.Get("refresh_token")) | ||
| } | ||
| if values.Get("client_id") != antigravityOAuthClientID { | ||
| t.Fatalf("unexpected client_id: %s", values.Get("client_id")) | ||
| } | ||
| if values.Get("client_secret") != antigravityOAuthClientSecret { | ||
| t.Fatalf("unexpected client_secret") | ||
| } | ||
|
|
||
| w.Header().Set("Content-Type", "application/json") | ||
| _ = json.NewEncoder(w).Encode(map[string]any{ | ||
| "access_token": "new-token", | ||
| "refresh_token": "rt2", | ||
| "expires_in": int64(3600), | ||
| "token_type": "Bearer", | ||
| }) | ||
| })) | ||
| t.Cleanup(srv.Close) | ||
|
|
||
| originalURL := antigravityOAuthTokenURL | ||
| antigravityOAuthTokenURL = srv.URL | ||
| t.Cleanup(func() { antigravityOAuthTokenURL = originalURL }) | ||
|
|
||
| store := &memoryAuthStore{} | ||
| manager := coreauth.NewManager(store, nil, nil) | ||
|
|
||
| auth := &coreauth.Auth{ | ||
| ID: "antigravity-test.json", | ||
| FileName: "antigravity-test.json", | ||
| Provider: "antigravity", | ||
| Metadata: map[string]any{ | ||
| "type": "antigravity", | ||
| "access_token": "old-token", | ||
| "refresh_token": "rt", | ||
| "expires_in": int64(3600), | ||
| "timestamp": time.Now().Add(-2 * time.Hour).UnixMilli(), | ||
| "expired": time.Now().Add(-1 * time.Hour).Format(time.RFC3339), | ||
| }, | ||
| } | ||
| if _, err := manager.Register(context.Background(), auth); err != nil { | ||
| t.Fatalf("register auth: %v", err) | ||
| } | ||
|
|
||
| h := &Handler{authManager: manager} | ||
| token, err := h.resolveTokenForAuth(context.Background(), auth) | ||
| if err != nil { | ||
| t.Fatalf("resolveTokenForAuth: %v", err) | ||
| } | ||
| if token != "new-token" { | ||
| t.Fatalf("expected refreshed token, got %q", token) | ||
| } | ||
| if callCount != 1 { | ||
| t.Fatalf("expected 1 refresh call, got %d", callCount) | ||
| } | ||
|
|
||
| updated, ok := manager.GetByID(auth.ID) | ||
| if !ok || updated == nil { | ||
| t.Fatalf("expected auth in manager after update") | ||
| } | ||
| if got := tokenValueFromMetadata(updated.Metadata); got != "new-token" { | ||
| t.Fatalf("expected manager metadata updated, got %q", got) | ||
| } | ||
| } | ||
|
|
||
| func TestResolveTokenForAuth_Antigravity_SkipsRefreshWhenTokenValid(t *testing.T) { | ||
| var callCount int | ||
| srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | ||
| callCount++ | ||
| w.WriteHeader(http.StatusInternalServerError) | ||
| })) | ||
| t.Cleanup(srv.Close) | ||
|
|
||
| originalURL := antigravityOAuthTokenURL | ||
| antigravityOAuthTokenURL = srv.URL | ||
| t.Cleanup(func() { antigravityOAuthTokenURL = originalURL }) | ||
|
|
||
| auth := &coreauth.Auth{ | ||
| ID: "antigravity-valid.json", | ||
| FileName: "antigravity-valid.json", | ||
| Provider: "antigravity", | ||
| Metadata: map[string]any{ | ||
| "type": "antigravity", | ||
| "access_token": "ok-token", | ||
| "expired": time.Now().Add(30 * time.Minute).Format(time.RFC3339), | ||
| }, | ||
| } | ||
| h := &Handler{} | ||
| token, err := h.resolveTokenForAuth(context.Background(), auth) | ||
| if err != nil { | ||
| t.Fatalf("resolveTokenForAuth: %v", err) | ||
| } | ||
| if token != "ok-token" { | ||
| t.Fatalf("expected existing token, got %q", token) | ||
| } | ||
| if callCount != 0 { | ||
| t.Fatalf("expected no refresh calls, got %d", callCount) | ||
| } | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hardcoding secrets like
antigravityOAuthClientIDandantigravityOAuthClientSecretis a major security risk. These credentials should be loaded from a secure configuration source (e.g., environment variables, a secret manager) and not be committed to the source code. Additionally, these constants appear to be duplicates of those ininternal/runtime/executor/antigravity_executor.go. This duplication should be avoided by centralizing the configuration and loading secrets securely.