通过 LimitedReader 防止未知响应体耗尽内存

当 Go 客户端面对未知长度、分块传输或不可信上游时，直接对 resp.Body 调用 io.ReadAll 会让内存占用跟着响应体增长。安全做法不是先相信 Content-Length，而是把真正读取的数据流包在 io.LimitReader 中，并多读一个哨兵字节：最多读取 maxBytes+1，若结果长度大于 maxBytes，就确定响应体已经超限。

官方文档：https://pkg.go.dev/io

这个“多读 1 字节”很关键。若只限制为 maxBytes，读取结果恰好等于上限时，无法判断原始响应是真的刚好结束，还是后面还有更多数据被 LimitedReader 截断。

实验目标与前置条件

本文实现一个适用于“小型 JSON、文本或错误响应”的客户端读取函数，满足以下约束：

• 不论上游是否提供 Content-Length，内存中的正文最多保留上限加一个哨兵字节。
• Content-Length 已知且明显超限时，可以在读取前快速拒绝。
• 读取失败与响应体超限使用不同错误，便于调用方决定重试或告警。
• 响应体始终关闭；正常小响应读取到 EOF，超限响应及时停止。

示例假设响应体应当完整放入内存。如果业务要下载大文件、流式解码或转存对象存储，应改用 io.Copy、json.Decoder 或分块处理，并在对应流上设置上限，而不是先 ReadAll。

为什么要读 max+1，而不是只读 max

io.LimitReader(r, n) 返回一个 Reader，最多从底层 Reader 返回 n 个字节；达到 n 后，它表现为 EOF。它能限制调用方从该包装器拿到的数据量，却不会主动告诉你“底层数据是否还有剩余”。

因此将 n 设置为 maxBytes+1：如果读到 0 到 maxBytes 字节，正文没有超过上限；如果读到 maxBytes+1 字节，最后一个字节就是超限证据。无论底层后面还有 1 字节还是 1 GB，程序都不会继续把它们读入内存。

图1：max+1 哨兵字节让读取边界可判定。

实现一个有界响应体读取函数

package safehttp

import (
    "errors"
    "fmt"
    "io"
    "net/http"
)

const MaxResponseBodyBytes int64 = 1  MaxResponseBodyBytes {
        return nil, fmt.Errorf(
            "%w: content-length=%d limit=%d",
            ErrResponseBodyTooLarge,
            resp.ContentLength,
            MaxResponseBodyBytes,
        )
    }

    // 多读一个哨兵字节，才能区分“刚好等于上限”和“超过上限”
    limited := io.LimitReader(resp.Body, MaxResponseBodyBytes+1)
    body, err := io.ReadAll(limited)
    if err != nil {
        return nil, fmt.Errorf("read response body: %w", err)
    }

    if int64(len(body)) > MaxResponseBodyBytes {
        return nil, fmt.Errorf(
            "%w: limit=%d",
            ErrResponseBodyTooLarge,
            MaxResponseBodyBytes,
        )
    }

    return body, nil
}

这里的上限控制的是返回给 io.ReadAll 的字节数，所以 ReadAll 的正文缓冲不会随恶意响应无限增长。实现仍会有切片容量、HTTP 栈和解压器等额外内存开销，因此“1 MiB 上限”不代表进程内存只增加 1 MiB；它代表响应正文这一项被严格限制在一个小范围内。

把函数接入带超时的 HTTP 客户端

内存上限不等于时间上限。一个对端可以非常缓慢地发送 maxBytes+1 字节，因此请求还需要上下文超时或客户端超时。

package safehttp

import (
    "context"
    "fmt"
    "net/http"
    "time"
)

func Fetch(ctx context.Context, client *http.Client, url string) ([]byte, error) {
    // 用上下文覆盖建连、发送请求、读取响应头和响应体的生命周期
    ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
    defer cancel()

    req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
    if err != nil {
        return nil, fmt.Errorf("build request: %w", err)
    }

    resp, err := client.Do(req)
    if err != nil {
        return nil, fmt.Errorf("send request: %w", err)
    }

    // ReadSmallBody 内部负责关闭 resp.Body
    return ReadSmallBody(resp)
}

