purgatory/Xray-core
1
0
Fork 0
mirror of https://github.com/XTLS/Xray-core.git synced 2025-04-06 07:13:33 +03:00
Code Issues Projects Releases Packages Wiki Activity

VLESS: Add rate limiting to fallback handling via token bucket

Browse source
This commit is contained in:
Meo597 2025-03-25 18:17:15 +08:00
parent 4afe2d0cff
commit 8d36fd7715
4 changed files with 41 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
common/buf/reader.go
View file

@ -3,6 +3,8 @@ package buf
import (
"io"
"github.com/juju/ratelimit"
"github.com/xtls/xray-core/common"
"github.com/xtls/xray-core/common/errors"
)
@ -172,3 +174,27 @@ func (r *PacketReader) ReadMultiBuffer() (MultiBuffer, error) {
}
return MultiBuffer{b}, nil
}
// Reader returns a reader that is rate limited by
// the given token bucket. Each token in the bucket
// represents one byte.
type RateLimitedReader struct {
Reader Reader
Bucket *ratelimit.Bucket
}
// ReadMultiBuffer implements Reader.
func (r *RateLimitedReader) ReadMultiBuffer() (MultiBuffer, error) {
b, err := r.Reader.ReadMultiBuffer()
if err != nil {
return nil, err
}
var total int64
for _, buf := range b {
total += int64(buf.Len())
}
r.Bucket.Wait(total)
return b, nil
}

1
go.mod
View file

@ -41,6 +41,7 @@ require (
github.com/go-task/slim-sprig/v3 v3.0.0 // indirect
github.com/google/btree v1.1.2 // indirect
github.com/google/pprof v0.0.0-20240528025155-186aa0362fba // indirect
github.com/juju/ratelimit v1.0.2
github.com/klauspost/compress v1.17.8 // indirect
github.com/klauspost/cpuid/v2 v2.2.7 // indirect
github.com/onsi/ginkgo/v2 v2.19.0 // indirect

2
go.sum
View file

@ -34,6 +34,8 @@ github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aN
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
github.com/h12w/go-socks5 v0.0.0-20200522160539-76189e178364 h1:5XxdakFhqd9dnXoAZy1Mb2R/DZ6D1e+0bGC/JhucGYI=
github.com/h12w/go-socks5 v0.0.0-20200522160539-76189e178364/go.mod h1:eDJQioIyy4Yn3MVivT7rv/39gAJTrA7lgmYr8EW950c=
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/klauspost/compress v1.17.8 h1:YcnTYrq7MikUT7k0Yb5eceMmALQPYBW/Xltxn0NAMnU=
github.com/klauspost/compress v1.17.8/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
github.com/klauspost/cpuid/v2 v2.2.7 h1:ZWSB3igEs+d0qvnxR/ZBzXVmxkgt8DdzP6m9pfuVLDM=

14
proxy/vless/inbound/inbound.go
View file

@ -11,6 +11,8 @@ import (
"time"
"unsafe"
"github.com/juju/ratelimit"
"github.com/xtls/xray-core/common"
"github.com/xtls/xray-core/common/buf"
"github.com/xtls/xray-core/common/errors"
@ -405,7 +407,11 @@ func (h *Handler) Process(ctx context.Context, network net.Network, connection s
return errors.New("failed to set PROXY protocol v", fb.Xver).Base(err).AtWarning()
}
}
if err := buf.Copy(reader, serverWriter, buf.UpdateActivity(timer)); err != nil {
rlReader := &buf.RateLimitedReader{
Reader: reader,
Bucket: ratelimit.NewBucketWithRate(128*1024, 512*1024),
}
if err := buf.Copy(rlReader, serverWriter, buf.UpdateActivity(timer)); err != nil {
return errors.New("failed to fallback request payload").Base(err).AtInfo()
}
return nil
@ -414,8 +420,12 @@ func (h *Handler) Process(ctx context.Context, network net.Network, connection s
writer := buf.NewWriter(connection)
getResponse := func() error {
rlServerReader := &buf.RateLimitedReader{
Reader: serverReader,
Bucket: ratelimit.NewBucketWithRate(128*1024, 512*1024),
}
defer timer.SetTimeout(sessionPolicy.Timeouts.UplinkOnly)
if err := buf.Copy(serverReader, writer, buf.UpdateActivity(timer)); err != nil {
if err := buf.Copy(rlServerReader, writer, buf.UpdateActivity(timer)); err != nil {
return errors.New("failed to deliver response payload").Base(err).AtInfo()
}
return nil