feat: Add rate limiting to fallback handling via token bucket (#12)

Co-authored-by: RPRX <63339210+RPRX@users.noreply.github.com>

README.en.md

@@ -45,7 +45,19 @@ TODO List: TODO
                     "shortIds": [ // Required, the acceptable shortId list, which can be used to distinguish different clients
                         "", // If there is this item, the client shortId can be empty
                         "0123456789abcdef" // 0 to f, the length is a multiple of 2, the maximum length is 16
-                    ]
+                    ],
+                    // These two limitations below are optional, for rate limiting fallback connections, bytesPerSec's default is 0, which means disabled
+                    // It's a detectable pattern, not recommended to be enabled, RANDOMIZE these parameters if you're a web-panel/one-click-script developer
+                    "limitFallbackUpload": {
+                        "afterBytes": 0, // Start throttling after (bytes)
+                        "bytesPerSec": 0, // Base speed (bytes/s)
+                        "burstBytesPerSec": 0 // Burst capacity (bytes/s), works only when it is larger than bytesPerSec
+                    },
+                    "limitFallbackDownload": {
+                        "afterBytes": 0, // Start throttling after (bytes)
+                        "bytesPerSec": 0, // Base speed (bytes/s)
+                        "burstBytesPerSec": 0 // Burst capacity (bytes/s), works only when it is larger than bytesPerSec
+                    }
                 }
             }
         }

README.md

@@ -45,7 +45,19 @@ TODO List: TODO
                     "shortIds": [ // 必填，客户端可用的 shortId 列表，可用于区分不同的客户端
                         "", // 若有此项，客户端 shortId 可为空
                         "0123456789abcdef" // 0 到 f，长度为 2 的倍数，长度上限为 16
-                    ]
+                    ],
+                    // 下列两个 limit 为选填，可对未通过验证的回落连接限速，bytesPerSec 默认为 0 即不启用
+                    // 回落限速是一种特征，不建议启用，如果您是面板/一键脚本开发者，务必让这些参数随机化
+                    "limitFallbackUpload": {
+                        "afterBytes": 0, // 传输指定字节后开始限速
+                        "bytesPerSec": 0, // 基准速率（字节/秒）
+                        "burstBytesPerSec": 0 // 突发速率（字节/秒），大于 bytesPerSec 时生效
+                    },
+                    "limitFallbackDownload": {
+                        "afterBytes": 0, // 传输指定字节后开始限速
+                        "bytesPerSec": 0, // 基准速率（字节/秒）
+                        "burstBytesPerSec": 0 // 突发速率（字节/秒），大于 bytesPerSec 时生效
+                    }
                 }
             }
         }

common.go

@@ -537,6 +537,12 @@ const (
 	RenegotiateFreelyAsClient
 )
 
+type LimitFallback struct {
+	AfterBytes       uint64
+	BytesPerSec      uint64
+	BurstBytesPerSec uint64
+}
+
 // A Config structure is used to configure a TLS client or server.
 // After one has been passed to a TLS function it must not be
 // modified. A Config may be reused; the tls package will also not
