Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 57 additions & 0 deletions controller/video_proxy.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,12 @@ package controller

import (
"context"
"encoding/base64"
"fmt"
"io"
"net/http"
"net/url"
"strings"
"time"

"github.com/QuantumNous/new-api/constant"
Expand Down Expand Up @@ -94,6 +96,13 @@ func VideoProxy(c *gin.Context) {
return
}
req.Header.Set("x-goog-api-key", apiKey)
case constant.ChannelTypeVertexAi:
videoURL, err = getVertexVideoURL(channel, task)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to resolve Vertex video URL for task %s: %s", taskID, err.Error()))
videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to resolve Vertex video URL")
return
}
case constant.ChannelTypeOpenAI, constant.ChannelTypeSora:
videoURL = fmt.Sprintf("%s/v1/videos/%s/content", baseURL, task.GetUpstreamTaskID())
req.Header.Set("Authorization", "Bearer "+channel.Key)
Expand All @@ -102,6 +111,21 @@ func VideoProxy(c *gin.Context) {
videoURL = task.GetResultURL()
}

videoURL = strings.TrimSpace(videoURL)
if videoURL == "" {
logger.LogError(c.Request.Context(), fmt.Sprintf("Video URL is empty for task %s", taskID))
videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to fetch video content")
return
}

if strings.HasPrefix(videoURL, "data:") {
if err := writeVideoDataURL(c, videoURL); err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to decode video data URL for task %s: %s", taskID, err.Error()))
videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to fetch video content")
}
return
}

req.URL, err = url.Parse(videoURL)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to parse URL %s: %s", videoURL, err.Error()))
Expand Down Expand Up @@ -136,3 +160,36 @@ func VideoProxy(c *gin.Context) {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to stream video content: %s", err.Error()))
}
}

