fix(api): 渠道请求头 - 补充自定义 Header 配置并应用到上游请求

Refs #39
2026-07-29 20:57:24 +08:00
parent 412399a473
commit 6de81049bf
21 changed files with 535 additions and 75 deletions
--- a/backend/internal/handler/auth.go
+++ b/backend/internal/handler/auth.go
@@ -659,6 +659,11 @@ func proxySystemRequest(c *gin.Context, svc *service.Service, user *model.User,
 		failService(c, err)
 		return
 	}
+	channelHeaders, err := service.ParseOutboundHeadersJSON(channel.HeadersJSON)
+	if err != nil {
+		failService(c, err)
+		return
+	}
 	// 同步代理与后台任务必须共享渠道槽位，否则两条入口会共同超过供应商并发上限。
 	releaseChannel, concurrencyLimit, err := svc.AcquireChannelSlot(c.Request.Context(), channel.ID, "", 36*time.Minute)
 	if err != nil {
@@ -694,6 +699,7 @@ func proxySystemRequest(c *gin.Context, svc *service.Service, user *model.User,
 	if accept := c.GetHeader("Accept"); accept != "" {
 		upstreamReq.Header.Set("Accept", accept)
 	}
+	service.ApplyOutboundHeaders(upstreamReq, channelHeaders)
 	service.ApplyDefaultOutboundHeaders(upstreamReq)
 	if protocol == model.ChannelInterfaceGeminiVeo {
 		upstreamReq.Header.Set("x-goog-api-key", channel.APIKey)
--- a/backend/internal/handler/custom_proxy.go
+++ b/backend/internal/handler/custom_proxy.go
@@ -64,6 +64,11 @@ func proxyCustomRelayRequest(c *gin.Context, policy service.RuntimeRequestPolicy
 		fail(c, http.StatusUnauthorized, err)
 		return
 	}
+	headers, err := service.DecodeRelayOutboundHeaders(c.GetHeader(service.CustomRelayHeadersHeader))
+	if err != nil {
+		failService(c, err)
+		return
+	}
 	requestLimit := policy.CustomRelayRequestMB << 20
 	if c.Request.ContentLength > requestLimit {
 		fail(c, http.StatusRequestEntityTooLarge, errors.New("自定义渠道请求超过配置上限"))
@@ -97,6 +102,7 @@ func proxyCustomRelayRequest(c *gin.Context, policy service.RuntimeRequestPolicy
 	} else {
 		upstreamReq.Header.Set("Accept", "application/json")
 	}
+	service.ApplyOutboundHeaders(upstreamReq, headers)
 	service.ApplyDefaultOutboundHeaders(upstreamReq)
 	if apiFormat == "gemini" {
 		upstreamReq.Header.Set("x-goog-api-key", apiKey)
@@ -110,10 +116,11 @@ func proxyCustomRelayRequest(c *gin.Context, policy service.RuntimeRequestPolicy
 		return
 	}
 	defer resp.Body.Close()
-	writeCustomRelayResponse(c, resp, apiKey, policy.CustomRelayResponseMB<<20)
+	allowBinary := customVideoContentPath.MatchString(target.EscapedPath()) || strings.HasSuffix(target.Path, "/audio/speech")
+	writeCustomRelayResponse(c, resp, apiKey, policy.CustomRelayResponseMB<<20, allowBinary)
 }

-func writeCustomRelayResponse(c *gin.Context, resp *http.Response, apiKey string, responseLimit int64) {
+func writeCustomRelayResponse(c *gin.Context, resp *http.Response, apiKey string, responseLimit int64, allowBinary bool) {
 	c.Header("Cache-Control", "no-store")
 	c.Header("X-Content-Type-Options", "nosniff")
 	mediaType, _, _ := mime.ParseMediaType(resp.Header.Get("Content-Type"))
@@ -129,6 +136,15 @@ func writeCustomRelayResponse(c *gin.Context, resp *http.Response, apiKey string
 		copyCustomRelayStream(c, resp.Body, apiKey, responseLimit)
 		return
 	}
+	if allowBinary && (strings.HasPrefix(mediaType, "video/") || strings.HasPrefix(mediaType, "audio/") || mediaType == "application/octet-stream") {
+		body, err := readLimitedRelayBody(resp.Body, responseLimit)
+		if err != nil {
+			fail(c, http.StatusBadGateway, errors.New("自定义渠道上游返回了过大的媒体文件"))
+			return
+		}
+		c.Data(resp.StatusCode, mediaType, body)
+		return
+	}
 	if mediaType != "application/json" && !strings.HasSuffix(mediaType, "+json") {
 		fail(c, http.StatusBadGateway, errors.New("自定义渠道上游返回了不支持的内容类型"))
 		return
--- a/backend/internal/handler/custom_proxy_test.go
+++ b/backend/internal/handler/custom_proxy_test.go
@@ -3,6 +3,7 @@ package handler
 import (
 	"bufio"
 	"bytes"
+	"encoding/base64"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -22,9 +23,12 @@ func TestCustomRelayForwardsOpenAIRequestWithoutBrowserHeaders(t *testing.T) {
 		if r.Header.Get("Authorization") != "Bearer "+apiKey {
 			t.Errorf("Authorization = %q", r.Header.Get("Authorization"))
 		}
-		if r.Header.Get("User-Agent") != service.DefaultOutboundUserAgent {
+		if r.Header.Get("User-Agent") != "Custom Relay Agent" {
 			t.Errorf("User-Agent = %q", r.Header.Get("User-Agent"))
 		}
+		if r.Header.Get("X-Gateway-Tenant") != "tenant-a" {
+			t.Errorf("X-Gateway-Tenant = %q", r.Header.Get("X-Gateway-Tenant"))
+		}
 		for _, name := range []string{"Cookie", "Origin", "Referer", "X-Canvas-Upstream-URL", "X-Forwarded-For"} {
 			if value := r.Header.Get(name); value != "" {
 				t.Errorf("upstream received %s = %q", name, value)
@@ -43,6 +47,7 @@ func TestCustomRelayForwardsOpenAIRequestWithoutBrowserHeaders(t *testing.T) {
 	request.Header.Set("Authorization", "Bearer "+apiKey)
 	request.Header.Set("X-Canvas-Upstream-URL", upstream.URL+"/v1/models")
 	request.Header.Set("X-Canvas-Upstream-Format", "openai")
+	request.Header.Set(service.CustomRelayHeadersHeader, base64.StdEncoding.EncodeToString([]byte(`[{"name":"User-Agent","value":"Custom Relay Agent"},{"name":"X-Gateway-Tenant","value":"tenant-a"}]`)))
 	request.Header.Set("Cookie", "browser=session")
 	request.Header.Set("Origin", "https://canvas.example.com")
 	response := httptest.NewRecorder()
@@ -61,6 +66,30 @@ func TestCustomRelayForwardsOpenAIRequestWithoutBrowserHeaders(t *testing.T) {
 	}
 }

+func TestCustomRelayReturnsVideoContent(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "video/mp4")
+		_, _ = w.Write([]byte("video-content"))
+	}))
+	defer upstream.Close()
+	useCustomRelayTestClient(t, upstream.Client())
+	t.Setenv("CANVAS_ALLOWED_PRIVATE_UPSTREAM_HOSTS", "127.0.0.1")
+
+	request := httptest.NewRequest(http.MethodGet, "/api/ai/custom", nil)
+	request.Header.Set("Authorization", "Bearer test-key")
+	request.Header.Set("X-Canvas-Upstream-URL", upstream.URL+"/v1/videos/task-1/content")
+	request.Header.Set("X-Canvas-Upstream-Format", "openai")
+	response := httptest.NewRecorder()
+	context, _ := gin.CreateTestContext(response)
+	context.Request = request
+
+	proxyCustomRelayRequest(context, defaultCustomRelayTestPolicy())
+	if response.Code != http.StatusOK || response.Header().Get("Content-Type") != "video/mp4" || response.Body.String() != "video-content" {
+		t.Fatalf("status = %d, content-type = %q, body = %q", response.Code, response.Header().Get("Content-Type"), response.Body.String())
+	}
+}
+
 func TestCustomRelayConvertsGeminiAuthentication(t *testing.T) {
 	gin.SetMode(gin.TestMode)
 	const apiKey = "gemini-secret-key"
--- a/backend/internal/handler/security.go
+++ b/backend/internal/handler/security.go
@@ -28,6 +28,7 @@ var (
 	customGeminiRelayPath     = regexp.MustCompile(`(?:^|/)models/[^/:]+:(generateContent|streamGenerateContent|predictLongRunning)$`)
 	customVideoTaskPath       = regexp.MustCompile(`(?:^|/)video/generations/[^/]+$`)
 	customXAIVideoTaskPath    = regexp.MustCompile(`(?:^|/)videos/[^/]+$`)
+	customVideoContentPath    = regexp.MustCompile(`(?:^|/)videos/[^/]+/content$`)
 	customArkVideoTaskPath    = regexp.MustCompile(`(?:^|/)contents/generations/tasks/[^/]+$`)
 	customGeminiOperationPath = regexp.MustCompile(`(?:^|/)(?:models/[^/]+/)?operations/[^/]+$`)
 	openAIPostEndpoints       = map[string]bool{
@@ -63,7 +64,7 @@ func authorizeCustomRelay(method string, target *url.URL, apiFormat string, cont
 	if method == http.MethodGet {
 		allowed := requestPath == "/models" || strings.HasSuffix(requestPath, "/models")
 		if apiFormat == "openai" {
-			allowed = allowed || customVideoTaskPath.MatchString(requestPath) || customXAIVideoTaskPath.MatchString(requestPath) || customArkVideoTaskPath.MatchString(requestPath)
+			allowed = allowed || customVideoTaskPath.MatchString(requestPath) || customXAIVideoTaskPath.MatchString(requestPath) || customVideoContentPath.MatchString(requestPath) || customArkVideoTaskPath.MatchString(requestPath)
 		} else {
 			allowed = allowed || customGeminiOperationPath.MatchString(requestPath)
 		}
@@ -72,19 +73,24 @@ func authorizeCustomRelay(method string, target *url.URL, apiFormat string, cont
 		}
 		return nil
 	}
-	mediaType, _, err := mime.ParseMediaType(contentType)
-	if err != nil || mediaType != "application/json" {
-		return errors.New("自定义渠道生成请求必须使用 application/json")
-	}
 	if method != http.MethodPost {
 		return errors.New("自定义渠道不允许使用该请求方法")
 	}
+	mediaType, _, err := mime.ParseMediaType(contentType)
+	if err != nil {
+		return errors.New("自定义渠道生成请求类型无效")
+	}
 	if apiFormat == "openai" {
-		if len(query) != 0 || (!strings.HasSuffix(requestPath, "/responses") && !strings.HasSuffix(requestPath, "/chat/completions") && !strings.HasSuffix(requestPath, "/images/generations") && !strings.HasSuffix(requestPath, "/audio/speech") && !strings.HasSuffix(requestPath, "/video/generations") && !strings.HasSuffix(requestPath, "/videos/generations") && !strings.HasSuffix(requestPath, "/contents/generations/tasks")) {
+		multipartAllowed := mediaType == "multipart/form-data" && (strings.HasSuffix(requestPath, "/images/edits") || strings.HasSuffix(requestPath, "/videos"))
+		jsonAllowed := mediaType == "application/json" && (strings.HasSuffix(requestPath, "/responses") || strings.HasSuffix(requestPath, "/chat/completions") || strings.HasSuffix(requestPath, "/images/generations") || strings.HasSuffix(requestPath, "/audio/speech") || strings.HasSuffix(requestPath, "/video/generations") || strings.HasSuffix(requestPath, "/videos/generations") || strings.HasSuffix(requestPath, "/videos") || strings.HasSuffix(requestPath, "/contents/generations/tasks"))
+		if len(query) != 0 || (!multipartAllowed && !jsonAllowed) {
 			return errors.New("自定义渠道不允许访问该上游接口")
 		}
 		return nil
 	}
+	if mediaType != "application/json" {
+		return errors.New("Gemini 自定义渠道生成请求必须使用 application/json")
+	}
 	if !customGeminiRelayPath.MatchString(requestPath) {
 		return errors.New("自定义渠道不允许访问该上游接口")
 	}
--- a/backend/internal/handler/security_test.go
+++ b/backend/internal/handler/security_test.go
@@ -29,6 +29,10 @@ func TestAuthorizeCustomRelayAllowsModelsAndAgentEndpoints(t *testing.T) {
 		{method: http.MethodPost, target: "https://api.example.com/v1/responses", apiFormat: "openai", contentType: "application/json"},
 		{method: http.MethodPost, target: "https://api.example.com/v1/chat/completions", apiFormat: "openai", contentType: "application/json; charset=utf-8"},
 		{method: http.MethodPost, target: "https://api.example.com/v1/audio/speech", apiFormat: "openai", contentType: "application/json"},
+		{method: http.MethodPost, target: "https://api.example.com/v1/images/edits", apiFormat: "openai", contentType: "multipart/form-data; boundary=test"},
+		{method: http.MethodPost, target: "https://api.example.com/v1/videos", apiFormat: "openai", contentType: "application/json"},
+		{method: http.MethodPost, target: "https://api.example.com/v1/videos", apiFormat: "openai", contentType: "multipart/form-data; boundary=test"},
+		{method: http.MethodGet, target: "https://api.example.com/v1/videos/task-1/content", apiFormat: "openai"},
 		{method: http.MethodPost, target: "https://api.example.com/v1/video/generations", apiFormat: "openai", contentType: "application/json"},
 		{method: http.MethodGet, target: "https://api.example.com/v1/video/generations/task-1", apiFormat: "openai"},
 		{method: http.MethodPost, target: "https://api.x.ai/v1/videos/generations", apiFormat: "openai", contentType: "application/json"},
@@ -62,6 +66,7 @@ func TestAuthorizeCustomRelayRejectsArbitraryRequestsAndCredentialQueries(t *tes
 		{method: http.MethodGet, target: "https://api.example.com/account", apiFormat: "openai"},
 		{method: http.MethodGet, target: "https://api.example.com/v1/models?api_key=secret", apiFormat: "openai"},
 		{method: http.MethodPost, target: "https://api.example.com/v1/responses", apiFormat: "openai", contentType: "text/plain"},
+		{method: http.MethodPost, target: "https://api.example.com/v1/account", apiFormat: "openai", contentType: "multipart/form-data; boundary=test"},
 		{method: http.MethodPost, target: "https://api.example.com/v1/../account/chat/completions", apiFormat: "openai", contentType: "application/json"},
 		{method: http.MethodPost, target: "https://api.example.com/v1/models/gemini:streamGenerateContent?alt=sse&token=secret", apiFormat: "gemini", contentType: "application/json"},
 	}