feat: support rate limit

examples/speed_limit.json

@@ -0,0 +1,14 @@
+{
+    "relay_configs": [
+        {
+            "listen": "127.0.0.1:1234",
+            "listen_type": "raw",
+            "transport_type": "raw",
+            "label": "iperf3",
+            "tcp_remotes": [
+                "0.0.0.0:5201"
+            ],
+            "max_read_rate_kbps": 10000
+        }
+    ]
+}

go.mod

@@ -65,6 +65,7 @@ require (
 	github.com/illumos/go-kstat v0.0.0-20210513183136-173c9b0a9973 // indirect
 	github.com/josharian/native v1.1.0 // indirect
 	github.com/jsimonetti/rtnetlink v1.4.2 // indirect
+	github.com/juju/ratelimit v1.0.2 // indirect
 	github.com/klauspost/compress v1.17.8 // indirect
 	github.com/klauspost/cpuid/v2 v2.2.7 // indirect
 	github.com/labstack/gommon v0.4.2 // indirect

go.sum

@@ -140,6 +140,8 @@ github.com/jsimonetti/rtnetlink v1.4.2 h1:Df9w9TZ3npHTyDn0Ev9e1uzmN2odmXd0QX+J5G
 github.com/jsimonetti/rtnetlink v1.4.2/go.mod h1:92s6LJdE+1iOrw+F2/RO7LYI2Qd8pPpFNNUYW06gcoM=
 github.com/json-iterator/go v1.1.6/go.mod h1:+SdeFBvtyEkXs7REEP0seUULqWtbJapLOCVDaaPEHmU=
 github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1:6v2b51hI/fHJwM22ozAgKL4VKDeJcHhJFhtBdhmNjmU=
+github.com/juju/ratelimit v1.0.2 h1:sRxmtRiajbvrcLQT7S+JbqU0ntsb9W2yhSdNN8tWfaI=
+github.com/juju/ratelimit v1.0.2/go.mod h1:qapgC/Gy+xNh9UxzV13HGGl/6UXNN+ct+vwSgWNm/qk=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
 github.com/klauspost/compress v1.17.8 h1:YcnTYrq7MikUT7k0Yb5eceMmALQPYBW/Xltxn0NAMnU=
 github.com/klauspost/compress v1.17.8/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=

internal/conn/limit_reader.go

@@ -0,0 +1,29 @@
+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)
+}

internal/relay/conf/cfg.go

@@ -33,6 +33,7 @@ type Config struct {
 
 	MaxConnection    int      `json:"max_connection,omitempty"`
 	BlockedProtocols []string `json:"blocked_protocols,omitempty"`
+	MaxReadRateKbps  int64    `json:"max_read_rate_kbps,omitempty"`
 
 	WSConfig *WSConfig `json:"ws_config,omitempty"`
 }

internal/transporter/base.go

@@ -94,6 +94,11 @@ func (b *baseTransporter) RelayTCPConn(c net.Conn, handshakeF TCPHandShakeF) err
 	}
 	defer rc.Close()
 
+	// rate limit
+	if b.cfg.MaxReadRateKbps > 0 {
+		c = conn.NewRateLimitedConn(c, b.cfg.MaxReadRateKbps)
+	}
+
 	b.l.Infof("RelayTCPConn from %s to %s", c.LocalAddr(), remote.Address)
 	relayConn := conn.NewRelayConn(
 		b.cfg.Label, c, rc, conn.WithHandshakeDuration(clonedRemote.HandShakeDuration))

pkg/bytes/utils.go

@@ -14,3 +14,13 @@ func PrettyByteSize(bf float64) string {
 	}
 	return fmt.Sprintf(" %.1fYiB ", bf)
 }
+
+func PrettyBitRate(bps float64) string {
+	for _, unit := range []string{"bps", "Kbps", "Mbps", "Gbps", "Tbps", "Pbps", "Ebps", "Zbps"} {
+		if math.Abs(bps) < 1000.0 {
+			return fmt.Sprintf(" %3.1f %s ", bps, unit)
+		}
+		bps /= 1000.0
+	}
+	return fmt.Sprintf(" %.1f Ybps ", bps)
+}