Go使用io.LimitedReader限制HTTP响应读取及返回io.Reader疑问

现有[]byte返回版本的正确性判断 #

你的实现不正确，存在3个核心问题：

• 多余的bufio.NewReader包装会引入预读副作用：bufio初始化时会从底层resp.Body预读一块数据到自身缓冲区，这部分数据可能已经超出了你设置的大小限制，同时后续defer逻辑是直接从resp.Body读取剩余内容丢弃，会漏掉bufio缓冲区中未被读取的缓存数据，导致响应体没有被完整读取，最终HTTP连接无法被复用，甚至出现后续请求响应串读的问题。
• 无法区分「响应体刚好等于限制大小正常结束」和「响应体超出限制被截断」两种场景：io.LimitReader读取到设定的上限字节数后会直接返回io.EOF，没有任何额外标识，你拿到截断后的内容会误以为是完整响应，埋下逻辑bug。
• 逻辑冗余：bufio.Reader在这里没有任何实际收益，平白增加了一次内存拷贝开销。

修正后的[]byte返回版本参考实现 #

import "errors"

var ErrResponseTooLarge = errors.New("http response size exceeds the set limit")

func getRequestData(req *http.Request, client *http.Client,
	responseSizeLimit int64) ([]byte, error) {
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer func() {
		// 完整丢弃剩余内容保证连接可复用
		io.Copy(ioutil.Discard, resp.Body)
		resp.Body.Close()
	}()

	// 直接包装原始响应Body，不需要bufio层，多读取1字节用于判断是否超限
	limitReader := io.LimitReader(resp.Body, responseSizeLimit + 1)
	data, err := ioutil.ReadAll(limitReader)
	if err != nil {
		return nil, err
	}
	// 读到的字节数超过限制，说明响应过大
	if int64(len(data)) > responseSizeLimit {
		return nil, ErrResponseTooLarge
	}
	return data, nil
}

该实现只需要多读取1个字节就能判定响应是否超限，不需要读取完整的超大响应，性能和正确性都能保证。

返回io.Reader类型的实现方案 #

完全可以基于http.Response.Body和限流逻辑实现，但是要注意不能在当前函数内用defer处理响应体的关闭和剩余内容丢弃——函数返回时调用方还没开始读取数据，提前关闭会导致读失败。正确的实现方式是自定义一个包装Reader，把资源回收逻辑内置到Close方法中，调用方只需要像处理普通io.ReadCloser一样读取、关闭即可。

参考实现：

type limitedResponseBody struct {
	resp   *http.Response
	reader io.Reader
	limit  int64
}

func (l *limitedResponseBody) Read(p []byte) (n int, err error) {
	n, err = l.reader.Read(p)
	l.limit -= int64(n)
	// 只要limit被减到负数，说明已经读到了超过限制的内容，返回明确错误
	if l.limit < 0 {
		return n, ErrResponseTooLarge
	}
	return n, err
}

func (l *limitedResponseBody) Close() error {
	// 关闭前自动丢弃剩余内容，保证连接可复用
	io.Copy(ioutil.Discard, l.resp.Body)
	return l.resp.Body.Close()
}

func GetRequestReader(req *http.Request, client *http.Client,
	responseSizeLimit int64) (io.ReadCloser, error) {
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	return &limitedResponseBody{
		resp:   resp,
		reader: io.LimitReader(resp.Body, responseSizeLimit+1),
		limit:  responseSizeLimit,
	}, nil
}

内容的提问来源于stack exchange，提问作者Vladimir Bershov