146 lines
5.0 KiB
Go
146 lines
5.0 KiB
Go
// Copyright (c) Tailscale Inc & contributors
|
|
// SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
// Package httpbody provides helpers for reading HTTP response bodies
|
|
// defensively, capping how much memory a response can make the reader
|
|
// buffer.
|
|
//
|
|
// A typical HTTP client applies the cap in its Do method, so every
|
|
// response it returns is capped unless the request opts out:
|
|
//
|
|
// res, err := c.Client.Do(req)
|
|
// httpbody.LimitSize(res)
|
|
package httpbody
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
|
|
"tailscale.com/util/ctxkey"
|
|
)
|
|
|
|
// ErrTooLarge is reported by a body wrapped with [LimitSize] once more
|
|
// than the configured number of bytes have been read. The error text
|
|
// includes the limit that was exceeded.
|
|
var ErrTooLarge = errors.New("response body too large")
|
|
|
|
// DefaultMaxSize is the default response body size cap, in bytes, used
|
|
// when a context carries no override set with [WithMaxSize]. HTTP clients
|
|
// that apply response body caps should honor it, so requests can tune or
|
|
// disable the cap uniformly across clients.
|
|
const DefaultMaxSize = 1 << 20
|
|
|
|
// maxSizeKey is the context key carrying the response body size cap.
|
|
var maxSizeKey = ctxkey.New[int64]("httpbody.maxSize", DefaultMaxSize)
|
|
|
|
// WithMaxSize returns a context that sets the response body size cap to
|
|
// max, instead of [DefaultMaxSize], for HTTP clients that apply response
|
|
// body caps. A max of zero or less disables the cap. Disable it for
|
|
// responses that stream an unbounded number of individually bounded
|
|
// messages, like a streaming map response.
|
|
func WithMaxSize(ctx context.Context, max int64) context.Context {
|
|
return maxSizeKey.WithValue(ctx, max)
|
|
}
|
|
|
|
// MaxSize returns the response body size cap in ctx: the value set with
|
|
// [WithMaxSize], or [DefaultMaxSize] if ctx carries none.
|
|
func MaxSize(ctx context.Context) int64 {
|
|
return maxSizeKey.Value(ctx)
|
|
}
|
|
|
|
// LimitSize replaces res.Body with a wrapper that fails reads with an
|
|
// error wrapping [ErrTooLarge] once more than the response body size
|
|
// limit in res.Request's context, per [WithMaxSize], has been read:
|
|
// [DefaultMaxSize] by default, or if res has no Request. A limit of zero
|
|
// or less imposes no limit, removing any limit a previous call set.
|
|
//
|
|
// Repeated calls to LimitSize or LimitSizeTo replace the previous limit
|
|
// rather than compounding it, so a later call can raise or remove the
|
|
// limit an earlier one set. If bytes have already been read through a
|
|
// previous limit, the new limit counts bytes read after it was applied.
|
|
//
|
|
// Unlike io.LimitReader, which silently truncates at its limit, the
|
|
// wrapper reports an oversize body so callers can distinguish a short
|
|
// body from a too-large one. A body of at most the limit, including one
|
|
// of exactly the limit, reads back without error.
|
|
func LimitSize(res *http.Response) {
|
|
max := int64(DefaultMaxSize)
|
|
if res.Request != nil {
|
|
max = MaxSize(res.Request.Context())
|
|
}
|
|
limitSizeTo(res, max)
|
|
}
|
|
|
|
// LimitSizeTo is like [LimitSize] but caps res.Body at exactly max bytes,
|
|
// ignoring any override in res.Request's context. A max of zero or less
|
|
// imposes no limit, removing any limit a previous call set. Repeated
|
|
// calls to LimitSize or LimitSizeTo replace the previous limit rather
|
|
// than compounding it.
|
|
func LimitSizeTo(res *http.Response, max int64) {
|
|
limitSizeTo(res, max)
|
|
}
|
|
|
|
func limitSizeTo(res *http.Response, max int64) {
|
|
// Unwrap any limit a previous call set instead of wrapping it again:
|
|
// nested limits would compound, with the smaller always winning, so
|
|
// re-limiting could never raise the cap. At most one wrapper can
|
|
// exist, because this always unwraps first.
|
|
if lb, ok := res.Body.(*limitedBody); ok {
|
|
res.Body = lb.body
|
|
}
|
|
if max <= 0 {
|
|
return
|
|
}
|
|
res.Body = &limitedBody{body: res.Body, max: max, remain: max}
|
|
}
|
|
|
|
// limitedBody is the io.ReadCloser installed by [LimitSize].
|
|
type limitedBody struct {
|
|
body io.ReadCloser
|
|
max int64 // the limit, for error reporting
|
|
// remain is the number of bytes still allowed. Once the limit is
|
|
// exhausted, a one-byte probe of body distinguishes a body of exactly
|
|
// max bytes (EOF, fine) from one with more to come. remain goes
|
|
// negative once oversize is detected, and stays there.
|
|
remain int64
|
|
}
|
|
|
|
func (b *limitedBody) Read(p []byte) (n int, err error) {
|
|
if len(p) == 0 {
|
|
return 0, nil
|
|
}
|
|
if b.remain < 0 {
|
|
return 0, b.tooLarge()
|
|
}
|
|
if b.remain == 0 {
|
|
// The limit is exhausted. A body of exactly max bytes is legal, so
|
|
// check whether the body has actually ended before declaring it
|
|
// too large. The probed byte, if any, is dropped: once a body is
|
|
// too large, callers only want the error.
|
|
var probe [1]byte
|
|
m, err := b.body.Read(probe[:])
|
|
if m == 0 && err != nil {
|
|
return 0, err // io.EOF, or the underlying error
|
|
}
|
|
b.remain = -1
|
|
return 0, b.tooLarge()
|
|
}
|
|
if int64(len(p)) > b.remain {
|
|
p = p[:b.remain]
|
|
}
|
|
n, err = b.body.Read(p)
|
|
b.remain -= int64(n)
|
|
return n, err
|
|
}
|
|
|
|
func (b *limitedBody) Close() error {
|
|
return b.body.Close()
|
|
}
|
|
|
|
func (b *limitedBody) tooLarge() error {
|
|
return fmt.Errorf("%w: limit %d bytes", ErrTooLarge, b.max)
|
|
}
|