mirror of
https://github.com/bolucat/Archive.git
synced 2025-10-05 08:08:03 +08:00
30 lines
622 B
Go
30 lines
622 B
Go
package conn
|
|
|
|
import (
|
|
"io"
|
|
"net"
|
|
|
|
"github.com/juju/ratelimit"
|
|
)
|
|
|
|
type RateLimitedConn struct {
|
|
net.Conn
|
|
bucket *ratelimit.Bucket
|
|
reader io.Reader
|
|
}
|
|
|
|
func NewRateLimitedConn(conn net.Conn, kbps int64) *RateLimitedConn {
|
|
bps := float64(kbps) * 1000 // Convert kbps to bps (1 kbps = 1000 bps)
|
|
rateBytesPerSec := bps / 8 // 1KB = 1024B, 1B = 8b
|
|
bucket := ratelimit.NewBucketWithRate(rateBytesPerSec, int64(rateBytesPerSec))
|
|
return &RateLimitedConn{
|
|
Conn: conn,
|
|
bucket: bucket,
|
|
reader: ratelimit.Reader(conn, bucket),
|
|
}
|
|
}
|
|
|
|
func (r *RateLimitedConn) Read(p []byte) (int, error) {
|
|
return r.reader.Read(p)
|
|
}
|