github.com/kelleygo/clashcore@v1.0.2/hub/updater/limitedreader.go (about)

     1  package updater
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  
     7  	"golang.org/x/exp/constraints"
     8  )
     9  
    10  // LimitReachedError records the limit and the operation that caused it.
    11  type LimitReachedError struct {
    12  	Limit int64
    13  }
    14  
    15  // Error implements the [error] interface for *LimitReachedError.
    16  //
    17  // TODO(a.garipov): Think about error string format.
    18  func (lre *LimitReachedError) Error() string {
    19  	return fmt.Sprintf("attempted to read more than %d bytes", lre.Limit)
    20  }
    21  
    22  // limitedReader is a wrapper for [io.Reader] limiting the input and dealing
    23  // with errors package.
    24  type limitedReader struct {
    25  	r     io.Reader
    26  	limit int64
    27  	n     int64
    28  }
    29  
    30  // Read implements the [io.Reader] interface.
    31  func (lr *limitedReader) Read(p []byte) (n int, err error) {
    32  	if lr.n == 0 {
    33  		return 0, &LimitReachedError{
    34  			Limit: lr.limit,
    35  		}
    36  	}
    37  
    38  	p = p[:Min(lr.n, int64(len(p)))]
    39  
    40  	n, err = lr.r.Read(p)
    41  	lr.n -= int64(n)
    42  
    43  	return n, err
    44  }
    45  
    46  // LimitReader wraps Reader to make it's Reader stop with ErrLimitReached after
    47  // n bytes read.
    48  func LimitReader(r io.Reader, n int64) (limited io.Reader, err error) {
    49  	if n < 0 {
    50  		return nil, &updateError{Message: "limit must be non-negative"}
    51  	}
    52  
    53  	return &limitedReader{
    54  		r:     r,
    55  		limit: n,
    56  		n:     n,
    57  	}, nil
    58  }
    59  
    60  // Min returns the smaller of x or y.
    61  func Min[T constraints.Integer | ~string](x, y T) (res T) {
    62  	if x < y {
    63  		return x
    64  	}
    65  
    66  	return y
    67  }