blob: a93464bdfcb1815b75064837ad53c0a7d57ee037 [file] [log] [blame]
Ivan Pilat8ecbed32015-01-13 21:30:05 -08001// LimitedWriter is an io.Writer wrapper that limits the total number of bytes
2// written to the underlying writer.
3//
4// All attempted writes count against the limit, regardless of whether they
5// succeed.
6// Not thread-safe.
7
8package lib
9
10import (
11 "errors"
12 "io"
13 "net/http"
14 "sync"
15)
16
17var ErrWriteLimitExceeded = errors.New("LimitedWriter: write limit exceeded")
18
19// Initialize using NewLimitedWriter.
20type LimitedWriter struct {
21 io.Writer
22 maxLen int
23 maxLenExceededCb func()
24 lenWritten int
25}
26
27func NewLimitedWriter(writer io.Writer, maxLen int, maxLenExceededCb func()) *LimitedWriter {
28 return &LimitedWriter{
29 Writer: writer,
30 maxLen: maxLen,
31 maxLenExceededCb: maxLenExceededCb,
32 }
33}
34
35func (t *LimitedWriter) Write(p []byte) (n int, err error) {
36 if t.lenWritten+len(p) > t.maxLen {
37 t.lenWritten = t.maxLen
38 if t.maxLenExceededCb != nil {
39 t.maxLenExceededCb()
40 }
41 return 0, ErrWriteLimitExceeded
42 }
43 if len(p) == 0 {
44 return 0, nil
45 }
46 t.lenWritten += len(p)
47 return t.Writer.Write(p)
48}
49
50var _ http.Flusher = (*LimitedWriter)(nil)
51
52func (t *LimitedWriter) Flush() {
53 if f, ok := t.Writer.(http.Flusher); ok {
54 f.Flush()
55 }
56}
57
58// Wraps a function to prevent it from executing more than once.
59func DoOnce(f func()) func() {
60 var once sync.Once
61 return func() {
62 once.Do(f)
63 }
64}