mirror of
https://github.com/infinite-iroha/touka.git
synced 2026-06-15 08:37:38 +08:00
64e2ad9e7b
- FileText: now respects the provided status code instead of defaulting to 200 OK
- Request body limits: prepareRequestBody() is now only called when MaxRequestBodySize > 0
- ShouldBindJSON, ShouldBindWANF, ShouldBindGOB, ShouldBindForm, GetReqBody, PostForm
all now use the original c.Request.Body path when no limit is configured
- maxBytesReader: fixed exact-limit boundary case where body size == limit was
incorrectly rejected
- Added regression tests for FileText status codes and body limit behavior
All existing tests pass, and new tests verify the corrected behavior.
114 lines
3.6 KiB
Go
114 lines
3.6 KiB
Go
// This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. If a copy of the MPL was not distributed with this
|
|
// file, You can obtain one at https://mozilla.org/MPL/2.0/.
|
|
// Copyright 2024 WJQSERVER. All rights reserved.
|
|
// All rights reserved by WJQSERVER, related rights can be exercised by the infinite-iroha organization.
|
|
package touka
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"sync/atomic"
|
|
)
|
|
|
|
// ErrBodyTooLarge 是当读取的字节数超过 MaxBytesReader 设置的限制时返回的错误.
|
|
// 将其定义为可导出的变量, 方便调用方使用 errors.Is 进行判断.
|
|
var ErrBodyTooLarge = fmt.Errorf("body too large")
|
|
|
|
// maxBytesReader 是一个实现了 io.ReadCloser 接口的结构体.
|
|
// 它包装了另一个 io.ReadCloser, 并限制了从其中读取的最大字节数.
|
|
type maxBytesReader struct {
|
|
// r 是底层的 io.ReadCloser.
|
|
r io.ReadCloser
|
|
// n 是允许读取的最大字节数.
|
|
n int64
|
|
// read 是一个原子计数器, 用于安全地在多个 goroutine 之间跟踪已读取的字节数.
|
|
read atomic.Int64
|
|
}
|
|
|
|
// NewMaxBytesReader 创建并返回一个 io.ReadCloser, 它从 r 读取数据,
|
|
// 但在读取的字节数超过 n 后会返回 ErrBodyTooLarge 错误.
|
|
//
|
|
// 如果 r 为 nil, 会 panic.
|
|
// 如果 n 小于 0, 则读取不受限制, 直接返回原始的 r.
|
|
func NewMaxBytesReader(r io.ReadCloser, n int64) io.ReadCloser {
|
|
if r == nil {
|
|
panic("NewMaxBytesReader called with a nil reader")
|
|
}
|
|
// 如果限制为负数, 意味着不限制, 直接返回原始的 ReadCloser.
|
|
if n < 0 {
|
|
return r
|
|
}
|
|
return &maxBytesReader{
|
|
r: r,
|
|
n: n,
|
|
}
|
|
}
|
|
|
|
// Read 方法从底层的 ReadCloser 读取数据, 同时检查是否超过了字节限制.
|
|
func (mbr *maxBytesReader) Read(p []byte) (int, error) {
|
|
if len(p) == 0 {
|
|
return 0, nil
|
|
}
|
|
|
|
// 在函数开始时只加载一次原子变量, 减少后续的原子操作开销.
|
|
readSoFar := mbr.read.Load()
|
|
|
|
if readSoFar > mbr.n {
|
|
return 0, ErrBodyTooLarge
|
|
}
|
|
|
|
// 当已恰好读满限制时, 需要探测底层是否还有额外数据.
|
|
// 如果下一次读取立即 EOF, 说明请求体大小恰好等于限制, 属于合法情况.
|
|
if readSoFar == mbr.n {
|
|
var probe [1]byte
|
|
n, err := mbr.r.Read(probe[:])
|
|
if n > 0 {
|
|
mbr.read.Add(int64(n))
|
|
return 0, ErrBodyTooLarge
|
|
}
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return 0, ErrBodyTooLarge
|
|
}
|
|
|
|
// 计算当前还可以读取多少字节.
|
|
remaining := mbr.n - readSoFar
|
|
|
|
// 如果请求读取的长度大于剩余可读长度, 我们需要限制本次读取的长度.
|
|
// 这样可以保证即使 p 很大, 我们也只读取到恰好达到 maxBytes 的字节数.
|
|
if int64(len(p)) > remaining {
|
|
p = p[:remaining]
|
|
}
|
|
|
|
// 从底层 Reader 读取数据.
|
|
n, err := mbr.r.Read(p)
|
|
|
|
// 如果实际读取到了数据, 更新原子计数器.
|
|
if n > 0 {
|
|
readSoFar = mbr.read.Add(int64(n))
|
|
}
|
|
|
|
// 如果底层 Read 返回错误 (例如 io.EOF).
|
|
if err != nil {
|
|
// 如果是 EOF, 并且我们还没有读满 n 个字节, 这是一个正常的结束.
|
|
// 如果已经读满了 n 个字节, 即使是 EOF, 也可以认为成功了.
|
|
return n, err
|
|
}
|
|
|
|
// 读后检查: 如果这次读取使得总字节数超过了限制, 返回超限错误.
|
|
// 这是处理"跨越"限制情况的关键.
|
|
if readSoFar > mbr.n {
|
|
// 返回实际读取的字节数 n, 并附上超限错误.
|
|
// 上层调用者知道已经有 n 字节被读入了缓冲区 p, 但流已因超限而关闭.
|
|
return n, ErrBodyTooLarge
|
|
}
|
|
|
|
// 一切正常, 返回读取的字节数和 nil 错误.
|
|
return n, nil
|
|
}
|
|
|
|
// Close 方法关闭底层的 ReadCloser, 保证资源释放.
|
|
func (mbr *maxBytesReader) Close() error {
|
|
return mbr.r.Close()
|
|
}
|