266 lines
6.5 KiB
Go
266 lines
6.5 KiB
Go
package antigravity
|
||
|
||
import (
|
||
"crypto/rand"
|
||
"crypto/sha256"
|
||
"encoding/base64"
|
||
"encoding/hex"
|
||
"fmt"
|
||
"net/url"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
const (
|
||
// Google OAuth 端点
|
||
AuthorizeURL = "https://accounts.google.com/o/oauth2/v2/auth"
|
||
TokenURL = "https://oauth2.googleapis.com/token"
|
||
UserInfoURL = "https://www.googleapis.com/oauth2/v2/userinfo"
|
||
|
||
// Antigravity OAuth 客户端凭证
|
||
ClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
|
||
ClientSecret = "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"
|
||
|
||
// 固定的 redirect_uri（用户需手动复制 code）
|
||
RedirectURI = "http://localhost:8085/callback"
|
||
|
||
// OAuth scopes
|
||
Scopes = "https://www.googleapis.com/auth/cloud-platform " +
|
||
"https://www.googleapis.com/auth/userinfo.email " +
|
||
"https://www.googleapis.com/auth/userinfo.profile " +
|
||
"https://www.googleapis.com/auth/cclog " +
|
||
"https://www.googleapis.com/auth/experimentsandconfigs"
|
||
|
||
// User-Agent（与 Antigravity-Manager 保持一致）
|
||
UserAgent = "antigravity/1.15.8 windows/amd64"
|
||
|
||
// Session 过期时间
|
||
SessionTTL = 30 * time.Minute
|
||
|
||
// URL 可用性 TTL（不可用 URL 的恢复时间）
|
||
URLAvailabilityTTL = 5 * time.Minute
|
||
)
|
||
|
||
// BaseURLs 定义 Antigravity API 端点（与 Antigravity-Manager 保持一致）
|
||
var BaseURLs = []string{
|
||
"https://cloudcode-pa.googleapis.com", // prod (优先)
|
||
"https://daily-cloudcode-pa.sandbox.googleapis.com", // daily sandbox (备用)
|
||
}
|
||
|
||
// BaseURL 默认 URL（保持向后兼容）
|
||
var BaseURL = BaseURLs[0]
|
||
|
||
// URLAvailability 管理 URL 可用性状态（带 TTL 自动恢复和动态优先级）
|
||
type URLAvailability struct {
|
||
mu sync.RWMutex
|
||
unavailable map[string]time.Time // URL -> 恢复时间
|
||
ttl time.Duration
|
||
lastSuccess string // 最近成功请求的 URL，优先使用
|
||
}
|
||
|
||
// DefaultURLAvailability 全局 URL 可用性管理器
|
||
var DefaultURLAvailability = NewURLAvailability(URLAvailabilityTTL)
|
||
|
||
// NewURLAvailability 创建 URL 可用性管理器
|
||
func NewURLAvailability(ttl time.Duration) *URLAvailability {
|
||
return &URLAvailability{
|
||
unavailable: make(map[string]time.Time),
|
||
ttl: ttl,
|
||
}
|
||
}
|
||
|
||
// MarkUnavailable 标记 URL 临时不可用
|
||
func (u *URLAvailability) MarkUnavailable(url string) {
|
||
u.mu.Lock()
|
||
defer u.mu.Unlock()
|
||
u.unavailable[url] = time.Now().Add(u.ttl)
|
||
}
|
||
|
||
// MarkSuccess 标记 URL 请求成功，将其设为优先使用
|
||
func (u *URLAvailability) MarkSuccess(url string) {
|
||
u.mu.Lock()
|
||
defer u.mu.Unlock()
|
||
u.lastSuccess = url
|
||
// 成功后清除该 URL 的不可用标记
|
||
delete(u.unavailable, url)
|
||
}
|
||
|
||
// IsAvailable 检查 URL 是否可用
|
||
func (u *URLAvailability) IsAvailable(url string) bool {
|
||
u.mu.RLock()
|
||
defer u.mu.RUnlock()
|
||
expiry, exists := u.unavailable[url]
|
||
if !exists {
|
||
return true
|
||
}
|
||
return time.Now().After(expiry)
|
||
}
|
||
|
||
// GetAvailableURLs 返回可用的 URL 列表
|
||
// 最近成功的 URL 优先，其他按默认顺序
|
||
func (u *URLAvailability) GetAvailableURLs() []string {
|
||
u.mu.RLock()
|
||
defer u.mu.RUnlock()
|
||
|
||
now := time.Now()
|
||
result := make([]string, 0, len(BaseURLs))
|
||
|
||
// 如果有最近成功的 URL 且可用，放在最前面
|
||
if u.lastSuccess != "" {
|
||
expiry, exists := u.unavailable[u.lastSuccess]
|
||
if !exists || now.After(expiry) {
|
||
result = append(result, u.lastSuccess)
|
||
}
|
||
}
|
||
|
||
// 添加其他可用的 URL（按默认顺序）
|
||
for _, url := range BaseURLs {
|
||
// 跳过已添加的 lastSuccess
|
||
if url == u.lastSuccess {
|
||
continue
|
||
}
|
||
expiry, exists := u.unavailable[url]
|
||
if !exists || now.After(expiry) {
|
||
result = append(result, url)
|
||
}
|
||
}
|
||
return result
|
||
}
|
||
|
||
// OAuthSession 保存 OAuth 授权流程的临时状态
|
||
type OAuthSession struct {
|
||
State string `json:"state"`
|
||
CodeVerifier string `json:"code_verifier"`
|
||
ProxyURL string `json:"proxy_url,omitempty"`
|
||
CreatedAt time.Time `json:"created_at"`
|
||
}
|
||
|
||
// SessionStore OAuth session 存储
|
||
type SessionStore struct {
|
||
mu sync.RWMutex
|
||
sessions map[string]*OAuthSession
|
||
stopCh chan struct{}
|
||
}
|
||
|
||
func NewSessionStore() *SessionStore {
|
||
store := &SessionStore{
|
||
sessions: make(map[string]*OAuthSession),
|
||
stopCh: make(chan struct{}),
|
||
}
|
||
go store.cleanup()
|
||
return store
|
||
}
|
||
|
||
func (s *SessionStore) Set(sessionID string, session *OAuthSession) {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
s.sessions[sessionID] = session
|
||
}
|
||
|
||
func (s *SessionStore) Get(sessionID string) (*OAuthSession, bool) {
|
||
s.mu.RLock()
|
||
defer s.mu.RUnlock()
|
||
session, ok := s.sessions[sessionID]
|
||
if !ok {
|
||
return nil, false
|
||
}
|
||
if time.Since(session.CreatedAt) > SessionTTL {
|
||
return nil, false
|
||
}
|
||
return session, true
|
||
}
|
||
|
||
func (s *SessionStore) Delete(sessionID string) {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
delete(s.sessions, sessionID)
|
||
}
|
||
|
||
func (s *SessionStore) Stop() {
|
||
select {
|
||
case <-s.stopCh:
|
||
return
|
||
default:
|
||
close(s.stopCh)
|
||
}
|
||
}
|
||
|
||
func (s *SessionStore) cleanup() {
|
||
ticker := time.NewTicker(5 * time.Minute)
|
||
defer ticker.Stop()
|
||
for {
|
||
select {
|
||
case <-s.stopCh:
|
||
return
|
||
case <-ticker.C:
|
||
s.mu.Lock()
|
||
for id, session := range s.sessions {
|
||
if time.Since(session.CreatedAt) > SessionTTL {
|
||
delete(s.sessions, id)
|
||
}
|
||
}
|
||
s.mu.Unlock()
|
||
}
|
||
}
|
||
}
|
||
|
||
func GenerateRandomBytes(n int) ([]byte, error) {
|
||
b := make([]byte, n)
|
||
_, err := rand.Read(b)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return b, nil
|
||
}
|
||
|
||
func GenerateState() (string, error) {
|
||
bytes, err := GenerateRandomBytes(32)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
return base64URLEncode(bytes), nil
|
||
}
|
||
|
||
func GenerateSessionID() (string, error) {
|
||
bytes, err := GenerateRandomBytes(16)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
return hex.EncodeToString(bytes), nil
|
||
}
|
||
|
||
func GenerateCodeVerifier() (string, error) {
|
||
bytes, err := GenerateRandomBytes(32)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
return base64URLEncode(bytes), nil
|
||
}
|
||
|
||
func GenerateCodeChallenge(verifier string) string {
|
||
hash := sha256.Sum256([]byte(verifier))
|
||
return base64URLEncode(hash[:])
|
||
}
|
||
|
||
func base64URLEncode(data []byte) string {
|
||
return strings.TrimRight(base64.URLEncoding.EncodeToString(data), "=")
|
||
}
|
||
|
||
// BuildAuthorizationURL 构建 Google OAuth 授权 URL
|
||
func BuildAuthorizationURL(state, codeChallenge string) string {
|
||
params := url.Values{}
|
||
params.Set("client_id", ClientID)
|
||
params.Set("redirect_uri", RedirectURI)
|
||
params.Set("response_type", "code")
|
||
params.Set("scope", Scopes)
|
||
params.Set("state", state)
|
||
params.Set("code_challenge", codeChallenge)
|
||
params.Set("code_challenge_method", "S256")
|
||
params.Set("access_type", "offline")
|
||
params.Set("prompt", "consent")
|
||
params.Set("include_granted_scopes", "true")
|
||
|
||
return fmt.Sprintf("%s?%s", AuthorizeURL, params.Encode())
|
||
}
|