Blocking large requests in Go using http library

If your server never expects over 100kb in a single request, you can block oversized requests to protect against abuse. You could block them in a reverse proxy, e.g. NGINX, but you can also block them in application code. Here I will talk about the latter, since this can be useful if you don’t have a reverse proxy, or you need fine control over the limiting.

Blocking a large request, by counting the bytes

Let’s say you want to block requests over 100kb. This is something I looked into for a recent project. My first thought when looking into this was to replace the body with a reader that closes after it consumes 100kb and throws an error. As it happens there is reader built into the http library called MaxBytesReader that does exactly this, so you can do something like this.

http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    maxBytesReader := http.MaxBytesReader(w, r.Body, maxBytes)
    ...
}

The great thing about doing it as a reader is you don’t really impact performance that much, you do the read you’d do anyway. This is as opposed to say reading the whole body into memory, then realizing it is 100mb, wasting memory, and then rejecting the request.

One thing to note is that the user of maxBytesReader needs to determine what to do when the reader returns the error reporting the max size has been reached. It will need to determine that if it is a too large error then return the response code. Like this:

var maxBytesError *http.MaxBytesError

body, err := io.ReadAll(maxBytesReader)
if err != nil {
	if errors.As(err, &maxBytesError) {
		w.WriteHeader(http.StatusRequestEntityTooLarge)
		fmt.Fprintf(w, "Content Too Large")
	}
} 

Next, let’s make the typical case more efficient. Most requests will send the size of their payload in the header, so we can block earlier if we know the request will be too big. Let’s look into that.

Blocking a large request using content-length header

The HTTP/1.1 content-length header is an highly recommended (but technically optional) field to indicate how much data you intend to send in a request or response.

We can use this header to block a request that is too large before even reading the body. However we can’t trust that it will be set, or if set that is is accurate, so you will want to combine this with counting the bytes.

Combining the two methods, looks like this:

http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {

	if r.ContentLength > maxBytes {
		w.WriteHeader(http.StatusRequestEntityTooLarge)
		fmt.Fprintf(w, "Content Too Large")
		return
	}

	maxBytesReader := http.MaxBytesReader(w, r.Body, maxBytes)

        ...
}

That is it. We are all done! Now let’s now visit a couple of things related to this of interest…

Additional Notes

Sometimes the http server will check body length, but don’t rely on that.

Go’s http server (see transfer.go) will limit the reader to only read a number of bytes equal to realLength which is usually set to content-length, but only if the content-length is supplied, and there are various conditions around whether it will do it or not. This is a good feature, but not 100% reliable for the purpose of blocking large requests to protect your service.

Here is the code inside the Go http library that does this:

