Update

vendor/tailscale.com/util/httpbody/httpbody.go

@@ -0,0 +1,145 @@
+// 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)
+}