mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Merge pull request #4543 from heathermhuang/codex/secure-protected-video-content-4498
fix(grok): securely proxy protected video content
This commit is contained in:
@@ -46,6 +46,11 @@ func (h *OpenAIGatewayHandler) GrokVideoStatus(c *gin.Context) {
|
||||
h.handleGrokMedia(c, service.GrokMediaEndpointVideoStatus, c.Param("request_id"))
|
||||
}
|
||||
|
||||
// GrokVideoContent proxies downloadable video content through the task's upstream account.
|
||||
func (h *OpenAIGatewayHandler) GrokVideoContent(c *gin.Context) {
|
||||
h.handleGrokMedia(c, service.GrokMediaEndpointVideoContent, c.Param("request_id"))
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.GrokMediaEndpoint, requestID string) {
|
||||
streamStarted := false
|
||||
defer h.recoverResponsesPanic(c, &streamStarted)
|
||||
@@ -100,7 +105,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||||
return
|
||||
}
|
||||
if endpoint == service.GrokMediaEndpointVideoStatus && strings.TrimSpace(requestID) == "" {
|
||||
if endpoint.IsVideoLookupRequest() && strings.TrimSpace(requestID) == "" {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "request_id is required")
|
||||
return
|
||||
}
|
||||
@@ -160,8 +165,17 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
sessionSeed = []byte(requestID)
|
||||
}
|
||||
sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, sessionSeed)
|
||||
if endpoint == service.GrokMediaEndpointVideoStatus {
|
||||
sessionHash = service.GrokMediaVideoRequestSessionHash(requestID)
|
||||
boundLookupAccountID := int64(0)
|
||||
if endpoint.IsVideoLookupRequest() {
|
||||
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{})
|
||||
@@ -238,6 +252,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),
|
||||
@@ -317,6 +339,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 {
|
||||
@@ -356,7 +382,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
continue
|
||||
}
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, grokMediaScheduleModel(account, routingModel, nil), false, nil)
|
||||
if c.Writer.Size() == writerSizeBeforeForward {
|
||||
if !service.IsResponseCommitted(c) && c.Writer.Size() == writerSizeBeforeForward {
|
||||
h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
|
||||
}
|
||||
reqLog.Warn("grok_media.forward_failed",
|
||||
@@ -368,7 +394,9 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, grokMediaScheduleModel(account, routingModel, result), 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),
|
||||
|
||||
@@ -52,6 +52,12 @@ func TestShouldRecordGrokMediaUsage(t *testing.T) {
|
||||
model: "",
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "video content skips usage",
|
||||
endpoint: service.GrokMediaEndpointVideoContent,
|
||||
model: "",
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "generation skips usage without model",
|
||||
endpoint: service.GrokMediaEndpointImagesGenerations,
|
||||
@@ -79,6 +85,7 @@ func TestGrokMediaRequiredCapability(t *testing.T) {
|
||||
{name: "video edit", endpoint: service.GrokMediaEndpointVideosEdits, want: service.OpenAIEndpointCapabilityGrokMediaGeneration},
|
||||
{name: "video extension", endpoint: service.GrokMediaEndpointVideosExtensions, want: service.OpenAIEndpointCapabilityGrokMediaGeneration},
|
||||
{name: "video status preserves lookup", endpoint: service.GrokMediaEndpointVideoStatus, want: ""},
|
||||
{name: "video content preserves lookup", endpoint: service.GrokMediaEndpointVideoContent, want: ""},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
|
||||
@@ -645,7 +645,8 @@ func TestGrokMedia429FailoverIsBounded(t *testing.T) {
|
||||
_, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "first_429")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/openai/v1/videos/request-1", nil)
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/videos/generations", bytes.NewBufferString(`{"model":"grok-imagine-video","prompt":"waves"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
@@ -657,7 +658,8 @@ func TestGrokMedia429FailoverIsBounded(t *testing.T) {
|
||||
_, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "all_429")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/openai/v1/videos/request-1", nil)
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/videos/generations", bytes.NewBufferString(`{"model":"grok-imagine-video","prompt":"waves"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
@@ -678,7 +680,7 @@ func TestGrokOAuthCredentialFailoverAcrossHTTPHandlers(t *testing.T) {
|
||||
{name: "messages", method: http.MethodPost, path: "/openai/v1/messages", body: `{"model":"grok","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`},
|
||||
{name: "chat completions", method: http.MethodPost, path: "/openai/v1/chat/completions", body: `{"model":"grok","messages":[{"role":"user","content":"hello"}],"stream":false}`},
|
||||
{name: "chat completions raw fallback", method: http.MethodPost, path: "/openai/v1/chat/completions", body: `{"model":"grok","messages":[{"role":"user","content":"hello"}],"stop":["END"],"stream":false}`},
|
||||
{name: "grok media", method: http.MethodGet, path: "/openai/v1/videos/request-1"},
|
||||
{name: "grok media", method: http.MethodPost, path: "/openai/v1/videos/generations", body: `{"model":"grok-imagine-video","prompt":"waves"}`},
|
||||
}
|
||||
|
||||
for _, endpoint := range endpoints {
|
||||
@@ -822,6 +824,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
"access_token": "expired", "refresh_token": "revoked-refresh",
|
||||
"expires_at": time.Now().Add(-time.Minute).UTC().Format(time.RFC3339),
|
||||
},
|
||||
Extra: map[string]any{service.GrokMediaEligibleExtraKey: true},
|
||||
},
|
||||
{
|
||||
ID: 802, Name: "healthy", Platform: service.PlatformGrok, Type: service.AccountTypeOAuth,
|
||||
@@ -830,6 +833,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
"access_token": "healthy-access", "refresh_token": "healthy-refresh",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
Extra: map[string]any{service.GrokMediaEligibleExtraKey: true},
|
||||
},
|
||||
}
|
||||
if mode == "postmap_cancel" || mode == "first_429" || mode == "all_429" || mode == "mixed_429_500" || mode == "mixed_500_429" || mode == "oauth_429_apikey_500" {
|
||||
@@ -843,6 +847,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
"access_token": "untried-healthy-access", "refresh_token": "untried-healthy-refresh",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
Extra: map[string]any{service.GrokMediaEligibleExtraKey: true},
|
||||
})
|
||||
}
|
||||
if mode == "oauth_429_apikey_500" {
|
||||
@@ -904,7 +909,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
apiKey := &service.APIKey{
|
||||
ID: 902, GroupID: &groupID,
|
||||
User: &service.User{ID: 903, Status: service.StatusActive},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, Status: service.StatusActive},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, Status: service.StatusActive, AllowImageGeneration: true},
|
||||
}
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
@@ -916,6 +921,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
router.GET("/openai/v1/responses", h.ResponsesWebSocket)
|
||||
router.POST("/openai/v1/messages", h.Messages)
|
||||
router.POST("/openai/v1/chat/completions", h.ChatCompletions)
|
||||
router.POST("/openai/v1/videos/generations", h.GrokVideoGeneration)
|
||||
router.GET("/openai/v1/videos/:request_id", h.GrokVideoStatus)
|
||||
handlerRefresherStarted.Store(router, refresher.started)
|
||||
cleanup := func() {
|
||||
|
||||
@@ -92,6 +92,19 @@ func RegisterGatewayRoutes(
|
||||
},
|
||||
})
|
||||
}
|
||||
videoContentHandler := func(c *gin.Context) {
|
||||
if getGroupPlatform(c) == service.PlatformGrok {
|
||||
h.OpenAIGateway.GrokVideoContent(c)
|
||||
return
|
||||
}
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
c.JSON(http.StatusNotFound, gin.H{
|
||||
"error": gin.H{
|
||||
"type": "not_found_error",
|
||||
"message": "Videos API is not supported for this platform",
|
||||
},
|
||||
})
|
||||
}
|
||||
videoEditHandler := func(c *gin.Context) {
|
||||
if getGroupPlatform(c) == service.PlatformGrok {
|
||||
h.OpenAIGateway.GrokVideoEdit(c)
|
||||
@@ -210,6 +223,7 @@ func RegisterGatewayRoutes(
|
||||
gateway.POST("/videos/edits", videoEditHandler)
|
||||
gateway.POST("/videos/extensions", videoExtensionHandler)
|
||||
gateway.GET("/videos/:request_id", videoStatusHandler)
|
||||
gateway.GET("/videos/:request_id/content", videoContentHandler)
|
||||
}
|
||||
|
||||
// Gemini 原生 API 兼容层(Gemini SDK/CLI 直连)
|
||||
@@ -283,6 +297,7 @@ func RegisterGatewayRoutes(
|
||||
r.POST("/videos/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoEditHandler)
|
||||
r.POST("/videos/extensions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoExtensionHandler)
|
||||
r.GET("/videos/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoStatusHandler)
|
||||
r.GET("/videos/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoContentHandler)
|
||||
|
||||
// Antigravity 模型列表
|
||||
r.GET("/antigravity/models", gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.Gateway.AntigravityModels)
|
||||
|
||||
@@ -160,6 +160,8 @@ func TestGatewayRoutesGrokImagesAndVideosPathsAreRegistered(t *testing.T) {
|
||||
for _, path := range []string{
|
||||
"/v1/videos/request-123",
|
||||
"/videos/request-123",
|
||||
"/v1/videos/request-123/content",
|
||||
"/videos/request-123/content",
|
||||
} {
|
||||
req := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
w := httptest.NewRecorder()
|
||||
@@ -186,6 +188,8 @@ func TestGatewayRoutesNonGrokVideosAreRejectedAtPlatformGate(t *testing.T) {
|
||||
{http.MethodPost, "/videos/extensions", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`},
|
||||
{http.MethodGet, "/v1/videos/request-123", ""},
|
||||
{http.MethodGet, "/videos/request-123", ""},
|
||||
{http.MethodGet, "/v1/videos/request-123/content", ""},
|
||||
{http.MethodGet, "/videos/request-123/content", ""},
|
||||
} {
|
||||
req := httptest.NewRequest(tc.method, tc.path, strings.NewReader(tc.body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
@@ -29,10 +29,15 @@ const (
|
||||
GrokMediaEndpointVideosEdits GrokMediaEndpoint = "videos_edits"
|
||||
GrokMediaEndpointVideosExtensions GrokMediaEndpoint = "videos_extensions"
|
||||
GrokMediaEndpointVideoStatus GrokMediaEndpoint = "video_status"
|
||||
GrokMediaEndpointVideoContent GrokMediaEndpoint = "video_content"
|
||||
)
|
||||
|
||||
func (e GrokMediaEndpoint) RequiresRequestBody() bool {
|
||||
return e != GrokMediaEndpointVideoStatus
|
||||
return !e.IsVideoLookupRequest()
|
||||
}
|
||||
|
||||
func (e GrokMediaEndpoint) IsVideoLookupRequest() bool {
|
||||
return e == GrokMediaEndpointVideoStatus || e == GrokMediaEndpointVideoContent
|
||||
}
|
||||
|
||||
func (e GrokMediaEndpoint) IsGenerationRequest() bool {
|
||||
@@ -97,7 +102,7 @@ func (r GrokMediaRequestInfo) ModerationBody() []byte {
|
||||
}
|
||||
|
||||
func (e GrokMediaEndpoint) httpMethod() string {
|
||||
if e == GrokMediaEndpointVideoStatus {
|
||||
if e.IsVideoLookupRequest() {
|
||||
return http.MethodGet
|
||||
}
|
||||
return http.MethodPost
|
||||
@@ -264,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(
|
||||
@@ -297,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
|
||||
@@ -386,6 +428,13 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
|
||||
}
|
||||
}
|
||||
}
|
||||
if endpoint == GrokMediaEndpointVideoStatus {
|
||||
respBody = rewriteGrokMediaVideoContentURLs(
|
||||
respBody,
|
||||
requestID,
|
||||
grokMediaContentProxyURL(c, requestID),
|
||||
)
|
||||
}
|
||||
writeGrokMediaResponse(c, resp, respBody, s.responseHeaderFilter)
|
||||
usage := grokMediaUsageFromResponse(endpoint, requestInfo, respBody)
|
||||
return &OpenAIForwardResult{
|
||||
@@ -407,6 +456,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")
|
||||
@@ -756,3 +938,140 @@ func writeGrokMediaResponse(c *gin.Context, resp *http.Response, body []byte, fi
|
||||
}
|
||||
c.Data(resp.StatusCode, contentType, body)
|
||||
}
|
||||
|
||||
func writeGrokMediaContentResponse(c *gin.Context, resp *http.Response) error {
|
||||
if c == nil || resp == nil || resp.Body == nil {
|
||||
return fmt.Errorf("grok media content response is incomplete")
|
||||
}
|
||||
|
||||
for _, name := range []string{
|
||||
"Content-Type",
|
||||
"Content-Length",
|
||||
"Content-Range",
|
||||
"Accept-Ranges",
|
||||
"Content-Disposition",
|
||||
} {
|
||||
if value := strings.TrimSpace(resp.Header.Get(name)); value != "" {
|
||||
c.Header(name, value)
|
||||
}
|
||||
}
|
||||
if strings.TrimSpace(c.Writer.Header().Get("Content-Length")) == "" && resp.ContentLength >= 0 {
|
||||
c.Header("Content-Length", strconv.FormatInt(resp.ContentLength, 10))
|
||||
}
|
||||
if strings.TrimSpace(c.Writer.Header().Get("Content-Type")) == "" {
|
||||
c.Header("Content-Type", "application/octet-stream")
|
||||
}
|
||||
c.Status(resp.StatusCode)
|
||||
MarkResponseCommitted(c)
|
||||
_, err := io.Copy(c.Writer, resp.Body)
|
||||
return err
|
||||
}
|
||||
|
||||
func rewriteGrokMediaVideoContentURLs(body []byte, requestID, proxyURL string) []byte {
|
||||
if len(body) == 0 || strings.TrimSpace(requestID) == "" || strings.TrimSpace(proxyURL) == "" || !gjson.ValidBytes(body) {
|
||||
return body
|
||||
}
|
||||
|
||||
decoder := json.NewDecoder(bytes.NewReader(body))
|
||||
decoder.UseNumber()
|
||||
var value any
|
||||
if err := decoder.Decode(&value); err != nil {
|
||||
return body
|
||||
}
|
||||
changed := rewriteGrokMediaKnownVideoURL(&value, proxyURL)
|
||||
if rewriteGrokMediaVideoContentURLValue(&value, requestID, proxyURL) {
|
||||
changed = true
|
||||
}
|
||||
if !changed {
|
||||
return body
|
||||
}
|
||||
rewritten, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return body
|
||||
}
|
||||
return rewritten
|
||||
}
|
||||
|
||||
func rewriteGrokMediaKnownVideoURL(value *any, proxyURL string) bool {
|
||||
if value == nil {
|
||||
return false
|
||||
}
|
||||
root, ok := (*value).(map[string]any)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
video, ok := root["video"].(map[string]any)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
rawURL, ok := video["url"].(string)
|
||||
if !ok || strings.TrimSpace(rawURL) == "" {
|
||||
return false
|
||||
}
|
||||
video["url"] = proxyURL
|
||||
return true
|
||||
}
|
||||
|
||||
func rewriteGrokMediaVideoContentURLValue(value *any, requestID, proxyURL string) bool {
|
||||
if value == nil {
|
||||
return false
|
||||
}
|
||||
switch typed := (*value).(type) {
|
||||
case map[string]any:
|
||||
changed := false
|
||||
for key, child := range typed {
|
||||
childValue := child
|
||||
if rewriteGrokMediaVideoContentURLValue(&childValue, requestID, proxyURL) {
|
||||
typed[key] = childValue
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
return changed
|
||||
case []any:
|
||||
changed := false
|
||||
for index, child := range typed {
|
||||
childValue := child
|
||||
if rewriteGrokMediaVideoContentURLValue(&childValue, requestID, proxyURL) {
|
||||
typed[index] = childValue
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
return changed
|
||||
case string:
|
||||
if isGrokMediaVideoContentURL(typed, requestID) {
|
||||
*value = proxyURL
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func isGrokMediaVideoContentURL(rawURL, requestID string) bool {
|
||||
parsed, err := url.Parse(strings.TrimSpace(rawURL))
|
||||
if err != nil || parsed.Path == "" {
|
||||
return false
|
||||
}
|
||||
segments := strings.Split(strings.Trim(parsed.EscapedPath(), "/"), "/")
|
||||
if len(segments) < 3 {
|
||||
return false
|
||||
}
|
||||
requestID = strings.Trim(requestID, "/")
|
||||
decodedID, err := url.PathUnescape(segments[len(segments)-2])
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return segments[len(segments)-3] == "videos" &&
|
||||
decodedID == requestID &&
|
||||
segments[len(segments)-1] == "content"
|
||||
}
|
||||
|
||||
func grokMediaContentProxyURL(c *gin.Context, requestID string) string {
|
||||
if c == nil || c.Request == nil || c.Request.URL == nil || strings.TrimSpace(requestID) == "" {
|
||||
return ""
|
||||
}
|
||||
pathPrefix := ""
|
||||
if strings.HasPrefix(c.Request.URL.Path, "/v1/") {
|
||||
pathPrefix = "/v1"
|
||||
}
|
||||
return pathPrefix + "/videos/" + url.PathEscape(strings.Trim(requestID, "/")) + "/content"
|
||||
}
|
||||
|
||||
@@ -0,0 +1,301 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type grokMediaContentUpstreamStub struct {
|
||||
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
|
||||
}
|
||||
|
||||
func (s *grokMediaContentUpstreamStub) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
|
||||
return s.Do(req, proxyURL, accountID, accountConcurrency)
|
||||
}
|
||||
|
||||
func grokMediaContentTestAccount() *Account {
|
||||
return &Account{
|
||||
ID: 9,
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "upstream-key",
|
||||
"base_url": "https://relay.example/v1",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func grokMediaContentTestContext(method, target string, headers map[string]string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(method, target, nil)
|
||||
for name, value := range headers {
|
||||
c.Request.Header.Set(name, value)
|
||||
}
|
||||
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{
|
||||
responses: []*http.Response{grokMediaContentStatusResponse(`{"status":"completed"}`), {
|
||||
StatusCode: http.StatusPartialContent,
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"video/mp4"},
|
||||
"Content-Length": []string{"13"},
|
||||
"Content-Range": []string{"bytes 0-12/100"},
|
||||
"Accept-Ranges": []string{"bytes"},
|
||||
"Content-Disposition": []string{
|
||||
`attachment; filename="task-1.mp4"`,
|
||||
},
|
||||
},
|
||||
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{
|
||||
"Range": "bytes=0-12",
|
||||
})
|
||||
|
||||
result, err := svc.ForwardGrokMedia(
|
||||
context.Background(), c, grokMediaContentTestAccount(),
|
||||
GrokMediaEndpointVideoContent, "task-1", nil, "",
|
||||
)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
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, "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"))
|
||||
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{
|
||||
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)
|
||||
|
||||
_, 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.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"))
|
||||
require.Empty(t, recorder.Header().Get("X-Upstream-Secret"))
|
||||
require.True(t, IsResponseCommitted(c))
|
||||
}
|
||||
|
||||
func TestForwardGrokMediaContentPreservesRangeNotSatisfiable(t *testing.T) {
|
||||
upstream := &grokMediaContentUpstreamStub{
|
||||
responses: []*http.Response{grokMediaContentStatusResponse(`{"status":"completed"}`), {
|
||||
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.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{
|
||||
response: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(statusBody)),
|
||||
},
|
||||
}
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
|
||||
c, recorder := grokMediaContentTestContext(http.MethodGet, "https://api.example/v1/videos/task-1", map[string]string{
|
||||
"X-Forwarded-Host": "malicious.invalid",
|
||||
"X-Forwarded-Proto": "https",
|
||||
})
|
||||
|
||||
_, err := svc.ForwardGrokMedia(
|
||||
context.Background(), c, grokMediaContentTestAccount(),
|
||||
GrokMediaEndpointVideoStatus, "task-1", nil, "",
|
||||
)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.Equal(t, "/v1/videos/task-1/content", gjson.Get(recorder.Body.String(), "url").String())
|
||||
require.Equal(t, "/v1/videos/task-1/content", gjson.Get(recorder.Body.String(), "download_url").String())
|
||||
require.Equal(t, "https://vidgen.x.ai/task-1.mp4", gjson.Get(recorder.Body.String(), "video_url").String())
|
||||
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())
|
||||
}
|
||||
|
||||
func TestRewriteGrokMediaVideoContentURLsRewritesSignedVideoURL(t *testing.T) {
|
||||
body := []byte(`{"status":"done","video":{"url":"https://vidgen.x.ai/signed-token/xai-video-request-1.mp4","duration":8}}`)
|
||||
|
||||
rewritten := rewriteGrokMediaVideoContentURLs(body, "request-1", "/v1/videos/request-1/content")
|
||||
|
||||
require.Equal(t, "/v1/videos/request-1/content", gjson.GetBytes(rewritten, "video.url").String())
|
||||
require.Equal(t, "8", gjson.GetBytes(rewritten, "video.duration").String())
|
||||
require.Equal(t, "done", gjson.GetBytes(rewritten, "status").String())
|
||||
}
|
||||
@@ -112,6 +112,12 @@ func buildGrokMediaURL(account *Account, cfg *config.Config, endpoint GrokMediaE
|
||||
return xai.BuildVideosExtensionsURLWithValidator(baseURL, validator)
|
||||
case GrokMediaEndpointVideoStatus:
|
||||
return xai.BuildVideoURLWithValidator(baseURL, requestID, validator)
|
||||
case GrokMediaEndpointVideoContent:
|
||||
videoURL, err := xai.BuildVideoURLWithValidator(baseURL, requestID, validator)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return videoURL + "/content", nil
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported grok media endpoint: %s", endpoint)
|
||||
}
|
||||
|
||||
@@ -35,6 +35,10 @@ func TestGrokAPIKeyURLPolicyFollowsGlobalSecurityConfig(t *testing.T) {
|
||||
mediaURL, err := buildGrokMediaURL(account, cfg, GrokMediaEndpointImagesGenerations, "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "http://grok.example.test/v1/images/generations", mediaURL)
|
||||
|
||||
contentURL, err := buildGrokMediaURL(account, cfg, GrokMediaEndpointVideoContent, "request 123")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "http://grok.example.test/v1/videos/request%20123/content", contentURL)
|
||||
})
|
||||
|
||||
t.Run("insecure HTTP disabled", func(t *testing.T) {
|
||||
|
||||
@@ -1034,19 +1034,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) {
|
||||
|
||||
Reference in New Issue
Block a user