V2bX/common/rate/writer.go
2023-07-29 18:47:47 +08:00

29 lines
549 B
Go

package rate
import (
"github.com/juju/ratelimit"
"github.com/xtls/xray-core/common"
"github.com/xtls/xray-core/common/buf"
)
type Writer struct {
writer buf.Writer
limiter *ratelimit.Bucket
}
func NewRateLimitWriter(writer buf.Writer, limiter *ratelimit.Bucket) buf.Writer {
return &Writer{
writer: writer,
limiter: limiter,
}
}
func (w *Writer) Close() error {
return common.Close(w.writer)
}
func (w *Writer) WriteMultiBuffer(mb buf.MultiBuffer) error {
w.limiter.Wait(int64(mb.Len()))
return w.writer.WriteMultiBuffer(mb)
}