diff --git a/common/init.go b/common/init.go index 4d4c62b27cac..27e00d674e5c 100644 --- a/common/init.go +++ b/common/init.go @@ -177,6 +177,7 @@ func initConstantEnv() { constant.StreamingTimeout = GetEnvOrDefault("STREAMING_TIMEOUT", 300) constant.DifyDebug = GetEnvOrDefaultBool("DIFY_DEBUG", true) constant.MaxFileDownloadMB = GetEnvOrDefault("MAX_FILE_DOWNLOAD_MB", 64) + constant.MaxRelayResponseMB = GetEnvOrDefault("MAX_RELAY_RESPONSE_MB", 64) constant.StreamScannerMaxBufferMB = GetEnvOrDefault("STREAM_SCANNER_MAX_BUFFER_MB", 128) // MaxRequestBodyMB 请求体最大大小（解压后），用于防止超大请求/zip bomb导致内存暴涨 constant.MaxRequestBodyMB = GetEnvOrDefault("MAX_REQUEST_BODY_MB", 128) diff --git a/common/page_info.go b/common/page_info.go index 2378a5d81fdb..2138cc43aa03 100644 --- a/common/page_info.go +++ b/common/page_info.go @@ -6,6 +6,16 @@ import ( "github.com/gin-gonic/gin" ) +// MaxPageSize caps the page_size accepted from query parameters. List +// endpoints use PageSize directly as SQL LIMIT; an unbounded value would let a +// single request pull every matching row into memory (F-37). GetPageQuery +// clamps every page size to this value. +const MaxPageSize = 100 + +// MaxPage caps the page number so (page-1)*pageSize cannot overflow into a +// negative SQL OFFSET (F-37 hardening). +const MaxPage = 100000 + type PageInfo struct { Page int `json:"page"` // page num 页码 PageSize int `json:"page_size"` // page size 页大小 @@ -73,9 +83,24 @@ func GetPageQuery(c *gin.Context) *PageInfo { pageInfo.PageSize = ItemsPerPage } } + // F-37: cap page size. List endpoints use PageSize directly as SQL LIMIT; + // an unbounded value (e.g. page_size=2000000000) makes a single request + // pull every matching row into memory (memory/DB exhaustion DoS). + if pageInfo.PageSize > MaxPageSize { + pageInfo.PageSize = MaxPageSize + } + if pageInfo.PageSize < 1 { + pageInfo.PageSize = 1 + } + if pageInfo.Page > MaxPage { + pageInfo.Page = MaxPage + } + if pageInfo.Page < 1 { + pageInfo.Page = 1 + } - if pageInfo.PageSize > 100 { - pageInfo.PageSize = 100 + if pageInfo.PageSize > MaxPageSize { + pageInfo.PageSize = MaxPageSize } return pageInfo diff --git a/constant/env.go b/constant/env.go index 512bfc31126b..4a01baec4243 100644 --- a/constant/env.go +++ b/constant/env.go @@ -3,6 +3,12 @@ package constant var StreamingTimeout int var DifyDebug bool var MaxFileDownloadMB int + +// MaxRelayResponseMB bounds the size of a non-stream upstream response body +// read fully into memory by relay handlers. A malicious or misbehaving +// upstream (e.g. free proxies) must not be able to OOM the gateway with an +// unbounded JSON body (F-36). +var MaxRelayResponseMB int var StreamScannerMaxBufferMB int var ForceStreamOption bool var CountToken bool diff --git a/controller/relay.go b/controller/relay.go index 8dccfe76dddd..148fe0ba3d88 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -230,6 +230,9 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { if newAPIError == nil { relayInfo.LastError = nil + if channel != nil { + service.ClearChannelErrorStreak(channel.Id) + } return } @@ -365,9 +368,12 @@ func processChannelError(c *gin.Context, channelError types.ChannelError, err *t // 不要使用context获取渠道信息，异步处理时可能会出现渠道信息不一致的情况 // do not use context to get channel info, there may be inconsistent channel info when processing asynchronously if service.ShouldDisableChannel(err) && channelError.AutoBan { - gopool.Go(func() { - service.DisableChannel(channelError, err.ErrorWithStatusCode()) - }) + userID := c.GetInt("id") + if service.RecordChannelErrorAndShouldDisable(channelError.ChannelId, userID) { + gopool.Go(func() { + service.DisableChannel(channelError, err.ErrorWithStatusCode()) + }) + } } if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) { @@ -556,6 +562,9 @@ func RelayTask(c *gin.Context) { result, taskErr = relay.RelayTaskSubmit(c, relayInfo) if taskErr == nil { + if channel != nil { + service.ClearChannelErrorStreak(channel.Id) + } break } diff --git a/go.mod b/go.mod index b0642f162db2..a060914995e3 100644 --- a/go.mod +++ b/go.mod @@ -90,7 +90,7 @@ require ( require ( github.com/DmitriyVTitov/size v1.5.0 // indirect github.com/anknown/darts v0.0.0-20151216065714-83ff685239e6 // indirect - github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 // indirect + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 // indirect github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 // indirect github.com/beorn7/perks v1.0.1 // indirect diff --git a/relay/channel/api_request.go b/relay/channel/api_request.go index 48241b14a5e9..91fc908c008a 100644 --- a/relay/channel/api_request.go +++ b/relay/channel/api_request.go @@ -12,6 +12,7 @@ import ( "time" common2 "github.com/QuantumNous/new-api/common" + rootconstant "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relay/constant" @@ -25,6 +26,32 @@ import ( "github.com/gorilla/websocket" ) +// limitReadCloser caps how many bytes can be read from an upstream response +// body and fails closed: reads past the limit return an error instead of a +// silent truncation, and Close is delegated to the wrapped body so transport +// resources are released (F-36). +type limitReadCloser struct { + rc io.ReadCloser + n int64 + max int64 +} + +func (l *limitReadCloser) Read(p []byte) (int, error) { + if l.n >= l.max { + return 0, fmt.Errorf("upstream response body exceeds %d bytes", l.max) + } + if int64(len(p)) > l.max-l.n { + p = p[:l.max-l.n] + } + n, err := l.rc.Read(p) + l.n += int64(n) + return n, err +} + +func (l *limitReadCloser) Close() error { + return l.rc.Close() +} + // ApplyUpstreamBodyMetadata restores metadata that net/http cannot infer from // a ReplayableBody. Callers must pass the original body because NewRequest // hides its dynamic type behind req.Body's io.ReadCloser wrapper. @@ -537,6 +564,20 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http if resp == nil { return nil, errors.New("resp is nil") } + // F-36: cap non-stream upstream response bodies. Handlers read them fully + // with io.ReadAll; an unbounded body from a misbehaving/malicious upstream + // (e.g. free proxies) would OOM the gateway. Streams are read + // incrementally and are not capped here. + if !info.IsStream && resp.Body != nil { + maxBytes := int64(rootconstant.MaxRelayResponseMB) << 20 + if maxBytes <= 0 { + maxBytes = 64 << 20 + } + // Fail closed on oversized bodies (ReadAll past the limit returns a + // distinct error instead of a silent truncation) and delegate Close to + // the original body so transport resources are released. + resp.Body = &limitReadCloser{rc: resp.Body, max: maxBytes} + } if common2.DebugEnabled { policy := service.NormalizeHTTPTransportPolicy(info.ChannelSetting) logger.LogDebug(c, fmt.Sprintf( diff --git a/relay/channel/openai/relay_realtime.go b/relay/channel/openai/relay_realtime.go index bea97780df1b..53be838d12ff 100644 --- a/relay/channel/openai/relay_realtime.go +++ b/relay/channel/openai/relay_realtime.go @@ -2,6 +2,7 @@ package openai import ( "fmt" + "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/logger" @@ -25,6 +26,20 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. clientConn := info.ClientWs targetConn := info.TargetWs + // F-39: an idle WebSocket (no client messages, no pings) must not hold the + // connection forever (slowloris resource exhaustion, uncharged idle + // sessions). Enforce a message-size cap and a read deadline refreshed on + // activity; gorilla auto-replies to client pings, and the pong handler + // refreshes the deadline so well-behaved keep-alive clients stay alive. + const realtimeReadLimit = 8 << 20 + const realtimeReadTimeout = 90 * time.Second + clientConn.SetReadLimit(realtimeReadLimit) + clientConn.SetPongHandler(func(string) error { + return clientConn.SetReadDeadline(time.Now().Add(realtimeReadTimeout)) + }) + targetConn.SetReadLimit(realtimeReadLimit) + targetConn.SetReadDeadline(time.Now().Add(realtimeReadTimeout)) + clientClosed := make(chan struct{}) targetClosed := make(chan struct{}) sendChan := make(chan []byte, 100) @@ -46,6 +61,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. case <-c.Done(): return default: + _ = clientConn.SetReadDeadline(time.Now().Add(realtimeReadTimeout)) _, message, err := clientConn.ReadMessage() if err != nil { if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { @@ -106,6 +122,7 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. case <-c.Done(): return default: + _ = targetConn.SetReadDeadline(time.Now().Add(realtimeReadTimeout)) _, message, err := targetConn.ReadMessage() if err != nil { if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { @@ -210,6 +227,11 @@ func OpenaiRealtimeHandler(c *gin.Context, info *relaycommon.RelayInfo) (*types. case <-c.Done(): } + // Close both connections so whichever reader is still blocked in + // ReadMessage unblocks promptly instead of waiting out the read deadline. + _ = clientConn.Close() + _ = targetConn.Close() + if usage.TotalTokens != 0 { _ = preConsumeUsage(c, info, usage, sumUsage) } diff --git a/service/channel.go b/service/channel.go index f348e081e2ed..f88e1cb0849f 100644 --- a/service/channel.go +++ b/service/channel.go @@ -3,6 +3,8 @@ package service import ( "fmt" "strings" + "sync" + "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" @@ -11,6 +13,67 @@ import ( "github.com/QuantumNous/new-api/setting/operation_setting" ) +// channelDisableWindow / channelDisableThreshold / channelDisableMinUsers +// implement a windowed, multi-user gate before an upstream error may +// auto-disable a channel. F-24: without a gate, a single request that hits an +// upstream quota/rate-limit message (e.g. "You exceeded your current quota") +// disables the whole channel for every user. Requiring several error events +// from at least two distinct users inside a short window makes accidental or +// single-account triggered disables materially harder while still catching +// genuinely broken channels. +const ( + channelDisableErrorWindow = 60 * time.Second + channelDisableErrorThreshold = 3 + channelDisableMinDistinctUsers = 2 +) + +type channelErrorStreak struct { + mu sync.Mutex + windowStart time.Time + users map[int]struct{} + count int +} + +var channelErrorStreaks sync.Map // channelID int -> *channelErrorStreak + +// RecordChannelErrorAndShouldDisable records an auto-ban-worthy error for a +// channel and reports whether the disable threshold has been crossed. +func RecordChannelErrorAndShouldDisable(channelID int, userID int) bool { + if userID <= 0 { + // Unauthenticated/background errors: treat as one distinct pseudo-user + // so a genuinely broken channel still converges with enough events. + userID = -1 + } + now := time.Now() + raw, _ := channelErrorStreaks.LoadOrStore(channelID, &channelErrorStreak{ + windowStart: now, + users: make(map[int]struct{}), + }) + s := raw.(*channelErrorStreak) + s.mu.Lock() + defer s.mu.Unlock() + if now.Sub(s.windowStart) > channelDisableErrorWindow { + s.windowStart = now + s.users = make(map[int]struct{}) + s.count = 0 + } + s.users[userID] = struct{}{} + s.count++ + if s.count >= channelDisableErrorThreshold && len(s.users) >= channelDisableMinDistinctUsers { + // Only delete the streak if it is still the live entry: a concurrent + // ClearChannelErrorStreak plus a new error may have replaced it, and + // deleting the new streak would discard fresh error events. + return channelErrorStreaks.CompareAndDelete(channelID, s) + } + return false +} + +// ClearChannelErrorStreak resets the disable gate for a channel after a +// successful request. +func ClearChannelErrorStreak(channelID int) { + channelErrorStreaks.Delete(channelID) +} + func formatNotifyType(channelId int, status int) string { return fmt.Sprintf("%s_%d_%d", dto.NotifyTypeChannelUpdate, channelId, status) } diff --git a/service/file_service.go b/service/file_service.go index 52652a3478a1..35796b24d9f8 100644 --- a/service/file_service.go +++ b/service/file_service.go @@ -549,7 +549,7 @@ func parseHEIFDimensions(data []byte) (int, int, bool) { if len(metaData) < 4 { return 0, 0, false } - return findISPE(metaData[4:]) + return findISPE(metaData[4:], 0) } offset += boxSize } @@ -558,7 +558,11 @@ func parseHEIFDimensions(data []byte) (int, int, bool) { // findISPE recursively searches for the ispe box within container boxes. // Path: meta -> iprp -> ipco -> ispe -func findISPE(data []byte) (int, int, bool) { +// 安全加固：限制递归深度，防止构造的嵌套 iprp/ipco 链导致无界递归栈耗尽 DoS (F-19)。 +func findISPE(data []byte, depth int) (int, int, bool) { + if depth > 16 { + return 0, 0, false + } offset := 0 size := len(data) for offset+8 <= size { @@ -570,7 +574,7 @@ func findISPE(data []byte) (int, int, bool) { content := data[offset+8 : offset+boxSize] switch boxType { case "iprp", "ipco": - if w, h, ok := findISPE(content); ok { + if w, h, ok := findISPE(content, depth+1); ok { return w, h, true } case "ispe":