github/sub2api
4
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-10-21 03:58:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(grok): fetch validated signed video content

Browse Source
This commit is contained in:
Heatherm Huang
2026-07-18 17:49:09 +08:00
parent f2e7a55703
commit 95d27f2dce
2 changed files with 248 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/service/grok_media.go
+138 -22
View File
@@ -336,6 +336,9 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
if err != nil {
return nil, err
}
if endpoint == GrokMediaEndpointVideoContent {
return s.forwardGrokMediaVideoContent(ctx, c, account, token, requestID, startTime)
}
targetURL, err := buildGrokMediaURL(account, s.cfg, endpoint, requestID)
if err != nil {
return nil, err
@@ -366,16 +369,7 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
return nil, err
}
upstreamReq.Header.Set("Authorization", "Bearer "+token)
if endpoint == GrokMediaEndpointVideoContent {
upstreamReq.Header.Set("Accept", "*/*")
if c != nil {
if rangeHeader := strings.TrimSpace(c.GetHeader("Range")); rangeHeader != "" {
upstreamReq.Header.Set("Range", rangeHeader)
}
}
} else {
upstreamReq.Header.Set("Accept", "application/json")
}
upstreamReq.Header.Set("Accept", "application/json")
if account.IsGrokOAuth() && isGrokCLIProxyTarget(targetURL) {
applyGrokCLIHeaders(upstreamReq.Header)
}
@@ -403,22 +397,11 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
requestIDHeader := firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id"))
requestModel := requestInfo.Model
preserveContentRangeError := endpoint == GrokMediaEndpointVideoContent && resp.StatusCode == http.StatusRequestedRangeNotSatisfiable
if resp.StatusCode >= 400 && !preserveContentRangeError {
if resp.StatusCode >= 400 {
return s.handleGrokMediaErrorResponse(ctx, resp, c, account, requestIDHeader, requestModel)
}
s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode)
if endpoint == GrokMediaEndpointVideoContent {
if err := writeGrokMediaContentResponse(c, resp); err != nil {
return nil, err
}
return &OpenAIForwardResult{
RequestID: requestIDHeader,
ResponseHeaders: resp.Header.Clone(),
Duration: time.Since(startTime),
}, nil
}
respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError)
if err != nil {
return nil, err
@@ -451,6 +434,139 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
}, nil
}
func (s *OpenAIGatewayService) forwardGrokMediaVideoContent(
ctx context.Context,
c *gin.Context,
account *Account,
token, requestID string,
startTime time.Time,
) (*OpenAIForwardResult, error) {
statusURL, err := buildGrokMediaURL(account, s.cfg, GrokMediaEndpointVideoStatus, requestID)
if err != nil {
return nil, err
}
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
defer releaseUpstreamCtx()
statusReq, err := http.NewRequestWithContext(
WithHTTPUpstreamRedirectsDisabled(upstreamCtx),
http.MethodGet,
statusURL,
nil,
)
if err != nil {
return nil, err
}
statusReq.Header.Set("Authorization", "Bearer "+token)
statusReq.Header.Set("Accept", "application/json")
if account.IsGrokOAuth() && isGrokCLIProxyTarget(statusURL) {
applyGrokCLIHeaders(statusReq.Header)
}
account.ApplyHeaderOverrides(statusReq.Header)
proxyURL := ""
if account.ProxyID != nil && account.Proxy != nil {
proxyURL = account.Proxy.URL()
}
upstreamStart := time.Now()
statusResp, err := s.httpUpstream.Do(statusReq, proxyURL, account.ID, account.Concurrency)
if err != nil {
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
}
statusRequestID := firstNonEmpty(statusResp.Header.Get("x-request-id"), statusResp.Header.Get("xai-request-id"))
if statusResp.StatusCode >= 300 {
defer func() { _ = statusResp.Body.Close() }()
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
if statusResp.StatusCode < 400 {
return nil, fmt.Errorf("grok media status redirect is not allowed")
}
return s.handleGrokMediaErrorResponse(ctx, statusResp, c, account, statusRequestID, "")
}
statusBody, err := ReadUpstreamResponseBody(statusResp.Body, s.cfg, c, openAITooLargeError)
_ = statusResp.Body.Close()
if err != nil {
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
return nil, err
}
contentURL, err := grokMediaSignedVideoContentURL(statusBody)
if err != nil {
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
return nil, err
}
signedContent := contentURL != ""
if !signedContent {
contentURL, err = buildGrokMediaURL(account, s.cfg, GrokMediaEndpointVideoContent, requestID)
if err != nil {
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
return nil, err
}
}
contentReq, err := http.NewRequestWithContext(
WithHTTPUpstreamRedirectsDisabled(upstreamCtx),
http.MethodGet,
contentURL,
nil,
)
if err != nil {
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
return nil, err
}
contentReq.Header.Set("Accept", "*/*")
if c != nil {
if rangeHeader := strings.TrimSpace(c.GetHeader("Range")); rangeHeader != "" {
contentReq.Header.Set("Range", rangeHeader)
}
}
if !signedContent {
contentReq.Header.Set("Authorization", "Bearer "+token)
if account.IsGrokOAuth() && isGrokCLIProxyTarget(contentURL) {
applyGrokCLIHeaders(contentReq.Header)
}
account.ApplyHeaderOverrides(contentReq.Header)
}
contentResp, err := s.httpUpstream.Do(contentReq, proxyURL, account.ID, account.Concurrency)
SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
if err != nil {
return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
}
defer func() { _ = contentResp.Body.Close() }()
contentRequestID := firstNonEmpty(contentResp.Header.Get("x-request-id"), contentResp.Header.Get("xai-request-id"), statusRequestID)
if contentResp.StatusCode >= 300 && contentResp.StatusCode < 400 {
return nil, fmt.Errorf("grok media signed content redirect is not allowed")
}
if contentResp.StatusCode >= 400 && contentResp.StatusCode != http.StatusRequestedRangeNotSatisfiable {
return s.handleGrokMediaErrorResponse(ctx, contentResp, c, account, contentRequestID, "")
}
s.updateGrokUsageFromResponse(ctx, account, contentResp.Header, contentResp.StatusCode)
if err := writeGrokMediaContentResponse(c, contentResp); err != nil {
return nil, err
}
return &OpenAIForwardResult{
RequestID: contentRequestID,
ResponseHeaders: contentResp.Header.Clone(),
Duration: time.Since(startTime),
}, nil
}
func grokMediaSignedVideoContentURL(body []byte) (string, error) {
rawURL := strings.TrimSpace(gjson.GetBytes(body, "video.url").String())
if rawURL == "" {
return "", nil
}
parsed, err := url.Parse(rawURL)
if err != nil || !strings.EqualFold(parsed.Scheme, "https") ||
!strings.EqualFold(parsed.Hostname(), "vidgen.x.ai") ||
(parsed.Port() != "" && parsed.Port() != "443") || parsed.User != nil {
return "", fmt.Errorf("grok media status returned an unsupported video content URL")
}
return parsed.String(), nil
}
func isGrokCLIProxyTarget(rawURL string) bool {
parsed, err := url.Parse(strings.TrimSpace(rawURL))
return err == nil && strings.EqualFold(parsed.Hostname(), "cli-chat-proxy.grok.com")
backend/internal/service/grok_media_content_test.go
+110 -14
View File
@@ -18,12 +18,20 @@ import (
)
type grokMediaContentUpstreamStub struct {
request *http.Request
response *http.Response
request *http.Request
requests []*http.Request
response *http.Response
responses []*http.Response
}
func (s *grokMediaContentUpstreamStub) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
s.request = req
s.requests = append(s.requests, req)
if len(s.responses) > 0 {
resp := s.responses[0]
s.responses = s.responses[1:]
return resp, nil
}
return s.response, nil
}
@@ -54,9 +62,17 @@ func grokMediaContentTestContext(method, target string, headers map[string]strin
return c, recorder
}
func grokMediaContentStatusResponse(body string) *http.Response {
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(body)),
}
}
func TestForwardGrokMediaContentUsesUpstreamCredentialAndStreamsRange(t *testing.T) {
upstream := &grokMediaContentUpstreamStub{
response: &http.Response{
responses: []*http.Response{grokMediaContentStatusResponse(`{"status":"completed"}`), {
StatusCode: http.StatusPartialContent,
Header: http.Header{
"Content-Type": []string{"video/mp4"},
@@ -68,7 +84,7 @@ func TestForwardGrokMediaContentUsesUpstreamCredentialAndStreamsRange(t *testing
},
},
Body: io.NopCloser(strings.NewReader("video-payload")),
},
}},
}
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
c, recorder := grokMediaContentTestContext(http.MethodGet, "https://api.example/v1/videos/task-1/content", map[string]string{
@@ -84,10 +100,13 @@ func TestForwardGrokMediaContentUsesUpstreamCredentialAndStreamsRange(t *testing
require.NotNil(t, result)
require.Equal(t, http.StatusPartialContent, recorder.Code)
require.Equal(t, "video-payload", recorder.Body.String())
require.Equal(t, "https://relay.example/v1/videos/task-1/content", upstream.request.URL.String())
require.Equal(t, "Bearer upstream-key", upstream.request.Header.Get("Authorization"))
require.Equal(t, "bytes=0-12", upstream.request.Header.Get("Range"))
require.Equal(t, "*/*", upstream.request.Header.Get("Accept"))
require.Len(t, upstream.requests, 2)
require.Equal(t, "https://relay.example/v1/videos/task-1", upstream.requests[0].URL.String())
require.Equal(t, "Bearer upstream-key", upstream.requests[0].Header.Get("Authorization"))
require.Equal(t, "https://relay.example/v1/videos/task-1/content", upstream.requests[1].URL.String())
require.Equal(t, "Bearer upstream-key", upstream.requests[1].Header.Get("Authorization"))
require.Equal(t, "bytes=0-12", upstream.requests[1].Header.Get("Range"))
require.Equal(t, "*/*", upstream.requests[1].Header.Get("Accept"))
require.Equal(t, "video/mp4", recorder.Header().Get("Content-Type"))
require.Equal(t, "13", recorder.Header().Get("Content-Length"))
require.Equal(t, "bytes 0-12/100", recorder.Header().Get("Content-Range"))
@@ -98,12 +117,12 @@ func TestForwardGrokMediaContentUsesUpstreamCredentialAndStreamsRange(t *testing
func TestForwardGrokMediaContentStreamsFullResponseWithSafeDefaults(t *testing.T) {
upstream := &grokMediaContentUpstreamStub{
response: &http.Response{
responses: []*http.Response{grokMediaContentStatusResponse(`{"status":"completed"}`), {
StatusCode: http.StatusOK,
Header: http.Header{"Set-Cookie": []string{"secret=upstream"}, "X-Upstream-Secret": []string{"hidden"}},
Body: io.NopCloser(strings.NewReader("full-video")),
ContentLength: -1,
},
}},
}
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
c, recorder := grokMediaContentTestContext(http.MethodGet, "https://api.example/v1/videos/task-1/content", nil)
@@ -116,7 +135,8 @@ func TestForwardGrokMediaContentStreamsFullResponseWithSafeDefaults(t *testing.T
require.NoError(t, err)
require.Equal(t, http.StatusOK, recorder.Code)
require.Equal(t, "full-video", recorder.Body.String())
require.Empty(t, upstream.request.Header.Get("Range"))
require.Len(t, upstream.requests, 2)
require.Empty(t, upstream.requests[1].Header.Get("Range"))
require.Equal(t, "application/octet-stream", recorder.Header().Get("Content-Type"))
require.Empty(t, recorder.Header().Get("Content-Length"))
require.Empty(t, recorder.Header().Get("Set-Cookie"))
@@ -126,7 +146,7 @@ func TestForwardGrokMediaContentStreamsFullResponseWithSafeDefaults(t *testing.T
func TestForwardGrokMediaContentPreservesRangeNotSatisfiable(t *testing.T) {
upstream := &grokMediaContentUpstreamStub{
response: &http.Response{
responses: []*http.Response{grokMediaContentStatusResponse(`{"status":"completed"}`), {
StatusCode: http.StatusRequestedRangeNotSatisfiable,
Header: http.Header{
"Content-Type": []string{"text/plain"},
@@ -135,7 +155,7 @@ func TestForwardGrokMediaContentPreservesRangeNotSatisfiable(t *testing.T) {
"Accept-Ranges": []string{"bytes"},
},
Body: io.NopCloser(strings.NewReader("bad-range!!")),
},
}},
}
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
c, recorder := grokMediaContentTestContext(http.MethodGet, "https://api.example/v1/videos/task-1/content", map[string]string{
@@ -150,12 +170,88 @@ func TestForwardGrokMediaContentPreservesRangeNotSatisfiable(t *testing.T) {
require.NoError(t, err)
require.Equal(t, http.StatusRequestedRangeNotSatisfiable, recorder.Code)
require.Equal(t, "bad-range!!", recorder.Body.String())
require.Equal(t, "bytes=500-600", upstream.request.Header.Get("Range"))
require.Len(t, upstream.requests, 2)
require.Equal(t, "bytes=500-600", upstream.requests[1].Header.Get("Range"))
require.Equal(t, "bytes */100", recorder.Header().Get("Content-Range"))
require.Equal(t, "bytes", recorder.Header().Get("Accept-Ranges"))
require.True(t, IsResponseCommitted(c))
}
func TestForwardGrokMediaContentFetchesValidatedSignedURLWithoutCredentials(t *testing.T) {
upstream := &grokMediaContentUpstreamStub{
responses: []*http.Response{
grokMediaContentStatusResponse(`{"status":"done","video":{"url":"https://vidgen.x.ai/signed-token/xai-video-task-1.mp4"}}`),
{
StatusCode: http.StatusPartialContent,
Header: http.Header{
"Content-Type": []string{"video/mp4"},
"Content-Length": []string{"13"},
"Content-Range": []string{"bytes 0-12/100"},
},
Body: io.NopCloser(strings.NewReader("video-payload")),
},
},
}
account := grokMediaContentTestAccount()
account.Credentials[credKeyHeaderOverrideEnabled] = true
account.Credentials[credKeyHeaderOverrides] = map[string]any{"user-agent": "private-agent"}
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
c, recorder := grokMediaContentTestContext(http.MethodGet, "https://api.example/v1/videos/task-1/content", map[string]string{
"Range": "bytes=0-12",
})
_, err := svc.ForwardGrokMedia(
context.Background(), c, account,
GrokMediaEndpointVideoContent, "task-1", nil, "",
)
require.NoError(t, err)
require.Equal(t, http.StatusPartialContent, recorder.Code)
require.Equal(t, "video-payload", recorder.Body.String())
require.Len(t, upstream.requests, 2)
require.Equal(t, "https://relay.example/v1/videos/task-1", upstream.requests[0].URL.String())
require.Equal(t, "Bearer upstream-key", upstream.requests[0].Header.Get("Authorization"))
require.Equal(t, "private-agent", upstream.requests[0].Header.Get("User-Agent"))
require.True(t, HTTPUpstreamRedirectsDisabled(upstream.requests[0].Context()))
require.Equal(t, "https://vidgen.x.ai/signed-token/xai-video-task-1.mp4", upstream.requests[1].URL.String())
require.Empty(t, upstream.requests[1].Header.Get("Authorization"))
require.Empty(t, upstream.requests[1].Header.Get("User-Agent"))
require.Equal(t, "bytes=0-12", upstream.requests[1].Header.Get("Range"))
require.True(t, HTTPUpstreamRedirectsDisabled(upstream.requests[1].Context()))
}
func TestForwardGrokMediaContentRejectsUntrustedSignedURL(t *testing.T) {
upstream := &grokMediaContentUpstreamStub{
responses: []*http.Response{
grokMediaContentStatusResponse(`{"status":"done","video":{"url":"http://169.254.169.254/latest/meta-data"}}`),
},
}
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
c, _ := grokMediaContentTestContext(http.MethodGet, "https://api.example/v1/videos/task-1/content", nil)
_, err := svc.ForwardGrokMedia(
context.Background(), c, grokMediaContentTestAccount(),
GrokMediaEndpointVideoContent, "task-1", nil, "",
)
require.ErrorContains(t, err, "unsupported video content URL")
require.Len(t, upstream.requests, 1)
}
func TestGrokMediaSignedVideoContentURLRejectsDeceptiveOrigins(t *testing.T) {
for _, rawURL := range []string{
"https://vidgen.x.ai.attacker.invalid/video.mp4",
"https://vidgen.x.ai@attacker.invalid/video.mp4",
"https://vidgen.x.ai:444/video.mp4",
"http://vidgen.x.ai/video.mp4",
} {
t.Run(rawURL, func(t *testing.T) {
_, err := grokMediaSignedVideoContentURL([]byte(`{"video":{"url":"` + rawURL + `"}}`))
require.ErrorContains(t, err, "unsupported video content URL")
})
}
}
func TestForwardGrokVideoStatusRewritesOnlyProtectedContentURL(t *testing.T) {
statusBody := `{"id":"task-1","status":"completed","url":"https://relay.example/v1/videos/task-1/content","download_url":"/v1/videos/task-1/content","video_url":"https://vidgen.x.ai/task-1.mp4","counter":9007199254740993}`
upstream := &grokMediaContentUpstreamStub{