forked from Mxmilu666/frp
37 lines
758 B
Go
37 lines
758 B
Go
package proxy
|
|
|
|
import "io"
|
|
|
|
type countingReadWriteCloser struct {
|
|
io.ReadWriteCloser
|
|
onRead func(int64)
|
|
onWrite func(int64)
|
|
}
|
|
|
|
func wrapCountingReadWriteCloser(rwc io.ReadWriteCloser, onRead, onWrite func(int64)) io.ReadWriteCloser {
|
|
if onRead == nil && onWrite == nil {
|
|
return rwc
|
|
}
|
|
return &countingReadWriteCloser{
|
|
ReadWriteCloser: rwc,
|
|
onRead: onRead,
|
|
onWrite: onWrite,
|
|
}
|
|
}
|
|
|
|
func (c *countingReadWriteCloser) Read(p []byte) (n int, err error) {
|
|
n, err = c.ReadWriteCloser.Read(p)
|
|
if n > 0 && c.onRead != nil {
|
|
c.onRead(int64(n))
|
|
}
|
|
return
|
|
}
|
|
|
|
func (c *countingReadWriteCloser) Write(p []byte) (n int, err error) {
|
|
n, err = c.ReadWriteCloser.Write(p)
|
|
if n > 0 && c.onWrite != nil {
|
|
c.onWrite(int64(n))
|
|
}
|
|
return
|
|
}
|