func writeVideoDataURL(c *gin.Context, dataURL string) error {
parts := strings.SplitN(dataURL, ",", 2)
if len(parts) != 2 {
return fmt.Errorf("invalid data url")
}

header := parts[0]
payload := parts[1]
if !strings.HasPrefix(header, "data:") || !strings.Contains(header, ";base64") {
return fmt.Errorf("unsupported data url")
}

mimeType := strings.TrimPrefix(header, "data:")
mimeType = strings.TrimSuffix(mimeType, ";base64")
if mimeType == "" {
mimeType = "video/mp4"
}

videoBytes, err := base64.StdEncoding.DecodeString(payload)
if err != nil {
videoBytes, err = base64.RawStdEncoding.DecodeString(payload)
if err != nil {
return err
}
}

c.Writer.Header().Set("Content-Type", mimeType)
c.Writer.Header().Set("Cache-Control", "public, max-age=86400")
c.Writer.WriteHeader(http.StatusOK)
_, err = c.Writer.Write(videoBytes)
return err
Comment on lines +164 to +194

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Guard decoded payload size before allocating full video bytes.

writeVideoDataURL currently decodes the whole base64 payload into memory. Large inline videos can cause high memory pressure or OOM under concurrency.

💡 Proposed fix
 func writeVideoDataURL(c *gin.Context, dataURL string) error {
+	const maxInlineVideoBytes = 64 << 20 // 64 MiB decoded payload cap
 	parts := strings.SplitN(dataURL, ",", 2)
 	if len(parts) != 2 {
 		return fmt.Errorf("invalid data url")
 	}
@@
 	mimeType := strings.TrimPrefix(header, "data:")
 	mimeType = strings.TrimSuffix(mimeType, ";base64")
 	if mimeType == "" {
 		mimeType = "video/mp4"
 	}
+
+	if base64.StdEncoding.DecodedLen(len(payload)) > maxInlineVideoBytes {
+		return fmt.Errorf("data url payload too large")
+	}
 
 	videoBytes, err := base64.StdEncoding.DecodeString(payload)
 	if err != nil {
 		videoBytes, err = base64.RawStdEncoding.DecodeString(payload)
 		if err != nil {
 			return err
 		}
 	}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
func writeVideoDataURL(c *gin.Context, dataURL string) error {
parts := strings.SplitN(dataURL, ",", 2)
if len(parts) != 2 {
return fmt.Errorf("invalid data url")
}
header := parts[0]
payload := parts[1]
if !strings.HasPrefix(header, "data:") || !strings.Contains(header, ";base64") {
return fmt.Errorf("unsupported data url")
}
mimeType := strings.TrimPrefix(header, "data:")
mimeType = strings.TrimSuffix(mimeType, ";base64")
if mimeType == "" {
mimeType = "video/mp4"
}
videoBytes, err := base64.StdEncoding.DecodeString(payload)
if err != nil {
videoBytes, err = base64.RawStdEncoding.DecodeString(payload)
if err != nil {
return err
}
}
c.Writer.Header().Set("Content-Type", mimeType)
c.Writer.Header().Set("Cache-Control", "public, max-age=86400")
c.Writer.WriteHeader(http.StatusOK)
_, err = c.Writer.Write(videoBytes)
return err
func writeVideoDataURL(c *gin.Context, dataURL string) error {
const maxInlineVideoBytes = 64 << 20 // 64 MiB decoded payload cap
parts := strings.SplitN(dataURL, ",", 2)
if len(parts) != 2 {
return fmt.Errorf("invalid data url")
}
header := parts[0]
payload := parts[1]
if !strings.HasPrefix(header, "data:") || !strings.Contains(header, ";base64") {
return fmt.Errorf("unsupported data url")
}
mimeType := strings.TrimPrefix(header, "data:")
mimeType = strings.TrimSuffix(mimeType, ";base64")
if mimeType == "" {
mimeType = "video/mp4"
}
if base64.StdEncoding.DecodedLen(len(payload)) > maxInlineVideoBytes {
return fmt.Errorf("data url payload too large")
}
videoBytes, err := base64.StdEncoding.DecodeString(payload)
if err != nil {
videoBytes, err = base64.RawStdEncoding.DecodeString(payload)
if err != nil {
return err
}
}
c.Writer.Header().Set("Content-Type", mimeType)
c.Writer.Header().Set("Cache-Control", "public, max-age=86400")
c.Writer.WriteHeader(http.StatusOK)
_, err = c.Writer.Write(videoBytes)
return err
}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@controller/video_proxy.go` around lines 164 - 194, writeVideoDataURL decodes
the entire base64 payload into memory (payload) which can OOM; before
allocating, estimate the decoded size from payload length (or set a maxBytes
constant) and reject/return an error if it exceeds the limit, and then stream
decode instead of decoding into []byte. Locate writeVideoDataURL and replace the
direct base64.StdEncoding.DecodeString(payload) usage with a guarded approach:
compute expectedDecodedSize := (len(payload) * 3) / 4 (adjust for padding) or
compare len(payload) against maxBase64Len, return an error if it exceeds allowed
size, and then use a streaming decoder (base64.NewDecoder) to copy into c.Writer
with an io.LimitedReader / io.CopyN to enforce the same limit rather than
allocating videoBytes. Ensure headers (Content-Type, Cache-Control) are set
before streaming and any decode errors are handled and returned.

}
128 changes: 128 additions & 0 deletions controller/video_proxy_gemini.go
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,134 @@ func extractGeminiVideoURLFromGeneratedSamples(gvr map[string]any) string {
return ""
}

func getVertexVideoURL(channel *model.Channel, task *model.Task) (string, error) {
if channel == nil || task == nil {
return "", fmt.Errorf("invalid channel or task")
}
if url := strings.TrimSpace(task.GetResultURL()); url != "" {
return url, nil
}
if url := extractVertexVideoURLFromTaskData(task); url != "" {
return url, nil
}

baseURL := constant.ChannelBaseURLs[channel.Type]
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}

adaptor := relay.GetTaskAdaptor(constant.TaskPlatform(strconv.Itoa(channel.Type)))
if adaptor == nil {
return "", fmt.Errorf("vertex task adaptor not found")
}

key := getVertexTaskKey(channel, task)
if key == "" {
return "", fmt.Errorf("vertex key not available for task")
}

resp, err := adaptor.FetchTask(baseURL, key, map[string]any{
"task_id": task.GetUpstreamTaskID(),
"action": task.Action,
}, channel.GetSetting().Proxy)
if err != nil {
return "", fmt.Errorf("fetch task failed: %w", err)
}
defer resp.Body.Close()

body, err := io.ReadAll(resp.Body)
if err != nil {
return "", fmt.Errorf("read task response failed: %w", err)
}

taskInfo, parseErr := adaptor.ParseTaskResult(body)
if parseErr == nil && taskInfo != nil && strings.TrimSpace(taskInfo.Url) != "" {
return taskInfo.Url, nil
Comment on lines +189 to +190

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Return the trimmed URL after validation.

You validate TrimSpace(taskInfo.Url) but return taskInfo.Url directly. Returning the trimmed value avoids avoidable URL-parse failures from surrounding whitespace.

💡 Proposed fix
-	if parseErr == nil && taskInfo != nil && strings.TrimSpace(taskInfo.Url) != "" {
-		return taskInfo.Url, nil
+	if parseErr == nil && taskInfo != nil && strings.TrimSpace(taskInfo.Url) != "" {
+		return strings.TrimSpace(taskInfo.Url), nil
 	}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if parseErr == nil && taskInfo != nil && strings.TrimSpace(taskInfo.Url) != "" {
return taskInfo.Url, nil
if parseErr == nil && taskInfo != nil && strings.TrimSpace(taskInfo.Url) != "" {
return strings.TrimSpace(taskInfo.Url), nil
}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@controller/video_proxy_gemini.go` around lines 184 - 185, The code validates
strings.TrimSpace(taskInfo.Url) but returns taskInfo.Url, so update the return
to use the trimmed URL: compute a trimmed variable (e.g., trimmedURL :=
strings.TrimSpace(taskInfo.Url)) after the parseErr check and return trimmedURL
(and nil error) instead of taskInfo.Url; adjust any subsequent use of
taskInfo.Url in this function (in controller/video_proxy_gemini.go) to use the
trimmed variable to avoid whitespace-related parse failures.

}
if url := extractVertexVideoURLFromPayload(body); url != "" {
return url, nil
}
if parseErr != nil {
return "", fmt.Errorf("parse task result failed: %w", parseErr)
}
return "", fmt.Errorf("vertex video url not found")
}

func getVertexTaskKey(channel *model.Channel, task *model.Task) string {
if task != nil {
if key := strings.TrimSpace(task.PrivateData.Key); key != "" {
return key
}
}
if channel == nil {
return ""
}
keys := channel.GetKeys()
for _, key := range keys {
key = strings.TrimSpace(key)
if key != "" {
return key
}
}
return strings.TrimSpace(channel.Key)
}

func extractVertexVideoURLFromTaskData(task *model.Task) string {
if task == nil || len(task.Data) == 0 {
return ""
}
return extractVertexVideoURLFromPayload(task.Data)
}

func extractVertexVideoURLFromPayload(body []byte) string {
var payload map[string]any
if err := common.Unmarshal(body, &payload); err != nil {
return ""
}
resp, ok := payload["response"].(map[string]any)
if !ok || resp == nil {
return ""
}

if videos, ok := resp["videos"].([]any); ok && len(videos) > 0 {
if video, ok := videos[0].(map[string]any); ok && video != nil {
if b64, _ := video["bytesBase64Encoded"].(string); strings.TrimSpace(b64) != "" {
mime, _ := video["mimeType"].(string)
enc, _ := video["encoding"].(string)
return buildVideoDataURL(mime, enc, b64)
}
}
Comment on lines +237 to +244

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

videos parsing currently misses common valid URL shapes.

The logic only checks videos[0] and only bytesBase64Encoded. If uri exists (or the valid item is later in the array), this returns empty and can incorrectly fail resolution.

💡 Proposed fix
-	if videos, ok := resp["videos"].([]any); ok && len(videos) > 0 {
-		if video, ok := videos[0].(map[string]any); ok && video != nil {
-			if b64, _ := video["bytesBase64Encoded"].(string); strings.TrimSpace(b64) != "" {
-				mime, _ := video["mimeType"].(string)
-				enc, _ := video["encoding"].(string)
-				return buildVideoDataURL(mime, enc, b64)
-			}
-		}
-	}
+	if videos, ok := resp["videos"].([]any); ok {
+		for _, item := range videos {
+			video, ok := item.(map[string]any)
+			if !ok || video == nil {
+				continue
+			}
+			if uri, _ := video["uri"].(string); strings.TrimSpace(uri) != "" {
+				return strings.TrimSpace(uri)
+			}
+			if b64, _ := video["bytesBase64Encoded"].(string); strings.TrimSpace(b64) != "" {
+				mime, _ := video["mimeType"].(string)
+				enc, _ := video["encoding"].(string)
+				return buildVideoDataURL(mime, enc, b64)
+			}
+		}
+	}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if videos, ok := resp["videos"].([]any); ok && len(videos) > 0 {
if video, ok := videos[0].(map[string]any); ok && video != nil {
if b64, _ := video["bytesBase64Encoded"].(string); strings.TrimSpace(b64) != "" {
mime, _ := video["mimeType"].(string)
enc, _ := video["encoding"].(string)
return buildVideoDataURL(mime, enc, b64)
}
}
if videos, ok := resp["videos"].([]any); ok {
for _, item := range videos {
video, ok := item.(map[string]any)
if !ok || video == nil {
continue
}
if uri, _ := video["uri"].(string); strings.TrimSpace(uri) != "" {
return strings.TrimSpace(uri)
}
if b64, _ := video["bytesBase64Encoded"].(string); strings.TrimSpace(b64) != "" {
mime, _ := video["mimeType"].(string)
enc, _ := video["encoding"].(string)
return buildVideoDataURL(mime, enc, b64)
}
}
}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@controller/video_proxy_gemini.go` around lines 213 - 220, The current parsing
only inspects videos[0] and only the bytesBase64Encoded field, so it misses
later entries and valid uri shapes; update the parsing of resp["videos"] in the
handler to iterate all items in the videos slice, for each item (map[string]any)
first check bytesBase64Encoded and if present build and return the data URL via
buildVideoDataURL(mimeType, encoding, bytesBase64Encoded), otherwise check for a
uri string and return that URI (ensuring it is non-empty); reference the resp
variable, the videos slice, each video map, and buildVideoDataURL when making
the change.

}
if b64, _ := resp["bytesBase64Encoded"].(string); strings.TrimSpace(b64) != "" {
enc, _ := resp["encoding"].(string)
return buildVideoDataURL("", enc, b64)
}
if video, _ := resp["video"].(string); strings.TrimSpace(video) != "" {
if strings.HasPrefix(video, "data:") || strings.HasPrefix(video, "http://") || strings.HasPrefix(video, "https://") {
return video
}
enc, _ := resp["encoding"].(string)
return buildVideoDataURL("", enc, video)
Comment on lines +250 to +255

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Normalize video before scheme checks and return.

Whitespace is stripped for emptiness only; scheme checks and return values should use the trimmed value too.

💡 Proposed fix
-	if video, _ := resp["video"].(string); strings.TrimSpace(video) != "" {
-		if strings.HasPrefix(video, "data:") || strings.HasPrefix(video, "http://") || strings.HasPrefix(video, "https://") {
-			return video
+	if video, _ := resp["video"].(string); strings.TrimSpace(video) != "" {
+		trimmedVideo := strings.TrimSpace(video)
+		if strings.HasPrefix(trimmedVideo, "data:") || strings.HasPrefix(trimmedVideo, "http://") || strings.HasPrefix(trimmedVideo, "https://") {
+			return trimmedVideo
 		}
 		enc, _ := resp["encoding"].(string)
-		return buildVideoDataURL("", enc, video)
+		return buildVideoDataURL("", enc, trimmedVideo)
 	}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if video, _ := resp["video"].(string); strings.TrimSpace(video) != "" {
if strings.HasPrefix(video, "data:") || strings.HasPrefix(video, "http://") || strings.HasPrefix(video, "https://") {
return video
}
enc, _ := resp["encoding"].(string)
return buildVideoDataURL("", enc, video)
if video, _ := resp["video"].(string); strings.TrimSpace(video) != "" {
trimmedVideo := strings.TrimSpace(video)
if strings.HasPrefix(trimmedVideo, "data:") || strings.HasPrefix(trimmedVideo, "http://") || strings.HasPrefix(trimmedVideo, "https://") {
return trimmedVideo
}
enc, _ := resp["encoding"].(string)
return buildVideoDataURL("", enc, trimmedVideo)
}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@controller/video_proxy_gemini.go` around lines 250 - 255, Trim the extracted
video string immediately and use that trimmed value for emptiness and scheme
checks and when returning or passing to buildVideoDataURL; specifically, in the
block that reads video from resp (the variable named video) replace usages of
the original untrimmed value with a trimmedVideo (or reassign video =
strings.TrimSpace(video)) so the HasPrefix checks and the returned value (or
call to buildVideoDataURL) use the normalized string, and ensure encoding is
read as before from resp["encoding"] for the buildVideoDataURL call.

}
return ""
}

func buildVideoDataURL(mimeType string, encoding string, base64Data string) string {
mime := strings.TrimSpace(mimeType)
if mime == "" {
enc := strings.TrimSpace(encoding)
if enc == "" {
enc = "mp4"
}
if strings.Contains(enc, "/") {
mime = enc
} else {
mime = "video/" + enc
}
Comment on lines +260 to +271

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Do not infer MIME as video/<encoding> for encoding markers like base64.

When encoding is a transport marker (for example, base64), this generates invalid content types such as video/base64, which can break playback/content handling.

💡 Proposed fix
 func buildVideoDataURL(mimeType string, encoding string, base64Data string) string {
 	mime := strings.TrimSpace(mimeType)
 	if mime == "" {
-		enc := strings.TrimSpace(encoding)
-		if enc == "" {
-			enc = "mp4"
-		}
-		if strings.Contains(enc, "/") {
-			mime = enc
-		} else {
-			mime = "video/" + enc
-		}
+		enc := strings.ToLower(strings.TrimSpace(encoding))
+		switch enc {
+		case "", "base64":
+			mime = "video/mp4"
+		default:
+			if strings.Contains(enc, "/") {
+				mime = enc
+			} else {
+				mime = "video/" + enc
+			}
+		}
 	}
 	return "data:" + mime + ";base64," + base64Data
 }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
func buildVideoDataURL(mimeType string, encoding string, base64Data string) string {
mime := strings.TrimSpace(mimeType)
if mime == "" {
enc := strings.TrimSpace(encoding)
if enc == "" {
enc = "mp4"
}
if strings.Contains(enc, "/") {
mime = enc
} else {
mime = "video/" + enc
}
func buildVideoDataURL(mimeType string, encoding string, base64Data string) string {
mime := strings.TrimSpace(mimeType)
if mime == "" {
enc := strings.ToLower(strings.TrimSpace(encoding))
switch enc {
case "", "base64":
mime = "video/mp4"
default:
if strings.Contains(enc, "/") {
mime = enc
} else {
mime = "video/" + enc
}
}
}
return "data:" + mime + ";base64," + base64Data
}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@controller/video_proxy_gemini.go` around lines 236 - 247, The
buildVideoDataURL function currently infers mime as "video/<encoding>" when
mimeType is empty, which incorrectly produces types like "video/base64"; update
buildVideoDataURL to treat encoding as a MIME subtype only if it looks like a
real mime (contains '/' ) or matches known video file extensions (e.g.,
"mp4","webm","mov","mkv","avi"); otherwise treat encodings like "base64" as
transport markers and fall back to a safe default MIME such as "video/mp4" (or
use encoding when it contains a slash), ensuring the logic around mimeType,
encoding and the enc variable in buildVideoDataURL is adjusted accordingly.

}
return "data:" + mime + ";base64," + base64Data
}

func ensureAPIKey(uri, key string) string {
if key == "" || uri == "" {
return uri
Expand Down
3 changes: 2 additions & 1 deletion model/task.go
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,8 @@ func InitTask(platform constant.TaskPlatform, relayInfo *commonRelay.RelayInfo)
properties := Properties{}
privateData := TaskPrivateData{}
if relayInfo != nil && relayInfo.ChannelMeta != nil {
if relayInfo.ChannelMeta.ChannelType == constant.ChannelTypeGemini {
if relayInfo.ChannelMeta.ChannelType == constant.ChannelTypeGemini ||
relayInfo.ChannelMeta.ChannelType == constant.ChannelTypeVertexAi {
privateData.Key = relayInfo.ChannelMeta.ApiKey
}
if relayInfo.UpstreamModelName != "" {
Expand Down