@@ -556,6 +562,9 @@ type Config struct {
 	MaxTimeDiff  time.Duration
 	ShortIds     map[[8]byte]bool
 
+	LimitFallbackUpload   LimitFallback
+	LimitFallbackDownload LimitFallback
+
 	// Rand provides the source of entropy for nonces and RSA blinding.
 	// If Rand is nil, TLS uses the cryptographic random reader in package
 	// crypto/rand.
@@ -913,7 +922,6 @@ type EncryptedClientHelloKey struct {
 	SendAsRetry bool
 }
 
-
 const (
 	// ticketKeyLifetime is how long a ticket key remains valid and can be used to
 	// resume a client connection.
@@ -971,6 +979,8 @@ func (c *Config) Clone() *Config {
 		MaxClientVer:                        c.MaxClientVer,
 		MaxTimeDiff:                         c.MaxTimeDiff,
 		ShortIds:                            c.ShortIds,
+		LimitFallbackUpload:                 c.LimitFallbackUpload,
+		LimitFallbackDownload:               c.LimitFallbackDownload,
 		Rand:                                c.Rand,
 		Time:                                c.Time,
 		Certificates:                        c.Certificates,
@@ -1793,4 +1803,4 @@ func fipsAllowChain(chain []*x509.Certificate) bool {
 	}
 
 	return true
-}
+}

go.mod

@@ -3,6 +3,7 @@ module github.com/xtls/reality
 go 1.24
 
 require (
+	github.com/juju/ratelimit v1.0.2
 	github.com/pires/go-proxyproto v0.8.1
 	github.com/refraction-networking/utls v1.7.3
 	golang.org/x/crypto v0.39.0

go.sum

@@ -1,3 +1,5 @@
+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/pires/go-proxyproto v0.8.1 h1:9KEixbdJfhrbtjpz/ZwCdWDD2Xem0NZ38qMYaASJgp0=
 github.com/pires/go-proxyproto v0.8.1/go.mod h1:ZKAAyp3cgy5Y5Mo4n9AlScrkCZwUy0g3Jf+slqQVcuU=
 github.com/refraction-networking/utls v1.7.3 h1:L0WRhHY7Oq1T0zkdzVZMR6zWZv+sXbHB9zcuvsAEqCo=

tls.go

@@ -50,6 +50,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/juju/ratelimit"
 	"github.com/pires/go-proxyproto"
 	"golang.org/x/crypto/curve25519"
 	"golang.org/x/crypto/hkdf"
@@ -100,6 +101,41 @@ func (c *MirrorConn) SetWriteDeadline(t time.Time) error {
 	return nil
 }
 
+type RatelimitedConn struct {
+	net.Conn
+	After  int64
+	Bucket *ratelimit.Bucket
+}
+
+func (c *RatelimitedConn) Read(b []byte) (int, error) {
+	n, err := c.Conn.Read(b)
+	if n != 0 {
+		if c.After > 0 {
+			c.After -= int64(n)
+		} else {
+			c.Bucket.Wait(int64(n))
+		}
+	}
+	return n, err
+}
+
+func NewRatelimitedConn(conn net.Conn, limit *LimitFallback) net.Conn {
+	if limit.BytesPerSec == 0 {
+		return conn
+	}
+
+	burstBytesPerSec := limit.BurstBytesPerSec
+	if burstBytesPerSec < limit.BytesPerSec {
+		burstBytesPerSec = limit.BytesPerSec
+	}
+
+	return &RatelimitedConn{
+		Conn:   conn,
+		After:  int64(limit.AfterBytes),
+		Bucket: ratelimit.NewBucketWithRate(float64(limit.BytesPerSec), int64(burstBytesPerSec)),
+	}
+}
+
 var (
 	size  = 8192
 	empty = make([]byte, size)
@@ -228,7 +264,7 @@ func Server(ctx context.Context, conn net.Conn, config *Config) (*Conn, error) {
 			if config.Show && hs.clientHello != nil {
 				fmt.Printf("REALITY remoteAddr: %v\tforwarded SNI: %v\n", remoteAddr, hs.clientHello.serverName)
 			}
-			io.Copy(target, underlying)
+			io.Copy(target, NewRatelimitedConn(underlying, &config.LimitFallbackUpload))
 		}
 		waitGroup.Done()
 	}()
@@ -359,12 +395,12 @@ func Server(ctx context.Context, conn net.Conn, config *Config) (*Conn, error) {
 			if hs.c.conn == conn { // if we processed the Client Hello successfully but the target did not
 				waitGroup.Add(1)
 				go func() {
-					io.Copy(target, underlying)
+					io.Copy(target, NewRatelimitedConn(underlying, &config.LimitFallbackUpload))
 					waitGroup.Done()
 				}()
 			}
 			conn.Write(s2cSaved)
-			io.Copy(underlying, target)
+			io.Copy(underlying, NewRatelimitedConn(target, &config.LimitFallbackDownload))
 			// Here is bidirectional direct forwarding:
 			// client ---underlying--- server ---target--- dest
 			// Call `underlying.CloseWrite()` once `io.Copy()` returned