mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-27 00:04:32 +00:00
feat(grok): 支持账号级自定义上游地址与请求头覆写
将账号级请求头覆写从 anthropic/openai 的 api_key 账号扩展到 Grok 的 api_key 与 oauth 账号,并放开 Grok OAuth 账号的自定义上游地址(仅作用于 转发端点,授权与 token 刷新链路不变)。 后端: - IsHeaderOverrideEligible 扩展到 Grok(api_key+oauth),禁止名单新增 x-grok-conv-id(逐请求会话路由头)。 - 所有 Grok 上游请求路径接线 ApplyHeaderOverrides(Responses/Chat 桥/ 媒体/配额探测/billing 探测/连通性测试),统一置于内置默认头之后。 - GetGrokBaseURL/GetGrokMediaBaseURL 放开 OAuth 自定义地址:官方地址 视同未定制回落官方网关,仅显式第三方 host 改发转发流量。 - 非官方 host 允许任意 path 前缀,官方 host 仍强制 /v1。 - OAuth base_url 校验按 host 判定官方/自定义,自定义 host 恒受运营方 URL 策略约束,不受 XAI_ALLOW_UNSAFE_URL_OVERRIDES 调试开关放宽。 - 给 GrokQuotaService 注入 config,使配额/billing 探测与转发共用同一 URL 策略。 前端: - Edit 模态为 Grok OAuth 账号新增「自定义上游地址」开关。 - 请求头表单新增 JSON 快速导入与按 JSON 一键复制(复用 useClipboard, 兼容非安全上下文 HTTP 页面)。 - 门控改为平台×类型判定,Bulk 批量 base_url 增加格式校验,zh/en 文案同步。
This commit is contained in:
@@ -183,7 +183,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
claudeUsageFetcher := repository.NewClaudeUsageFetcher(httpUpstream)
|
||||
antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository)
|
||||
grokQuotaFetcher := service.NewGrokQuotaFetcher()
|
||||
grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream, usageLogRepository)
|
||||
grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream, configConfig, usageLogRepository)
|
||||
openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory, openAIGatewayService)
|
||||
usageCache := service.NewUsageCache()
|
||||
accountUsageService := service.ProvideAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, grokQuotaService, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService, openAIGatewayService)
|
||||
|
||||
@@ -113,7 +113,7 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
|
||||
},
|
||||
}}
|
||||
upstream := &grokQuotaHandlerUpstream{}
|
||||
quotaService := service.NewGrokQuotaService(repo, nil, service.NewGrokTokenProvider(repo, nil), upstream)
|
||||
quotaService := service.NewGrokQuotaService(repo, nil, service.NewGrokTokenProvider(repo, nil), upstream, nil)
|
||||
handler := NewGrokOAuthHandler(nil, nil, quotaService, nil)
|
||||
|
||||
router := gin.New()
|
||||
@@ -151,7 +151,7 @@ func TestGrokOAuthHandlerResetQuotaReturnsUnsupported(t *testing.T) {
|
||||
Platform: service.PlatformGrok,
|
||||
Type: service.AccountTypeOAuth,
|
||||
}}
|
||||
quotaService := service.NewGrokQuotaService(repo, nil, nil, nil)
|
||||
quotaService := service.NewGrokQuotaService(repo, nil, nil, nil, nil)
|
||||
handler := NewGrokOAuthHandler(nil, nil, quotaService, nil)
|
||||
|
||||
router := gin.New()
|
||||
|
||||
@@ -94,6 +94,21 @@ func BuildBillingURL(formatCredits bool) string {
|
||||
return base + BillingMonthlyPath
|
||||
}
|
||||
|
||||
// BuildBillingURLWithValidator builds the weekly or monthly billing URL against
|
||||
// the caller-resolved base URL, applying the caller's outbound URL trust policy
|
||||
// first. Accounts forwarding through a custom upstream keep their billing
|
||||
// probes on the same upstream.
|
||||
func BuildBillingURLWithValidator(baseURL string, formatCredits bool, validator BaseURLValidator) (string, error) {
|
||||
validatedBaseURL, err := validatedBaseURLWithValidator(baseURL, validator)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid base url: %w", err)
|
||||
}
|
||||
if formatCredits {
|
||||
return validatedBaseURL + BillingWeeklyPath, nil
|
||||
}
|
||||
return validatedBaseURL + BillingMonthlyPath, nil
|
||||
}
|
||||
|
||||
// ApplyCLIBillingHeaders sets Authorization + CLI identity headers for billing GETs.
|
||||
func ApplyCLIBillingHeaders(req *http.Request, accessToken string) {
|
||||
if req == nil {
|
||||
|
||||
@@ -14,6 +14,20 @@ func TestBuildBillingURL(t *testing.T) {
|
||||
require.Equal(t, "https://cli-chat-proxy.grok.com/v1/billing", BuildBillingURL(false))
|
||||
}
|
||||
|
||||
func TestBuildBillingURLWithValidator(t *testing.T) {
|
||||
t.Parallel()
|
||||
weeklyURL, err := BuildBillingURLWithValidator(DefaultCLIBaseURL, true, ValidateTrustedBaseURL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://cli-chat-proxy.grok.com/v1/billing?format=credits", weeklyURL)
|
||||
|
||||
monthlyURL, err := BuildBillingURLWithValidator("https://relay.example.test/v1", false, ValidateBaseURL)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://relay.example.test/v1/billing", monthlyURL)
|
||||
|
||||
_, err = BuildBillingURLWithValidator("https://relay.example.test/v1", true, ValidateTrustedBaseURL)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestApplyCLIBillingHeaders(t *testing.T) {
|
||||
t.Parallel()
|
||||
req, err := http.NewRequest(http.MethodGet, BuildBillingURL(true), nil)
|
||||
|
||||
@@ -298,6 +298,12 @@ func ValidateTrustedBaseURL(raw string) (string, error) {
|
||||
return normalizeKnownBaseURLPath(normalized)
|
||||
}
|
||||
|
||||
// normalizeKnownBaseURLPath 规范化 base URL 的 path 部分:
|
||||
// - 官方主机固定使用 /v1 前缀(空 path 自动补齐,其余 path 拒绝);
|
||||
// - 其他主机保留管理员配置的任意 path 前缀(第三方转发地址常见
|
||||
// /xxx/v1 之类的路由前缀),空 path 仍按惯例补 /v1。
|
||||
//
|
||||
// 所有主机统一禁止 userinfo/query/fragment,并去除尾部斜杠。
|
||||
func normalizeKnownBaseURLPath(raw string) (string, error) {
|
||||
parsed, err := url.Parse(raw)
|
||||
if err != nil || parsed.Scheme == "" || parsed.Host == "" {
|
||||
@@ -318,7 +324,7 @@ func normalizeKnownBaseURLPath(raw string) (string, error) {
|
||||
parsed.RawPath = ""
|
||||
return strings.TrimRight(parsed.String(), "/"), nil
|
||||
}
|
||||
if path != "/v1" {
|
||||
if path != "/v1" && IsOfficialBaseURLHost(parsed.Hostname()) {
|
||||
return "", fmt.Errorf("base URL path must be /v1")
|
||||
}
|
||||
parsed.Path = path
|
||||
@@ -326,6 +332,32 @@ func normalizeKnownBaseURLPath(raw string) (string, error) {
|
||||
return strings.TrimRight(parsed.String(), "/"), nil
|
||||
}
|
||||
|
||||
// IsOfficialBaseURLHost 报告 host 是否属于官方 API / CLI 网关主机。
|
||||
func IsOfficialBaseURLHost(host string) bool {
|
||||
host = strings.ToLower(strings.TrimSpace(host))
|
||||
for _, allowed := range baseURLAllowedHosts {
|
||||
if host == allowed {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// IsOfficialBaseURL 报告 raw 是否指向官方主机(api.x.ai 或 CLI 网关),
|
||||
// 容忍存量凭证中的历史变体(大小写、显式 443 端口、百分号编码 path 等)。
|
||||
// 无法解析的值一并视为官方,调用方据此回落默认端点而不是把流量发往未定义目标。
|
||||
func IsOfficialBaseURL(raw string) bool {
|
||||
trimmed := strings.TrimSpace(raw)
|
||||
if trimmed == "" {
|
||||
return true
|
||||
}
|
||||
parsed, err := url.Parse(trimmed)
|
||||
if err != nil || parsed.Host == "" {
|
||||
return true
|
||||
}
|
||||
return IsOfficialBaseURLHost(parsed.Hostname())
|
||||
}
|
||||
|
||||
func AllowUnsafeURLOverrides() bool {
|
||||
return envBool(EnvAllowUnsafeURLOverrides)
|
||||
}
|
||||
|
||||
@@ -166,6 +166,57 @@ func TestValidateBaseURLAllowsPublicThirdPartyGrokAPI(t *testing.T) {
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestValidateBaseURLPathPrefixPolicy(t *testing.T) {
|
||||
// 非官方主机保留管理员配置的任意 path 前缀。
|
||||
prefixed, err := ValidateBaseURL("https://relay.example.test/xai/v1/")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://relay.example.test/xai/v1", prefixed)
|
||||
|
||||
deepPrefixed, err := ValidateBaseURL("https://relay.example.test/tenant-a/proxy")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://relay.example.test/tenant-a/proxy", deepPrefixed)
|
||||
|
||||
// 空 path 仍按惯例补 /v1,保持既有配置兼容。
|
||||
rootOnly, err := ValidateBaseURL("https://relay.example.test")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://relay.example.test/v1", rootOnly)
|
||||
|
||||
// 官方主机固定 /v1 前缀。
|
||||
_, err = ValidateBaseURL("https://api.x.ai/xai/v1")
|
||||
require.Error(t, err)
|
||||
_, err = ValidateBaseURL("https://cli-chat-proxy.grok.com/other")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestIsOfficialBaseURL(t *testing.T) {
|
||||
official := []string{
|
||||
"",
|
||||
" ",
|
||||
DefaultBaseURL,
|
||||
DefaultCLIBaseURL,
|
||||
"https://api.x.ai",
|
||||
"HTTPS://API.X.AI:443/",
|
||||
"https://api.x.ai:0443/v1",
|
||||
"https://api.x.ai/%76%31",
|
||||
"https://api.x.ai:8443/v1",
|
||||
"HTTPS://CLI-CHAT-PROXY.GROK.COM:443/%76%31/",
|
||||
"::invalid::url", // 无法解析的值按官方处理,回落默认端点
|
||||
}
|
||||
for _, raw := range official {
|
||||
require.True(t, IsOfficialBaseURL(raw), "expected official: %q", raw)
|
||||
}
|
||||
|
||||
custom := []string{
|
||||
"https://relay.example.test/v1",
|
||||
"https://relay.example.test/xai/v1",
|
||||
"http://relay.example.test/v1",
|
||||
"https://grok.com.evil.example.test/v1",
|
||||
}
|
||||
for _, raw := range custom {
|
||||
require.False(t, IsOfficialBaseURL(raw), "expected custom: %q", raw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateBaseURLsRejectEmptyQueryDelimiter(t *testing.T) {
|
||||
_, err := ValidateBaseURL("https://grok.example.test/v1?")
|
||||
require.Error(t, err)
|
||||
|
||||
@@ -1263,17 +1263,25 @@ func (a *Account) GetOpenAIRefreshToken() string {
|
||||
// GetGrokBaseURL selects the upstream used by Grok text and Responses traffic.
|
||||
// Grok media traffic has a different transport contract and must use
|
||||
// GetGrokMediaBaseURL instead.
|
||||
//
|
||||
// The stored base_url only rewrites forwarding endpoints. Credential lifecycle
|
||||
// traffic (OAuth authorization and token refresh) always uses the official
|
||||
// auth endpoints regardless of this value.
|
||||
func (a *Account) GetGrokBaseURL() string {
|
||||
if !a.IsGrok() {
|
||||
return ""
|
||||
}
|
||||
baseURL := strings.TrimSpace(a.GetCredential("base_url"))
|
||||
if a.IsGrokOAuth() {
|
||||
// OAuth bearer credentials are subscription credentials and may only be
|
||||
// sent to the supported CLI gateway. Stored base_url values and unsafe
|
||||
// development overrides apply exclusively to API-key accounts.
|
||||
return xai.DefaultCLIBaseURL
|
||||
// Subscription traffic defaults to the supported CLI gateway. Stored
|
||||
// official-host values (written by credential creation/refresh, or
|
||||
// legacy variants) mean "not customized"; only an explicit custom-host
|
||||
// forwarding address redirects traffic.
|
||||
if baseURL == "" || xai.IsOfficialBaseURL(baseURL) {
|
||||
return xai.DefaultCLIBaseURL
|
||||
}
|
||||
return baseURL
|
||||
}
|
||||
baseURL := a.GetCredential("base_url")
|
||||
if baseURL != "" {
|
||||
return baseURL
|
||||
}
|
||||
@@ -1281,17 +1289,12 @@ func (a *Account) GetGrokBaseURL() string {
|
||||
}
|
||||
|
||||
// GetGrokMediaBaseURL selects the upstream used by Grok Imagine APIs.
|
||||
//
|
||||
// OAuth media credentials have the same trust boundary as OAuth text traffic:
|
||||
// they are pinned to the supported CLI gateway even for large request bodies.
|
||||
// API-key accounts retain their configured public/custom upstream behavior.
|
||||
// It currently resolves the same way as text traffic; the separate accessor
|
||||
// preserves the media/text distinction at call sites.
|
||||
func (a *Account) GetGrokMediaBaseURL() string {
|
||||
if !a.IsGrok() {
|
||||
return ""
|
||||
}
|
||||
if a.IsGrokOAuth() {
|
||||
return xai.DefaultCLIBaseURL
|
||||
}
|
||||
return a.GetGrokBaseURL()
|
||||
}
|
||||
|
||||
|
||||
@@ -266,7 +266,7 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) {
|
||||
expected: xai.DefaultCLIBaseURL,
|
||||
},
|
||||
{
|
||||
name: "oauth explicit custom base_url stays pinned to CLI proxy by default",
|
||||
name: "oauth explicit custom base_url redirects forwarding traffic",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformGrok,
|
||||
@@ -274,7 +274,18 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) {
|
||||
"base_url": "https://custom.example.com/v1",
|
||||
},
|
||||
},
|
||||
expected: xai.DefaultCLIBaseURL,
|
||||
expected: "https://custom.example.com/v1",
|
||||
},
|
||||
{
|
||||
name: "oauth custom base_url with path prefix redirects forwarding traffic",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformGrok,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "https://relay.example.com/xai/v1",
|
||||
},
|
||||
},
|
||||
expected: "https://relay.example.com/xai/v1",
|
||||
},
|
||||
{
|
||||
name: "API key without base_url uses official credit-backed API",
|
||||
@@ -294,7 +305,7 @@ func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetGrokBaseURLPinsOAuthWhenUnsafeOverridesEnabled(t *testing.T) {
|
||||
func TestGetGrokBaseURLHonorsOAuthCustomRegardlessOfUnsafeOverrides(t *testing.T) {
|
||||
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
|
||||
account := Account{
|
||||
Type: AccountTypeOAuth,
|
||||
@@ -304,7 +315,7 @@ func TestGetGrokBaseURLPinsOAuthWhenUnsafeOverridesEnabled(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
require.Equal(t, xai.DefaultCLIBaseURL, account.GetGrokBaseURL())
|
||||
require.Equal(t, "https://custom.example.com/v1", account.GetGrokBaseURL())
|
||||
}
|
||||
|
||||
func TestGetGrokMediaBaseURLPinsOAuthMediaToCLIProxy(t *testing.T) {
|
||||
@@ -356,7 +367,7 @@ func TestGetGrokMediaBaseURLPinsOAuthMediaToCLIProxy(t *testing.T) {
|
||||
expected: xai.DefaultCLIBaseURL,
|
||||
},
|
||||
{
|
||||
name: "oauth untrusted custom base_url is pinned to CLI proxy",
|
||||
name: "oauth custom base_url redirects media traffic",
|
||||
account: Account{
|
||||
Type: AccountTypeOAuth,
|
||||
Platform: PlatformGrok,
|
||||
@@ -364,7 +375,7 @@ func TestGetGrokMediaBaseURLPinsOAuthMediaToCLIProxy(t *testing.T) {
|
||||
"base_url": "https://custom.example.com/v1",
|
||||
},
|
||||
},
|
||||
expected: xai.DefaultCLIBaseURL,
|
||||
expected: "https://custom.example.com/v1",
|
||||
},
|
||||
{
|
||||
name: "API key retains its configured media API",
|
||||
@@ -395,7 +406,7 @@ func TestGetGrokMediaBaseURLPinsOAuthMediaToCLIProxy(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetGrokMediaBaseURLPinsOAuthWhenUnsafeOverridesEnabled(t *testing.T) {
|
||||
func TestGetGrokMediaBaseURLHonorsOAuthCustomRegardlessOfUnsafeOverrides(t *testing.T) {
|
||||
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
|
||||
account := Account{
|
||||
Type: AccountTypeOAuth,
|
||||
@@ -405,5 +416,5 @@ func TestGetGrokMediaBaseURLPinsOAuthWhenUnsafeOverridesEnabled(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
require.Equal(t, xai.DefaultCLIBaseURL, account.GetGrokMediaBaseURL())
|
||||
require.Equal(t, "https://custom.example.com/v1", account.GetGrokMediaBaseURL())
|
||||
}
|
||||
|
||||
@@ -9,7 +9,8 @@ import (
|
||||
"golang.org/x/net/http/httpguts"
|
||||
)
|
||||
|
||||
// 请求头覆写(header override):仅对 Anthropic / OpenAI 平台的 api_key 账号生效。
|
||||
// 请求头覆写(header override):对 Anthropic / OpenAI 平台的 api_key 账号
|
||||
// 以及 Grok 平台的 api_key / oauth 账号生效。
|
||||
// 管理员在账号上配置一组 header name -> value,转发到上游前用配置值覆盖同名请求头
|
||||
// (匹配不区分大小写);value 为空的条目视为"未填写",不参与覆盖。
|
||||
const (
|
||||
@@ -28,7 +29,8 @@ const (
|
||||
// - authorization/x-api-key/cookie 等:上游认证头由账号凭据统一注入,禁止通过覆写篡改或重新引入;
|
||||
// - accept-encoding:强制压缩会破坏网关对上游流式响应(SSE/usage)的解析;
|
||||
// - sec-websocket-*:WebSocket 握手头由拨号器管理(OpenAI WS 模式);
|
||||
// - session_id/x-claude-code-session-id 等:逐请求会话隔离头,固定值会造成会话串扰。
|
||||
// - session_id/x-claude-code-session-id/x-grok-conv-id 等:逐请求会话隔离头,
|
||||
// 固定值会造成会话串扰。
|
||||
var headerOverrideBlockedNames = map[string]struct{}{
|
||||
"host": {},
|
||||
"content-length": {},
|
||||
@@ -59,6 +61,7 @@ var headerOverrideBlockedNames = map[string]struct{}{
|
||||
"chatgpt-account-id": {},
|
||||
"x-claude-code-session-id": {},
|
||||
"x-client-request-id": {},
|
||||
"x-grok-conv-id": {},
|
||||
}
|
||||
|
||||
func isHeaderOverrideBlockedName(lowerName string) bool {
|
||||
@@ -67,12 +70,20 @@ func isHeaderOverrideBlockedName(lowerName string) bool {
|
||||
}
|
||||
|
||||
// IsHeaderOverrideEligible 报告账号类型是否支持请求头覆写。
|
||||
// 目前仅开放 Anthropic / OpenAI 两个平台的 api_key 账号。
|
||||
// Anthropic / OpenAI 仅开放 api_key 账号;Grok 额外开放 oauth 账号——
|
||||
// 订阅流量改发自定义转发地址时,通常需要补充中间层要求的准入头。
|
||||
func (a *Account) IsHeaderOverrideEligible() bool {
|
||||
if a == nil || a.Type != AccountTypeAPIKey {
|
||||
if a == nil {
|
||||
return false
|
||||
}
|
||||
switch a.Platform {
|
||||
case PlatformAnthropic, PlatformOpenAI:
|
||||
return a.Type == AccountTypeAPIKey
|
||||
case PlatformGrok:
|
||||
return a.Type == AccountTypeAPIKey || a.Type == AccountTypeOAuth
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return a.Platform == PlatformAnthropic || a.Platform == PlatformOpenAI
|
||||
}
|
||||
|
||||
// IsHeaderOverrideEnabled 报告账号是否启用了请求头覆写。
|
||||
|
||||
@@ -30,7 +30,8 @@ func TestIsHeaderOverrideEligible(t *testing.T) {
|
||||
{"anthropic oauth", PlatformAnthropic, AccountTypeOAuth, false},
|
||||
{"openai oauth", PlatformOpenAI, AccountTypeOAuth, false},
|
||||
{"gemini apikey", PlatformGemini, AccountTypeAPIKey, false},
|
||||
{"grok apikey", PlatformGrok, AccountTypeAPIKey, false},
|
||||
{"grok apikey", PlatformGrok, AccountTypeAPIKey, true},
|
||||
{"grok oauth", PlatformGrok, AccountTypeOAuth, true},
|
||||
{"antigravity apikey", PlatformAntigravity, AccountTypeAPIKey, false},
|
||||
{"anthropic bedrock", PlatformAnthropic, AccountTypeBedrock, false},
|
||||
}
|
||||
|
||||
@@ -758,6 +758,8 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *
|
||||
if account.IsGrokOAuth() {
|
||||
applyGrokCLIHeaders(req.Header)
|
||||
}
|
||||
// 连通性测试与真实转发保持同一套账号级请求头覆写。
|
||||
account.ApplyHeaderOverrides(req.Header)
|
||||
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
|
||||
@@ -329,6 +329,8 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
|
||||
}
|
||||
upstreamReq.Header.Set("Content-Type", contentType)
|
||||
}
|
||||
// 账号级请求头覆写最后应用,配置值优先于内置默认头。
|
||||
account.ApplyHeaderOverrides(upstreamReq.Header)
|
||||
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"golang.org/x/sync/singleflight"
|
||||
@@ -52,6 +53,7 @@ type GrokQuotaService struct {
|
||||
tokenProvider *GrokTokenProvider
|
||||
httpUpstream HTTPUpstream
|
||||
usageLogRepo UsageLogRepository
|
||||
cfg *config.Config
|
||||
probeFlight singleflight.Group
|
||||
}
|
||||
|
||||
@@ -60,6 +62,7 @@ func NewGrokQuotaService(
|
||||
proxyRepo ProxyRepository,
|
||||
tokenProvider *GrokTokenProvider,
|
||||
httpUpstream HTTPUpstream,
|
||||
cfg *config.Config,
|
||||
usageLogRepos ...UsageLogRepository,
|
||||
) *GrokQuotaService {
|
||||
var usageLogRepo UsageLogRepository
|
||||
@@ -72,6 +75,7 @@ func NewGrokQuotaService(
|
||||
tokenProvider: tokenProvider,
|
||||
httpUpstream: httpUpstream,
|
||||
usageLogRepo: usageLogRepo,
|
||||
cfg: cfg,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -135,7 +139,7 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_PROBE_BODY_ERROR", "failed to build probe body: %v", err)
|
||||
}
|
||||
targetURL, err := buildGrokResponsesURL(account, nil)
|
||||
targetURL, err := buildGrokResponsesURL(account, s.cfg)
|
||||
if err != nil {
|
||||
return nil, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_BASE_URL_INVALID", "invalid Grok base_url: %v", err)
|
||||
}
|
||||
@@ -152,6 +156,8 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr
|
||||
if account.IsGrokOAuth() {
|
||||
applyGrokCLIHeaders(req.Header)
|
||||
}
|
||||
// 探测请求与真实转发保持同一套账号级请求头覆写,避免探测通过但转发失败。
|
||||
account.ApplyHeaderOverrides(req.Header)
|
||||
|
||||
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 1))
|
||||
if err != nil {
|
||||
@@ -306,11 +312,17 @@ func (s *GrokQuotaService) fetchBilling(
|
||||
proxyURL string,
|
||||
weekly bool,
|
||||
) (*xai.BillingSummary, int, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, xai.BuildBillingURL(weekly), nil)
|
||||
billingURL, err := buildGrokBillingURL(account, s.cfg, weekly)
|
||||
if err != nil {
|
||||
return nil, 0, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_BASE_URL_INVALID", "invalid Grok base_url: %v", err)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, billingURL, nil)
|
||||
if err != nil {
|
||||
return nil, 0, infraerrors.Newf(http.StatusInternalServerError, "GROK_QUOTA_PROBE_REQUEST_BUILD_FAILED", "failed to build billing request: %v", err)
|
||||
}
|
||||
xai.ApplyCLIBillingHeaders(req, token)
|
||||
// billing 探测与真实转发保持同一套账号级请求头覆写。
|
||||
account.ApplyHeaderOverrides(req.Header)
|
||||
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 2))
|
||||
if err != nil {
|
||||
return nil, 0, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed: %v", err)
|
||||
|
||||
@@ -227,7 +227,7 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) {
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
|
||||
}}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil)
|
||||
|
||||
result, err := svc.ProbeUsage(context.Background(), 42)
|
||||
require.NoError(t, err)
|
||||
@@ -269,7 +269,7 @@ func TestGrokQuotaServiceProbeUsageIgnoresAccountGrokMapping(t *testing.T) {
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
|
||||
}}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil)
|
||||
|
||||
result, err := svc.ProbeUsage(context.Background(), 47)
|
||||
require.NoError(t, err)
|
||||
@@ -292,7 +292,7 @@ func TestGrokQuotaServiceProbeUsageReportsProbeModelOnUpstreamError(t *testing.T
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(`{"code":"invalid-argument","error":"Model not found"}`)),
|
||||
}}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil)
|
||||
|
||||
_, err := svc.ProbeUsage(context.Background(), 48)
|
||||
require.Error(t, err)
|
||||
@@ -320,6 +320,7 @@ func TestGrokQuotaServiceProbeUsageRedactsUpstreamErrorBodyFromErrorAndLogs(t *t
|
||||
nil,
|
||||
NewGrokTokenProvider(repo, nil),
|
||||
upstream,
|
||||
nil,
|
||||
)
|
||||
|
||||
var logs bytes.Buffer
|
||||
@@ -365,7 +366,7 @@ func TestGrokQuotaServiceProbeUsageLoadsProxyWhenAccountEdgeMissing(t *testing.T
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
|
||||
}}
|
||||
svc := NewGrokQuotaService(repo, proxyRepo, NewGrokTokenProvider(repo, nil), upstream)
|
||||
svc := NewGrokQuotaService(repo, proxyRepo, NewGrokTokenProvider(repo, nil), upstream, nil)
|
||||
|
||||
_, err := svc.ProbeUsage(context.Background(), 46)
|
||||
require.NoError(t, err)
|
||||
@@ -392,7 +393,7 @@ func TestGrokQuotaServiceProbeUsageStoresNoHeadersState(t *testing.T) {
|
||||
Header: http.Header{},
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
|
||||
}}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil)
|
||||
|
||||
result, err := svc.ProbeUsage(context.Background(), 45)
|
||||
require.NoError(t, err)
|
||||
@@ -427,7 +428,7 @@ func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) {
|
||||
Header: http.Header{"Retry-After": []string{"45"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
|
||||
}}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil)
|
||||
|
||||
result, err := svc.ProbeUsage(context.Background(), 43)
|
||||
require.NoError(t, err)
|
||||
@@ -450,7 +451,7 @@ func TestGrokQuotaServiceQueryQuotaFreeFallsBackToGrok45(t *testing.T) {
|
||||
}}
|
||||
upstream := &grokHybridUpstream{}
|
||||
usageRepo := &grokQuotaUsageLogRepo{stats: &usagestats.AccountStats{Tokens: 1_000_000}}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, usageRepo)
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil, usageRepo)
|
||||
|
||||
result, err := svc.QueryQuota(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
@@ -492,7 +493,7 @@ func TestGrokQuotaServiceQueryQuotaPaidBillingSkipsActiveProbe(t *testing.T) {
|
||||
usagePercent := 25.0
|
||||
upstream := &grokHybridUpstream{weeklyUsagePercent: &usagePercent}
|
||||
usageRepo := &grokQuotaUsageLogRepo{stats: &usagestats.AccountStats{Tokens: 1_000_000}}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, usageRepo)
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil, usageRepo)
|
||||
|
||||
result, err := svc.QueryQuota(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
@@ -519,7 +520,7 @@ func TestGrokQuotaServiceQueryQuotaCustomPaidMonthlyLimitSkipsActiveProbe(t *tes
|
||||
}}
|
||||
monthlyLimit := 25_000.0
|
||||
upstream := &grokHybridUpstream{monthlyLimitCents: &monthlyLimit}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil)
|
||||
|
||||
result, err := svc.QueryQuota(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
@@ -652,7 +653,7 @@ func TestAccountUsageServiceGrokRefreshUsesBillingOnly(t *testing.T) {
|
||||
}}
|
||||
upstream := &grokHybridUpstream{}
|
||||
usageRepo := &grokQuotaUsageLogRepo{stats: &usagestats.AccountStats{Tokens: 750_000}}
|
||||
quotaService := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, usageRepo)
|
||||
quotaService := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil, usageRepo)
|
||||
usageService := &AccountUsageService{
|
||||
grokQuotaFetcher: NewGrokQuotaFetcher(),
|
||||
grokQuotaService: quotaService,
|
||||
@@ -688,7 +689,7 @@ func TestGrokQuotaServiceProbeFlightsDeduplicateBillingAndSeparateActive(t *test
|
||||
billingStarted := make(chan struct{})
|
||||
billingRelease := make(chan struct{})
|
||||
upstream := &grokHybridUpstream{billingStarted: billingStarted, billingRelease: billingRelease}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil)
|
||||
|
||||
type probeOutcome struct {
|
||||
result *GrokQuotaProbeResult
|
||||
@@ -745,7 +746,7 @@ func TestGrokQuotaServiceBilling429DoesNotPauseModelScheduling(t *testing.T) {
|
||||
billingStatus: http.StatusTooManyRequests,
|
||||
billingHeaders: http.Header{"Retry-After": []string{"45"}},
|
||||
}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil)
|
||||
|
||||
result, err := svc.ProbeBilling(context.Background(), account.ID)
|
||||
|
||||
@@ -765,7 +766,7 @@ func TestGrokQuotaServiceQueryQuotaFree429PersistsLimitAndKeepsBilling(t *testin
|
||||
activeStatus: http.StatusTooManyRequests,
|
||||
activeHeaders: http.Header{"Retry-After": []string{"45"}},
|
||||
}
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream)
|
||||
svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil)
|
||||
|
||||
result, err := svc.QueryQuota(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
@@ -791,7 +792,7 @@ func TestGrokQuotaServiceResetQuotaUnsupported(t *testing.T) {
|
||||
accountsByID: map[int64]*Account{44: account},
|
||||
},
|
||||
}
|
||||
svc := NewGrokQuotaService(repo, nil, nil, nil)
|
||||
svc := NewGrokQuotaService(repo, nil, nil, nil, nil)
|
||||
|
||||
_, err := svc.ResetQuota(context.Background(), 44)
|
||||
require.Error(t, err)
|
||||
|
||||
@@ -15,30 +15,48 @@ func grokBaseURLValidator(account *Account, cfg *config.Config) (xai.BaseURLVali
|
||||
}
|
||||
switch account.Type {
|
||||
case AccountTypeOAuth:
|
||||
// Subscription credentials are never governed by the operator's API-key
|
||||
// URL policy. They stay pinned to the supported CLI gateway.
|
||||
return redactedGrokBaseURLValidator(xai.ValidateTrustedBaseURL), nil
|
||||
case AccountTypeAPIKey:
|
||||
if cfg == nil {
|
||||
return redactedGrokBaseURLValidator(xai.ValidateBaseURL), nil
|
||||
}
|
||||
if !cfg.Security.URLAllowlist.Enabled {
|
||||
return redactedGrokBaseURLValidator(func(raw string) (string, error) {
|
||||
return urlvalidator.ValidateURLFormat(raw, cfg.Security.URLAllowlist.AllowInsecureHTTP)
|
||||
}), nil
|
||||
}
|
||||
// Official gateway hosts are always trusted and always usable, even when
|
||||
// the operator enables a restrictive URL allowlist. A custom forwarding
|
||||
// host is vetted by the same operator policy as API-key accounts.
|
||||
//
|
||||
// The official-vs-custom decision is made on the host, not via
|
||||
// ValidateTrustedBaseURL: that validator relaxes to accept-any under the
|
||||
// XAI_ALLOW_UNSAFE_URL_OVERRIDES debug switch, which must never let an
|
||||
// OAuth bearer token reach an arbitrary custom host.
|
||||
policyValidator := grokOperatorPolicyValidator(cfg)
|
||||
return redactedGrokBaseURLValidator(func(raw string) (string, error) {
|
||||
return urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{
|
||||
AllowedHosts: cfg.Security.URLAllowlist.UpstreamHosts,
|
||||
RequireAllowlist: true,
|
||||
AllowPrivate: cfg.Security.URLAllowlist.AllowPrivateHosts,
|
||||
})
|
||||
if xai.IsOfficialBaseURL(raw) {
|
||||
return xai.ValidateTrustedBaseURL(raw)
|
||||
}
|
||||
return policyValidator(raw)
|
||||
}), nil
|
||||
case AccountTypeAPIKey:
|
||||
return redactedGrokBaseURLValidator(grokOperatorPolicyValidator(cfg)), nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported grok account type: %s", account.Type)
|
||||
}
|
||||
}
|
||||
|
||||
// grokOperatorPolicyValidator 按全局出站 URL 安全策略校验自定义 base_url:
|
||||
// 白名单开启时强制 UpstreamHosts;关闭时仅做格式校验(HTTP 允许与否跟随配置)。
|
||||
func grokOperatorPolicyValidator(cfg *config.Config) xai.BaseURLValidator {
|
||||
if cfg == nil {
|
||||
return xai.ValidateBaseURL
|
||||
}
|
||||
if !cfg.Security.URLAllowlist.Enabled {
|
||||
return func(raw string) (string, error) {
|
||||
return urlvalidator.ValidateURLFormat(raw, cfg.Security.URLAllowlist.AllowInsecureHTTP)
|
||||
}
|
||||
}
|
||||
return func(raw string) (string, error) {
|
||||
return urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{
|
||||
AllowedHosts: cfg.Security.URLAllowlist.UpstreamHosts,
|
||||
RequireAllowlist: true,
|
||||
AllowPrivate: cfg.Security.URLAllowlist.AllowPrivateHosts,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func redactedGrokBaseURLValidator(validator xai.BaseURLValidator) xai.BaseURLValidator {
|
||||
return func(raw string) (string, error) {
|
||||
validated, err := validator(raw)
|
||||
@@ -65,6 +83,16 @@ func buildGrokChatCompletionsURL(account *Account, cfg *config.Config) (string,
|
||||
return xai.BuildChatCompletionsURLWithValidator(account.GetGrokBaseURL(), validator)
|
||||
}
|
||||
|
||||
// buildGrokBillingURL 解析 billing 探测端点:跟随账号的转发 base_url,
|
||||
// 未定制的账号仍指向官方 CLI 网关。
|
||||
func buildGrokBillingURL(account *Account, cfg *config.Config, weekly bool) (string, error) {
|
||||
validator, err := grokBaseURLValidator(account, cfg)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return xai.BuildBillingURLWithValidator(account.GetGrokBaseURL(), weekly, validator)
|
||||
}
|
||||
|
||||
func buildGrokMediaURL(account *Account, cfg *config.Config, endpoint GrokMediaEndpoint, requestID string) (string, error) {
|
||||
validator, err := grokBaseURLValidator(account, cfg)
|
||||
if err != nil {
|
||||
|
||||
@@ -104,19 +104,182 @@ func TestGrokAPIKeyURLPolicyRedactsMalformedConfiguredURL(t *testing.T) {
|
||||
require.NotContains(t, err.Error(), "secret")
|
||||
}
|
||||
|
||||
func TestGrokOAuthURLPolicyIgnoresAPIKeyOverrides(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "http://attacker.example.test/v1",
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.Enabled = false
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
|
||||
func TestGrokOAuthURLPolicy(t *testing.T) {
|
||||
t.Run("default CLI gateway always allowed under restrictive allowlist", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.Enabled = true
|
||||
cfg.Security.URLAllowlist.UpstreamHosts = []string{"other.example.test"}
|
||||
|
||||
target, err := buildGrokResponsesURL(account, cfg)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", target)
|
||||
target, err := buildGrokResponsesURL(account, cfg)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", target)
|
||||
})
|
||||
|
||||
t.Run("stored official-host variant stays on CLI gateway", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "HTTPS://API.X.AI:443/",
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
|
||||
target, err := buildGrokResponsesURL(account, cfg)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", target)
|
||||
})
|
||||
|
||||
t.Run("custom forwarding address follows operator policy", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "https://relay.example.test/v1",
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.Enabled = false
|
||||
|
||||
target, err := buildGrokResponsesURL(account, cfg)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://relay.example.test/v1/responses", target)
|
||||
})
|
||||
|
||||
t.Run("custom path prefix is preserved", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "https://relay.example.test/xai/v1",
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
|
||||
target, err := buildGrokResponsesURL(account, cfg)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://relay.example.test/xai/v1/responses", target)
|
||||
})
|
||||
|
||||
t.Run("custom forwarding address rejected by allowlist", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "https://relay.example.test/v1",
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.Enabled = true
|
||||
cfg.Security.URLAllowlist.UpstreamHosts = []string{"other.example.test"}
|
||||
|
||||
_, err := buildGrokResponsesURL(account, cfg)
|
||||
require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy")
|
||||
})
|
||||
|
||||
t.Run("insecure HTTP custom address requires operator opt-in", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "http://relay.example.test/v1",
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.Enabled = false
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = false
|
||||
|
||||
_, err := buildGrokResponsesURL(account, cfg)
|
||||
require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy")
|
||||
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
|
||||
target, err := buildGrokResponsesURL(account, cfg)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "http://relay.example.test/v1/responses", target)
|
||||
})
|
||||
|
||||
t.Run("unsafe override switch does not relax the operator allowlist for custom hosts", func(t *testing.T) {
|
||||
// XAI_ALLOW_UNSAFE_URL_OVERRIDES relaxes the trusted-host validator to
|
||||
// accept-any; a custom OAuth forwarding host must still be governed by
|
||||
// the operator allowlist so the bearer token cannot reach arbitrary hosts.
|
||||
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.Enabled = true
|
||||
cfg.Security.URLAllowlist.UpstreamHosts = []string{"cli-chat-proxy.grok.com"}
|
||||
|
||||
custom := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "http://10.0.0.1/v1",
|
||||
},
|
||||
}
|
||||
_, err := buildGrokResponsesURL(custom, cfg)
|
||||
require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy")
|
||||
|
||||
// The official gateway still resolves even under the restrictive allowlist.
|
||||
official := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{},
|
||||
}
|
||||
target, err := buildGrokResponsesURL(official, cfg)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", target)
|
||||
})
|
||||
}
|
||||
|
||||
func TestGrokBillingURLFollowsAccountBaseURL(t *testing.T) {
|
||||
t.Run("oauth default stays on CLI gateway", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{},
|
||||
}
|
||||
|
||||
weeklyURL, err := buildGrokBillingURL(account, nil, true)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/billing?format=credits", weeklyURL)
|
||||
|
||||
monthlyURL, err := buildGrokBillingURL(account, nil, false)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/billing", monthlyURL)
|
||||
})
|
||||
|
||||
t.Run("oauth custom forwarding address carries billing probes", func(t *testing.T) {
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "https://relay.example.test/v1",
|
||||
},
|
||||
}
|
||||
|
||||
weeklyURL, err := buildGrokBillingURL(account, nil, true)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://relay.example.test/v1/billing?format=credits", weeklyURL)
|
||||
})
|
||||
|
||||
t.Run("billing probe honors the operator allowlist like forwarding", func(t *testing.T) {
|
||||
// Probe paths must share the forwarding URL policy so a custom host the
|
||||
// allowlist rejects cannot receive the OAuth bearer via a billing probe.
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "https://relay.example.test/v1",
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{}
|
||||
cfg.Security.URLAllowlist.Enabled = true
|
||||
cfg.Security.URLAllowlist.UpstreamHosts = []string{"cli-chat-proxy.grok.com"}
|
||||
|
||||
_, err := buildGrokBillingURL(account, cfg, true)
|
||||
require.EqualError(t, err, "invalid base url: base URL rejected by URL security policy")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -195,14 +195,15 @@ func (s *OpenAIGatewayService) sendCCUpstreamRequest(
|
||||
upstreamReq.Header.Set("user-agent", userAgent)
|
||||
}
|
||||
|
||||
// 账号级请求头覆写(仅 openai api_key 账号启用时生效)
|
||||
account.ApplyHeaderOverrides(upstreamReq.Header)
|
||||
if account.Platform == PlatformGrok {
|
||||
if account.IsGrokOAuth() {
|
||||
applyGrokCLIHeaders(upstreamReq.Header)
|
||||
}
|
||||
applyGrokCacheHeaders(upstreamReq.Header, grokCacheIdentity)
|
||||
}
|
||||
// 账号级请求头覆写:放在所有内置默认头(含 Grok CLI 身份头)之后应用,
|
||||
// 使配置值获得除共享传输层强制头之外的最高优先级。
|
||||
account.ApplyHeaderOverrides(upstreamReq.Header)
|
||||
|
||||
proxyURL := ""
|
||||
if account.Proxy != nil {
|
||||
|
||||
@@ -769,6 +769,9 @@ func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Acc
|
||||
req.Header.Set("OpenAI-Beta", v)
|
||||
}
|
||||
}
|
||||
// 账号级请求头覆写最后应用,使配置值优先于上面的内置默认头;
|
||||
// 打到官方 CLI 网关时身份头仍由共享传输层最终强制。
|
||||
account.ApplyHeaderOverrides(req.Header)
|
||||
return req, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -259,7 +259,7 @@ func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T)
|
||||
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "isolated-cache-id", nil)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.MethodPost, req.Method)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", req.URL.String())
|
||||
require.Equal(t, "https://xai.test/v1/responses", req.URL.String())
|
||||
require.Equal(t, "Bearer access-token", req.Header.Get("Authorization"))
|
||||
require.Equal(t, "application/json", req.Header.Get("Content-Type"))
|
||||
require.Contains(t, req.Header.Get("Accept"), "text/event-stream")
|
||||
@@ -288,14 +288,14 @@ func TestBuildGrokResponsesRequestAllowsPublicAPIKeyBaseURLByDefault(t *testing.
|
||||
require.NotEqual(t, grokUpstreamUserAgent, req.Header.Get("User-Agent"))
|
||||
}
|
||||
|
||||
func TestBuildGrokResponsesRequestPinsOAuthCustomBaseURLByDefault(t *testing.T) {
|
||||
func TestBuildGrokResponsesRequestPinsOAuthOfficialVariantBaseURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "https://xai.test/v1",
|
||||
"base_url": "HTTPS://API.X.AI:443/",
|
||||
},
|
||||
}
|
||||
|
||||
@@ -304,6 +304,58 @@ func TestBuildGrokResponsesRequestPinsOAuthCustomBaseURLByDefault(t *testing.T)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", req.URL.String())
|
||||
}
|
||||
|
||||
func TestBuildGrokResponsesRequestAppliesHeaderOverridesLast(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"base_url": "https://relay.example.test/v1",
|
||||
"header_override_enabled": true,
|
||||
"header_overrides": map[string]any{
|
||||
"User-Agent": "relay-client/2.0",
|
||||
"X-Grok-Client-Version": "9.9.9",
|
||||
"X-Relay-Token": "relay-secret",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "conv-1", nil)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://relay.example.test/v1/responses", req.URL.String())
|
||||
// 覆写值优先于内置 CLI 身份头。名字不在 wire casing 映射中的覆写头
|
||||
// 以小写键直写(HTTP/2 线上语义),需按写入形态断言。
|
||||
require.Equal(t, "relay-client/2.0", req.Header.Get("User-Agent"))
|
||||
require.Equal(t, []string{"9.9.9"}, req.Header["x-grok-client-version"])
|
||||
require.Empty(t, req.Header.Get("X-Grok-Client-Version"))
|
||||
require.Equal(t, []string{"relay-secret"}, req.Header["x-relay-token"])
|
||||
// 会话路由头与认证头不受覆写影响。
|
||||
require.Equal(t, "conv-1", req.Header.Get(grokConversationIDHeader))
|
||||
require.Equal(t, "Bearer access-token", req.Header.Get("Authorization"))
|
||||
}
|
||||
|
||||
func TestBuildGrokResponsesRequestIgnoresBlockedHeaderOverrides(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
account := &Account{
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"header_override_enabled": true,
|
||||
"header_overrides": map[string]any{
|
||||
"Authorization": "Bearer stolen",
|
||||
"x-grok-conv-id": "pinned-conversation",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "conv-2", nil)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "Bearer api-key", req.Header.Get("Authorization"))
|
||||
require.Equal(t, "conv-2", req.Header.Get(grokConversationIDHeader))
|
||||
}
|
||||
|
||||
func TestGrokMediaGenerationGateCoversImagesAndVideo(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -2129,4 +2181,3 @@ func TestIsGrokImageGenerationModel(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -199,9 +199,10 @@ func ProvideGrokQuotaService(
|
||||
proxyRepo ProxyRepository,
|
||||
tokenProvider *GrokTokenProvider,
|
||||
httpUpstream HTTPUpstream,
|
||||
cfg *config.Config,
|
||||
usageLogRepo UsageLogRepository,
|
||||
) *GrokQuotaService {
|
||||
return NewGrokQuotaService(accountRepo, proxyRepo, tokenProvider, httpUpstream, usageLogRepo)
|
||||
return NewGrokQuotaService(accountRepo, proxyRepo, tokenProvider, httpUpstream, cfg, usageLogRepo)
|
||||
}
|
||||
|
||||
// ProvideGeminiTokenProvider creates GeminiTokenProvider with OAuthRefreshAPI injection
|
||||
|
||||
Reference in New Issue
Block a user