不要同时在多个层级随意关闭同一个 Body，也不要在返回 resp 给上层后又由下层提前关闭。选择一个清晰的所有权规则：谁消费完整正文，谁负责关闭。

Content-Length 为什么不能代替流式上限

resp.ContentLength 为 -1 时表示长度未知。分块传输常见于动态响应，无法在读取前获知最终大小。即使长度已知，它也描述协议层看到的长度，不能替代对实际消费数据流的限制。

Go 的默认 Transport 在自动处理 gzip 时，调用方读取的 resp.Body 可能是解压后的数据流，同时 ContentLength 会变为未知。压缩包很小不代表解压结果也小，所以限制应包在“应用最终读取的流”上。若程序自行设置 Accept-Encoding 并自行解压，则应在解压器之后再套 LimitedReader。

图2：内存安全优先，同时显式处理响应体和连接复用。

用 httptest 覆盖边界条件

边界测试至少要覆盖“小于上限、刚好等于上限、超过上限、已知 Content-Length 超限”四种情况。下面用较小上限演示核心判定，正式代码可把上限做成参数或配置。

package safehttp

import (
    "errors"
    "io"
    "net/http"
    "strings"
    "testing"
)

func responseOf(body string) *http.Response {
    // 用内存 Reader 构造可控响应，测试不访问网络
    return &http.Response{
        Body:          io.NopCloser(strings.NewReader(body)),
        ContentLength: int64(len(body)),
    }
}

func TestReadSmallBody(t *testing.T) {
    // 生产常量为 1 MiB，这里直接构造等于和超过上限的数据
    okBody := strings.Repeat("a", int(MaxResponseBodyBytes))
    body, err := ReadSmallBody(responseOf(okBody))
    if err != nil {
        t.Fatalf("exact limit should succeed: %v", err)
    }
    if int64(len(body)) != MaxResponseBodyBytes {
        t.Fatalf("unexpected length: %d", len(body))
    }

    tooLarge := strings.Repeat("b", int(MaxResponseBodyBytes)+1)
    _, err = ReadSmallBody(responseOf(tooLarge))
    if !errors.Is(err, ErrResponseBodyTooLarge) {
        t.Fatalf("want size error, got %v", err)
    }
}

生产项目更适合让函数接收 maxBytes，测试就不必分配 1 MiB 数据。传参时要拒绝负数，并避免对 math.MaxInt64 做 +1 导致溢出。固定且合理的业务上限则更简单，也更不容易被错误配置放大。

压缩、分块传输与连接复用

Go 官方文档说明，HTTP 客户端调用方应关闭响应体。如果 Body 没有读到 EOF 且没有关闭，底层 Transport 可能无法复用持久连接。对于正常且未超限的响应，上面的实现通过 ReadAll 读到 EOF，再由 defer 关闭，通常有利于连接复用。

超限时情况不同：LimitReader 在 maxBytes+1 处停止，底层响应体可能仍有大量未读数据。此时应优先保护内存、带宽和请求时间，及时关闭 Body，并接受该连接可能不能复用。不要为了复用连接而对一个未知巨大响应执行无界 io.Copy(io.Discard, resp.Body)，那会把内存风险换成带宽和时间风险。

如果业务确实希望尝试排空少量剩余数据，可以再加一个很小的排空上限和独立超时，但这属于连接池优化，不能破坏响应体总读取上限，也不能延长故障请求的生命周期。

常见错误对照

相关问题

LimitedReader 会关闭底层 Reader 吗？不会。它只是 Reader 包装器，没有替调用方关闭 resp.Body，所以仍需显式 Close。

服务器接收请求体也用同一方案吗？服务器 Handler 限制传入请求体时，可优先使用 http.MaxBytesReader，它专门面向服务端请求体，并能向 ResponseWriter 传递超限语义。

读取 JSON 时还需要限制吗？需要。即使改用 json.Decoder 流式解析，也应把 Decoder 的输入包在 LimitedReader 后，避免巨型字符串、数组或持续输入绕过正文边界。

最终原则很简单：协议头可以帮助快速拒绝，但真正的安全边界必须落在应用实际消费的数据流上。用 max+1 的哨兵字节建立可判定上限，再配合超时、错误分类和明确的 Body 所有权，就能让未知响应体从“潜在无界分配”变成可测试、可监控的工程约束。