diff --git a/adapter/outbound/base.go b/adapter/outbound/base.go index 01e0079846..47c0d52479 100644 --- a/adapter/outbound/base.go +++ b/adapter/outbound/base.go @@ -142,6 +142,11 @@ func (b *Base) Unwrap(metadata *C.Metadata, touch bool) C.Proxy { return nil } +// Bandwidth implements C.ProxyAdapter. Standalone proxies have no bandwidth limit. +func (b *Base) Bandwidth() uint64 { + return 0 +} + // DialOptions return []dialer.Option from struct func (b *Base) DialOptions() (opts []dialer.Option) { if b.iface != "" { diff --git a/adapter/outboundgroup/fallback.go b/adapter/outboundgroup/fallback.go index 98aaf20a9d..85cf9040c3 100644 --- a/adapter/outboundgroup/fallback.go +++ b/adapter/outboundgroup/fallback.go @@ -48,6 +48,9 @@ func (f *Fallback) DialContext(ctx context.Context, metadata *C.Metadata) (C.Con }) } + if err == nil { + c = f.LimitConn(c) + } return c, err } @@ -57,6 +60,7 @@ func (f *Fallback) ListenPacketContext(ctx context.Context, metadata *C.Metadata pc, err := proxy.ListenPacketContext(ctx, metadata) if err == nil { pc.AppendToChains(f) + pc = f.LimitPacketConn(pc) } return pc, err @@ -159,7 +163,7 @@ func (f *Fallback) Proxies() []C.Proxy { return f.GetProxies(false) } -func NewFallback(option GroupCommonOption, fallbackOption FallbackOption, emptyFallback C.Proxy, providers []P.ProxyProvider) (*Fallback, error) { +func NewFallback(option GroupCommonOption, fallbackOption FallbackOption, emptyFallback C.Proxy, providers []P.ProxyProvider, bandwidth uint64) (*Fallback, error) { return &Fallback{ GroupBase: NewGroupBase(GroupBaseOption{ Name: option.Name, @@ -173,6 +177,7 @@ func NewFallback(option GroupCommonOption, fallbackOption FallbackOption, emptyF MaxFailedTimes: option.MaxFailedTimes, EmptyFallback: emptyFallback, Providers: providers, + Bandwidth: bandwidth, }), disableUDP: option.DisableUDP, testUrl: option.URL, diff --git a/adapter/outboundgroup/groupbase.go b/adapter/outboundgroup/groupbase.go index eed45e0d61..44808eb727 100644 --- a/adapter/outboundgroup/groupbase.go +++ b/adapter/outboundgroup/groupbase.go @@ -11,6 +11,7 @@ import ( "github.com/metacubex/mihomo/adapter/outbound" "github.com/metacubex/mihomo/common/atomic" "github.com/metacubex/mihomo/common/utils" + "github.com/metacubex/mihomo/component/ratelimit" C "github.com/metacubex/mihomo/constant" P "github.com/metacubex/mihomo/constant/provider" "github.com/metacubex/mihomo/log" @@ -34,6 +35,7 @@ type GroupBase struct { testTimeout int maxFailedTimes int emptyFallback C.Proxy + limiter *ratelimit.Limiter // shared group bandwidth limit, nil = unlimited // for GetProxies getProxiesMutex sync.Mutex @@ -53,6 +55,7 @@ type GroupBaseOption struct { MaxFailedTimes int EmptyFallback C.Proxy Providers []P.ProxyProvider + Bandwidth uint64 // bits per second, 0 = unlimited } func NewGroupBase(opt GroupBaseOption) *GroupBase { @@ -89,6 +92,7 @@ func NewGroupBase(opt GroupBaseOption) *GroupBase { testTimeout: opt.TestTimeout, maxFailedTimes: opt.MaxFailedTimes, emptyFallback: opt.EmptyFallback, + limiter: ratelimit.NewLimiter(opt.Bandwidth), } if gb.testTimeout == 0 { @@ -109,6 +113,20 @@ func (gb *GroupBase) Icon() string { return gb.icon } +func (gb *GroupBase) Bandwidth() uint64 { + return gb.limiter.Rate() +} + +// LimitConn applies the group's shared bandwidth limiter to conn. +func (gb *GroupBase) LimitConn(conn C.Conn) C.Conn { + return gb.limiter.WrapCConn(conn) +} + +// LimitPacketConn applies the group's shared bandwidth limiter to a packet conn. +func (gb *GroupBase) LimitPacketConn(pc C.PacketConn) C.PacketConn { + return gb.limiter.WrapPacketConn(pc) +} + func (gb *GroupBase) EmptyFallback() C.Proxy { return gb.emptyFallback } diff --git a/adapter/outboundgroup/loadbalance.go b/adapter/outboundgroup/loadbalance.go index 88d4674a03..d6cb6ff68a 100644 --- a/adapter/outboundgroup/loadbalance.go +++ b/adapter/outboundgroup/loadbalance.go @@ -101,6 +101,9 @@ func (lb *LoadBalance) DialContext(ctx context.Context, metadata *C.Metadata) (c }) } + if err == nil { + c = lb.LimitConn(c) + } return } @@ -109,6 +112,7 @@ func (lb *LoadBalance) ListenPacketContext(ctx context.Context, metadata *C.Meta defer func() { if err == nil { pc.AppendToChains(lb) + pc = lb.LimitPacketConn(pc) } }() @@ -247,7 +251,7 @@ func (lb *LoadBalance) Now() string { return "" } -func NewLoadBalance(option GroupCommonOption, loadBalanceOption LoadBalanceOption, emptyFallback C.Proxy, providers []P.ProxyProvider) (lb *LoadBalance, err error) { +func NewLoadBalance(option GroupCommonOption, loadBalanceOption LoadBalanceOption, emptyFallback C.Proxy, providers []P.ProxyProvider, bandwidth uint64) (lb *LoadBalance, err error) { var strategyFn strategyFn switch loadBalanceOption.Strategy { case "", "consistent-hashing": @@ -272,6 +276,7 @@ func NewLoadBalance(option GroupCommonOption, loadBalanceOption LoadBalanceOptio MaxFailedTimes: option.MaxFailedTimes, EmptyFallback: emptyFallback, Providers: providers, + Bandwidth: bandwidth, }), strategyFn: strategyFn, disableUDP: option.DisableUDP, diff --git a/adapter/outboundgroup/parser.go b/adapter/outboundgroup/parser.go index e97f8da04b..27182ba385 100644 --- a/adapter/outboundgroup/parser.go +++ b/adapter/outboundgroup/parser.go @@ -10,6 +10,7 @@ import ( "github.com/metacubex/mihomo/adapter/provider" "github.com/metacubex/mihomo/common/structure" "github.com/metacubex/mihomo/common/utils" + "github.com/metacubex/mihomo/component/ratelimit" C "github.com/metacubex/mihomo/constant" P "github.com/metacubex/mihomo/constant/provider" "github.com/metacubex/mihomo/log" @@ -43,6 +44,7 @@ type GroupCommonOption struct { IncludeAllProviders bool `group:"include-all-providers,omitempty"` Hidden bool `group:"hidden,omitempty"` Icon string `group:"icon,omitempty"` + Bandwidth string `group:"bandwidth,omitempty"` } func ParseProxyGroup(config map[string]any, proxyMap map[string]C.Proxy, providersMap map[string]P.ProxyProvider, AllProxies []string, AllProviders []string) (ProxyGroup, error) { @@ -55,6 +57,8 @@ func ParseProxyGroup(config map[string]any, proxyMap map[string]C.Proxy, provide return nil, errFormat } + bandwidth := ratelimit.ParseBandwidth(groupOption.Bandwidth) + if groupOption.Type == "" || groupOption.Name == "" { return nil, errFormat } @@ -191,28 +195,28 @@ func ParseProxyGroup(config map[string]any, proxyMap map[string]C.Proxy, provide if err != nil { return nil, err } - return NewURLTest(groupOption, opt, emptyFallback, providers) + return NewURLTest(groupOption, opt, emptyFallback, providers, bandwidth) case "select": opt := SelectorOption{} err = decoder.Decode(config, &opt) if err != nil { return nil, err } - return NewSelector(groupOption, opt, emptyFallback, providers) + return NewSelector(groupOption, opt, emptyFallback, providers, bandwidth) case "fallback": opt := FallbackOption{} err = decoder.Decode(config, &opt) if err != nil { return nil, err } - return NewFallback(groupOption, opt, emptyFallback, providers) + return NewFallback(groupOption, opt, emptyFallback, providers, bandwidth) case "load-balance": opt := LoadBalanceOption{} err = decoder.Decode(config, &opt) if err != nil { return nil, err } - return NewLoadBalance(groupOption, opt, emptyFallback, providers) + return NewLoadBalance(groupOption, opt, emptyFallback, providers, bandwidth) case "relay": return nil, fmt.Errorf("%w: The group [%s] with relay type was removed, please using dialer-proxy instead", errType, groupName) default: diff --git a/adapter/outboundgroup/selector.go b/adapter/outboundgroup/selector.go index 2fea75032a..9ee598dfb6 100644 --- a/adapter/outboundgroup/selector.go +++ b/adapter/outboundgroup/selector.go @@ -25,6 +25,7 @@ func (s *Selector) DialContext(ctx context.Context, metadata *C.Metadata) (C.Con c, err := s.selectedProxy(true).DialContext(ctx, metadata) if err == nil { c.AppendToChains(s) + c = s.LimitConn(c) } return c, err } @@ -34,6 +35,7 @@ func (s *Selector) ListenPacketContext(ctx context.Context, metadata *C.Metadata pc, err := s.selectedProxy(true).ListenPacketContext(ctx, metadata) if err == nil { pc.AppendToChains(s) + pc = s.LimitPacketConn(pc) } return pc, err } @@ -119,7 +121,7 @@ func (s *Selector) Proxies() []C.Proxy { return s.GetProxies(false) } -func NewSelector(option GroupCommonOption, selectorOption SelectorOption, emptyFallback C.Proxy, providers []P.ProxyProvider) (*Selector, error) { +func NewSelector(option GroupCommonOption, selectorOption SelectorOption, emptyFallback C.Proxy, providers []P.ProxyProvider, bandwidth uint64) (*Selector, error) { return &Selector{ GroupBase: NewGroupBase(GroupBaseOption{ Name: option.Name, @@ -133,6 +135,7 @@ func NewSelector(option GroupCommonOption, selectorOption SelectorOption, emptyF MaxFailedTimes: option.MaxFailedTimes, EmptyFallback: emptyFallback, Providers: providers, + Bandwidth: bandwidth, }), selected: selectorOption.DefaultSelected, disableUDP: option.DisableUDP, diff --git a/adapter/outboundgroup/urltest.go b/adapter/outboundgroup/urltest.go index 64e4def0de..fc8b51b8f1 100644 --- a/adapter/outboundgroup/urltest.go +++ b/adapter/outboundgroup/urltest.go @@ -73,6 +73,9 @@ func (u *URLTest) DialContext(ctx context.Context, metadata *C.Metadata) (c C.Co }) } + if err == nil { + c = u.LimitConn(c) + } return c, err } @@ -82,6 +85,7 @@ func (u *URLTest) ListenPacketContext(ctx context.Context, metadata *C.Metadata) pc, err := proxy.ListenPacketContext(ctx, metadata) if err == nil { pc.AppendToChains(u) + pc = u.LimitPacketConn(pc) } else { u.onDialFailed(proxy.Type(), err, u.healthCheck) } @@ -192,7 +196,7 @@ func (u *URLTest) URLTest(ctx context.Context, url string, expectedStatus utils. return u.GroupBase.URLTest(ctx, u.testUrl, expectedStatus) } -func NewURLTest(option GroupCommonOption, urlTestOption URLTestOption, emptyFallback C.Proxy, providers []P.ProxyProvider) (*URLTest, error) { +func NewURLTest(option GroupCommonOption, urlTestOption URLTestOption, emptyFallback C.Proxy, providers []P.ProxyProvider, bandwidth uint64) (*URLTest, error) { if emptyFallback == nil { return nil, errors.New("empty fallback proxy not exist") } @@ -209,6 +213,7 @@ func NewURLTest(option GroupCommonOption, urlTestOption URLTestOption, emptyFall MaxFailedTimes: option.MaxFailedTimes, EmptyFallback: emptyFallback, Providers: providers, + Bandwidth: bandwidth, }), fastSingle: singledo.NewSingle[C.Proxy](time.Second * 10), disableUDP: option.DisableUDP, diff --git a/component/ratelimit/ratelimit.go b/component/ratelimit/ratelimit.go new file mode 100644 index 0000000000..7475c9e0f2 --- /dev/null +++ b/component/ratelimit/ratelimit.go @@ -0,0 +1,417 @@ +package ratelimit + +import ( + "context" + "io" + "net" + "strconv" + "strings" + "sync" + "time" + "unicode" + + "github.com/metacubex/mihomo/common/buf" + C "github.com/metacubex/mihomo/constant" +) + +const ( + bitsPerByte = 8 + rateLimitCycle = 10 * time.Millisecond + maxRateLimitBurstBytes = 64 * 1024 +) + +// ParseBandwidth converts human-readable bandwidth to bits per second. +// Accepted: "500k", "1m", "2g", "500kbps", "1Mbps", bare number as Mbps. +// Empty or "0" → 0 (unlimited). +func ParseBandwidth(s string) uint64 { + s = strings.TrimSpace(s) + if s == "" || s == "0" { + return 0 + } + + // Bare integer → Mbps (same convention as utils.StringToBps). + if v, err := strconv.ParseUint(s, 10, 64); err == nil { + return v * 1_000_000 + } + + s = strings.ToLower(strings.ReplaceAll(s, " ", "")) + s = strings.TrimSuffix(s, "ps") // kbps/mbps → kb/mb + s = strings.TrimSuffix(s, "b") // kb/mb → k/m (or bare bit unit) + + var numStr strings.Builder + unit := "" + for i, r := range s { + if unicode.IsDigit(r) { + numStr.WriteRune(r) + continue + } + unit = s[i:] + break + } + if numStr.Len() == 0 { + return 0 + } + v, err := strconv.ParseUint(numStr.String(), 10, 64) + if err != nil || v == 0 { + return 0 + } + + switch unit { + case "", "b": + return v + case "k": + return v * 1_000 + case "m": + return v * 1_000_000 + case "g": + return v * 1_000_000_000 + case "t": + return v * 1_000_000_000_000 + default: + return 0 + } +} + +func burstFor(rateBps uint64) int { + burst := rateBps / bitsPerByte / uint64(time.Second/rateLimitCycle) + if burst == 0 { + return 1 + } + if burst > maxRateLimitBurstBytes { + return maxRateLimitBurstBytes + } + return int(burst) +} + +// bitRateLimiter is a token-bucket style rate limiter measured in bits per second. +type bitRateLimiter struct { + mu sync.Mutex + rateBps uint64 + next time.Time +} + +func (l *bitRateLimiter) WaitN(ctx context.Context, n int) error { + delay := l.reserveN(time.Now(), n) + if delay <= 0 { + return nil + } + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (l *bitRateLimiter) reserveN(now time.Time, n int) time.Duration { + interval := time.Duration(uint64(n) * bitsPerByte * uint64(time.Second) / l.rateBps) + + l.mu.Lock() + ready := l.next + if ready.Before(now) { + ready = now + } + delay := ready.Sub(now) + l.next = ready.Add(interval) + l.mu.Unlock() + return delay +} + +// Limiter is a shared bidirectional bandwidth limiter. +// Connections wrapped by the same Limiter share one rate budget. +type Limiter struct { + rateBps uint64 + burst int + readLimiter *bitRateLimiter + writeLimiter *bitRateLimiter +} + +// NewLimiter creates a shared limiter at rateBps bits per second. +// rateBps 0 returns nil (unlimited). +func NewLimiter(rateBps uint64) *Limiter { + if rateBps == 0 { + return nil + } + return &Limiter{ + rateBps: rateBps, + burst: burstFor(rateBps), + readLimiter: &bitRateLimiter{rateBps: rateBps}, + writeLimiter: &bitRateLimiter{rateBps: rateBps}, + } +} + +// Rate returns the configured limit in bits per second. Nil-safe. +func (l *Limiter) Rate() uint64 { + if l == nil { + return 0 + } + return l.rateBps +} + +// WrapConn wraps a net.Conn with this shared limiter. +func (l *Limiter) WrapConn(conn net.Conn) net.Conn { + if l == nil { + return conn + } + limitCtx, cancel := context.WithCancel(context.Background()) + return &RateLimitedConn{ + Conn: conn, + ctx: limitCtx, + cancel: cancel, + readLimiter: l.readLimiter, + writeLimiter: l.writeLimiter, + burst: l.burst, + } +} + +// WrapCConn wraps a C.Conn with this shared limiter. +func (l *Limiter) WrapCConn(conn C.Conn) C.Conn { + if l == nil { + return conn + } + limitCtx, cancel := context.WithCancel(context.Background()) + return &limitedCConn{ + Conn: conn, + ctx: limitCtx, + cancel: cancel, + readLimiter: l.readLimiter, + writeLimiter: l.writeLimiter, + burst: l.burst, + } +} + +// RateLimitedConn wraps a net.Conn with per-direction rate limiting. +type RateLimitedConn struct { + net.Conn + ctx context.Context + cancel context.CancelFunc + readLimiter *bitRateLimiter + writeLimiter *bitRateLimiter + burst int +} + +// NewRateLimitedConn wraps conn with a private per-connection limit. +// Prefer Limiter.WrapConn for group-wide shared limits. +func NewRateLimitedConn(conn net.Conn, rateBps uint64) net.Conn { + return NewLimiter(rateBps).WrapConn(conn) +} + +func (c *RateLimitedConn) Read(p []byte) (n int, err error) { + if len(p) > c.burst { + p = p[:c.burst] + } + n, err = c.Conn.Read(p) + if n > 0 { + if limitErr := c.readLimiter.WaitN(c.ctx, n); err == nil { + err = limitErr + } + } + return +} + +func (c *RateLimitedConn) Write(p []byte) (n int, err error) { + for len(p) > 0 { + chunkSize := len(p) + if chunkSize > c.burst { + chunkSize = c.burst + } + if err = c.writeLimiter.WaitN(c.ctx, chunkSize); err != nil { + return n, err + } + var written int + written, err = c.Conn.Write(p[:chunkSize]) + n += written + p = p[written:] + if err != nil { + return n, err + } + if written != chunkSize { + return n, io.ErrShortWrite + } + } + return n, nil +} + +func (c *RateLimitedConn) Close() error { + c.cancel() + return c.Conn.Close() +} + +func (c *RateLimitedConn) CloseWrite() error { + if conn, ok := c.Conn.(interface{ CloseWrite() error }); ok { + return conn.CloseWrite() + } + return c.Close() +} + +// limitedCConn wraps a C.Conn with shared rate limiting. +// It overrides Read/Write AND ReadBuffer/WriteBuffer so bufio.Copy/Relay +// cannot bypass the limiter via ExtendedConn methods promoted by embedding. +// +// Upload rate limiting: WaitN BEFORE write reserves tokens first, then +// the kernel absorbs the write instantly. This limits the rate at which +// data enters the kernel buffer, which controls network throughput. +// (After-write WaitN is useless because TCP kernel buffers absorb instantly.) +// +// Download rate limiting: WaitN AFTER read works because Read blocks when +// kernel buffer is empty, creating backpressure that slows the TCP sender. +type limitedCConn struct { + C.Conn + ctx context.Context + cancel context.CancelFunc + readLimiter *bitRateLimiter + writeLimiter *bitRateLimiter + burst int +} + +// NewLimitedCConn wraps a C.Conn with a private per-connection limit. +func NewLimitedCConn(conn C.Conn, rateBps uint64) C.Conn { + return NewLimiter(rateBps).WrapCConn(conn) +} + +func (c *limitedCConn) Read(p []byte) (n int, err error) { + if len(p) > c.burst { + p = p[:c.burst] + } + n, err = c.Conn.Read(p) + if n > 0 { + if limitErr := c.readLimiter.WaitN(c.ctx, n); err == nil { + err = limitErr + } + } + return +} + +func (c *limitedCConn) Write(p []byte) (n int, err error) { + for len(p) > 0 { + chunkSize := len(p) + if chunkSize > c.burst { + chunkSize = c.burst + } + // Reserve token BEFORE write so data enters kernel at rate-limited pace. + if wErr := c.writeLimiter.WaitN(c.ctx, chunkSize); wErr != nil { + return n, wErr + } + var written int + written, err = c.Conn.Write(p[:chunkSize]) + n += written + p = p[written:] + if err != nil { + return n, err + } + if written != chunkSize { + return n, io.ErrShortWrite + } + } + return n, nil +} + +// ReadBuffer rate-limits ExtendedConn buffer reads used by Relay/bufio.Copy. +func (c *limitedCConn) ReadBuffer(buffer *buf.Buffer) error { + before := buffer.Len() + err := c.Conn.ReadBuffer(buffer) + n := buffer.Len() - before + if n > 0 { + if limitErr := c.readLimiter.WaitN(c.ctx, n); err == nil { + err = limitErr + } + } + return err +} + +// WriteBuffer rate-limits ExtendedConn buffer writes used by Relay/bufio.Copy. +func (c *limitedCConn) WriteBuffer(buffer *buf.Buffer) error { + // Match ExtendedWriterWrapper: write then release the buffer. + defer buffer.Release() + _, err := c.Write(buffer.Bytes()) + return err +} + +// ReaderReplaceable reports that Relay must not unwrap past this limiter. +func (c *limitedCConn) ReaderReplaceable() bool { return false } + +// WriterReplaceable reports that Relay must not unwrap past this limiter. +func (c *limitedCConn) WriterReplaceable() bool { return false } + +func (c *limitedCConn) Upstream() any { return c.Conn } + +func (c *limitedCConn) Close() error { + c.cancel() + return c.Conn.Close() +} + +func (c *limitedCConn) CloseWrite() error { + if conn, ok := c.Conn.(interface{ CloseWrite() error }); ok { + return conn.CloseWrite() + } + return c.Close() +} + +// WrapPacketConn wraps a C.PacketConn with shared rate limiting. +func (l *Limiter) WrapPacketConn(pc C.PacketConn) C.PacketConn { + if l == nil { + return pc + } + limitCtx, cancel := context.WithCancel(context.Background()) + return &limitedPacketConn{ + PacketConn: pc, + ctx: limitCtx, + cancel: cancel, + readLimiter: l.readLimiter, + writeLimiter: l.writeLimiter, + } +} + +type limitedPacketConn struct { + C.PacketConn + ctx context.Context + cancel context.CancelFunc + readLimiter *bitRateLimiter + writeLimiter *bitRateLimiter +} + +func (c *limitedPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) { + n, addr, err = c.PacketConn.ReadFrom(p) + if n > 0 { + if limitErr := c.readLimiter.WaitN(c.ctx, n); err == nil { + err = limitErr + } + } + return +} + +// WaitReadFrom overrides the zero-copy UDP read path so rate limiting applies. +func (c *limitedPacketConn) WaitReadFrom() (data []byte, put func(), addr net.Addr, err error) { + type waitReader interface { + WaitReadFrom() (data []byte, put func(), addr net.Addr, err error) + } + if wr, ok := c.PacketConn.(waitReader); ok { + data, put, addr, err = wr.WaitReadFrom() + } + if len(data) > 0 { + if limitErr := c.readLimiter.WaitN(c.ctx, len(data)); err == nil { + err = limitErr + } + } + return +} + +func (c *limitedPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) { + n, err = c.PacketConn.WriteTo(p, addr) + if n > 0 { + if limitErr := c.writeLimiter.WaitN(c.ctx, n); err == nil { + err = limitErr + } + } + return +} + +func (c *limitedPacketConn) Close() error { + c.cancel() + return c.PacketConn.Close() +} + +func (c *limitedPacketConn) ReaderReplaceable() bool { return false } +func (c *limitedPacketConn) WriterReplaceable() bool { return false } diff --git a/component/ratelimit/ratelimit_test.go b/component/ratelimit/ratelimit_test.go new file mode 100644 index 0000000000..a9b635a94c --- /dev/null +++ b/component/ratelimit/ratelimit_test.go @@ -0,0 +1,201 @@ +package ratelimit + +import ( + "net" + "testing" + "time" +) + +func TestParseBandwidth(t *testing.T) { + cases := []struct { + in string + want uint64 + }{ + {"", 0}, + {"0", 0}, + {"500k", 500_000}, + {"500K", 500_000}, + {"500kbps", 500_000}, + {"500Kbps", 500_000}, + {"1m", 1_000_000}, + {"1M", 1_000_000}, + {"1Mbps", 1_000_000}, + {"2g", 2_000_000_000}, + {"1", 1_000_000}, // bare number = Mbps + {" 500k ", 500_000}, + {"bogus", 0}, + } + for _, tc := range cases { + got := ParseBandwidth(tc.in) + if got != tc.want { + t.Errorf("ParseBandwidth(%q) = %d, want %d", tc.in, got, tc.want) + } + } +} + +func TestNewRateLimitedConn_ZeroRate(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + conn := NewRateLimitedConn(client, 0) + if conn != client { + t.Error("zero rate should return original connection") + } +} + +func TestNewRateLimitedConn_WithRate(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + // 1 Mbps = 1,000,000 bits per second + conn := NewRateLimitedConn(client, 1_000_000) + if conn == client { + t.Error("non-zero rate should return wrapped connection") + } + + rlc, ok := conn.(*RateLimitedConn) + if !ok { + t.Fatal("should return *RateLimitedConn") + } + if rlc.burst == 0 { + t.Error("burst should be non-zero") + } +} + +func TestBitRateLimiter_Reservation(t *testing.T) { + limiter := &bitRateLimiter{rateBps: 800} + now := time.Unix(0, 0) + if delay := limiter.reserveN(now, 1); delay != 0 { + t.Fatalf("initial reservation delay = %s, want 0", delay) + } + if delay := limiter.reserveN(now, 1); delay != 10*time.Millisecond { + t.Fatalf("second reservation delay = %s, want 10ms", delay) + } +} + +func TestRateLimitedConn_ReadWrite(t *testing.T) { + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + // 8 Mbps = 1 MB/s → 1 byte every 1µs theoretically; use high enough for test + limited := NewRateLimitedConn(client, 8_000_000) + + go func() { + buf := make([]byte, 100) + _, _ = server.Read(buf) + _, _ = server.Write([]byte("pong")) + }() + + n, err := limited.Write([]byte("ping")) + if err != nil { + t.Fatalf("write: %v", err) + } + if n != 4 { + t.Fatalf("wrote %d, want 4", n) + } + + buf := make([]byte, 4) + n, err = limited.Read(buf) + if err != nil { + t.Fatalf("read: %v", err) + } + if string(buf[:n]) != "pong" { + t.Fatalf("read %q, want pong", buf[:n]) + } +} + +func TestRateLimitedConn_BurstLimits(t *testing.T) { + // 800 bps → burst = 800/8/(1000/10) = 1 + conn := NewRateLimitedConn(&net.TCPConn{}, 800) + rlc := conn.(*RateLimitedConn) + if rlc.burst != 1 { + t.Fatalf("burst = %d, want 1", rlc.burst) + } +} + +func TestRateLimitedConn_HighRateBurstCap(t *testing.T) { + // Very high rate should cap burst at 64KB + conn := NewRateLimitedConn(&net.TCPConn{}, 100_000_000_000) + rlc := conn.(*RateLimitedConn) + if rlc.burst != maxRateLimitBurstBytes { + t.Fatalf("burst = %d, want %d", rlc.burst, maxRateLimitBurstBytes) + } +} + +func TestLimiter_SharedAcrossConns(t *testing.T) { + lim := NewLimiter(800) // 800 bps = 100 B/s + if lim == nil { + t.Fatal("expected non-nil limiter") + } + if lim.Rate() != 800 { + t.Fatalf("rate = %d, want 800", lim.Rate()) + } + + // Both wraps must share the same underlying limiters. + s1, c1 := net.Pipe() + defer s1.Close() + defer c1.Close() + s2, c2 := net.Pipe() + defer s2.Close() + defer c2.Close() + + w1 := lim.WrapConn(c1).(*RateLimitedConn) + w2 := lim.WrapConn(c2).(*RateLimitedConn) + if w1.readLimiter != w2.readLimiter || w1.writeLimiter != w2.writeLimiter { + t.Fatal("wrapped conns must share the same limiters") + } +} + +func TestNewLimiter_Zero(t *testing.T) { + if NewLimiter(0) != nil { + t.Fatal("zero rate should return nil limiter") + } +} + +func TestRateLimitedConn_Throughput(t *testing.T) { + // 16_000 bps = 2_000 B/s. 4_000 bytes should take ~2s. + const rateBps = 16_000 + const payload = 4_000 + + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + limited := NewRateLimitedConn(client, rateBps) + + go func() { + _, _ = server.Write(make([]byte, payload)) + _ = server.Close() + }() + + start := time.Now() + buf := make([]byte, 512) + var total int + for { + n, err := limited.Read(buf) + total += n + if err != nil { + break + } + } + elapsed := time.Since(start) + if total != payload { + t.Fatalf("read %d, want %d", total, payload) + } + if elapsed < time.Second { + t.Fatalf("read too fast: %s for %d bytes at %d bps", elapsed, payload, rateBps) + } +} + +func TestLimitedCConn_NotReplaceable(t *testing.T) { + c := &limitedCConn{} + if c.ReaderReplaceable() { + t.Fatal("ReaderReplaceable must be false") + } + if c.WriterReplaceable() { + t.Fatal("WriterReplaceable must be false") + } +} diff --git a/config/config.go b/config/config.go index 8959dda922..8007ae52e1 100644 --- a/config/config.go +++ b/config/config.go @@ -977,6 +977,7 @@ func parseProxies(cfg *RawConfig) (proxies map[string]C.Proxy, providersMap map[ outboundgroup.SelectorOption{}, proxies["COMPATIBLE"], []P.ProxyProvider{pd}, + 0, ) if err != nil { return nil, nil, fmt.Errorf("new GLOBAL proxy group error: %w", err) diff --git a/constant/adapters.go b/constant/adapters.go index ddcdbead1f..668adbf256 100644 --- a/constant/adapters.go +++ b/constant/adapters.go @@ -145,6 +145,9 @@ type ProxyAdapter interface { // Unwrap extracts the proxy from a proxy-group. It returns nil when nothing to extract. Unwrap(metadata *Metadata, touch bool) Proxy + // Bandwidth returns the bandwidth limit in bits per second. 0 means unlimited. + Bandwidth() uint64 + // Close releasing associated resources Close() error } diff --git a/transport/jls/jls.go b/transport/jls/jls.go index 1ca20fa590..b6a298b5b6 100644 --- a/transport/jls/jls.go +++ b/transport/jls/jls.go @@ -5,23 +5,18 @@ import ( "context" "errors" "fmt" - "io" "net" - "sync" - "time" N "github.com/metacubex/mihomo/common/net" "github.com/metacubex/mihomo/component/ca" + "github.com/metacubex/mihomo/component/ratelimit" "github.com/metacubex/mihomo/ntp" tls "github.com/metacubex/jls-tls" ) const ( - Mode = "jls" - bitsPerByte = 8 - rateLimitCycle = 10 * time.Millisecond - maxRateLimitBurstBytes = 64 * 1024 + Mode = "jls" ) var ( @@ -216,124 +211,13 @@ func relayFallback(ctx context.Context, inbound net.Conn, prefix []byte, config return err } inbound = N.NewCachedConn(inbound, prefix) - upstream = newRateLimitedConn(upstream, config.RateLimit) + upstream = ratelimit.NewRateLimitedConn(upstream, config.RateLimit) if err = N.RelayContext(ctx, inbound, upstream); err != nil { return err } return ErrFallbackCompleted } -type rateLimitedConn struct { - net.Conn - ctx context.Context - cancel context.CancelFunc - readLimiter *bitRateLimiter - writeLimiter *bitRateLimiter - burst int -} - -func newRateLimitedConn(conn net.Conn, rateBps uint64) net.Conn { - if rateBps == 0 { - return conn - } - burst := rateBps / bitsPerByte / uint64(time.Second/rateLimitCycle) - if burst == 0 { - burst = 1 - } else if burst > maxRateLimitBurstBytes { - burst = maxRateLimitBurstBytes - } - limitCtx, cancel := context.WithCancel(context.Background()) - return &rateLimitedConn{ - Conn: conn, - ctx: limitCtx, - cancel: cancel, - readLimiter: &bitRateLimiter{rateBps: rateBps}, - writeLimiter: &bitRateLimiter{rateBps: rateBps}, - burst: int(burst), - } -} - -func (c *rateLimitedConn) Read(p []byte) (n int, err error) { - if len(p) > c.burst { - p = p[:c.burst] - } - n, err = c.Conn.Read(p) - if n > 0 { - if limitErr := c.readLimiter.WaitN(c.ctx, n); err == nil { - err = limitErr - } - } - return -} - -func (c *rateLimitedConn) Write(p []byte) (n int, err error) { - for len(p) > 0 { - chunkSize := len(p) - if chunkSize > c.burst { - chunkSize = c.burst - } - if err = c.writeLimiter.WaitN(c.ctx, chunkSize); err != nil { - return n, err - } - var written int - written, err = c.Conn.Write(p[:chunkSize]) - n += written - p = p[written:] - if err != nil { - return n, err - } - if written != chunkSize { - return n, io.ErrShortWrite - } - } - return n, nil -} - -func (c *rateLimitedConn) Close() error { - c.cancel() - return c.Conn.Close() -} - -func (c *rateLimitedConn) CloseWrite() error { - if conn, ok := c.Conn.(interface{ CloseWrite() error }); ok { - return conn.CloseWrite() - } - return c.Close() -} - -type bitRateLimiter struct { - mu sync.Mutex - rateBps uint64 - next time.Time -} - -func (l *bitRateLimiter) WaitN(ctx context.Context, n int) error { - delay := l.reserveN(time.Now(), n) - if delay <= 0 { - return nil - } - timer := time.NewTimer(delay) - defer timer.Stop() - select { - case <-timer.C: - return nil - case <-ctx.Done(): - return ctx.Err() - } -} - -func (l *bitRateLimiter) reserveN(now time.Time, n int) time.Duration { - interval := time.Duration(uint64(n) * bitsPerByte * uint64(time.Second) / l.rateBps) - - l.mu.Lock() - ready := l.next - if ready.Before(now) { - ready = now - } - l.next = ready.Add(interval) - l.mu.Unlock() - return ready.Sub(now) -} type handshakeRecorderConn struct { net.Conn diff --git a/transport/jls/jls_test.go b/transport/jls/jls_test.go index c6dab56d8c..057b52bdaf 100644 --- a/transport/jls/jls_test.go +++ b/transport/jls/jls_test.go @@ -8,7 +8,6 @@ import ( "net" "sync/atomic" "testing" - "time" N "github.com/metacubex/mihomo/common/net" "github.com/metacubex/mihomo/component/ca" @@ -478,16 +477,6 @@ func TestJLSServerFallbackReplaysRejectedTLS(t *testing.T) { } } -func TestBitRateLimiterReservations(t *testing.T) { - limiter := &bitRateLimiter{rateBps: 800} - now := time.Unix(0, 0) - if delay := limiter.reserveN(now, 1); delay != 0 { - t.Fatalf("initial reservation delay = %s, want 0", delay) - } - if delay := limiter.reserveN(now, 1); delay != 10*time.Millisecond { - t.Fatalf("second reservation delay = %s, want 10ms", delay) - } -} func newTestTLSServerConfig(t *testing.T, version uint16) *tls.Config { t.Helper()