feat(admin): 资源与策略 - 后台统一维护账号配额与运行时请求限制 (#37)

Fixes #35

Co-authored-by: ddCat <ddcat666@88.com>
This commit is contained in:
ddcat666
2026-07-23 21:19:07 +08:00
committed by GitHub
parent 78e27507c4
commit d3dfede48f
40 changed files with 1202 additions and 488 deletions

View File

@@ -4,6 +4,7 @@ import (
"bytes"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
@@ -33,7 +34,8 @@ func RegisterAuthRoutes(r *gin.RouterGroup, svc *service.Service) {
fail(c, http.StatusBadRequest, err)
return
}
if !enforceRateLimit(c, "register:"+c.ClientIP(), 5, time.Hour) {
policy, available := loadRuntimePolicy(c, svc)
if !available || !enforceRateLimit(c, "register:"+c.ClientIP(), policy.Request.RegisterPerHour, time.Hour) {
return
}
result, err := svc.Register(req)
@@ -53,7 +55,8 @@ func RegisterAuthRoutes(r *gin.RouterGroup, svc *service.Service) {
fail(c, http.StatusBadRequest, err)
return
}
if !enforceRateLimit(c, "email-code:"+c.ClientIP(), 10, time.Hour) {
policy, available := loadRuntimePolicy(c, svc)
if !available || !enforceRateLimit(c, "email-code:"+c.ClientIP(), policy.Request.EmailCodePerHour, time.Hour) {
return
}
if err := svc.SendRegistrationEmailCode(req.Email); err != nil {
@@ -69,10 +72,11 @@ func RegisterAuthRoutes(r *gin.RouterGroup, svc *service.Service) {
fail(c, http.StatusBadRequest, err)
return
}
if !enforceRateLimit(c, "login-ip:"+c.ClientIP(), 50, 10*time.Minute) {
policy, available := loadRuntimePolicy(c, svc)
if !available || !enforceRateLimit(c, "login-ip:"+c.ClientIP(), policy.Request.LoginIPPerTenMinutes, 10*time.Minute) {
return
}
if !enforceRateLimit(c, "login:"+c.ClientIP()+":"+strings.ToLower(strings.TrimSpace(req.Username)), 10, 10*time.Minute) {
if !enforceRateLimit(c, "login:"+c.ClientIP()+":"+strings.ToLower(strings.TrimSpace(req.Username)), policy.Request.LoginAccountPerTenMinutes, 10*time.Minute) {
return
}
result, err := svc.Login(req)
@@ -112,7 +116,12 @@ func RegisterAuthRoutes(r *gin.RouterGroup, svc *service.Service) {
return
}
channels, _ := svc.PublicSystemChannels()
ok(c, gin.H{"user": publicUser, "systemChannels": channels})
limits, err := svc.PublicRuntimeLimits()
if err != nil {
failService(c, err)
return
}
ok(c, gin.H{"user": publicUser, "systemChannels": channels, "runtimeLimits": limits})
})
r.GET("/channels/system", func(c *gin.Context) {
if _, err := currentUser(c, svc); err != nil {
@@ -439,31 +448,58 @@ func RegisterAdminRoutes(r *gin.RouterGroup, svc *service.Service) {
}
ok(c, gin.H{"setting": setting})
})
r.GET("/admin/settings/concurrency", func(c *gin.Context) {
r.GET("/admin/settings/runtime-policy", func(c *gin.Context) {
user, err := currentUser(c, svc)
if err != nil {
failService(c, err)
return
}
setting, err := svc.AdminRuntimeConcurrencySetting(user)
setting, err := svc.AdminRuntimePolicySetting(user)
if err != nil {
failService(c, err)
return
}
ok(c, gin.H{"setting": setting})
})
r.PATCH("/admin/settings/concurrency", func(c *gin.Context) {
r.GET("/admin/settings/runtime-policy/self-use", func(c *gin.Context) {
user, err := currentUser(c, svc)
if err != nil {
failService(c, err)
return
}
var req service.RuntimeConcurrencySettingRequest
setting, err := svc.AdminSelfUseRuntimePolicy(user)
if err != nil {
failService(c, err)
return
}
ok(c, gin.H{"setting": setting})
})
r.PUT("/admin/settings/runtime-policy", func(c *gin.Context) {
user, err := currentUser(c, svc)
if err != nil {
failService(c, err)
return
}
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 64<<10)
var req service.RuntimePolicySetting
if err := c.ShouldBindJSON(&req); err != nil {
fail(c, http.StatusBadRequest, err)
return
}
setting, err := svc.UpdateRuntimeConcurrencySetting(user, req)
setting, err := svc.UpdateRuntimePolicySetting(user, req)
if err != nil {
failService(c, err)
return
}
ok(c, gin.H{"setting": setting})
})
r.DELETE("/admin/settings/runtime-policy", func(c *gin.Context) {
user, err := currentUser(c, svc)
if err != nil {
failService(c, err)
return
}
setting, err := svc.ResetRuntimePolicySetting(user)
if err != nil {
failService(c, err)
return
@@ -536,14 +572,15 @@ func RegisterSystemProxyRoutes(r *gin.RouterGroup, svc *service.Service) {
func proxySystemRequest(c *gin.Context, svc *service.Service, user *model.User, channel *model.ModelChannel) {
startedAt := time.Now()
if !enforceRateLimit(c, "system-proxy:"+user.ID, 120, time.Minute) {
policy, available := loadRuntimePolicy(c, svc)
if !available || !enforceRateLimit(c, "system-proxy:"+user.ID, policy.Request.SystemRelayPerMinute, time.Minute) {
return
}
path := c.Param("path")
if path == "" {
path = "/"
}
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, maxSystemProxyBodyBytes)
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, policy.Request.SystemRelayRequestMB<<20)
body, err := io.ReadAll(c.Request.Body)
if err != nil {
fail(c, http.StatusBadRequest, err)
@@ -633,7 +670,8 @@ func proxySystemRequest(c *gin.Context, svc *service.Service, user *model.User,
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
status = model.ApiCallStatusFailed
}
responseBody, readErr := io.ReadAll(io.LimitReader(resp.Body, (128<<20)+1))
responseLimit := policy.Request.SystemRelayResponseMB << 20
responseBody, readErr := io.ReadAll(io.LimitReader(resp.Body, responseLimit+1))
if readErr != nil {
status = model.ApiCallStatusFailed
errorText = readErr.Error()
@@ -642,9 +680,9 @@ func proxySystemRequest(c *gin.Context, svc *service.Service, user *model.User,
fail(c, http.StatusBadGateway, errors.New("系统渠道响应读取失败"))
return
}
if len(responseBody) > 128<<20 {
if int64(len(responseBody)) > responseLimit {
_ = svc.MarkBillingUncertain(billingOrderID, "上游已响应但响应体超过限制,费用状态待核对")
fail(c, http.StatusBadGateway, errors.New("系统渠道响应超过 128MB 限制"))
fail(c, http.StatusBadGateway, fmt.Errorf("系统渠道响应超过 %dMB 限制", policy.Request.SystemRelayResponseMB))
return
}
if status == model.ApiCallStatusSucceeded {

View File

@@ -8,7 +8,6 @@ import (
"mime"
"net/http"
"strings"
"sync"
"time"
"infinite-canvas/backend/internal/service"
@@ -16,23 +15,9 @@ import (
"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
)
const maxCustomRelayErrorResponseBytes int64 = 64 << 10
var (
customRelayClient = service.CustomRelayHTTPClient
customRelayActive = struct {
sync.Mutex
users map[string]int
}{users: map[string]int{}}
)
var customRelayClient = service.CustomRelayHTTPClient
func RegisterCustomRelayRoutes(r *gin.RouterGroup, svc *service.Service) {
r.Any("/ai/custom", func(c *gin.Context) {
@@ -41,19 +26,26 @@ func RegisterCustomRelayRoutes(r *gin.RouterGroup, svc *service.Service) {
failService(c, err)
return
}
if !enforceRateLimit(c, "custom-relay:"+user.ID, 120, time.Minute) {
policy, available := loadRuntimePolicy(c, svc)
if !available || !enforceRateLimit(c, "custom-relay:"+user.ID, policy.Request.CustomRelayPerMinute, time.Minute) {
return
}
if !acquireCustomRelaySlot(user.ID) {
ttl := time.Duration(policy.Request.CustomRelayTimeoutMinutes+1) * time.Minute
release, acquired, err := svc.AcquireCustomRelaySlot(c.Request.Context(), user.ID, policy.Request.CustomRelayConcurrency, ttl)
if err != nil {
fail(c, http.StatusServiceUnavailable, errors.New("自定义渠道并发协调服务不可用"))
return
}
if !acquired {
fail(c, http.StatusTooManyRequests, errors.New("自定义渠道并发请求过多,请等待已有请求完成"))
return
}
defer releaseCustomRelaySlot(user.ID)
proxyCustomRelayRequest(c)
defer release()
proxyCustomRelayRequest(c, policy.Request)
})
}
func proxyCustomRelayRequest(c *gin.Context) {
func proxyCustomRelayRequest(c *gin.Context, policy service.RuntimeRequestPolicy) {
target, err := service.ValidateCustomRelayURL(c.GetHeader("X-Canvas-Upstream-URL"))
if err != nil {
failService(c, err)
@@ -72,16 +64,17 @@ func proxyCustomRelayRequest(c *gin.Context) {
fail(c, http.StatusUnauthorized, err)
return
}
if c.Request.ContentLength > maxCustomRelayBodyBytes {
fail(c, http.StatusRequestEntityTooLarge, errors.New("自定义渠道请求不能超过 32MB"))
requestLimit := policy.CustomRelayRequestMB << 20
if c.Request.ContentLength > requestLimit {
fail(c, http.StatusRequestEntityTooLarge, errors.New("自定义渠道请求超过配置上限"))
return
}
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, maxCustomRelayBodyBytes)
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, requestLimit)
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"))
fail(c, http.StatusRequestEntityTooLarge, errors.New("自定义渠道请求超过配置上限"))
return
}
fail(c, http.StatusBadRequest, errors.New("读取自定义渠道请求失败"))
@@ -111,16 +104,16 @@ func proxyCustomRelayRequest(c *gin.Context) {
upstreamReq.Header.Set("Authorization", "Bearer "+apiKey)
}
resp, err := customRelayClient(customRelayHTTPTimeout).Do(upstreamReq)
resp, err := customRelayClient(time.Duration(policy.CustomRelayTimeoutMinutes) * time.Minute).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)
writeCustomRelayResponse(c, resp, apiKey, policy.CustomRelayResponseMB<<20)
}
func writeCustomRelayResponse(c *gin.Context, resp *http.Response, apiKey string, modelList bool) {
func writeCustomRelayResponse(c *gin.Context, resp *http.Response, apiKey string, responseLimit int64) {
c.Header("Cache-Control", "no-store")
c.Header("X-Content-Type-Options", "nosniff")
mediaType, _, _ := mime.ParseMediaType(resp.Header.Get("Content-Type"))
@@ -133,17 +126,14 @@ func writeCustomRelayResponse(c *gin.Context, resp *http.Response, apiKey string
c.Header("X-Accel-Buffering", "no")
c.Status(resp.StatusCode)
c.Writer.WriteHeaderNow()
copyCustomRelayStream(c, resp.Body, apiKey)
copyCustomRelayStream(c, resp.Body, apiKey, responseLimit)
return
}
if mediaType != "application/json" && !strings.HasSuffix(mediaType, "+json") {
fail(c, http.StatusBadGateway, errors.New("自定义渠道上游返回了不支持的内容类型"))
return
}
limit := maxCustomRelayJSONResponseBytes
if modelList {
limit = maxCustomRelayModelResponseBytes
}
limit := responseLimit
body, err := readLimitedRelayBody(resp.Body, limit)
if err != nil || !json.Valid(body) {
fail(c, http.StatusBadGateway, errors.New("自定义渠道上游返回无效或过大的 JSON"))
@@ -167,14 +157,14 @@ func writeCustomRelayError(c *gin.Context, resp *http.Response, apiKey string, m
fail(c, resp.StatusCode, errors.New("自定义渠道上游请求失败"))
}
func copyCustomRelayStream(c *gin.Context, source io.Reader, apiKey string) {
func copyCustomRelayStream(c *gin.Context, source io.Reader, apiKey string, maxBytes int64) {
redactor := newRelayStreamRedactor(apiKey)
buffer := make([]byte, 32<<10)
var written int64
for written < maxCustomRelayStreamBytes {
for written < maxBytes {
read, err := source.Read(buffer)
if read > 0 {
remaining := maxCustomRelayStreamBytes - written
remaining := maxBytes - written
if int64(read) > remaining {
read = int(remaining)
}
@@ -217,26 +207,6 @@ func customRelayAPIKey(value string) (string, error) {
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

View File

@@ -10,6 +10,8 @@ import (
"testing"
"time"
"infinite-canvas/backend/internal/service"
"github.com/gin-gonic/gin"
)
@@ -44,7 +46,7 @@ func TestCustomRelayForwardsOpenAIRequestWithoutBrowserHeaders(t *testing.T) {
context, _ := gin.CreateTestContext(response)
context.Request = request
proxyCustomRelayRequest(context)
proxyCustomRelayRequest(context, defaultCustomRelayTestPolicy())
if response.Code != http.StatusOK {
t.Fatalf("status = %d, body = %s", response.Code, response.Body.String())
}
@@ -82,7 +84,7 @@ func TestCustomRelayConvertsGeminiAuthentication(t *testing.T) {
context, _ := gin.CreateTestContext(response)
context.Request = request
proxyCustomRelayRequest(context)
proxyCustomRelayRequest(context, defaultCustomRelayTestPolicy())
if response.Code != http.StatusOK {
t.Fatalf("status = %d, body = %s", response.Code, response.Body.String())
}
@@ -106,7 +108,7 @@ func TestCustomRelayStreamsBeforeUpstreamCompletes(t *testing.T) {
proxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
context, _ := gin.CreateTestContext(w)
context.Request = r
proxyCustomRelayRequest(context)
proxyCustomRelayRequest(context, defaultCustomRelayTestPolicy())
}))
defer proxy.Close()
request, _ := http.NewRequest(http.MethodPost, proxy.URL, strings.NewReader(`{"model":"test"}`))
@@ -157,7 +159,7 @@ func TestCustomRelayRejectsOversizedDeclaredBodyBeforeConnecting(t *testing.T) {
t.Cleanup(func() { customRelayClient = previous })
request := httptest.NewRequest(http.MethodPost, "/api/ai/custom", strings.NewReader(`{"model":"test"}`))
request.ContentLength = maxCustomRelayBodyBytes + 1
request.ContentLength = (defaultCustomRelayTestPolicy().CustomRelayRequestMB << 20) + 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")
@@ -166,7 +168,7 @@ func TestCustomRelayRejectsOversizedDeclaredBodyBeforeConnecting(t *testing.T) {
context, _ := gin.CreateTestContext(response)
context.Request = request
proxyCustomRelayRequest(context)
proxyCustomRelayRequest(context, defaultCustomRelayTestPolicy())
if response.Code != http.StatusRequestEntityTooLarge {
t.Fatalf("status = %d, body = %s", response.Code, response.Body.String())
}
@@ -175,28 +177,15 @@ func TestCustomRelayRejectsOversizedDeclaredBodyBeforeConnecting(t *testing.T) {
}
}
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 })
}
func defaultCustomRelayTestPolicy() service.RuntimeRequestPolicy {
return service.RuntimeRequestPolicy{
CustomRelayRequestMB: 32, CustomRelayResponseMB: 32, CustomRelayTimeoutMinutes: 10,
}
}

View File

@@ -17,7 +17,8 @@ func RegisterTaskRoutes(r *gin.RouterGroup, svc *service.Service) {
failService(c, err)
return
}
if !enforceRateLimit(c, "tasks:"+user.ID, 30, time.Minute) {
policy, available := loadRuntimePolicy(c, svc)
if !available || !enforceRateLimit(c, "tasks:"+user.ID, policy.Request.TaskCreatePerMinute, time.Minute) {
return
}
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 16<<20)
@@ -112,7 +113,8 @@ func RegisterSessionRoutes(r *gin.RouterGroup, svc *service.Service) {
failService(c, err)
return
}
if !enforceRateLimit(c, "sessions:"+user.ID, 20, time.Minute) {
policy, available := loadRuntimePolicy(c, svc)
if !available || !enforceRateLimit(c, "sessions:"+user.ID, policy.Request.SessionCreatePerMinute, time.Minute) {
return
}
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 16<<20)
@@ -147,10 +149,11 @@ func RegisterSessionRoutes(r *gin.RouterGroup, svc *service.Service) {
failService(c, err)
return
}
if !enforceRateLimit(c, "session-files:"+user.ID, 30, time.Minute) {
policy, available := loadRuntimePolicy(c, svc)
if !available || !enforceRateLimit(c, "session-files:"+user.ID, policy.Request.SessionFilePerMinute, time.Minute) {
return
}
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, service.MaxSessionUploadBytes+(1<<20))
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, (policy.Resource.SessionUploadMB<<20)+(1<<20))
file, err := c.FormFile("file")
if err != nil {
fail(c, http.StatusBadRequest, err)

View File

@@ -20,8 +20,6 @@ import (
"github.com/gin-gonic/gin"
)
const maxSystemProxyBodyBytes int64 = 64 << 20
var (
runtimeService *service.Service
geminiGeneratePath = regexp.MustCompile(`^/models/([^/:]+):(generateContent|streamGenerateContent)$`)
@@ -120,6 +118,15 @@ func enforceRateLimit(c *gin.Context, key string, limit int, window time.Duratio
return false
}
func loadRuntimePolicy(c *gin.Context, svc *service.Service) (service.RuntimePolicySetting, bool) {
policy, err := svc.RuntimePolicy()
if err != nil {
fail(c, http.StatusServiceUnavailable, errors.New("读取运行时策略失败:"+err.Error()))
return service.RuntimePolicySetting{}, false
}
return policy, true
}
func authorizeSystemProxy(channel *model.ModelChannel, method string, requestPath string, contentType string, body []byte) error {
requestPath, err := normalizedProxyPath(requestPath)
if err != nil {

View File

@@ -68,10 +68,11 @@ func RegisterUserDataRoutes(r *gin.RouterGroup, svc *service.Service) {
failService(c, err)
return
}
if !enforceRateLimit(c, "resources-upload:"+user.ID, 30, time.Minute) {
policy, available := loadRuntimePolicy(c, svc)
if !available || !enforceRateLimit(c, "resources-upload:"+user.ID, policy.Request.ResourceUploadPerMinute, time.Minute) {
return
}
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, service.MaxResourceUploadBytes+(1<<20))
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, (policy.Resource.ResourceUploadMB<<20)+(1<<20))
file, err := c.FormFile("file")
if err != nil {
fail(c, http.StatusBadRequest, err)
@@ -93,7 +94,8 @@ func RegisterUserDataRoutes(r *gin.RouterGroup, svc *service.Service) {
failService(c, err)
return
}
if !enforceRateLimit(c, "resources-import:"+user.ID, 30, time.Minute) {
policy, available := loadRuntimePolicy(c, svc)
if !available || !enforceRateLimit(c, "resources-import:"+user.ID, policy.Request.ResourceImportPerMinute, time.Minute) {
return
}
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 64<<10)
@@ -247,7 +249,8 @@ func RegisterUserDataRoutes(r *gin.RouterGroup, svc *service.Service) {
failService(c, err)
return
}
if !enforceRateLimit(c, "assets-write:"+user.ID, 120, time.Minute) {
policy, available := loadRuntimePolicy(c, svc)
if !available || !enforceRateLimit(c, "assets-write:"+user.ID, policy.Request.AssetWritePerMinute, time.Minute) {
return
}
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 5<<20)
@@ -316,7 +319,8 @@ func RegisterUserDataRoutes(r *gin.RouterGroup, svc *service.Service) {
failService(c, err)
return
}
if !enforceRateLimit(c, "canvas-write:"+user.ID, 120, time.Minute) {
policy, available := loadRuntimePolicy(c, svc)
if !available || !enforceRateLimit(c, "canvas-write:"+user.ID, policy.Request.CanvasWritePerMinute, time.Minute) {
return
}
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 5<<20)