I was working on a process to automate publishing this very blog, and needed to read the payload from an HTTP POST request. My first instinct was to use Golang’s io.ReadAll. This function will take a reader and read all the bytes until an EOF. That gave me pause.

What if an unscrupulous actor decided to send a really, really large payload. That could be a problem. A bit of searching revealed the io.LimitReader wrapper. This wraps a reader and a limit of the number of bytes to read.

At first glance, this appeared to solve my problems. It would guarantee that the process wouldn’t allocate an excessive amount of memory. Once the LimitReader reached the number of bytes read from the wrapped reader, it would simply return an EOF.

As I started to use this in my implementation, it occurred to me that I could not know if I reached the real EOF, or if the Reader simply had too many bytes to read.

I would really like to know if I successfully read the entire payload, or if it was truncated. For my use case, I would be receiving JSON and if the data was truncated, the un-marshalling would likely result in an error and could be detected that way. But what if the body was plain text? In that scenario one would not be able to discern between a successful read and a truncated one.

I looked up the source for the stock LimitReader. It is pretty straightforward, and would be easy to implement a version that would signal an error if the limit was reached before the true EOF.

Original Code

func LimitReader(r Reader, n int64) Reader { 
  return &LimitedReader{r, n}
}

type LimitedReader struct {
	R Reader // underlying reader
	N int64  // max bytes remaining
}

func (l *LimitedReader) Read(p []byte) (n int, err error) {
	if l.N <= 0 {
		return 0, EOF
	}
	if int64(len(p)) > l.N {
		p = p[0:l.N]
	}
	n, err = l.R.Read(p)
	l.N -= int64(n)
	return
}

The original code is pretty straight forward. The LimitedReader struct keeps track of the underlying Reader (R) as well as the remaining number of bytes to read (N).

Each Read call checks to see if we have reached the limit and simply returns an EOF. If not, it will resize the buffer to be no larger than the remaining bytes to read, this ensures that no more than N bytes are read.

Read is called on the underlying Reader with the resized buffer. N is decremented by the number of bytes read, and both the number of bytes and any error will be returned to the caller.

This guarantees that N will either reduce to 0, or an EOF will be reached from the underlying Reader. Any subsequent calls to Read will return EOF in either situation (or possibly some other i/o error).

New Code

type limitedReader struct {
	r        io.Reader
	n        int64
	complete error
}

func New(r io.Reader, n int64) io.Reader {
	return &limitedReader{r, n, nil}
}

func (l *limitedReader) Read(p []byte) (int, error) {
	if l.complete != nil {
		return 0, l.complete
	}

	if int64(len(p)) > l.n {
		p = p[0:l.n]
	}
	n, err := l.r.Read(p)
	l.n -= int64(n)

	if err != nil {
		l.complete = err
	} else {
		if len(p) == 0 {
			err = ReaderBoundsExceededError{}
		}
	}
	return n, err
}

The new implementation is essentially the same1 but it also keeps track of the error it must return to the caller as it reads from the underlying Reader (the complete field of the struct).

If the complete field is not nil, it is simply returned to the caller.

Otherwise, the buffer is resized, Read is called on the underlying Reader, and n is decremented as before.

Where this implementation differs is that a check is performed on the error from the Read call. If there is an error returned, we set the complete field to this error. This ensures that if an EOF is reached (or any other i/o error) it will be returned on any subsequent calls.

If there is no error, the length of the resized buffer is checked. If the buffer length is 0 it means that we have read the all the bytes to which the reader was limited. In this case we will set the complete field to ReaderBoundsExceededError to signal to the caller that the limit has been reached without receiving the EOF.

Example Usage

r := strings.NewReader(someText)
lr := New(r, 512)
buf, err := io.ReadAll(lr)
if err != nil {
	// Handle error case
	// io.ReadAll will never return an EOF as an error 
	// so we know in this example that the original 
	// Reader had more bytes than we are willing to process
} else {
	// Handle the success case here
}

The source for this little utility can be found here. If you would like to use this in your projects simply add the following to your go.mod file:

require github.com/luciddev13/limit_reader v1.0.1

I hope you find this useful.

Thanks for visiting.

  1. The most obvious is that I made the struct itself as well as it’s fields private. This, to me was really internal state that should not be exported to the caller. ↩