package blobgen

import (
	"errors"
	"io"
)

// ErrOutputTooLarge is returned by a reader from LimitReader once it has
// been asked for more than its limit. It bounds how far an untrusted
// compressed stream may expand, so a small, highly compressible object
// from the store cannot decompress without limit.
var ErrOutputTooLarge = errors.New("output exceeds size limit")

// LimitReader returns a reader that yields at most limit bytes from r and
// then fails with ErrOutputTooLarge. Unlike io.LimitReader, which reports
// a silent io.EOF at the limit (indistinguishable from a stream that
// simply ended), this fails, so a caller decoding or copying the stream
// sees an error rather than a truncated value. A stream of exactly limit
// bytes reads back cleanly to EOF; the first byte beyond it is the error.
func LimitReader(r io.Reader, limit int64) io.Reader {
	// remaining counts down from limit+1: the extra byte is the one that,
	// if it ever arrives, proves the stream is longer than the limit.
	return &limitReader{r: r, remaining: limit + 1}
}

type limitReader struct {
	r         io.Reader
	remaining int64
}

func (l *limitReader) Read(p []byte) (int, error) {
	if l.remaining <= 0 {
		return 0, ErrOutputTooLarge
	}

	if int64(len(p)) > l.remaining {
		p = p[:l.remaining]
	}

	n, err := l.r.Read(p)
	l.remaining -= int64(n)

	if l.remaining <= 0 {
		// The (limit+1)th byte was just read: the stream is too long.
		return n, ErrOutputTooLarge
	}

	return n, err
}