repos / pico

pico services mono repo
git clone https://github.com/picosh/pico.git

pico / pkg / send / utils
Antonio Mika  ·  2025-03-12

limit_reader.go

 1package utils
 2
 3import (
 4	"io"
 5	"sync"
 6)
 7
 8func NewLimitReader(r io.Reader, limit int) io.Reader {
 9	return &LimitReader{
10		r:    r,
11		left: limit,
12	}
13}
14
15type LimitReader struct {
16	r io.Reader
17
18	lock sync.Mutex
19	left int
20}
21
22func (r *LimitReader) Read(b []byte) (int, error) {
23	r.lock.Lock()
24	defer r.lock.Unlock()
25
26	if r.left <= 0 {
27		return 0, io.EOF
28	}
29	if len(b) > r.left {
30		b = b[0:r.left]
31	}
32	n, err := r.r.Read(b)
33	r.left -= n
34	return n, err
35}