优化

internal/server/handler/media.go

@@ -1,6 +1,8 @@
 package handler
 
 import (
+	"errors"
+	"net/http"
 	"strconv"
 
 	"kra/internal/server/httpx"
@@ -15,13 +17,40 @@ type Media struct{ service *service.MediaService }
 
 func NewMedia(service *service.MediaService) *Media { return &Media{service: service} }
 
+func (h *Media) limitMultipartBody(c *gin.Context) int64 {
+	limit := h.service.MediaConfig().EffectiveMaxFileSize()
+	c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, limit+(1<<20))
+	return limit
+}
+
+func rejectMediaTooLarge(c *gin.Context, err error, message string) bool {
+	var tooLarge *http.MaxBytesError
+	if !errors.As(err, &tooLarge) {
+		return false
+	}
+	c.AbortWithStatusJSON(http.StatusRequestEntityTooLarge, httpx.Response{Code: httpx.CodeError, Data: gin.H{}, Msg: message})
+	return true
+}
+
 func (h *Media) Upload(c *gin.Context) {
 	claims := servermiddleware.Claims(c)
-	header, err := c.FormFile("file")
-	if claims == nil || err != nil {
+	if claims == nil {
 		httpx.Fail(c, "接收文件失败")
 		return
 	}
+	limit := h.limitMultipartBody(c)
+	header, err := c.FormFile("file")
+	if err != nil {
+		if rejectMediaTooLarge(c, err, "文件超过大小上限") {
+			return
+		}
+		httpx.Fail(c, "接收文件失败")
+		return
+	}
+	if header.Size > limit {
+		c.AbortWithStatusJSON(http.StatusRequestEntityTooLarge, httpx.Response{Code: httpx.CodeError, Data: gin.H{}, Msg: "文件超过大小上限"})
+		return
+	}
 	opened, err := header.Open()
 	if err != nil {
 		// Let the upload service surface this as the generic upload
@@ -191,13 +220,33 @@ func (h *Media) InitUpload(c *gin.Context) {
 }
 func (h *Media) SaveChunk(c *gin.Context) {
 	claims := servermiddleware.Claims(c)
-	uploadID, _ := strconv.ParseUint(c.PostForm("uploadId"), 10, 64)
-	index, _ := strconv.Atoi(c.PostForm("chunkIndex"))
-	header, err := c.FormFile("chunk")
-	if claims == nil || err != nil {
+	if claims == nil {
 		httpx.Fail(c, "接收分片失败")
 		return
 	}
+	limit := h.limitMultipartBody(c)
+	header, err := c.FormFile("chunk")
+	if err != nil {
+		if rejectMediaTooLarge(c, err, "分片超过大小上限") {
+			return
+		}
+		httpx.Fail(c, "接收分片失败")
+		return
+	}
+	if header.Size > limit {
+		c.AbortWithStatusJSON(http.StatusRequestEntityTooLarge, httpx.Response{Code: httpx.CodeError, Data: gin.H{}, Msg: "分片超过大小上限"})
+		return
+	}
+	uploadID, parseErr := strconv.ParseUint(c.PostForm("uploadId"), 10, 64)
+	if parseErr != nil || uploadID == 0 {
+		httpx.Fail(c, "上传会话 ID 非法")
+		return
+	}
+	index, parseErr := strconv.Atoi(c.PostForm("chunkIndex"))
+	if parseErr != nil || index < 0 {
+		httpx.Fail(c, "分片序号非法")
+		return
+	}
 	opened, err := header.Open()
 	if err != nil {
 		httpx.Fail(c, "分片读取失败")

internal/server/handler/public.go

@@ -34,12 +34,20 @@ func (h *Public) captchaConfig() (int, int, int) {
 
 func (h *Public) Captcha(c *gin.Context) {
 	keyLong, width, height := h.captchaConfig()
-	security, _ := h.settings.CurrentSecurity(c.Request.Context())
+	security, err := h.settings.CurrentSecurity(c.Request.Context())
+	if err != nil || security == nil {
+		httpx.Fail(c, "安全服务暂不可用")
+		return
+	}
 	openCaptcha := true
 	if security != nil {
 		keyLong, width, height = security.KeyLong, security.ImgWidth, security.ImgHeight
 		ttl := time.Duration(security.CaptchaTimeout) * time.Second
-		failures, _ := h.settings.EnsureLoginIPCounter(c.Request.Context(), c.ClientIP(), ttl)
+		failures, counterErr := h.settings.EnsureLoginIPCounter(c.Request.Context(), c.ClientIP(), ttl)
+		if counterErr != nil {
+			httpx.Fail(c, "安全服务暂不可用")
+			return
+		}
 		openCaptcha = security.CaptchaOpen == 0 || failures > security.CaptchaOpen
 	}
 	driver := base64Captcha.NewDriverDigit(height, width, keyLong, 0.7, 80)