yanistyle/sing-box-extended
1
0
Fork 0
mirror of https://github.com/shtorm-7/sing-box-extended.git synced 2026-08-08 19:45:18 +03:00
Code Issues Packages Projects Releases Wiki Activity

Fix limiters

Browse Source
This commit is contained in:
Shtorm
2026-08-08 10:11:10 +03:00
parent 6c81ddd00f
commit 3310474137
4 changed files with 77 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
protocol/limiter/bandwidth/outbound.go
View File

@@ -106,7 +106,12 @@ func (h *Outbound) DialContext(ctx context.Context, network string, destination
if err != nil {
return nil, err
}
return h.strategy.wrapConn(ctx, conn, adapter.ContextFrom(ctx), true)
wrappedConn, err := h.strategy.wrapConn(ctx, conn, adapter.ContextFrom(ctx), true)
if err != nil {
conn.Close()
return nil, err
}
return wrappedConn, nil
}
func (h *Outbound) ListenPacket(ctx context.Context, destination M.Socksaddr) (net.PacketConn, error) {
@@ -114,12 +119,17 @@ func (h *Outbound) ListenPacket(ctx context.Context, destination M.Socksaddr) (n
if err != nil {
return nil, err
}
return h.strategy.wrapPacketConn(ctx, conn, adapter.ContextFrom(ctx), true)
wrappedConn, err := h.strategy.wrapPacketConn(ctx, conn, adapter.ContextFrom(ctx), true)
if err != nil {
conn.Close()
return nil, err
}
return wrappedConn, nil
}
func (h *Outbound) NewConnectionEx(ctx context.Context, conn net.Conn, metadata adapter.InboundContext, onClose N.CloseHandlerFunc) {
ctx = adapter.WithContext(ctx, &metadata)
conn, err := h.strategy.wrapConn(ctx, conn, &metadata, false)
wrappedConn, err := h.strategy.wrapConn(ctx, conn, &metadata, false)
if err != nil {
h.logger.ErrorContext(ctx, err)
N.CloseOnHandshakeFailure(conn, onClose, err)
@@ -127,7 +137,7 @@ func (h *Outbound) NewConnectionEx(ctx context.Context, conn net.Conn, metadata
}
metadata.Inbound = h.Tag()
metadata.InboundType = h.Type()
h.router.RouteConnectionEx(ctx, conn, metadata, onClose)
h.router.RouteConnectionEx(ctx, wrappedConn, metadata, onClose)
return
}

23
protocol/limiter/bandwidth/strategy.go
View File

@@ -2,6 +2,7 @@ package bandwidth
import (
"context"
"io"
"net"
"strconv"
"sync"
@@ -191,7 +192,7 @@ func (s *UsersBandwidthStrategy) getStrategy(ctx context.Context, metadata *adap
}
type bwConnEntry struct {
conn net.Conn
conn io.Closer
}
type ManagerBandwidthStrategy struct {
@@ -251,7 +252,25 @@ func (s *ManagerBandwidthStrategy) wrapPacketConn(ctx context.Context, conn net.
if !ok {
return nil, E.New("user strategy not found: ", user)
}
return strategy.wrapPacketConn(ctx, conn, metadata, reverse)
wrapped, err := strategy.wrapPacketConn(ctx, conn, metadata, reverse)
if err != nil {
return nil, err
}
entry := &bwConnEntry{conn: conn}
s.mtx.Lock()
s.conns[user] = append(s.conns[user], entry)
s.mtx.Unlock()
return onclose.NewPacketConn(wrapped, func() {
s.mtx.Lock()
entries := s.conns[user]
for i, e := range entries {
if e == entry {
s.conns[user] = append(entries[:i], entries[i+1:]...)
break
}
}
s.mtx.Unlock()
}), nil
}
func (s *ManagerBandwidthStrategy) UpdateStrategies(strategies map[string]BandwidthStrategy) {

18
protocol/limiter/traffic/outbound.go
View File

@@ -82,7 +82,12 @@ func (h *Outbound) DialContext(ctx context.Context, network string, destination
if err != nil {
return nil, err
}
return h.strategy.wrapConn(ctx, conn, adapter.ContextFrom(ctx), true)
wrappedConn, err := h.strategy.wrapConn(ctx, conn, adapter.ContextFrom(ctx), true)
if err != nil {
conn.Close()
return nil, err
}
return wrappedConn, nil
}
func (h *Outbound) ListenPacket(ctx context.Context, destination M.Socksaddr) (net.PacketConn, error) {
@@ -90,11 +95,16 @@ func (h *Outbound) ListenPacket(ctx context.Context, destination M.Socksaddr) (n
if err != nil {
return nil, err
}
return h.strategy.wrapPacketConn(ctx, conn, adapter.ContextFrom(ctx), true)
wrappedConn, err := h.strategy.wrapPacketConn(ctx, conn, adapter.ContextFrom(ctx), true)
if err != nil {
conn.Close()
return nil, err
}
return wrappedConn, nil
}
func (h *Outbound) NewConnectionEx(ctx context.Context, conn net.Conn, metadata adapter.InboundContext, onClose N.CloseHandlerFunc) {
conn, err := h.strategy.wrapConn(ctx, conn, &metadata, false)
wrappedConn, err := h.strategy.wrapConn(ctx, conn, &metadata, false)
if err != nil {
if err.Error() != "traffic limit exceeded" {
h.logger.ErrorContext(ctx, err)
@@ -104,7 +114,7 @@ func (h *Outbound) NewConnectionEx(ctx context.Context, conn net.Conn, metadata
}
metadata.Inbound = h.Tag()
metadata.InboundType = h.Type()
h.router.RouteConnectionEx(ctx, conn, metadata, onClose)
h.router.RouteConnectionEx(ctx, wrappedConn, metadata, onClose)
}
func (h *Outbound) NewPacketConnectionEx(ctx context.Context, conn N.PacketConn, metadata adapter.InboundContext, onClose N.CloseHandlerFunc) {

37
protocol/limiter/traffic/strategy.go
View File

@@ -2,6 +2,7 @@ package traffic
import (
"context"
"io"
"net"
"sync"
@@ -37,11 +38,11 @@ func NewDefaultWrapStrategy(limiterStrategy TrafficLimiterStrategy, connWrapper
func (s *DefaultWrapStrategy) wrapConn(ctx context.Context, conn net.Conn, metadata *adapter.InboundContext, reverse bool) (net.Conn, error) {
limiter, err := s.limiterStrategy.getLimiter(ctx, metadata)
if err != nil {
return conn, err
return nil, err
}
_, err = limiter.Reserve(0)
if err != nil {
return conn, err
return nil, err
}
return s.connWrapper(ctx, conn, limiter, reverse), nil
}
@@ -49,11 +50,11 @@ func (s *DefaultWrapStrategy) wrapConn(ctx context.Context, conn net.Conn, metad
func (s *DefaultWrapStrategy) wrapPacketConn(ctx context.Context, conn net.PacketConn, metadata *adapter.InboundContext, reverse bool) (net.PacketConn, error) {
limiter, err := s.limiterStrategy.getLimiter(ctx, metadata)
if err != nil {
return conn, err
return nil, err
}
_, err = limiter.Reserve(0)
if err != nil {
return conn, err
return nil, err
}
return s.packetConnWrapper(ctx, conn, limiter, reverse), nil
}
@@ -73,7 +74,7 @@ func (s *GlobalTrafficStrategy) getLimiter(ctx context.Context, metadata *adapte
}
type connEntry struct {
conn net.Conn
conn io.Closer
}
type ManagerTrafficStrategy struct {
@@ -94,7 +95,7 @@ func (s *ManagerTrafficStrategy) wrapConn(ctx context.Context, conn net.Conn, me
if err != nil {
return nil, err
}
wrapped, err := strategy.wrapConn(ctx, conn, metadata, reverse)
wrappedConn, err := strategy.wrapConn(ctx, conn, metadata, reverse)
if err != nil {
return nil, err
}
@@ -102,7 +103,7 @@ func (s *ManagerTrafficStrategy) wrapConn(ctx context.Context, conn net.Conn, me
s.mtx.Lock()
s.conns[user] = append(s.conns[user], entry)
s.mtx.Unlock()
return onclose.NewConn(wrapped, func() {
return onclose.NewConn(wrappedConn, func() {
s.mtx.Lock()
entries := s.conns[user]
for i, e := range entries {
@@ -116,11 +117,29 @@ func (s *ManagerTrafficStrategy) wrapConn(ctx context.Context, conn net.Conn, me
}
func (s *ManagerTrafficStrategy) wrapPacketConn(ctx context.Context, conn net.PacketConn, metadata *adapter.InboundContext, reverse bool) (net.PacketConn, error) {
strategy, _, err := s.getStrategy(ctx, metadata)
strategy, user, err := s.getStrategy(ctx, metadata)
if err != nil {
return nil, err
}
return strategy.wrapPacketConn(ctx, conn, metadata, reverse)
wrappedConn, err := strategy.wrapPacketConn(ctx, conn, metadata, reverse)
if err != nil {
return nil, err
}
entry := &connEntry{conn: conn}
s.mtx.Lock()
s.conns[user] = append(s.conns[user], entry)
s.mtx.Unlock()
return onclose.NewPacketConn(wrappedConn, func() {
s.mtx.Lock()
entries := s.conns[user]
for i, e := range entries {
if e == entry {
s.conns[user] = append(entries[:i], entries[i+1:]...)
break
}
}
s.mtx.Unlock()
}), nil
}
func (s *ManagerTrafficStrategy) getStrategy(ctx context.Context, metadata *adapter.InboundContext) (TrafficStrategy, string, error) {