89 lines
1.5 KiB
Go
89 lines
1.5 KiB
Go
package httpserver
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
)
|
|
|
|
// LimitBuffer buf with limit
|
|
type LimitBuffer interface {
|
|
io.Writer
|
|
String() string
|
|
Remain() int
|
|
}
|
|
|
|
// limitBuffer buf with limit
|
|
type limitBuffer struct {
|
|
buf *bytes.Buffer
|
|
limit int
|
|
remain int
|
|
length int
|
|
}
|
|
|
|
// NewLimitBuffer ..
|
|
func NewLimitBuffer(limit int) LimitBuffer {
|
|
buf := bytes.NewBuffer(make([]byte, 0, limit))
|
|
return &limitBuffer{buf: buf, limit: limit, remain: limit}
|
|
}
|
|
|
|
func (b *limitBuffer) Write(p []byte) (n int, err error) {
|
|
defer func() {
|
|
b.length += n
|
|
}()
|
|
|
|
// discord
|
|
if b.remain <= 0 {
|
|
return len(p), nil
|
|
}
|
|
|
|
// write remain
|
|
if len(p) > b.remain {
|
|
n, err = b.buf.Write(p[:b.remain])
|
|
b.remain -= n
|
|
|
|
// alway return all writed length
|
|
n = len(p)
|
|
return
|
|
}
|
|
|
|
// write all
|
|
n, err = b.buf.Write(p)
|
|
b.remain -= n
|
|
return
|
|
}
|
|
|
|
// String ..
|
|
func (b *limitBuffer) String() string {
|
|
if b.length > b.limit {
|
|
return b.buf.String() + fmt.Sprintf("...(total %dB)", b.length)
|
|
}
|
|
|
|
return b.buf.String()
|
|
}
|
|
|
|
// Remain 剩余多少个字节没有写满
|
|
func (b *limitBuffer) Remain() int {
|
|
return b.remain
|
|
}
|
|
|
|
type teeReadCloser struct {
|
|
tee io.Reader
|
|
r io.Closer
|
|
}
|
|
|
|
// Read implement io.Read interface
|
|
func (t *teeReadCloser) Read(p []byte) (n int, err error) {
|
|
return t.tee.Read(p)
|
|
}
|
|
|
|
// Close implement io.Close interface
|
|
func (t *teeReadCloser) Close() error {
|
|
return t.r.Close()
|
|
}
|
|
|
|
// TeeReadCloser TeeReader with limit
|
|
func TeeReadCloser(r io.ReadCloser, w io.Writer) io.ReadCloser {
|
|
return &teeReadCloser{r: r, tee: io.TeeReader(r, w)}
|
|
}
|