diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 359ae32df9..5844257497 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -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") diff --git a/backend/internal/service/grok_media_content_test.go b/backend/internal/service/grok_media_content_test.go index 45bd2b7890..a967c88bd5 100644 --- a/backend/internal/service/grok_media_content_test.go +++ b/backend/internal/service/grok_media_content_test.go @@ -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{