diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index f51167c6e6..77c3e74abc 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -164,8 +164,17 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. sessionSeed = []byte(requestID) } sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, sessionSeed) + boundLookupAccountID := int64(0) if endpoint.IsVideoLookupRequest() { - sessionHash = service.GrokMediaVideoRequestSessionHash(requestID) + sessionHash = service.GrokMediaVideoRequestSessionHash(requestID, subject.UserID, apiKey.ID) + boundLookupAccountID, err = h.gatewayService.ResolveGrokMediaVideoRequestAccount( + c.Request.Context(), apiKey.GroupID, requestID, subject.UserID, apiKey.ID, + ) + if err != nil || boundLookupAccountID <= 0 { + reqLog.Info("grok_media.video_lookup_owner_binding_missing", zap.Error(err)) + h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found") + return + } } requestCtx := c.Request.Context() failedAccountIDs := make(map[int64]struct{}) @@ -240,6 +249,14 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. h.errorResponse(c, cls.Status, cls.ErrType, cls.Message) return } + if boundLookupAccountID > 0 && selection.Account.ID != boundLookupAccountID { + reqLog.Warn("grok_media.video_lookup_bound_account_unavailable", + zap.Int64("bound_account_id", boundLookupAccountID), + zap.Int64("selected_account_id", selection.Account.ID), + ) + h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found") + return + } reqLog.Debug("grok_media.account_schedule_decision", zap.String("layer", scheduleDecision.Layer), @@ -300,6 +317,10 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. h.handleFailoverExhausted(c, failoverErr, false) return } + if endpoint.IsVideoLookupRequest() { + h.handleFailoverExhausted(c, failoverErr, false) + return + } if failoverErr.RetryableOnSameAccount { retryLimit := account.GetPoolModeRetryCount() if sameAccountRetryCount[account.ID] < retryLimit { @@ -351,7 +372,9 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), true, nil) if endpoint.IsGenerationRequest() && strings.TrimSpace(result.ResponseID) != "" { - if err := h.gatewayService.BindGrokMediaVideoRequestAccount(requestCtx, apiKey.GroupID, result.ResponseID, account.ID); err != nil { + if err := h.gatewayService.BindGrokMediaVideoRequestAccount( + requestCtx, apiKey.GroupID, result.ResponseID, subject.UserID, apiKey.ID, account.ID, + ); err != nil { reqLog.Warn("grok_media.bind_video_request_account_failed", zap.Int64("account_id", account.ID), zap.String("request_id", result.ResponseID), diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index a107bd72f3..36bd80f728 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -269,16 +269,50 @@ func parseGrokMediaMultipartRequest(contentType string, body []byte, info *GrokM } } -func GrokMediaVideoRequestSessionHash(requestID string) string { +func GrokMediaVideoRequestSessionHash(requestID string, userID, apiKeyID int64) string { requestID = strings.TrimSpace(requestID) - if requestID == "" { + if requestID == "" || userID <= 0 || apiKeyID <= 0 { return "" } - return "grok-video:" + DeriveSessionHashFromSeed(requestID) + ownerSeed := fmt.Sprintf("%d:%d:%s", userID, apiKeyID, requestID) + return "grok-video:" + DeriveSessionHashFromSeed(ownerSeed) } -func (s *OpenAIGatewayService) BindGrokMediaVideoRequestAccount(ctx context.Context, groupID *int64, requestID string, accountID int64) error { - return s.BindStickySession(ctx, groupID, GrokMediaVideoRequestSessionHash(requestID), accountID) +func (s *OpenAIGatewayService) BindGrokMediaVideoRequestAccount( + ctx context.Context, + groupID *int64, + requestID string, + userID, apiKeyID, accountID int64, +) error { + if s == nil || s.cache == nil { + return fmt.Errorf("grok video request binding cache is unavailable") + } + sessionHash := GrokMediaVideoRequestSessionHash(requestID, userID, apiKeyID) + cacheKey := s.openAISessionCacheKey(sessionHash) + if cacheKey == "" || accountID <= 0 { + return fmt.Errorf("grok video request binding is invalid") + } + ttl := openaiStickySessionTTL + if s.cfg != nil && s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds > 0 { + ttl = time.Duration(s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds) * time.Second + } + return s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), cacheKey, accountID, ttl) +} + +func (s *OpenAIGatewayService) ResolveGrokMediaVideoRequestAccount( + ctx context.Context, + groupID *int64, + requestID string, + userID, apiKeyID int64, +) (int64, error) { + if s == nil || s.cache == nil { + return 0, fmt.Errorf("grok video request binding cache is unavailable") + } + cacheKey := s.openAISessionCacheKey(GrokMediaVideoRequestSessionHash(requestID, userID, apiKeyID)) + if cacheKey == "" { + return 0, fmt.Errorf("grok video request binding is invalid") + } + return s.cache.GetSessionAccountID(ctx, derefGroupID(groupID), cacheKey) } func (s *OpenAIGatewayService) ForwardGrokMedia( @@ -369,7 +403,8 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( requestIDHeader := firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")) requestModel := requestInfo.Model - if resp.StatusCode >= 400 { + preserveContentRangeError := endpoint == GrokMediaEndpointVideoContent && resp.StatusCode == http.StatusRequestedRangeNotSatisfiable + if resp.StatusCode >= 400 && !preserveContentRangeError { return s.handleGrokMediaErrorResponse(ctx, resp, c, account, requestIDHeader, requestModel) } @@ -859,7 +894,7 @@ func isGrokMediaVideoContentURL(rawURL, requestID string) bool { if err != nil || parsed.Path == "" { return false } - segments := strings.Split(strings.Trim(parsed.Path, "/"), "/") + segments := strings.Split(strings.Trim(parsed.EscapedPath(), "/"), "/") if len(segments) < 3 { return false } diff --git a/backend/internal/service/grok_media_content_test.go b/backend/internal/service/grok_media_content_test.go index 7e52b7e2f0..839bc01b68 100644 --- a/backend/internal/service/grok_media_content_test.go +++ b/backend/internal/service/grok_media_content_test.go @@ -93,6 +93,67 @@ func TestForwardGrokMediaContentUsesUpstreamCredentialAndStreamsRange(t *testing require.Equal(t, "bytes 0-12/100", recorder.Header().Get("Content-Range")) require.Equal(t, "bytes", recorder.Header().Get("Accept-Ranges")) require.Equal(t, `attachment; filename="task-1.mp4"`, recorder.Header().Get("Content-Disposition")) + require.True(t, IsResponseCommitted(c)) +} + +func TestForwardGrokMediaContentStreamsFullResponseWithSafeDefaults(t *testing.T) { + upstream := &grokMediaContentUpstreamStub{ + response: &http.Response{ + 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) + + _, err := svc.ForwardGrokMedia( + context.Background(), c, grokMediaContentTestAccount(), + GrokMediaEndpointVideoContent, "task-1", nil, "", + ) + + 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.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")) + require.Empty(t, recorder.Header().Get("X-Upstream-Secret")) + require.True(t, IsResponseCommitted(c)) +} + +func TestForwardGrokMediaContentPreservesRangeNotSatisfiable(t *testing.T) { + upstream := &grokMediaContentUpstreamStub{ + response: &http.Response{ + StatusCode: http.StatusRequestedRangeNotSatisfiable, + Header: http.Header{ + "Content-Type": []string{"text/plain"}, + "Content-Length": []string{"11"}, + "Content-Range": []string{"bytes */100"}, + "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{ + "Range": "bytes=500-600", + }) + + _, err := svc.ForwardGrokMedia( + context.Background(), c, grokMediaContentTestAccount(), + GrokMediaEndpointVideoContent, "task-1", nil, "", + ) + + 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.Equal(t, "bytes */100", recorder.Header().Get("Content-Range")) + require.Equal(t, "bytes", recorder.Header().Get("Accept-Ranges")) + require.True(t, IsResponseCommitted(c)) } func TestForwardGrokVideoStatusRewritesOnlyProtectedContentURL(t *testing.T) { @@ -123,3 +184,12 @@ func TestForwardGrokVideoStatusRewritesOnlyProtectedContentURL(t *testing.T) { require.Equal(t, "9007199254740993", gjson.Get(recorder.Body.String(), "counter").String()) require.NotContains(t, recorder.Body.String(), "malicious.invalid") } + +func TestRewriteGrokMediaVideoContentURLsPreservesOtherIDsAndHandlesNestedEscapedID(t *testing.T) { + body := []byte(`{"nested":[{"url":"https://relay.example/v1/videos/task%2Fone/content"},{"url":"https://relay.example/v1/videos/task-two/content"}]}`) + + rewritten := rewriteGrokMediaVideoContentURLs(body, "task/one", "/v1/videos/task%2Fone/content") + + require.Equal(t, "/v1/videos/task%2Fone/content", gjson.GetBytes(rewritten, "nested.0.url").String()) + require.Equal(t, "https://relay.example/v1/videos/task-two/content", gjson.GetBytes(rewritten, "nested.1.url").String()) +} diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 67ad712147..6ed35d4112 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -878,19 +878,34 @@ func TestForwardGrokMediaVideoMutationEndpoints(t *testing.T) { } } -func TestBindGrokMediaVideoRequestAccountUsesRequestIDStickyHash(t *testing.T) { - ctx := context.Background() +func TestGrokMediaVideoRequestBindingIsScopedToUserAndAPIKey(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/videos/video-request-123", nil) + c.Request.Header.Set("session_id", "shared-client-session") groupID := int64(7) cache := &stubGatewayCache{} svc := &OpenAIGatewayService{cache: cache} + const userID int64 = 41 + const apiKeyID int64 = 51 + require.NotEmpty(t, svc.GenerateExplicitSessionHash(c, nil)) + ctx := c.Request.Context() - hash := GrokMediaVideoRequestSessionHash("video-request-123") + hash := GrokMediaVideoRequestSessionHash("video-request-123", userID, apiKeyID) require.NotEmpty(t, hash) - require.NoError(t, svc.BindGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", 63)) + require.NoError(t, svc.BindGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", userID, apiKeyID, 63)) - accountID, err := svc.getStickySessionAccountID(ctx, &groupID, hash) + accountID, err := svc.ResolveGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", userID, apiKeyID) require.NoError(t, err) require.Equal(t, int64(63), accountID) + + accountID, err = svc.ResolveGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", userID+1, apiKeyID) + require.Error(t, err) + require.Zero(t, accountID) + + accountID, err = svc.ResolveGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", userID, apiKeyID+1) + require.Error(t, err) + require.Zero(t, accountID) } func TestForwardGrokMedia429ReconcilesRateLimitBeforeCustomErrorBypass(t *testing.T) {