func readTransfer(msg any, r *bufio.Reader) (err error) {

//...
	switch {
//...
	case realLength > 0:
		t.Body = &body{src: io.LimitReader(r, realLength), closing: t.Close}
//...

LimitReader vs. MaxBytesReader

There are two similar readers for limiting request size, io.LimitReader and http.MaxBytesReader. Here is what the documentation says:

MaxBytesReader is similar to io.LimitReader but is intended for limiting the size of incoming request bodies. In contrast to io.LimitReader, MaxBytesReader’s result is a ReadCloser, returns a non-nil error of type *MaxBytesError for a Read beyond the limit, and closes the underlying reader when its Close method is called.

MaxBytesReader prevents clients from accidentally or maliciously sending a large request and wasting server resources. If possible, it tells the ResponseWriter to close the connection after the limit has been reached.

And for the complete picture, the docs for LimitReader:

LimitReader returns a Reader that reads from r but stops with EOF after n bytes. The underlying implementation is a *LimitedReader.

If we used a LimitReader, we wouldn’t know that the client had sent too much data, and instead continued processing the request, likely hitting an error later because the payload is truncated.

Unit test mocks in Go

Recently I have been using Go for the first time seriously. Go is famously easy to learn for programmers from other languages, and I found this to be true. However, performing mocking in Go unit tests required a bit of my head scratching, so I’ve decided to write about it!

This post explores different ways you can mock in Go, and the tradeoffs involved between them.

What is test mocking?

It’s good to go over the definitions of mocks. I am going to use these definitions, taken from Martin Fowler’s article on mocks vs. stubs, which are in turn taken from Gerard Meszaros’s book xUnit Test Patterns, with my minor edits for brevity. Don’t worry about remembering all of this—I’ll recap as I use them:

Test Double is the generic term for any kind of pretend object used in place of a real object for testing purpose, can be one of:

  • Dummy objects are passed around but never actually used. Usually they are just used to fill parameter lists.
  • Fake objects actually have working implementations, but usually take some shortcut which makes them not suitable for production (an in memory database is a good example).
  • Stubs provide canned answers to calls made during the test, usually not responding at all to anything outside what’s programmed in for the test.
  • Spies are stubs that also record some information based on how they were called. One form of this might be an email service that records how many messages it was sent.
  • Mocks are what we are talking about here: objects pre-programmed with expectations which form a specification of the calls they are expected to receive.

I am going to show how to do various test doubles in Go, some of which are mocks.

Test doubles help reduce the computational cost of a test, and make setup simpler (no need to create a DB, for example). They tend to help make more reliable and focused tests.

Download-a-cool-RSS-feed example program.

We will get started with testing, but we need something to test. The program we will use for our examples downloads the feed of (perhaps) your favorite programming blog. Feel free to change the URL!

It’s simple enough, but there is enough complexity here to demonstrate different types of testing. Here is the code:

package main

import (
	“io”
	“log”
	“net/http”
	"os"
)

func main() {
	err := Download("https://martincapodici.com/feed", "feed.xml")
	if err != nil {
		log.Fatalln(err)
		os.Exit(1)
	}
	log.Println("Download complete")
}

func Download(url string, fileName string) error {
	resp, err := http.Get(url)
	if err != nil {
		return err
	}

	defer resp.Body.Close()

	file, err := os.Create(fileName)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(file, resp.Body)
	return err
}

You can try this now, by copying this into main.go and running the command:

go run main.go

This code has a function Download(...) that performs IO, both to make an HTTP call, and to save the result to disk. I am looking to unit test this without doing any real IO, so I need to change this a bit so it can be run with or without IO. My first job is to modify the function so that mocked IO is an option.

I will create an interface called IO that will expose CreateFile() and GetUrl() methods. I can then create mock implementations of that interface.

First let’s add the interface that lets this method do the IO things it needs to do:

type IO interface {
	GetUrl(url string) (resp *http.Response, err error)
	CreateFile(fileName string) (file io.ReadWriteCloser, err error)
}

The file is passed using the io library interface ReadWriteCloser, to make it swappable for a test object.

You might notice the absence of an interface for the io.Copy call. Since io.Copy works with interfaces, it is already test-ready.

Here is the io.Copy documentation to prove it works with Reader and Writer interfaces:

// Copy copies from src to dst until either EOF is reached
// on src or an error occurs. It returns the number of bytes
// copied and the first error encountered while copying, if any.
// ...
func Copy(dst Writer, src Reader) (written int64, err error) {

Anyway, lets get this done! We now will implement the real (i.e. not test) methods for the interface. We need a “receiver” object to attach the methods to. Since there is no state we need to keep, this is an empty struct that we will call RealIO:

type RealIO struct{}


func NewRealIO() IO {
	return &RealIO{}
}

Here are its methods:



func (r *RealIO) GetUrl(url string) (resp *http.Response, err error) {
	return http.Get(url)
}

func (r *RealIO) CreateFile(fileName string) (file io.ReadWriteCloser, err error) {
	return os.Create(fileName)
}

The download function can now use the interface instead of direct library calls:

func Download(testIO IO, url string, fileName string) error {
	resp, err := testIO.GetUrl(url)
	if err != nil {
		return err
	}

	defer resp.Body.Close()

	file, err := testIO.CreateFile(fileName)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(file, resp.Body)
	return err
}

Finally, we can update the main program to provide this. A one line change:

err := Download(NewRealIO(), "https://martincapodici.com/feed", "feed.xml")

Testing the app, using mocks (and friends)

We will now look at a few ways we can test the code above. We’ll start with the pure Go solution, then use a library, and then touch on code generation options.

Testing with vanilla Go: Or, just use the interface

Without importing any external libraries we can create a test double that implements the IO interface, providing test data and recording calls.

However, before we get to that, lets create a fake file as well need that for constructing the test and for the interface implementation.

I have chosen to fake the file using a buffer (bytes.Buffer). A buffer is basically bytes in memory that happens to have the same Read and Write methods as a file—so swapping in a buffer for a file should be easy enough.

Buffer + Close() = What we want

A buffer however doesn’t have a Close method, so we need to add one. But how do we do this when bytes.Buffer is a standard library object we can’t modify?

To make a buffer that is a ReadWriteCloser the following code uses a technique called struct embedding. Struct embedding is syntax sugar, meaning it’s short hand for something that would normally need more code.

Struct embedding lets you to make a struct a field of another struct and have all of the functions defined on that struct be part of the outer struct.

If this is confusing, it should become clear when you see the code or try this code out for yourself and experiment.

With struct embedding we can embed the buffer in a new struct, called InMemoryFile that gets given the Read and Write methods from the buffer. We then add the additional Close method so that InMemoryFile implements the entire ReadWriteCloser interface:

type InMemoryFile struct {
	bytes.Buffer
}

func (f *InMemoryFile) Close() error { return nil }

// + Struct Embedding Magic! InMemoryFile's Read() and Write() functions are implemented behind the
// scenes by Go so we don't have to write them by hand. What do they do? They just call
// the Read() and Write() functions of the buffer.

We will now create a TestIO struct, that will be used to sort-of act like the real IO, and also help us make assertions for the test.

Lets define the struct and a function to create a new instance. We have an InMemoryFile in our struct so we can fake the file:

type TestIO struct {
	file InMemoryFile
}

func NewTestIO() *TestIO {
	file := InMemoryFile{*bytes.NewBuffer(nil)}
	return &TestIO{file}
}

The interface for IO has two methods that we need to implement. The first method is GetUrl, which is set up to return an OK response for a test URL, and otherwise an error:

func (f *TestIO) GetUrl(url string) (resp *http.Response, err error) {
	if url == "http://testurl" {
		resp = &http.Response{
			StatusCode: http.StatusOK,
			Body:       io.NopCloser(strings.NewReader("OK")),
		}
		return resp, nil
	}
	return nil, errors.New("error")
}

The CreateFile method is next, and will return the fake file, or an error:

func (f *TestIO) CreateFile(fileName string) (file io.ReadWriteCloser, err error) {
	if fileName == "testfile" {
		return &f.file, nil
	}
	return nil, errors.New("error")
}

This is a test double object that is pretty hard wired for a specific kind of test. Hint hint! (there is a quiz question later!)

So we now have enough scaffold to do a basic test that simulates a file download, and checks that the mocked file has the correct contents:

func TestDownload(t *testing.T) {
	testIO:= NewTestIO()
	err := Download(testIO, "http://testurl", "testfile")
	if err != nil {
		t.Error(err)
	}
	b := testIO.file.Buffer.String()
	if string(b) != "OK" {
		t.Errorf("expected file to contain 'OK', got %s", string(b))
	}
}

This is only tests one scenario, and doesn’t test different kinds of IO errors. To do that we’d need to extend the test double to do more, and possibly record values that are passed in to them.

Doing this by hand gets tedious, so we will look into how to make this less work with mocking libraries. There are two good options I have seen, and we can weigh up pros and cons after.

Quick quiz

What kind of test double is TestIO? (Revisit the definitions under What is test mocking if needed.)

Use the stretchr/testify/mock library

The stretchr/testify/mock library is straightforward to use. It’s hard to beat the example given on https://pkg.go.dev/github.com/stretchr/testify/mock#hdr-Example_Usage for a description of how this library works and how to use it. Please read that first.

In short though the idea is your mock struct embeds the mock.Mock object, which you then use in each method you want to implement in the mock, to pass in the given parameters and return mock values that were set up in the test.

Here is how to use it for our example. First install it:

go get github.com/stretchr/testify/mock

Then in a new test file, say main2_test.go, we can create the mock, along with the interface methods that delegate the work to the mock. Each method passes it’s arguments into Called to tell the mock it was called, and then the mock returns the result that can be deconstructed. Nothing specific is set up yet; we get to control the behavior from the test code.

type MockIO struct {
	mock.Mock
}

func (m *MockIO) GetUrl(url string) (resp *http.Response, err error) {
	args := m.Called(url)
	return args.Get(0).(*http.Response), args.Error(1)
}

func (m *MockIO) CreateFile(fileName string) (file io.ReadWriteCloser, err error) {
	args := m.Called(fileName)
	return args.Get(0).(io.ReadWriteCloser), args.Error(1)
}

I also like add a line that asserts that MockIO has implemented the IO interface correctly. This gives you immediate feedback via the IDE if the interfaces is not properly implemented, plus it prevents the code from compiling if such an issue exists:

var _ IO = (*MockIO)(nil)

With this in the place we can now have a succinct test that defines how the mock should behave, then tests the function with that behavior. The rather haphazard fake implementation used earlier is replaced with just two testIO.On statements:


func TestDownloadUsingMock(t *testing.T) {
	testIO := new(MockIO)

	file := &InMemoryFile{*bytes.NewBuffer(nil)}

	resp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader("OK")),
	}

	testIO.On("GetUrl", "http://testurl").Return(resp, nil)
	testIO.On("CreateFile", "testfile").Return(file, nil)

	err := Download(testIO, "http://testurl", "testfile")
	if err != nil {
		t.Error(err)
	}

}

Hopefully you can see that it is now much easier to test other scenarios without needing to change the mock implementation. In addition to the On/Return pattern shown here there a multitude of behaviors available in the library, as well as the ability to bake in expectations e.g. the function should be called exactly once. There is also a MatchedBy argument you can use to specify a function that determines if an argument, for example, this could be used in the test too:

	testIO.On("GetUrl", mock.MatchedBy(func(url string) bool { return strings.Contains(url, "testurl") })).Return(resp, nil)

These features make it strictly better than building your own fake implementations in my opinion, which some rare exceptions.

An example of where a fake might be preferred is you want the fake to act like the real object by maintaining internal state. For example where the real object talks to a Redis cache, and the fake object uses a crude in-memory cache, and consumers want to test their code integrated with that cache. In a test where the function you are testing is making heavy use of the cache, it might be less attractive to set up all of the possible interactions in a mock.

Mockery

I’ll give a quick shout out to mockery, a tool that generates your stretchr/testify/mock mocks. https://vektra.github.io/mockery/latest

The benefit of this is not typing out a lot of boilerplate, but the downside is perhaps having mocks for things you never need, and adding bloat to your code base.

Check out their documentation if you want to use it. I haven’t yet decided if I prefer Mockery, or hand writing mock objects as needed. If you have an opinion or know of other interesting Go mocking libraries, please let me know, e.g. leave a comment.

Image by Bianca Van Dijk from Pixabay

Human-made Content