43 lines
605 B
Go
43 lines
605 B
Go
package limit
|
|
|
|
import (
|
|
"io"
|
|
|
|
frpNet "github.com/fatedier/frp/utils/net"
|
|
)
|
|
|
|
const (
|
|
B uint64 = 1 << (10 * (iota))
|
|
KB
|
|
MB
|
|
GB
|
|
TB
|
|
PB
|
|
EB
|
|
)
|
|
|
|
const burstLimit = 1024 * 1024 * 1024
|
|
|
|
type LimitConn struct {
|
|
frpNet.Conn
|
|
|
|
lr io.Reader
|
|
lw io.Writer
|
|
}
|
|
|
|
func NewLimitConn(maxread, maxwrite uint64, c frpNet.Conn) LimitConn {
|
|
return LimitConn{
|
|
lr: NewReaderWithLimit(c, maxread*KB),
|
|
lw: NewWriterWithLimit(c, maxwrite*KB),
|
|
Conn: c,
|
|
}
|
|
}
|
|
|
|
func (c LimitConn) Read(p []byte) (n int, err error) {
|
|
return c.lr.Read(p)
|
|
}
|
|
|
|
func (c LimitConn) Write(p []byte) (n int, err error) {
|
|
return c.lw.Write(p)
|
|
}
|