TIL: 18th July 2023 — reader limiting in Go

Go has standard library functionality that can limit the amount of bytes read from some io.Reader stream. This can be easily used to e.g. prevent reading an unexpectedly large payload into the memory or pick a subset of an input that has unknown size.

io.LimitReader§

One of the methods to do so is the io.LimitReader function, which works on any implementation of the io.Reader interface (a file, a TCP stream, stdin, etc.)

A snippet below uses this function to generate a random 16 byte stream (e.g. a password) and copy it to stdout:

package main

import (
	"crypto/rand"
	"io"
	"os"
)

func main() {
	limitedRandom := io.LimitReader(rand.Reader, 16)

	if _, err := io.Copy(os.Stdout, limitedRandom); err != nil {
		panic(err)
	}
}

http.MaxBytesReader§

There's also a more specialised http.MaxBytesReader function that works well with the Go standard library HTTP server interfaces that are widely used in the ecosystem. As the documentation says, it has a benefit of closing the underlying reader as well as telling the writer to close the connection after the limit has been reached so that keep-alive would be overridden, forcing the client to make a new connection.

The ideal use case is to limit the HTTP request size in order to prevent excessive load on the infrastructure. The snippet below demonstrates its usage as a middleware function:

package main

import (
	"io"
	"net/http"
)

const MaxBodySize = 1_00

func echo(w http.ResponseWriter, r *http.Request) {
	io.Copy(w, r.Body)
}

func maxSizeMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		r.Body = http.MaxBytesReader(w, r.Body, MaxBodySize)

		next.ServeHTTP(w, r)
	})
}

func main() {
	mux := http.NewServeMux()
	
	mux.Handle("/echo", maxSizeMiddleware(http.HandlerFunc(echo)))
	
	if err := http.ListenAndServe(":8080", mux); err != nil {
		panic(err)
	}
}

This implements a trivial echo handler that writes the request's Body into the response, but adds a middleware that limits the request's Body reader to 100 bytes.

If you run this server and send a payload larger than 100 bytes to the /echo endpoint the response will contain this payload truncated accordingly:

❯ xh http://localhost:8080/echo body=dgdfskdjfhsjkdhfjksdhfjksdhfkjsdhfjkkjhjkhjkdfgdfgdfsgdfffffffffffffffffffffffffffffffffffgdfgdfg
HTTP/1.1 200 OK
Connection: close
Content-Length: 100
Content-Type: text/plain; charset=utf-8
Date: Tue, 18 Jul 2023 20:12:44 GMT

{"body":"dgdfskdjfhsjkdhfjksdhfjksdhfkjsdhfjkkjhjkhjkdfgdfgdfsgdfffffffffffffffffffffffffffffffffffg

Note how the connection was also hinted as close in the response headers instead of being keep-alive as well as the content length being exactly 100 bytes.

This is a trivial example of course, and error handling should be added to return the right status code (http.StatusRequestEntityTooLarge) with a descriptive error message, which will be returned by the reader when it reaches the limit.

In conclusion, the two functions have a similar purpose (limit the reader stream), but different side effects and therefore use cases.