-
Notifications
You must be signed in to change notification settings - Fork 47
/
limit_bytes.go
66 lines (52 loc) · 1.27 KB
/
limit_bytes.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
/*
Copyright (c) 2021 - Present. Blend Labs, Inc. All rights reserved
Use of this source code is governed by a MIT license that can be found in the LICENSE file.
*/
package sh
import (
"io"
"sync"
"github.com/blend/go-sdk/ex"
)
// Errors
const (
ErrMaxBytesWriterCapacityLimit ex.Class = "write failed; maximum capacity reached or would be exceede"
)
// LimitBytes returns a new max bytes writer.
func LimitBytes(maxBytes int, inner io.Writer) *MaxBytesWriter {
return &MaxBytesWriter{
max: maxBytes,
inner: inner,
}
}
// MaxBytesWriter returns a maximum bytes writer.
type MaxBytesWriter struct {
sync.Mutex
max int
count int
inner io.Writer
}
// Max returns the maximum bytes allowed.
func (mbw *MaxBytesWriter) Max() int {
return int(mbw.max)
}
// Count returns the current written bytes..
func (mbw *MaxBytesWriter) Count() int {
return int(mbw.count)
}
func (mbw *MaxBytesWriter) Write(contents []byte) (int, error) {
mbw.Lock()
defer mbw.Unlock()
if mbw.count >= mbw.max {
return 0, ex.New(ErrMaxBytesWriterCapacityLimit)
}
if len(contents)+mbw.count >= mbw.max {
return 0, ex.New(ErrMaxBytesWriterCapacityLimit)
}
written, err := mbw.inner.Write(contents)
mbw.count += written
if err != nil {
return written, ex.New(err)
}
return written, nil
}