Introduce core.ErrOnLimitReader and DefaultMaxRead

Author: ezekiel

bdns/dns.go

@@ -17,6 +17,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 
+	"github.com/letsencrypt/boulder/core"
 	blog "github.com/letsencrypt/boulder/log"
 	"github.com/letsencrypt/boulder/metrics"
 )
@@ -375,7 +376,9 @@ func (d *dohExchanger) ExchangeContext(ctx context.Context, query *dns.Msg, serv
 		return nil, d.clk.Since(start), fmt.Errorf("doh: http status %d", resp.StatusCode)
 	}
 
-	b, err := io.ReadAll(&io.LimitedReader{R: resp.Body, N: 300_000})
+	// DNS response over 65535 is malformed https://datatracker.ietf.org/doc/html/rfc8484#section-6
+	lmr := core.ErrOnLimitReader(resp.Body, 65_535)
+	b, err := io.ReadAll(lmr)
 	if err != nil {
 		return nil, d.clk.Since(start), fmt.Errorf("doh: reading response body: %w", err)
 	}

cmd/crl-checker/main.go

@@ -28,7 +28,8 @@ func downloadShard(url string) (*x509.RevocationList, error) {
 		return nil, fmt.Errorf("downloading crl: http status %d", resp.StatusCode)
 	}
 
-	crlBytes, err := io.ReadAll(&io.LimitedReader{R: resp.Body, N: 1_000_000_000})
+	lmr := core.ErrOnLimitReader(resp.Body, 1_000_000_000)
+	crlBytes, err := io.ReadAll(lmr)
 	if err != nil {
 		return nil, fmt.Errorf("reading CRL bytes: %w", err)
 	}

cmd/shell.go

