oauth.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746
  1. package xai
  2. import (
  3. "context"
  4. "crypto/rand"
  5. "crypto/sha256"
  6. "encoding/base64"
  7. "encoding/hex"
  8. "errors"
  9. "fmt"
  10. "log/slog"
  11. "net/url"
  12. "os"
  13. "strings"
  14. "sync"
  15. "time"
  16. "github.com/Wei-Shaw/sub2api/internal/pkg/redissession"
  17. "github.com/Wei-Shaw/sub2api/internal/util/logredact"
  18. "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
  19. "github.com/redis/go-redis/v9"
  20. )
  21. const (
  22. OAuthIssuer = "https://auth.x.ai"
  23. DiscoveryURL = OAuthIssuer + "/.well-known/openid-configuration"
  24. DefaultAuthorizeURL = OAuthIssuer + "/oauth2/authorize"
  25. DefaultTokenURL = OAuthIssuer + "/oauth2/token"
  26. DefaultBaseURL = "https://api.x.ai/v1"
  27. DefaultCLIBaseURL = "https://cli-chat-proxy.grok.com/v1"
  28. DefaultUSEast1BaseURL = "https://us-east-1.api.x.ai/v1"
  29. DefaultUSWest2BaseURL = "https://us-west-2.api.x.ai/v1"
  30. DefaultEUWest1BaseURL = "https://eu-west-1.api.x.ai/v1"
  31. DefaultClientID = "b1a00492-073a-47ea-816f-4c329264a828"
  32. DefaultScope = "openid profile email offline_access grok-cli:access api:access"
  33. DefaultRedirectURI = "http://127.0.0.1:56121/callback"
  34. SessionTTL = 30 * time.Minute
  35. EnvAuthorizeURL = "XAI_OAUTH_AUTHORIZE_URL"
  36. EnvTokenURL = "XAI_OAUTH_TOKEN_URL"
  37. EnvClientID = "XAI_OAUTH_CLIENT_ID"
  38. EnvScope = "XAI_OAUTH_SCOPE"
  39. EnvRedirectURI = "XAI_OAUTH_REDIRECT_URI"
  40. EnvBaseURL = "XAI_BASE_URL"
  41. EnvAllowUnsafeURLOverrides = "XAI_ALLOW_UNSAFE_URL_OVERRIDES"
  42. EnvUnsafeAllowHighConcurrency = "XAI_GROK_UNSAFE_ALLOW_CONCURRENCY_GT_ONE"
  43. )
  44. var (
  45. oauthEndpointAllowedHosts = []string{"x.ai", "*.x.ai"}
  46. // *.api.x.ai 覆盖 xAI 区域端点(us-east-1/us-west-2/eu-west-1 等),
  47. // 运营方可在端点间手动切换以规避单点不可用。
  48. baseURLAllowedHosts = []string{"api.x.ai", "*.api.x.ai", "cli-chat-proxy.grok.com"}
  49. )
  50. // OAuthSession stores one PKCE OAuth flow.
  51. type OAuthSession struct {
  52. State string `json:"state"`
  53. CodeVerifier string `json:"code_verifier"`
  54. CodeChallenge string `json:"code_challenge"`
  55. ClientID string `json:"client_id,omitempty"`
  56. Scope string `json:"scope,omitempty"`
  57. ProxyURL string `json:"proxy_url,omitempty"`
  58. RedirectURI string `json:"redirect_uri"`
  59. CreatedAt time.Time `json:"created_at"`
  60. mu sync.Mutex
  61. consumed bool
  62. }
  63. func (s *OAuthSession) TryConsume() bool {
  64. if s == nil {
  65. return false
  66. }
  67. s.mu.Lock()
  68. defer s.mu.Unlock()
  69. if s.consumed {
  70. return false
  71. }
  72. s.consumed = true
  73. return true
  74. }
  75. // SessionStore manages xAI OAuth sessions with an optional Redis backend.
  76. type SessionStore struct {
  77. mu sync.RWMutex
  78. sessions map[string]*OAuthSession
  79. localOnly map[string]struct{}
  80. stopOnce sync.Once
  81. stopCh chan struct{}
  82. remote *redissession.Store
  83. }
  84. type oauthSessionDTO struct {
  85. State string `json:"state"`
  86. CodeVerifier string `json:"code_verifier"`
  87. CodeChallenge string `json:"code_challenge"`
  88. ClientID string `json:"client_id,omitempty"`
  89. Scope string `json:"scope,omitempty"`
  90. ProxyURL string `json:"proxy_url,omitempty"`
  91. RedirectURI string `json:"redirect_uri"`
  92. CreatedAt time.Time `json:"created_at"`
  93. }
  94. func NewSessionStore() *SessionStore {
  95. store := &SessionStore{
  96. sessions: make(map[string]*OAuthSession),
  97. localOnly: make(map[string]struct{}),
  98. stopCh: make(chan struct{}),
  99. }
  100. go store.cleanup()
  101. return store
  102. }
  103. func NewRedisSessionStore(rdb *redis.Client) *SessionStore {
  104. store := NewSessionStore()
  105. if rdb != nil {
  106. store.remote = redissession.New(rdb, "oauth:session:xai", SessionTTL)
  107. }
  108. return store
  109. }
  110. func (s *SessionStore) Set(sessionID string, session *OAuthSession) {
  111. if session == nil {
  112. return
  113. }
  114. var remoteErr error
  115. if s != nil && s.remote != nil {
  116. remoteErr = s.remote.Set(context.Background(), sessionID, oauthSessionDTO{
  117. State: session.State, CodeVerifier: session.CodeVerifier, CodeChallenge: session.CodeChallenge,
  118. ClientID: session.ClientID, Scope: session.Scope, ProxyURL: session.ProxyURL,
  119. RedirectURI: session.RedirectURI, CreatedAt: session.CreatedAt,
  120. })
  121. }
  122. s.mu.Lock()
  123. defer s.mu.Unlock()
  124. s.sessions[sessionID] = session
  125. if remoteErr != nil {
  126. s.localOnly[sessionID] = struct{}{}
  127. slog.Warn("xai oauth session Redis write failed; using process-local fallback", "error", remoteErr)
  128. } else {
  129. delete(s.localOnly, sessionID)
  130. }
  131. }
  132. func (s *SessionStore) Get(sessionID string) (*OAuthSession, bool) {
  133. if s.isLocalOnly(sessionID) {
  134. return s.getMemory(sessionID)
  135. }
  136. if s != nil && s.remote != nil {
  137. var dto oauthSessionDTO
  138. ok, err := s.remote.Get(context.Background(), sessionID, &dto)
  139. if err != nil || !ok || time.Since(dto.CreatedAt) > SessionTTL {
  140. return nil, false
  141. }
  142. session := &OAuthSession{
  143. State: dto.State, CodeVerifier: dto.CodeVerifier, CodeChallenge: dto.CodeChallenge,
  144. ClientID: dto.ClientID, Scope: dto.Scope, ProxyURL: dto.ProxyURL,
  145. RedirectURI: dto.RedirectURI, CreatedAt: dto.CreatedAt,
  146. }
  147. s.mu.Lock()
  148. s.sessions[sessionID] = session
  149. s.mu.Unlock()
  150. return session, true
  151. }
  152. return s.getMemory(sessionID)
  153. }
  154. func (s *SessionStore) getMemory(sessionID string) (*OAuthSession, bool) {
  155. s.mu.RLock()
  156. defer s.mu.RUnlock()
  157. session, ok := s.sessions[sessionID]
  158. if !ok {
  159. return nil, false
  160. }
  161. if time.Since(session.CreatedAt) > SessionTTL {
  162. return nil, false
  163. }
  164. return session, true
  165. }
  166. func (s *SessionStore) Delete(sessionID string) {
  167. if s != nil && s.remote != nil {
  168. _ = s.remote.Delete(context.Background(), sessionID)
  169. }
  170. s.mu.Lock()
  171. defer s.mu.Unlock()
  172. delete(s.sessions, sessionID)
  173. delete(s.localOnly, sessionID)
  174. }
  175. func (s *SessionStore) TryConsumeSession(sessionID string) bool {
  176. if s == nil {
  177. return false
  178. }
  179. if s.isLocalOnly(sessionID) {
  180. return s.tryConsumeMemory(sessionID)
  181. }
  182. if s.remote != nil {
  183. ok, err := s.remote.TryConsume(context.Background(), sessionID)
  184. return err == nil && ok
  185. }
  186. return s.tryConsumeMemory(sessionID)
  187. }
  188. func (s *SessionStore) isLocalOnly(sessionID string) bool {
  189. s.mu.RLock()
  190. defer s.mu.RUnlock()
  191. _, ok := s.localOnly[sessionID]
  192. return ok
  193. }
  194. func (s *SessionStore) tryConsumeMemory(sessionID string) bool {
  195. session, ok := s.getMemory(sessionID)
  196. return ok && session.TryConsume()
  197. }
  198. func (s *SessionStore) Stop() {
  199. s.stopOnce.Do(func() {
  200. close(s.stopCh)
  201. })
  202. }
  203. func (s *SessionStore) cleanup() {
  204. ticker := time.NewTicker(5 * time.Minute)
  205. defer ticker.Stop()
  206. for {
  207. select {
  208. case <-s.stopCh:
  209. return
  210. case <-ticker.C:
  211. s.mu.Lock()
  212. for id, session := range s.sessions {
  213. if time.Since(session.CreatedAt) > SessionTTL {
  214. delete(s.sessions, id)
  215. delete(s.localOnly, id)
  216. }
  217. }
  218. s.mu.Unlock()
  219. }
  220. }
  221. }
  222. func EffectiveAuthorizeURL() string {
  223. return envOrDefault(EnvAuthorizeURL, DefaultAuthorizeURL)
  224. }
  225. func ValidatedAuthorizeURL() (string, error) {
  226. return ValidateOAuthEndpointURL(EffectiveAuthorizeURL())
  227. }
  228. func EffectiveTokenURL() string {
  229. return envOrDefault(EnvTokenURL, DefaultTokenURL)
  230. }
  231. func ValidatedTokenURL() (string, error) {
  232. return ValidateOAuthEndpointURL(EffectiveTokenURL())
  233. }
  234. func EffectiveClientID() string {
  235. return envOrDefault(EnvClientID, DefaultClientID)
  236. }
  237. func EffectiveScope() string {
  238. return envOrDefault(EnvScope, DefaultScope)
  239. }
  240. func EffectiveRedirectURI(override string) string {
  241. if trimmed := strings.TrimSpace(override); trimmed != "" {
  242. return trimmed
  243. }
  244. return envOrDefault(EnvRedirectURI, DefaultRedirectURI)
  245. }
  246. func EffectiveBaseURL(override string) string {
  247. if trimmed := strings.TrimSpace(override); trimmed != "" {
  248. return strings.TrimRight(trimmed, "/")
  249. }
  250. return strings.TrimRight(envOrDefault(EnvBaseURL, DefaultBaseURL), "/")
  251. }
  252. func ValidatedBaseURL(override string) (string, error) {
  253. return ValidateBaseURL(EffectiveBaseURL(override))
  254. }
  255. // BaseURLValidator applies the caller's outbound URL trust policy before xAI
  256. // endpoint paths are appended. The service layer uses this for API-key accounts
  257. // so the global security.url_allowlist policy remains the single source of
  258. // truth; OAuth callers keep using the strict trusted-host validator.
  259. type BaseURLValidator func(string) (string, error)
  260. func validatedBaseURLWithValidator(override string, validator BaseURLValidator) (string, error) {
  261. if validator == nil {
  262. return ValidatedBaseURL(override)
  263. }
  264. raw := EffectiveBaseURL(override)
  265. validated, err := validator(raw)
  266. if err != nil {
  267. return "", err
  268. }
  269. return normalizeKnownBaseURLPath(validated)
  270. }
  271. type RuntimeSanityCheck struct {
  272. Value string `json:"value"`
  273. Valid bool `json:"valid"`
  274. Error string `json:"error,omitempty"`
  275. IsDefault bool `json:"is_default,omitempty"`
  276. }
  277. type RuntimeSanityReport struct {
  278. BaseURL RuntimeSanityCheck `json:"base_url"`
  279. OAuthAuthorizeURL RuntimeSanityCheck `json:"oauth_authorize_url"`
  280. OAuthTokenURL RuntimeSanityCheck `json:"oauth_token_url"`
  281. OAuthRedirectURI RuntimeSanityCheck `json:"oauth_redirect_uri"`
  282. UnsafeURLOverrides bool `json:"unsafe_url_overrides"`
  283. UnsafeHighConcurrency bool `json:"unsafe_high_concurrency"`
  284. PublicGatewayScope string `json:"public_gateway_scope"`
  285. ProxyPolicy string `json:"proxy_policy"`
  286. }
  287. func RuntimeSanity() RuntimeSanityReport {
  288. return RuntimeSanityReport{
  289. BaseURL: runtimeSanityCheck(EffectiveBaseURL(""), EnvBaseURL, ValidatedBaseURL),
  290. OAuthAuthorizeURL: runtimeSanityCheck(EffectiveAuthorizeURL(), EnvAuthorizeURL, func(string) (string, error) { return ValidatedAuthorizeURL() }),
  291. OAuthTokenURL: runtimeSanityCheck(EffectiveTokenURL(), EnvTokenURL, func(string) (string, error) { return ValidatedTokenURL() }),
  292. OAuthRedirectURI: runtimeSanityCheck(EffectiveRedirectURI(""), EnvRedirectURI, validateRedirectURI),
  293. UnsafeURLOverrides: AllowUnsafeURLOverrides(),
  294. UnsafeHighConcurrency: AllowUnsafeHighConcurrency(),
  295. PublicGatewayScope: "responses_only",
  296. ProxyPolicy: "account_proxy_optional; OAuth URLs use trusted-host allowlists; API-key base URLs require public HTTPS unless unsafe overrides are enabled",
  297. }
  298. }
  299. func runtimeSanityCheck(value string, envKey string, validate func(string) (string, error)) RuntimeSanityCheck {
  300. normalized, err := validate(value)
  301. check := RuntimeSanityCheck{
  302. Value: sanitizeRuntimeURLValue(normalized),
  303. Valid: err == nil,
  304. IsDefault: strings.TrimSpace(os.Getenv(envKey)) == "",
  305. }
  306. if err != nil {
  307. check.Value = sanitizeRuntimeURLValue(value)
  308. check.Error = sanitizeRuntimeError(err.Error(), value)
  309. }
  310. return check
  311. }
  312. func validateRedirectURI(raw string) (string, error) {
  313. return urlvalidator.ValidateURLFormat(raw, true)
  314. }
  315. func sanitizeRuntimeURLValue(raw string) string {
  316. trimmed := strings.TrimSpace(raw)
  317. if trimmed == "" {
  318. return ""
  319. }
  320. parsed, err := url.Parse(trimmed)
  321. if err != nil || parsed.Scheme == "" || parsed.Host == "" {
  322. return trimmed
  323. }
  324. parsed.User = nil
  325. parsed.RawQuery = ""
  326. parsed.Fragment = ""
  327. return strings.TrimRight(parsed.String(), "/")
  328. }
  329. func sanitizeRuntimeError(rawErr string, rawValue string) string {
  330. redacted := logredact.RedactText(rawErr)
  331. trimmedValue := strings.TrimSpace(rawValue)
  332. if trimmedValue == "" {
  333. return redacted
  334. }
  335. sanitizedValue := sanitizeRuntimeURLValue(trimmedValue)
  336. redacted = strings.ReplaceAll(redacted, trimmedValue, sanitizedValue)
  337. redacted = strings.ReplaceAll(redacted, logredact.RedactText(trimmedValue), sanitizedValue)
  338. return redacted
  339. }
  340. func ValidateOAuthEndpointURL(raw string) (string, error) {
  341. if AllowUnsafeURLOverrides() {
  342. return urlvalidator.ValidateURLFormat(raw, true)
  343. }
  344. return urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{
  345. AllowedHosts: oauthEndpointAllowedHosts,
  346. RequireAllowlist: true,
  347. AllowPrivate: false,
  348. })
  349. }
  350. func ValidateBaseURL(raw string) (string, error) {
  351. if AllowUnsafeURLOverrides() {
  352. return urlvalidator.ValidateURLFormat(raw, true)
  353. }
  354. normalized, err := urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{
  355. AllowPrivate: false,
  356. })
  357. if err != nil {
  358. return "", err
  359. }
  360. return normalizeKnownBaseURLPath(normalized)
  361. }
  362. func ValidateTrustedBaseURL(raw string) (string, error) {
  363. if AllowUnsafeURLOverrides() {
  364. return urlvalidator.ValidateURLFormat(raw, true)
  365. }
  366. normalized, err := urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{
  367. AllowedHosts: baseURLAllowedHosts,
  368. RequireAllowlist: true,
  369. AllowPrivate: false,
  370. })
  371. if err != nil {
  372. return "", err
  373. }
  374. return normalizeKnownBaseURLPath(normalized)
  375. }
  376. // normalizeKnownBaseURLPath 规范化 base URL 的 path 部分:
  377. // - 官方主机固定使用 /v1 前缀(空 path 自动补齐,其余 path 拒绝);
  378. // - 其他主机保留管理员配置的任意 path 前缀(第三方转发地址常见
  379. // /xxx/v1 之类的路由前缀),空 path 仍按惯例补 /v1。
  380. //
  381. // 所有主机统一禁止 userinfo/query/fragment,并去除尾部斜杠。
  382. func normalizeKnownBaseURLPath(raw string) (string, error) {
  383. parsed, err := url.Parse(raw)
  384. if err != nil || parsed.Scheme == "" || parsed.Host == "" {
  385. return "", errors.New("invalid base URL")
  386. }
  387. if parsed.User != nil {
  388. return "", errors.New("base URL must not include userinfo")
  389. }
  390. if parsed.ForceQuery || parsed.RawQuery != "" {
  391. return "", errors.New("base URL must not include a query")
  392. }
  393. if parsed.Fragment != "" {
  394. return "", errors.New("base URL must not include a fragment")
  395. }
  396. path := strings.TrimRight(parsed.Path, "/")
  397. if path == "" {
  398. parsed.Path = "/v1"
  399. parsed.RawPath = ""
  400. return strings.TrimRight(parsed.String(), "/"), nil
  401. }
  402. if path != "/v1" && IsOfficialBaseURLHost(parsed.Hostname()) {
  403. return "", fmt.Errorf("base URL path must be /v1")
  404. }
  405. parsed.Path = path
  406. parsed.RawPath = ""
  407. return strings.TrimRight(parsed.String(), "/"), nil
  408. }
  409. // IsOfficialBaseURLHost 报告 host 是否属于官方 API / 区域 API / CLI 网关主机。
  410. func IsOfficialBaseURLHost(host string) bool {
  411. host = strings.ToLower(strings.TrimSpace(host))
  412. for _, allowed := range baseURLAllowedHosts {
  413. if strings.HasPrefix(allowed, "*.") {
  414. suffix := strings.TrimPrefix(allowed, "*.")
  415. if host == suffix || strings.HasSuffix(host, "."+suffix) {
  416. return true
  417. }
  418. continue
  419. }
  420. if host == allowed {
  421. return true
  422. }
  423. }
  424. return false
  425. }
  426. // IsParseableBaseURL 报告 raw 是否能解析出 host。
  427. // 供读取路径判定存量脏数据:无法解析的值应回落默认端点,而不是把流量发往未定义目标。
  428. func IsParseableBaseURL(raw string) bool {
  429. trimmed := strings.TrimSpace(raw)
  430. if trimmed == "" {
  431. return false
  432. }
  433. parsed, err := url.Parse(trimmed)
  434. return err == nil && parsed.Host != ""
  435. }
  436. // IsOfficialBaseURL 报告 raw 是否指向官方主机(api.x.ai / *.api.x.ai 区域端点 / CLI 网关),
  437. // 容忍存量凭证中的历史变体(大小写、显式 443 端口、百分号编码 path 等)。
  438. // 无法解析的值一并视为官方,调用方据此回落默认端点而不是把流量发往未定义目标。
  439. func IsOfficialBaseURL(raw string) bool {
  440. trimmed := strings.TrimSpace(raw)
  441. if trimmed == "" {
  442. return true
  443. }
  444. parsed, err := url.Parse(trimmed)
  445. if err != nil || parsed.Host == "" {
  446. return true
  447. }
  448. return IsOfficialBaseURLHost(parsed.Hostname())
  449. }
  450. func AllowUnsafeURLOverrides() bool {
  451. return envBool(EnvAllowUnsafeURLOverrides)
  452. }
  453. func AllowUnsafeHighConcurrency() bool {
  454. return envBool(EnvUnsafeAllowHighConcurrency)
  455. }
  456. func envOrDefault(key, fallback string) string {
  457. if value := strings.TrimSpace(os.Getenv(key)); value != "" {
  458. return value
  459. }
  460. return fallback
  461. }
  462. func envBool(key string) bool {
  463. switch strings.ToLower(strings.TrimSpace(os.Getenv(key))) {
  464. case "1", "true", "yes", "y", "on":
  465. return true
  466. default:
  467. return false
  468. }
  469. }
  470. func GenerateRandomBytes(n int) ([]byte, error) {
  471. b := make([]byte, n)
  472. if _, err := rand.Read(b); err != nil {
  473. return nil, err
  474. }
  475. return b, nil
  476. }
  477. func GenerateState() (string, error) {
  478. bytes, err := GenerateRandomBytes(32)
  479. if err != nil {
  480. return "", err
  481. }
  482. return hex.EncodeToString(bytes), nil
  483. }
  484. func GenerateNonce() (string, error) {
  485. bytes, err := GenerateRandomBytes(16)
  486. if err != nil {
  487. return "", err
  488. }
  489. return hex.EncodeToString(bytes), nil
  490. }
  491. func GenerateSessionID() (string, error) {
  492. bytes, err := GenerateRandomBytes(16)
  493. if err != nil {
  494. return "", err
  495. }
  496. return hex.EncodeToString(bytes), nil
  497. }
  498. func GenerateCodeVerifier() (string, error) {
  499. bytes, err := GenerateRandomBytes(32)
  500. if err != nil {
  501. return "", err
  502. }
  503. return base64URLEncode(bytes), nil
  504. }
  505. func GenerateCodeChallenge(verifier string) string {
  506. hash := sha256.Sum256([]byte(verifier))
  507. return base64URLEncode(hash[:])
  508. }
  509. func base64URLEncode(data []byte) string {
  510. return strings.TrimRight(base64.URLEncoding.EncodeToString(data), "=")
  511. }
  512. func BuildAuthorizationURL(state, codeChallenge, redirectURI, nonce string) (string, error) {
  513. redirectURI = EffectiveRedirectURI(redirectURI)
  514. authorizeURL, err := ValidatedAuthorizeURL()
  515. if err != nil {
  516. return "", fmt.Errorf("invalid authorize url: %w", err)
  517. }
  518. params := url.Values{}
  519. params.Set("response_type", "code")
  520. params.Set("client_id", EffectiveClientID())
  521. params.Set("redirect_uri", redirectURI)
  522. params.Set("scope", EffectiveScope())
  523. params.Set("state", state)
  524. params.Set("nonce", nonce)
  525. params.Set("code_challenge", codeChallenge)
  526. params.Set("code_challenge_method", "S256")
  527. params.Set("plan", "generic")
  528. params.Set("referrer", "sub2api")
  529. return fmt.Sprintf("%s?%s", authorizeURL, params.Encode()), nil
  530. }
  531. // AuthorizationInput is a parsed manual OAuth callback input.
  532. type AuthorizationInput struct {
  533. Code string
  534. State string
  535. RequiresState bool
  536. }
  537. // ParseAuthorizationInput accepts a full callback URL, query string, or bare code.
  538. func ParseAuthorizationInput(raw string) AuthorizationInput {
  539. trimmed := strings.TrimSpace(raw)
  540. if trimmed == "" {
  541. return AuthorizationInput{}
  542. }
  543. if parsed, err := url.Parse(trimmed); err == nil && parsed != nil {
  544. values := parsed.Query()
  545. if code := strings.TrimSpace(values.Get("code")); code != "" {
  546. return AuthorizationInput{
  547. Code: code,
  548. State: strings.TrimSpace(values.Get("state")),
  549. RequiresState: true,
  550. }
  551. }
  552. }
  553. queryCandidate := strings.TrimPrefix(trimmed, "?")
  554. if strings.Contains(queryCandidate, "=") {
  555. if values, err := url.ParseQuery(queryCandidate); err == nil {
  556. if code := strings.TrimSpace(values.Get("code")); code != "" {
  557. return AuthorizationInput{
  558. Code: code,
  559. State: strings.TrimSpace(values.Get("state")),
  560. RequiresState: true,
  561. }
  562. }
  563. }
  564. }
  565. return AuthorizationInput{Code: trimmed}
  566. }
  567. func BuildResponsesURL(baseURL string) (string, error) {
  568. return BuildResponsesURLWithValidator(baseURL, nil)
  569. }
  570. func BuildResponsesURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
  571. validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
  572. if err != nil {
  573. return "", fmt.Errorf("invalid base url: %w", err)
  574. }
  575. return validatedBaseURL + "/responses", nil
  576. }
  577. func BuildChatCompletionsURL(baseURL string) (string, error) {
  578. return BuildChatCompletionsURLWithValidator(baseURL, nil)
  579. }
  580. func BuildChatCompletionsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
  581. validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
  582. if err != nil {
  583. return "", fmt.Errorf("invalid base url: %w", err)
  584. }
  585. return validatedBaseURL + "/chat/completions", nil
  586. }
  587. func BuildImagesGenerationsURL(baseURL string) (string, error) {
  588. return BuildImagesGenerationsURLWithValidator(baseURL, nil)
  589. }
  590. func BuildImagesGenerationsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
  591. validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
  592. if err != nil {
  593. return "", fmt.Errorf("invalid base url: %w", err)
  594. }
  595. return validatedBaseURL + "/images/generations", nil
  596. }
  597. func BuildImagesEditsURL(baseURL string) (string, error) {
  598. return BuildImagesEditsURLWithValidator(baseURL, nil)
  599. }
  600. func BuildImagesEditsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
  601. validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
  602. if err != nil {
  603. return "", fmt.Errorf("invalid base url: %w", err)
  604. }
  605. return validatedBaseURL + "/images/edits", nil
  606. }
  607. func BuildVideosGenerationsURL(baseURL string) (string, error) {
  608. return BuildVideosGenerationsURLWithValidator(baseURL, nil)
  609. }
  610. func BuildVideosGenerationsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
  611. validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
  612. if err != nil {
  613. return "", fmt.Errorf("invalid base url: %w", err)
  614. }
  615. return validatedBaseURL + "/videos/generations", nil
  616. }
  617. func BuildVideosEditsURL(baseURL string) (string, error) {
  618. return BuildVideosEditsURLWithValidator(baseURL, nil)
  619. }
  620. func BuildVideosEditsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
  621. validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
  622. if err != nil {
  623. return "", fmt.Errorf("invalid base url: %w", err)
  624. }
  625. return validatedBaseURL + "/videos/edits", nil
  626. }
  627. func BuildVideosExtensionsURL(baseURL string) (string, error) {
  628. return BuildVideosExtensionsURLWithValidator(baseURL, nil)
  629. }
  630. func BuildVideosExtensionsURLWithValidator(baseURL string, validator BaseURLValidator) (string, error) {
  631. validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
  632. if err != nil {
  633. return "", fmt.Errorf("invalid base url: %w", err)
  634. }
  635. return validatedBaseURL + "/videos/extensions", nil
  636. }
  637. func BuildVideoURL(baseURL, requestID string) (string, error) {
  638. return BuildVideoURLWithValidator(baseURL, requestID, nil)
  639. }
  640. func BuildVideoURLWithValidator(baseURL, requestID string, validator BaseURLValidator) (string, error) {
  641. validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
  642. if err != nil {
  643. return "", fmt.Errorf("invalid base url: %w", err)
  644. }
  645. requestID = strings.TrimSpace(requestID)
  646. if requestID == "" {
  647. return "", fmt.Errorf("request id is required")
  648. }
  649. // requestID 由客户端提供并拼进上游 URL 的 path。PathEscape 之外再要求它不是
  650. // 纯点片段、不含控制字符,保证它只能是一个普通的路径片段。
  651. if requestID == "." || requestID == ".." || strings.ContainsAny(requestID, "\x00\r\n") {
  652. return "", fmt.Errorf("invalid request id")
  653. }
  654. return validatedBaseURL + "/videos/" + url.PathEscape(requestID), nil
  655. }
  656. // TokenResponse represents xAI OAuth token responses.
  657. type TokenResponse struct {
  658. AccessToken string `json:"access_token"`
  659. RefreshToken string `json:"refresh_token,omitempty"`
  660. IDToken string `json:"id_token,omitempty"`
  661. TokenType string `json:"token_type,omitempty"`
  662. ExpiresIn int64 `json:"expires_in,omitempty"`
  663. Scope string `json:"scope,omitempty"`
  664. }