fix(api): 自定义渠道 - 移除浏览器 CORS 依赖

2026-07-22 21:44:44 +08:00
parent 180baf4be8
commit a5bc71c194
9 changed files with 820 additions and 28 deletions
--- a/backend/internal/handler/custom_proxy.go
+++ b/backend/internal/handler/custom_proxy.go
@@ -0,0 +1,287 @@
+package handler
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"io"
+	"mime"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"infinite-canvas/backend/internal/service"
+
+	"github.com/gin-gonic/gin"
+)
+
+const (
+	maxCustomRelayBodyBytes          int64 = 32 << 20
+	maxCustomRelayModelResponseBytes int64 = 4 << 20
+	maxCustomRelayJSONResponseBytes  int64 = 16 << 20
+	maxCustomRelayErrorResponseBytes int64 = 64 << 10
+	maxCustomRelayStreamBytes        int64 = 32 << 20
+	maxCustomRelayConcurrency              = 4
+	customRelayHTTPTimeout                 = 10 * time.Minute
+)
+
+var (
+	customRelayClient = service.CustomRelayHTTPClient
+	customRelayActive = struct {
+		sync.Mutex
+		users map[string]int
+	}{users: map[string]int{}}
+)
+
+func RegisterCustomRelayRoutes(r *gin.RouterGroup, svc *service.Service) {
+	r.Any("/ai/custom", func(c *gin.Context) {
+		user, err := currentUser(c, svc)
+		if err != nil {
+			failService(c, err)
+			return
+		}
+		if !enforceRateLimit(c, "custom-relay:"+user.ID, 120, time.Minute) {
+			return
+		}
+		if !acquireCustomRelaySlot(user.ID) {
+			fail(c, http.StatusTooManyRequests, errors.New("自定义渠道并发请求过多，请等待已有请求完成"))
+			return
+		}
+		defer releaseCustomRelaySlot(user.ID)
+		proxyCustomRelayRequest(c)
+	})
+}
+
+func proxyCustomRelayRequest(c *gin.Context) {
+	target, err := service.ValidateCustomRelayURL(c.GetHeader("X-Canvas-Upstream-URL"))
+	if err != nil {
+		failService(c, err)
+		return
+	}
+	apiFormat := strings.ToLower(strings.TrimSpace(c.GetHeader("X-Canvas-Upstream-Format")))
+	if apiFormat == "" {
+		apiFormat = "openai"
+	}
+	if err := authorizeCustomRelay(c.Request.Method, target, apiFormat, c.GetHeader("Content-Type")); err != nil {
+		fail(c, http.StatusForbidden, err)
+		return
+	}
+	apiKey, err := customRelayAPIKey(c.GetHeader("Authorization"))
+	if err != nil {
+		fail(c, http.StatusUnauthorized, err)
+		return
+	}
+	if c.Request.ContentLength > maxCustomRelayBodyBytes {
+		fail(c, http.StatusRequestEntityTooLarge, errors.New("自定义渠道请求不能超过 32MB"))
+		return
+	}
+	c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, maxCustomRelayBodyBytes)
+	body, err := io.ReadAll(c.Request.Body)
+	if err != nil {
+		var maxBytesError *http.MaxBytesError
+		if errors.As(err, &maxBytesError) {
+			fail(c, http.StatusRequestEntityTooLarge, errors.New("自定义渠道请求不能超过 32MB"))
+			return
+		}
+		fail(c, http.StatusBadRequest, errors.New("读取自定义渠道请求失败"))
+		return
+	}
+	if c.Request.Method == http.MethodGet && len(body) != 0 {
+		fail(c, http.StatusBadRequest, errors.New("模型列表请求不允许携带请求体"))
+		return
+	}
+	upstreamReq, err := http.NewRequestWithContext(c.Request.Context(), c.Request.Method, target.String(), bytes.NewReader(body))
+	if err != nil {
+		fail(c, http.StatusBadRequest, errors.New("构造自定义渠道请求失败"))
+		return
+	}
+	if contentType := c.GetHeader("Content-Type"); contentType != "" {
+		upstreamReq.Header.Set("Content-Type", contentType)
+	}
+	if strings.Contains(strings.ToLower(c.GetHeader("Accept")), "text/event-stream") {
+		upstreamReq.Header.Set("Accept", "text/event-stream")
+	} else {
+		upstreamReq.Header.Set("Accept", "application/json")
+	}
+	upstreamReq.Header.Set("User-Agent", "InfiniteCanvas/custom-channel-relay")
+	if apiFormat == "gemini" {
+		upstreamReq.Header.Set("x-goog-api-key", apiKey)
+	} else {
+		upstreamReq.Header.Set("Authorization", "Bearer "+apiKey)
+	}
+
+	resp, err := customRelayClient(customRelayHTTPTimeout).Do(upstreamReq)
+	if err != nil {
+		fail(c, http.StatusBadGateway, errors.New("自定义渠道上游连接失败"))
+		return
+	}
+	defer resp.Body.Close()
+	writeCustomRelayResponse(c, resp, apiKey, c.Request.Method == http.MethodGet)
+}
+
+func writeCustomRelayResponse(c *gin.Context, resp *http.Response, apiKey string, modelList bool) {
+	c.Header("Cache-Control", "no-store")
+	c.Header("X-Content-Type-Options", "nosniff")
+	mediaType, _, _ := mime.ParseMediaType(resp.Header.Get("Content-Type"))
+	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
+		writeCustomRelayError(c, resp, apiKey, mediaType)
+		return
+	}
+	if mediaType == "text/event-stream" {
+		c.Header("Content-Type", "text/event-stream; charset=utf-8")
+		c.Header("X-Accel-Buffering", "no")
+		c.Status(resp.StatusCode)
+		c.Writer.WriteHeaderNow()
+		copyCustomRelayStream(c, resp.Body, apiKey)
+		return
+	}
+	if mediaType != "application/json" && !strings.HasSuffix(mediaType, "+json") {
+		fail(c, http.StatusBadGateway, errors.New("自定义渠道上游返回了不支持的内容类型"))
+		return
+	}
+	limit := maxCustomRelayJSONResponseBytes
+	if modelList {
+		limit = maxCustomRelayModelResponseBytes
+	}
+	body, err := readLimitedRelayBody(resp.Body, limit)
+	if err != nil || !json.Valid(body) {
+		fail(c, http.StatusBadGateway, errors.New("自定义渠道上游返回无效或过大的 JSON"))
+		return
+	}
+	body = redactRelaySecret(body, apiKey)
+	c.Data(resp.StatusCode, "application/json; charset=utf-8", body)
+}
+
+func writeCustomRelayError(c *gin.Context, resp *http.Response, apiKey string, mediaType string) {
+	body, err := readLimitedRelayBody(resp.Body, maxCustomRelayErrorResponseBytes)
+	if err != nil {
+		fail(c, http.StatusBadGateway, errors.New("自定义渠道上游请求失败"))
+		return
+	}
+	body = redactRelaySecret(body, apiKey)
+	if (mediaType == "application/json" || strings.HasSuffix(mediaType, "+json")) && json.Valid(body) {
+		c.Data(resp.StatusCode, "application/json; charset=utf-8", body)
+		return
+	}
+	fail(c, resp.StatusCode, errors.New("自定义渠道上游请求失败"))
+}
+
+func copyCustomRelayStream(c *gin.Context, source io.Reader, apiKey string) {
+	redactor := newRelayStreamRedactor(apiKey)
+	buffer := make([]byte, 32<<10)
+	var written int64
+	for written < maxCustomRelayStreamBytes {
+		read, err := source.Read(buffer)
+		if read > 0 {
+			remaining := maxCustomRelayStreamBytes - written
+			if int64(read) > remaining {
+				read = int(remaining)
+			}
+			chunk := redactor.Push(buffer[:read], false)
+			if len(chunk) > 0 {
+				if _, writeErr := c.Writer.Write(chunk); writeErr != nil {
+					return
+				}
+				c.Writer.Flush()
+			}
+			written += int64(read)
+		}
+		if err != nil {
+			break
+		}
+	}
+	if tail := redactor.Push(nil, true); len(tail) > 0 {
+		_, _ = c.Writer.Write(tail)
+		c.Writer.Flush()
+	}
+}
+
+func readLimitedRelayBody(body io.Reader, limit int64) ([]byte, error) {
+	data, err := io.ReadAll(io.LimitReader(body, limit+1))
+	if err != nil {
+		return nil, err
+	}
+	if int64(len(data)) > limit {
+		return nil, errors.New("response body is too large")
+	}
+	return data, nil
+}
+
+func customRelayAPIKey(value string) (string, error) {
+	scheme, apiKey, found := strings.Cut(strings.TrimSpace(value), " ")
+	apiKey = strings.TrimSpace(apiKey)
+	if !found || !strings.EqualFold(scheme, "Bearer") || apiKey == "" || len(apiKey) > 512 || strings.ContainsAny(apiKey, "\r\n") {
+		return "", errors.New("自定义渠道 API Key 无效")
+	}
+	return apiKey, nil
+}
+
+func acquireCustomRelaySlot(userID string) bool {
+	customRelayActive.Lock()
+	defer customRelayActive.Unlock()
+	if customRelayActive.users[userID] >= maxCustomRelayConcurrency {
+		return false
+	}
+	customRelayActive.users[userID]++
+	return true
+}
+
+func releaseCustomRelaySlot(userID string) {
+	customRelayActive.Lock()
+	defer customRelayActive.Unlock()
+	if customRelayActive.users[userID] <= 1 {
+		delete(customRelayActive.users, userID)
+		return
+	}
+	customRelayActive.users[userID]--
+}
+
+func redactRelaySecret(body []byte, apiKey string) []byte {
+	if apiKey == "" {
+		return body
+	}
+	return bytes.ReplaceAll(body, []byte(apiKey), []byte("[REDACTED]"))
+}
+
+type relayStreamRedactor struct {
+	secret  []byte
+	pending []byte
+}
+
+func newRelayStreamRedactor(secret string) *relayStreamRedactor {
+	return &relayStreamRedactor{secret: []byte(secret)}
+}
+
+func (r *relayStreamRedactor) Push(chunk []byte, final bool) []byte {
+	r.pending = append(r.pending, chunk...)
+	if len(r.secret) == 0 {
+		result := append([]byte(nil), r.pending...)
+		r.pending = r.pending[:0]
+		return result
+	}
+	r.pending = bytes.ReplaceAll(r.pending, r.secret, []byte("[REDACTED]"))
+	if final {
+		result := append([]byte(nil), r.pending...)
+		r.pending = r.pending[:0]
+		return result
+	}
+	keep := relaySecretPrefixSuffixLength(r.pending, r.secret)
+	cut := len(r.pending) - keep
+	result := append([]byte(nil), r.pending[:cut]...)
+	r.pending = append(r.pending[:0], r.pending[cut:]...)
+	return result
+}
+
+func relaySecretPrefixSuffixLength(data []byte, secret []byte) int {
+	limit := len(secret) - 1
+	if len(data) < limit {
+		limit = len(data)
+	}
+	for length := limit; length > 0; length-- {
+		if bytes.Equal(data[len(data)-length:], secret[:length]) {
+			return length
+		}
+	}
+	return 0
+}
--- a/backend/internal/handler/custom_proxy_test.go
+++ b/backend/internal/handler/custom_proxy_test.go
@@ -0,0 +1,202 @@
+package handler
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/gin-gonic/gin"
+)
+
+func TestCustomRelayForwardsOpenAIRequestWithoutBrowserHeaders(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	const apiKey = "relay-secret-key"
+	upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("Authorization") != "Bearer "+apiKey {
+			t.Errorf("Authorization = %q", r.Header.Get("Authorization"))
+		}
+		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)
+			}
+		}
+		w.Header().Set("Content-Type", "application/json")
+		w.Header().Set("Set-Cookie", "upstream=unsafe")
+		_, _ = io.WriteString(w, `{"data":[]}`)
+	}))
+	defer upstream.Close()
+	useCustomRelayTestClient(t, upstream.Client())
+	t.Setenv("CANVAS_ALLOW_PRIVATE_UPSTREAMS", "true")
+	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 "+apiKey)
+	request.Header.Set("X-Canvas-Upstream-URL", upstream.URL+"/v1/models")
+	request.Header.Set("X-Canvas-Upstream-Format", "openai")
+	request.Header.Set("Cookie", "browser=session")
+	request.Header.Set("Origin", "https://canvas.example.com")
+	response := httptest.NewRecorder()
+	context, _ := gin.CreateTestContext(response)
+	context.Request = request
+
+	proxyCustomRelayRequest(context)
+	if response.Code != http.StatusOK {
+		t.Fatalf("status = %d, body = %s", response.Code, response.Body.String())
+	}
+	if response.Header().Get("Set-Cookie") != "" {
+		t.Fatal("upstream Set-Cookie should not be forwarded")
+	}
+	if strings.Contains(response.Body.String(), apiKey) {
+		t.Fatal("response leaked API key")
+	}
+}
+
+func TestCustomRelayConvertsGeminiAuthentication(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	const apiKey = "gemini-secret-key"
+	upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("x-goog-api-key") != apiKey {
+			t.Errorf("x-goog-api-key = %q", r.Header.Get("x-goog-api-key"))
+		}
+		if r.Header.Get("Authorization") != "" {
+			t.Errorf("Authorization should not be forwarded, got %q", r.Header.Get("Authorization"))
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = io.WriteString(w, `{"candidates":[]}`)
+	}))
+	defer upstream.Close()
+	useCustomRelayTestClient(t, upstream.Client())
+	t.Setenv("CANVAS_ALLOWED_PRIVATE_UPSTREAM_HOSTS", "127.0.0.1")
+
+	request := httptest.NewRequest(http.MethodPost, "/api/ai/custom", strings.NewReader(`{"contents":[]}`))
+	request.Header.Set("Authorization", "Bearer "+apiKey)
+	request.Header.Set("Content-Type", "application/json")
+	request.Header.Set("X-Canvas-Upstream-URL", upstream.URL+"/v1beta/models/gemini-test:generateContent")
+	request.Header.Set("X-Canvas-Upstream-Format", "gemini")
+	response := httptest.NewRecorder()
+	context, _ := gin.CreateTestContext(response)
+	context.Request = request
+
+	proxyCustomRelayRequest(context)
+	if response.Code != http.StatusOK {
+		t.Fatalf("status = %d, body = %s", response.Code, response.Body.String())
+	}
+}
+
+func TestCustomRelayStreamsBeforeUpstreamCompletes(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	const apiKey = "stream-secret"
+	release := make(chan struct{})
+	upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+		_, _ = io.WriteString(w, "data: first\n\n")
+		w.(http.Flusher).Flush()
+		<-release
+		_, _ = io.WriteString(w, "data: second\n\n")
+	}))
+	defer upstream.Close()
+	useCustomRelayTestClient(t, upstream.Client())
+	t.Setenv("CANVAS_ALLOWED_PRIVATE_UPSTREAM_HOSTS", "127.0.0.1")
+
+	proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		context, _ := gin.CreateTestContext(w)
+		context.Request = r
+		proxyCustomRelayRequest(context)
+	}))
+	defer proxy.Close()
+	request, _ := http.NewRequest(http.MethodPost, proxy.URL, strings.NewReader(`{"model":"test"}`))
+	request.Header.Set("Authorization", "Bearer "+apiKey)
+	request.Header.Set("Content-Type", "application/json")
+	request.Header.Set("Accept", "text/event-stream")
+	request.Header.Set("X-Canvas-Upstream-URL", upstream.URL+"/v1/responses")
+	request.Header.Set("X-Canvas-Upstream-Format", "openai")
+	client := &http.Client{Timeout: 3 * time.Second}
+	response, err := client.Do(request)
+	if err != nil {
+		close(release)
+		t.Fatal(err)
+	}
+	reader := bufio.NewReader(response.Body)
+	line, err := reader.ReadString('\n')
+	if err != nil {
+		close(release)
+		_ = response.Body.Close()
+		t.Fatal(err)
+	}
+	if line != "data: first\n" {
+		close(release)
+		_ = response.Body.Close()
+		t.Fatalf("first streamed line = %q", line)
+	}
+	close(release)
+	_ = response.Body.Close()
+}
+
+func TestRelayStreamRedactorHandlesSplitSecret(t *testing.T) {
+	redactor := newRelayStreamRedactor("split-secret")
+	output := append(redactor.Push([]byte("before split-"), false), redactor.Push([]byte("secret after"), true)...)
+	if bytes.Contains(output, []byte("split-secret")) || !bytes.Contains(output, []byte("[REDACTED]")) {
+		t.Fatalf("redacted output = %q", output)
+	}
+}
+
+func TestCustomRelayRejectsOversizedDeclaredBodyBeforeConnecting(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	t.Setenv("CANVAS_ALLOWED_PRIVATE_UPSTREAM_HOSTS", "127.0.0.1")
+	connected := false
+	previous := customRelayClient
+	customRelayClient = func(time.Duration) *http.Client {
+		connected = true
+		return http.DefaultClient
+	}
+	t.Cleanup(func() { customRelayClient = previous })
+
+	request := httptest.NewRequest(http.MethodPost, "/api/ai/custom", strings.NewReader(`{"model":"test"}`))
+	request.ContentLength = maxCustomRelayBodyBytes + 1
+	request.Header.Set("Authorization", "Bearer test-key")
+	request.Header.Set("Content-Type", "application/json")
+	request.Header.Set("X-Canvas-Upstream-URL", "https://127.0.0.1/v1/responses")
+	request.Header.Set("X-Canvas-Upstream-Format", "openai")
+	response := httptest.NewRecorder()
+	context, _ := gin.CreateTestContext(response)
+	context.Request = request
+
+	proxyCustomRelayRequest(context)
+	if response.Code != http.StatusRequestEntityTooLarge {
+		t.Fatalf("status = %d, body = %s", response.Code, response.Body.String())
+	}
+	if connected {
+		t.Fatal("oversized request should not create an upstream client")
+	}
+}
+
+func TestCustomRelayConcurrencyLimitReleasesSlots(t *testing.T) {
+	const userID = "relay-limit-test-user"
+	for index := 0; index < maxCustomRelayConcurrency; index++ {
+		if !acquireCustomRelaySlot(userID) {
+			t.Fatalf("slot %d should be available", index)
+		}
+	}
+	if acquireCustomRelaySlot(userID) {
+		t.Fatal("request above the concurrency limit should be rejected")
+	}
+	for index := 0; index < maxCustomRelayConcurrency; index++ {
+		releaseCustomRelaySlot(userID)
+	}
+	if !acquireCustomRelaySlot(userID) {
+		t.Fatal("released slot should be reusable")
+	}
+	releaseCustomRelaySlot(userID)
+}
+
+func useCustomRelayTestClient(t *testing.T, client *http.Client) {
+	t.Helper()
+	previous := customRelayClient
+	customRelayClient = func(time.Duration) *http.Client { return client }
+	t.Cleanup(func() { customRelayClient = previous })
+}
--- a/backend/internal/handler/security.go
+++ b/backend/internal/handler/security.go
@@ -23,9 +23,10 @@ import (
 const maxSystemProxyBodyBytes int64 = 64 << 20

 var (
-	runtimeService      *service.Service
-	geminiGeneratePath  = regexp.MustCompile(`^/models/([^/:]+):(generateContent|streamGenerateContent)$`)
-	openAIPostEndpoints = map[string]bool{
+	runtimeService        *service.Service
+	geminiGeneratePath    = regexp.MustCompile(`^/models/([^/:]+):(generateContent|streamGenerateContent)$`)
+	customGeminiRelayPath = regexp.MustCompile(`(?:^|/)models/[^/:]+:(generateContent|streamGenerateContent)$`)
+	openAIPostEndpoints   = map[string]bool{
 		"/responses": true, "/chat/completions": true, "/images/generations": true, "/images/edits": true,
 		"/audio/speech": true,
 	}
@@ -35,6 +36,72 @@ func ConfigureRuntime(svc *service.Service) {
 	runtimeService = svc
 }

+func authorizeCustomRelay(method string, target *url.URL, apiFormat string, contentType string) error {
+	requestPath, err := normalizedCustomRelayPath(target.EscapedPath())
+	if err != nil {
+		return err
+	}
+	query, err := url.ParseQuery(target.RawQuery)
+	if err != nil {
+		return errors.New("自定义渠道查询参数无效")
+	}
+	for key := range query {
+		switch strings.ToLower(strings.TrimSpace(key)) {
+		case "key", "api_key", "access_token", "token":
+			return errors.New("自定义渠道地址不允许在查询参数中携带密钥")
+		}
+	}
+
+	apiFormat = strings.ToLower(strings.TrimSpace(apiFormat))
+	if apiFormat != "openai" && apiFormat != "gemini" {
+		return errors.New("自定义渠道调用格式无效")
+	}
+	if method == http.MethodGet {
+		if len(query) != 0 || (requestPath != "/models" && !strings.HasSuffix(requestPath, "/models")) {
+			return errors.New("自定义渠道不允许访问该上游接口")
+		}
+		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("自定义渠道不允许使用该请求方法")
+	}
+	if apiFormat == "openai" {
+		if len(query) != 0 || (!strings.HasSuffix(requestPath, "/responses") && !strings.HasSuffix(requestPath, "/chat/completions")) {
+			return errors.New("自定义渠道不允许访问该上游接口")
+		}
+		return nil
+	}
+	if !customGeminiRelayPath.MatchString(requestPath) {
+		return errors.New("自定义渠道不允许访问该上游接口")
+	}
+	if len(query) == 0 {
+		return nil
+	}
+	if len(query) == 1 && len(query["alt"]) == 1 && query.Get("alt") == "sse" && strings.HasSuffix(requestPath, ":streamGenerateContent") {
+		return nil
+	}
+	return errors.New("自定义渠道不允许使用该查询参数")
+}
+
+func normalizedCustomRelayPath(value string) (string, error) {
+	if len(value) > 2048 {
+		return "", errors.New("自定义渠道请求路径过长")
+	}
+	decoded, err := url.PathUnescape(value)
+	if err != nil || strings.Contains(decoded, "\\") || strings.Contains(decoded, "\x00") {
+		return "", errors.New("自定义渠道请求路径无效")
+	}
+	cleaned := path.Clean("/" + strings.TrimPrefix(decoded, "/"))
+	if cleaned != decoded && cleaned != "/"+strings.TrimPrefix(decoded, "/") {
+		return "", errors.New("自定义渠道请求路径无效")
+	}
+	return cleaned, nil
+}
+
 func enforceRateLimit(c *gin.Context, key string, limit int, window time.Duration) bool {
 	if runtimeService == nil {
 		fail(c, http.StatusServiceUnavailable, errors.New("请求协调器尚未初始化"))
--- a/backend/internal/handler/security_test.go
+++ b/backend/internal/handler/security_test.go
@@ -3,6 +3,7 @@ package handler
 import (
 	"mime/multipart"
 	"net/http"
+	"net/url"
 	"strings"
 	"testing"

@@ -17,6 +18,54 @@ func TestAuthorizeSystemProxyAllowsConfiguredGenerationModel(t *testing.T) {
 	}
 }

+func TestAuthorizeCustomRelayAllowsModelsAndAgentEndpoints(t *testing.T) {
+	tests := []struct {
+		method      string
+		target      string
+		apiFormat   string
+		contentType string
+	}{
+		{method: http.MethodGet, target: "https://api.example.com/v1/models", apiFormat: "openai"},
+		{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://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse", apiFormat: "gemini", contentType: "application/json"},
+	}
+	for _, test := range tests {
+		target, err := url.Parse(test.target)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := authorizeCustomRelay(test.method, target, test.apiFormat, test.contentType); err != nil {
+			t.Fatalf("authorizeCustomRelay(%s %s) error = %v", test.method, test.target, err)
+		}
+	}
+}
+
+func TestAuthorizeCustomRelayRejectsArbitraryRequestsAndCredentialQueries(t *testing.T) {
+	tests := []struct {
+		method      string
+		target      string
+		apiFormat   string
+		contentType string
+	}{
+		{method: http.MethodDelete, target: "https://api.example.com/v1/models", apiFormat: "openai"},
+		{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/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"},
+	}
+	for _, test := range tests {
+		target, err := url.Parse(test.target)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := authorizeCustomRelay(test.method, target, test.apiFormat, test.contentType); err == nil {
+			t.Fatalf("authorizeCustomRelay(%s %s) should fail", test.method, test.target)
+		}
+	}
+}
+
 func TestAuthorizeSystemProxyRejectsArbitraryPathAndModel(t *testing.T) {
 	channel := &model.ModelChannel{APIFormat: "openai", ModelsJSON: `["gpt-image-1"]`}
 	if err := authorizeSystemProxy(channel, http.MethodDelete, "/account", "application/json", nil); err == nil {