badguardhome/internal/aghio/limitedreader.go

// Package aghio contains extensions for io package's types and methods
package aghio

import (
	"fmt"
	"io"
)

// LimitReachedError records the limit and the operation that caused it.
type LimitReachedError struct {
	Limit int64
}

// Error implements the error interface for LimitReachedError.
//
// TODO(a.garipov): Think about error string format.
func (lre *LimitReachedError) Error() string {
	return fmt.Sprintf("attempted to read more than %d bytes", lre.Limit)
}

// limitedReader is a wrapper for io.Reader with limited reader and dealing with
// errors package.
type limitedReader struct {
	r     io.Reader
	limit int64
	n     int64
}

// Read implements Reader interface.
func (lr *limitedReader) Read(p []byte) (n int, err error) {
	if lr.n == 0 {
		return 0, &LimitReachedError{
			Limit: lr.limit,
		}
	}

	if int64(len(p)) > lr.n {
		p = p[0:lr.n]
	}

	n, err = lr.r.Read(p)
	lr.n -= int64(n)

	return n, err
}

// LimitReader wraps Reader to make it's Reader stop with ErrLimitReached after
// n bytes read.
func LimitReader(r io.Reader, n int64) (limited io.Reader, err error) {
	if n < 0 {
		return nil, fmt.Errorf("aghio: invalid n in LimitReader: %d", n)
	}

	return &limitedReader{
		r:     r,
		limit: n,
		n:     n,
	}, nil
}