Something went wrong. Try again.
sloproxy
Something went wrong. Try again.
9.8 kB · 320 lines
Go
at commit e8e39526
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321package executor
import ( "context" "encoding/json" "fmt" "io" "net/http" "net/url" "strconv" "strings" "time"
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps" sdkAuth "github.com/router-for-me/CLIProxyAPI/v7/sdk/auth" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth" log "github.com/sirupsen/logrus")
// Refresh refreshes the authentication credentials using the refresh token.func (e *AntigravityExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled { return refreshed, err } if auth == nil { return auth, nil } updated, errRefresh := e.refreshToken(ctx, auth.Clone()) if errRefresh != nil { return nil, errRefresh } return updated, nil}
func (e *AntigravityExecutor) ShouldPrepareRequestAuth(auth *cliproxyauth.Auth) bool { return antigravityProjectIDFromAuth(auth) == ""}
func (e *AntigravityExecutor) PrepareRequestAuth(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { if auth == nil || !e.ShouldPrepareRequestAuth(auth) { return nil, nil }
updated := auth.Clone() token, refreshedAuth, errToken := e.ensureAccessToken(ctx, updated) if errToken != nil { return nil, errToken } if refreshedAuth != nil { updated = refreshedAuth } if antigravityProjectIDFromAuth(updated) != "" { return updated, nil }
projectID, errProject := e.fetchAntigravityProjectID(ctx, updated, token) if errProject != nil { return nil, missingAntigravityProjectIDError(errProject) } if projectID == "" { return nil, missingAntigravityProjectIDError(nil) } if updated.Metadata == nil { updated.Metadata = make(map[string]any) } updated.Metadata["project_id"] = projectID return updated, nil}
func (e *AntigravityExecutor) ensureAccessToken(ctx context.Context, auth *cliproxyauth.Auth) (string, *cliproxyauth.Auth, error) { if auth == nil { return "", nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"} } accessToken := metaStringValue(auth.Metadata, "access_token") expiry := tokenExpiry(auth.Metadata) if accessToken != "" && expiry.After(time.Now().Add(refreshSkew)) { e.maybeRefreshAntigravityCreditsHint(ctx, auth, accessToken) return accessToken, nil, nil } refreshCtx := context.Background() if ctx != nil { if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil { refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt) } } if refreshed, handled, err := helps.RefreshAuthViaHome(refreshCtx, e.cfg, auth); handled { if err != nil { return "", nil, err } token := metaStringValue(refreshed.Metadata, "access_token") if strings.TrimSpace(token) == "" { return "", nil, statusErr{code: http.StatusUnauthorized, msg: "missing access token"} } e.maybeRefreshAntigravityCreditsHint(ctx, refreshed, token) return token, refreshed, nil }
updated, errRefresh := e.refreshToken(refreshCtx, auth.Clone()) if errRefresh != nil { return "", nil, errRefresh } return metaStringValue(updated.Metadata, "access_token"), updated, nil}
func (e *AntigravityExecutor) refreshToken(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) { if auth == nil { return nil, statusErr{code: http.StatusUnauthorized, msg: "missing auth"} } refreshToken := metaStringValue(auth.Metadata, "refresh_token") if refreshToken == "" { return auth, statusErr{code: http.StatusUnauthorized, msg: "missing refresh token"} } if ctx == nil { ctx = context.Background() } refreshToken = strings.TrimSpace(refreshToken)
result, errRefresh, _ := antigravityRefreshGroup.Do(refreshToken, func() (interface{}, error) { return e.refreshTokenSingleFlight(context.WithoutCancel(ctx), auth, refreshToken) }) if errRefresh != nil { return auth, errRefresh } tokenResp, ok := result.(*antigravityTokenRefreshData) if !ok || tokenResp == nil { return auth, fmt.Errorf("antigravity token refresh failed: invalid single-flight result") }
if auth.Metadata == nil { auth.Metadata = make(map[string]any) } auth.Metadata["access_token"] = tokenResp.AccessToken if tokenResp.RefreshToken != "" { auth.Metadata["refresh_token"] = tokenResp.RefreshToken } auth.Metadata["expires_in"] = tokenResp.ExpiresIn now := time.Now() auth.Metadata["timestamp"] = now.UnixMilli() auth.Metadata["expired"] = now.Add(time.Duration(tokenResp.ExpiresIn) * time.Second).Format(time.RFC3339) auth.Metadata["type"] = antigravityAuthType if errProject := e.ensureAntigravityProjectID(ctx, auth, tokenResp.AccessToken); errProject != nil { log.Warnf("antigravity executor: ensure project id failed: %v", errProject) } e.updateAntigravityCreditsBalance(ctx, auth, tokenResp.AccessToken) return auth, nil}
func (e *AntigravityExecutor) refreshTokenSingleFlight(ctx context.Context, auth *cliproxyauth.Auth, refreshToken string) (*antigravityTokenRefreshData, error) { form := url.Values{} form.Set("client_id", antigravityClientID) form.Set("client_secret", antigravityClientSecret) form.Set("grant_type", "refresh_token") form.Set("refresh_token", refreshToken)
httpReq, errReq := http.NewRequestWithContext(ctx, http.MethodPost, "https://oauth2.googleapis.com/token", strings.NewReader(form.Encode())) if errReq != nil { return nil, errReq } httpReq.Header.Set("Host", "oauth2.googleapis.com") httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") // Real Antigravity uses Go's default User-Agent for OAuth token refresh httpReq.Header.Set("User-Agent", "Go-http-client/2.0")
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) httpResp, errDo := httpClient.Do(httpReq) if errDo != nil { return nil, errDo } defer func() { if errClose := httpResp.Body.Close(); errClose != nil { log.Errorf("antigravity executor: close response body error: %v", errClose) } }()
bodyBytes, errRead := io.ReadAll(httpResp.Body) if errRead != nil { return nil, errRead }
if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { sErr := statusErr{code: httpResp.StatusCode, msg: string(bodyBytes)} if httpResp.StatusCode == http.StatusTooManyRequests { if retryAfter, parseErr := helps.ParseRetryDelay(bodyBytes); parseErr == nil && retryAfter != nil { sErr.retryAfter = retryAfter } } return nil, sErr }
var tokenResp antigravityTokenRefreshData if errUnmarshal := json.Unmarshal(bodyBytes, &tokenResp); errUnmarshal != nil { return nil, errUnmarshal }
return &tokenResp, nil}
func (e *AntigravityExecutor) ensureAntigravityProjectID(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) error { if auth == nil { return nil }
if antigravityProjectIDFromAuth(auth) != "" { return nil }
projectID, errFetch := e.fetchAntigravityProjectID(ctx, auth, accessToken) if errFetch != nil { return errFetch } if projectID == "" { return nil } if auth.Metadata == nil { auth.Metadata = make(map[string]any) } auth.Metadata["project_id"] = projectID
return nil}
func (e *AntigravityExecutor) fetchAntigravityProjectID(ctx context.Context, auth *cliproxyauth.Auth, accessToken string) (string, error) { token := strings.TrimSpace(accessToken) if token == "" { token = metaStringValue(auth.Metadata, "access_token") } if token == "" { return "", nil }
httpClient := newAntigravityHTTPClient(ctx, e.cfg, auth, 0) projectID, errFetch := sdkAuth.FetchAntigravityProjectID(ctx, token, httpClient) if errFetch != nil { return "", errFetch } return strings.TrimSpace(projectID), nil}
func (e *AntigravityExecutor) projectIDForRequest(_ context.Context, auth *cliproxyauth.Auth, _ string) (string, error) { if projectID := antigravityProjectIDFromAuth(auth); projectID != "" { return projectID, nil } return "", missingAntigravityProjectIDError(nil)}
func antigravityProjectIDFromAuth(auth *cliproxyauth.Auth) string { if auth == nil || auth.Metadata == nil { return "" } if pid, ok := auth.Metadata["project_id"].(string); ok { return strings.TrimSpace(pid) } return ""}
func missingAntigravityProjectIDError(cause error) statusErr { msg := "antigravity auth missing project_id" if cause != nil { msg = fmt.Sprintf("%s: %v", msg, cause) } return statusErr{code: http.StatusBadRequest, msg: msg}}
func tokenExpiry(metadata map[string]any) time.Time { if metadata == nil { return time.Time{} } if expStr, ok := metadata["expired"].(string); ok { expStr = strings.TrimSpace(expStr) if expStr != "" { if parsed, errParse := time.Parse(time.RFC3339, expStr); errParse == nil { return parsed } } } expiresIn, hasExpires := int64Value(metadata["expires_in"]) tsMs, hasTimestamp := int64Value(metadata["timestamp"]) if hasExpires && hasTimestamp { return time.Unix(0, tsMs*int64(time.Millisecond)).Add(time.Duration(expiresIn) * time.Second) } return time.Time{}}
func metaStringValue(metadata map[string]any, key string) string { if metadata == nil { return "" } if v, ok := metadata[key]; ok { switch typed := v.(type) { case string: return strings.TrimSpace(typed) case []byte: return strings.TrimSpace(string(typed)) } } return ""}
func int64Value(value any) (int64, bool) { switch typed := value.(type) { case int: return int64(typed), true case int64: return typed, true case float64: return int64(typed), true case json.Number: if i, errParse := typed.Int64(); errParse == nil { return i, true } case string: if strings.TrimSpace(typed) == "" { return 0, false } if i, errParse := strconv.ParseInt(strings.TrimSpace(typed), 10, 64); errParse == nil { return i, true } } return 0, false}