Fix limiters

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

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
}

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) {