@@ -504,7 +504,8 @@ func ValidateYAMLConfig(cv *ConfigValidator, in io.Reader) error {
 	// Register custom types for use with existing validation tags.
 	validate.RegisterCustomTypeFunc(config.DurationCustomTypeFunc, config.Duration{})
 
-	inBytes, err := io.ReadAll(&io.LimitedReader{R: in, N: 300_000})
+	lmr := core.ErrOnLimitReader(in, core.DefaultMaxRead)
+	inBytes, err := io.ReadAll(lmr)
 	if err != nil {
 		return err
 	}

core/core_test.go

@@ -3,6 +3,8 @@ package core
 import (
 	"encoding/base64"
 	"encoding/json"
+	"io"
+	"strings"
 	"testing"
 
 	"github.com/go-jose/go-jose/v4"
@@ -80,3 +82,38 @@ func TestFingerprint(t *testing.T) {
 		t.Errorf("Incorrect SHA-256 fingerprint: %v", digest)
 	}
 }
+
+func TestErrOnLimitReader(t *testing.T) {
+	// test a read where the limit is larger than the source
+	strReader := strings.NewReader("foo bar baz bot")
+	lmr := ErrOnLimitReader(strReader, 21)
+
+	strOut01, err := io.ReadAll(lmr)
+
+	// should return no err, and the full source
+	test.AssertNotError(t, err, "unexpectedly errored")
+	test.AssertEquals(t, string(strOut01), "foo bar baz bot")
+
+	// test a read where the limit exactly matches source size
+	strReader = strings.NewReader("foo bar baz bot")
+	lmr = ErrOnLimitReader(strReader, 15)
+
+	strOut02, err := io.ReadAll(lmr)
+
+	// should return an error, and the full source
+	test.AssertError(t, err, "unexpectedly succeeded")
+	test.AssertEquals(t, err, ErrReaderLimitReached)
+	test.AssertEquals(t, string(strOut02), "foo bar baz bot")
+
+	// test a read where the limit is lower than the source
+	strReader = strings.NewReader("foo bar baz bot")
+	lmr = ErrOnLimitReader(strReader, 9)
+
+	strOut03, err := io.ReadAll(lmr)
+
+	// should return an error, and a partial read result
+	test.AssertError(t, err, "unexpectedly succeeded")
+	test.AssertEquals(t, err, ErrReaderLimitReached)
+	test.AssertEquals(t, string(strOut03), "foo bar b")
+
+}

core/util.go

@@ -51,6 +51,14 @@ var BuildHost string
 // BuildTime is set by the compiler and is used by GetBuildTime
 var BuildTime string
 
+// DefaultMaxRead is for use by ErrOnLimitReader when to conveniently limit
+// reads to less than half a MB, which should be most of the time
+var DefaultMaxRead int64 = 300_000
+
+// ErrReaderLimitReached as an exported error type allows callers to check for
+// this error type after Read
+var ErrReaderLimitReached error = errors.New("reader size limit reached")
+
 func init() {
 	expvar.NewString("BuildID").Set(BuildID)
 	expvar.NewString("BuildTime").Set(BuildTime)
@@ -431,3 +439,39 @@ func NormalizeIssuerDomainName(name string) (string, error) {
 	}
 	return name, nil
 }
+
+// An ErrOnLimitedReader reads from R but limits the amount of data returned to
+// just N bytes. Each call to Read updates N to reflect the new amount
+// remaining.
+// Read returns ErrReaderLimitExceeded when N <= 0, or EOF when the underlying R
+// returns EOF.
+type ErrOnLimitedReader struct {
+	R io.Reader
+	N int64
+}
+
+// ErrOnLimitReader returns a Reader that reads from r but stops with
+// ErrReaderLimitExceeded after n bytes.
+// The underlying implementation is a *ErrOnLimitedReader.
+//
+// If LimitedReader gets an Err field, we can switch to use that
+// https://github.com/golang/go/issues/51115
+func ErrOnLimitReader(r io.Reader, n int64) io.Reader {
+	return &ErrOnLimitedReader{r, n}
+}
+
+func (ltdR *ErrOnLimitedReader) Read(b []byte) (n int, err error) {
+	// Important: when an ErrOnLimitedReader reads _exactly_ the limit amount,
+	// it results in the LimitReached Error
+	if ltdR.N <= 0 {
+		return 0, ErrReaderLimitReached
+	}
+
+	if int64(len(b)) > ltdR.N {
+		b = b[0:ltdR.N]
+	}
+
+	n, err = ltdR.R.Read(b)
+	ltdR.N -= int64(n)
+	return
+}

linter/pkimetal/client.go

@@ -15,6 +15,8 @@ import (
 	"time"
 
 	"github.com/zmap/zlint/v3/lint"
+
+	"github.com/letsencrypt/boulder/core"
 )
 
 // Config holds configuration for linting both certs and CRLs using PKIMetal.
@@ -80,7 +82,8 @@ func (pkim *Client) Execute(endpoint string, der []byte) (*lint.LintResult, erro
 		return nil, fmt.Errorf("got status %d (%s) from pkimetal API", resp.StatusCode, resp.Status)
 	}
 
-	resJSON, err := io.ReadAll(&io.LimitedReader{R: resp.Body, N: 300_000})
+	lmr := core.ErrOnLimitReader(resp.Body, core.DefaultMaxRead)
+	resJSON, err := io.ReadAll(lmr)
 	if err != nil {
 		return nil, fmt.Errorf("reading response from pkimetal API: %s", err)
 	}

observer/probers/aia/aia.go

@@ -8,6 +8,8 @@ import (
 	"net/http"
 
 	"github.com/prometheus/client_golang/prometheus"
+
+	"github.com/letsencrypt/boulder/core"
 )
 
 // AIAProbe is the exported 'Prober' object for monitors configured to
@@ -48,7 +50,8 @@ func (p AIAProbe) Probe(ctx context.Context) error {
 		return fmt.Errorf("certificate Content-Type is %q but want application/pkix-cert", contentType)
 	}
 
-	body, err := io.ReadAll(&io.LimitedReader{R: resp.Body, N: 300_000})
+	lmr := core.ErrOnLimitReader(resp.Body, core.DefaultMaxRead)
+	body, err := io.ReadAll(lmr)
 	if err != nil {
 		return err
 	}

observer/probers/ccadb/retryhttp.go

@@ -6,6 +6,8 @@ import (
 	"io"
 	"net/http"
 	"time"
+
+	"github.com/letsencrypt/boulder/core"
 )
 
 func getBody(ctx context.Context, url string) ([]byte, error) {
@@ -21,19 +23,24 @@ func getBody(ctx context.Context, url string) ([]byte, error) {
 	}
 	defer resp.Body.Close()
 
-	body, err := io.ReadAll(&io.LimitedReader{R: resp.Body, N: 1_000_000_000})
-	if err != nil {
-		return nil, err
-	}
-
 	if resp.StatusCode != http.StatusOK {
-		// Truncate the response body in case it's too big to be useful in logs.
-		if len(body) > 400 {
-			body = body[:400]
+		// Read up to 400 bytes of response body to include in error logs
+		lmr := core.ErrOnLimitReader(resp.Body, 400)
+		body, err := io.ReadAll(lmr)
+		if err != nil && err != core.ErrReaderLimitReached {
+			return nil, err
 		}
+
 		return nil, fmt.Errorf("http status %d for %q: %s", resp.StatusCode, url, string(body))
 	}
 
+	// Read up to ~1G of response body to return to caller
+	lmr := core.ErrOnLimitReader(resp.Body, 1_000_000_000)
+	body, err := io.ReadAll(lmr)
+	if err != nil {
+		return nil, err
+	}
+
 	return body, nil
 }
 

observer/probers/crl/crl.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/prometheus/client_golang/prometheus"
 
+	"github.com/letsencrypt/boulder/core"
 	"github.com/letsencrypt/boulder/crl/idp"
 	"github.com/letsencrypt/boulder/observer/obsclient"
 )
@@ -47,7 +48,8 @@ func (p CRLProbe) Probe(ctx context.Context) error {
 	}
 	defer resp.Body.Close()
 
-	body, err := io.ReadAll(&io.LimitedReader{R: resp.Body, N: 1_000_000_000})
+	lmr := core.ErrOnLimitReader(resp.Body, 1_000_000_000)
+	body, err := io.ReadAll(lmr)
 	if err != nil {
 		return err
 	}

observer/probers/tls/tls.go

@@ -16,6 +16,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/crypto/ocsp"
 
+	"github.com/letsencrypt/boulder/core"
 	"github.com/letsencrypt/boulder/observer/obsclient"
 )
 
@@ -86,7 +87,8 @@ func checkOCSP(ctx context.Context, cert, issuer *x509.Certificate, want int) (b
 	}
 	defer res.Body.Close()
 
-	output, err := io.ReadAll(&io.LimitedReader{R: res.Body, N: 300_000})
+	lmr := core.ErrOnLimitReader(res.Body, core.DefaultMaxRead)
+	output, err := io.ReadAll(lmr)
 	if err != nil {
 		return false, err
 	}
@@ -115,7 +117,8 @@ func checkCRL(ctx context.Context, cert, issuer *x509.Certificate, want int) (bo
 	}
 	defer resp.Body.Close()
 
-	der, err := io.ReadAll(&io.LimitedReader{R: resp.Body, N: 1_000_000_000})
+	lmr := core.ErrOnLimitReader(resp.Body, 1_000_000_000)
+	der, err := io.ReadAll(lmr)
 	if err != nil {
 		return false, fmt.Errorf("reading CRL: %w", err)
 	}