feat: improve proxy connection handling and add bandwidth throttling
Kim committed
Apr 15, 2026 at 15:45 UTC
3691964124e753262e3de30c31d30a831d16d921
5 files changed
+183
-54
portal/api_server.go
+23
-7
@@ -163,8 +163,8 @@ func (s *Server) signedRelayDescriptor(now time.Time) (types.RelayDescriptor, er
163
} else {
164
now = now.UTC()
165
}
166
- activeConns := float64(s.proxy.ActiveConns())
167
- tcpTrafficBPS := s.proxy.CurrentTCPBPS(now)
166
+ activeConns := float64(s.proxy.activeConnectionCount())
167
+ tcpTrafficBPS := s.proxy.currentTCPBPS(now)
168
ingressAddr := s.identity.Name
169
if s.cfg.SNIPort != 0 && s.cfg.SNIPort != 443 {
170
ingressAddr = fmt.Sprintf("%s:%d", ingressAddr, s.cfg.SNIPort)
@@ -263,10 +263,24 @@ func (s *Server) handleRelayDiscoveryAnnounce(w http.ResponseWriter, r *http.Req
263
// Self-announce guard: the relay's own URL is established locally, not
264
// gossiped through the announce endpoint. Reject loopback / own-host
265
// announces to prevent self-amplification or misconfiguration loops.
266
- announceURL, parseErr := url.Parse(desc.APIHTTPSAddr)
267
- if parseErr == nil && announceURL != nil {
268
- if utils.IsLocalRelayHost(announceURL.Hostname()) {
269
- utils.WriteAPIError(w, http.StatusBadRequest, types.APIErrorCodeInvalidRequest, "self-announce rejected")
266
+ announceURL, err := url.Parse(strings.TrimSpace(desc.APIHTTPSAddr))
267
+ if err == nil && announceURL != nil {
268
+ host := utils.NormalizeHostname(announceURL.Hostname())
269
+ if utils.IsLocalRelayHost(host) {
270
+ utils.WriteAPIError(w, http.StatusBadRequest, types.APIErrorCodeInvalidRequest,
271
+ fmt.Sprintf("self-announce rejected: host %q is local-only", host))
272
+ return
273
+ }
274
+ if selfURL, err := utils.NormalizeRelayURL(s.cfg.PortalURL); err == nil {
275
+ if announceRelayURL, err := utils.NormalizeRelayURL(desc.APIHTTPSAddr); err == nil && announceRelayURL == selfURL {
276
+ utils.WriteAPIError(w, http.StatusBadRequest, types.APIErrorCodeInvalidRequest,
277
+ fmt.Sprintf("self-announce rejected: %q matches receiving relay url", announceRelayURL))
278
+ return
279
+ }
280
+ }
281
+ if host != "" && host == utils.NormalizeHostname(s.identity.Name) {
282
+ utils.WriteAPIError(w, http.StatusBadRequest, types.APIErrorCodeInvalidRequest,
283
+ fmt.Sprintf("self-announce rejected: host %q matches receiving relay host", host))
284
return
285
}
286
}
@@ -684,7 +698,9 @@ func (s *Server) registerLease(req types.RegisterChallengeRequest, clientIP, rep
698
}
699
return types.RegisterResponse{}, err
700
}
687
- record.tcpPort = transport.NewRelayTCPPort(identityKey, port, stream)
701
+ record.tcpPort = transport.NewRelayTCPPort(identityKey, port, stream, func(left, right net.Conn) {
702
+ s.proxy.bridge(left, right, identityKey, s.registry.policy.BPSManager())
703
+ })
704
record.tcpPorts = s.tcpPorts
705
}
706
portal/policy/bps_manager.go
+101
-3
@@ -3,16 +3,19 @@ package policy
3
import (
4
"maps"
5
"sync"
6
+ "time"
7
)
8
9
type BPSManager struct {
9
- identityBPS map[string]int64
10
- mu sync.RWMutex
10
+ identityBPS map[string]int64
11
+ identityLimiters map[string]*bpsLimiter
12
+ mu sync.RWMutex
13
}
14
15
func NewBPSManager() *BPSManager {
16
return &BPSManager{
15
- identityBPS: make(map[string]int64),
17
+ identityBPS: make(map[string]int64),
18
+ identityLimiters: make(map[string]*bpsLimiter),
19
}
20
}
21
@@ -35,6 +38,7 @@ func (m *BPSManager) SetIdentityBPS(key string, bps int64) {
38
defer m.mu.Unlock()
39
if bps <= 0 {
40
delete(m.identityBPS, key)
41
+ delete(m.identityLimiters, key)
42
return
43
}
44
m.identityBPS[key] = bps
@@ -48,6 +52,7 @@ func (m *BPSManager) DeleteIdentityBPS(key string) {
52
m.mu.Lock()
53
defer m.mu.Unlock()
54
delete(m.identityBPS, key)
55
+ delete(m.identityLimiters, key)
56
}
57
58
func (m *BPSManager) IdentityBPSLimits() map[string]int64 {
@@ -78,5 +83,98 @@ func (m *BPSManager) SetIdentityBPSLimits(limits map[string]int64) {
83
84
m.mu.Lock()
85
m.identityBPS = next
86
+ m.identityLimiters = make(map[string]*bpsLimiter)
87
m.mu.Unlock()
88
}
89
+
90
+func (m *BPSManager) ThrottleIdentityBPS(key string, maxBytes int) int {
91
+ if m == nil || key == "" || maxBytes <= 0 {
92
+ return maxBytes
93
+ }
94
+
95
+ for {
96
+ bps, limiter := m.identityLimiter(key)
97
+ if bps <= 0 || limiter == nil {
98
+ return maxBytes
99
+ }
100
+ chunkSize := bpsChunkSize(maxBytes, bps)
101
+ if wait := limiter.reserve(float64(chunkSize), float64(bps)); wait > 0 {
102
+ time.Sleep(wait)
103
+ continue
104
+ }
105
+ return chunkSize
106
+ }
107
+}
108
+
109
+func (m *BPSManager) identityLimiter(key string) (int64, *bpsLimiter) {
110
+ m.mu.RLock()
111
+ bps := m.identityBPS[key]
112
+ limiter := m.identityLimiters[key]
113
+ m.mu.RUnlock()
114
+ if bps <= 0 || limiter != nil {
115
+ return bps, limiter
116
+ }
117
+
118
+ m.mu.Lock()
119
+ defer m.mu.Unlock()
120
+
121
+ bps = m.identityBPS[key]
122
+ if bps <= 0 {
123
+ return 0, nil
124
+ }
125
+ if m.identityLimiters == nil {
126
+ m.identityLimiters = make(map[string]*bpsLimiter)
127
+ }
128
+ limiter = m.identityLimiters[key]
129
+ if limiter == nil {
130
+ limiter = &bpsLimiter{}
131
+ m.identityLimiters[key] = limiter
132
+ }
133
+ return bps, limiter
134
+}
135
+
136
+func bpsChunkSize(length int, bps int64) int {
137
+ if bps <= 0 {
138
+ return length
139
+ }
140
+ chunk := bps / 10
141
+ if chunk < 1 {
142
+ chunk = 1
143
+ }
144
+ if chunk > int64(length) {
145
+ chunk = int64(length)
146
+ }
147
+ return int(chunk)
148
+}
149
+
150
+type bpsLimiter struct {
151
+ mu sync.Mutex
152
+ tokens float64
153
+ updatedAt time.Time
154
+}
155
+
156
+func (l *bpsLimiter) reserve(bytes, bps float64) time.Duration {
157
+ l.mu.Lock()
158
+ defer l.mu.Unlock()
159
+
160
+ now := time.Now()
161
+ if l.updatedAt.IsZero() {
162
+ l.updatedAt = now
163
+ } else if elapsed := now.Sub(l.updatedAt).Seconds(); elapsed > 0 {
164
+ l.tokens += elapsed * bps
165
+ l.updatedAt = now
166
+ }
167
+ if l.tokens > bps {
168
+ l.tokens = bps
169
+ }
170
+
171
+ if l.tokens >= bytes {
172
+ l.tokens -= bytes
173
+ return 0
174
+ }
175
+
176
+ missing := bytes - l.tokens
177
+ l.tokens = 0
178
+ l.updatedAt = now
179
+ return time.Duration(missing / bps * float64(time.Second))
180
+}
portal/proxy.go
+48
-5
@@ -8,6 +8,8 @@ import (
8
"time"
9
10
"golang.org/x/sync/errgroup"
11
+
12
+ "github.com/gosuda/portal-tunnel/v2/portal/policy"
13
)
14
15
type proxy struct {
@@ -18,32 +20,33 @@ type proxy struct {
20
tcpLoadBytes int64
21
}
22
21
-func (p *proxy) Bridge(left, right net.Conn) {
23
+func (p *proxy) bridge(left, right net.Conn, identityKey string, bpsManager *policy.BPSManager) {
24
p.activeConns.Add(1)
25
defer p.activeConns.Add(-1)
26
27
defer left.Close()
28
defer right.Close()
29
30
+ throttled := bpsManager != nil && bpsManager.IdentityBPS(identityKey) > 0
31
var group errgroup.Group
32
group.Go(func() error {
30
- _, err := io.Copy(&countingConn{Conn: right, bytes: &p.tcpBytes}, left)
33
+ err := p.copy(right, left, identityKey, bpsManager, throttled)
34
closeWrite(right)
35
return err
36
})
37
group.Go(func() error {
35
- _, err := io.Copy(&countingConn{Conn: left, bytes: &p.tcpBytes}, right)
38
+ err := p.copy(left, right, identityKey, bpsManager, throttled)
39
closeWrite(left)
40
return err
41
})
42
_ = group.Wait()
43
}
44
42
-func (p *proxy) ActiveConns() int64 {
45
+func (p *proxy) activeConnectionCount() int64 {
46
return p.activeConns.Load()
47
}
48
46
-func (p *proxy) CurrentTCPBPS(now time.Time) float64 {
49
+func (p *proxy) currentTCPBPS(now time.Time) float64 {
50
totalTCPBytes := p.tcpBytes.Load()
51
52
p.tcpLoadMu.Lock()
@@ -65,6 +68,46 @@ func (p *proxy) CurrentTCPBPS(now time.Time) float64 {
68
return 0
69
}
70
71
+func (p *proxy) copy(dst, src net.Conn, identityKey string, bpsManager *policy.BPSManager, throttled bool) error {
72
+ // fast path
73
+ if !throttled {
74
+ _, err := io.Copy(&countingConn{Conn: dst, bytes: &p.tcpBytes}, src)
75
+ return err
76
+ }
77
+
78
+ buf := make([]byte, 32*1024)
79
+ for {
80
+ nr, readErr := src.Read(buf)
81
+ if nr > 0 {
82
+ data := buf[:nr]
83
+ for len(data) > 0 {
84
+ chunkSize := len(data)
85
+ if bpsManager != nil {
86
+ chunkSize = bpsManager.ThrottleIdentityBPS(identityKey, chunkSize)
87
+ }
88
+
89
+ n, err := dst.Write(data[:chunkSize])
90
+ if n > 0 {
91
+ p.tcpBytes.Add(int64(n))
92
+ data = data[n:]
93
+ }
94
+ if err != nil {
95
+ return err
96
+ }
97
+ if n == 0 {
98
+ return io.ErrShortWrite
99
+ }
100
+ }
101
+ }
102
+ if readErr != nil {
103
+ if readErr == io.EOF {
104
+ return nil
105
+ }
106
+ return readErr
107
+ }
108
+ }
109
+}
110
+
111
type countingConn struct {
112
net.Conn
113
bytes *atomic.Int64
portal/server.go
+3
-3
@@ -459,12 +459,12 @@ func (s *Server) runSNIListener(ctx context.Context) error {
459
_ = wrappedConn.Close()
460
return
461
}
462
- s.proxy.Bridge(wrappedConn, upstream)
462
+ s.proxy.bridge(wrappedConn, upstream, "", nil)
463
return
464
}
465
466
record, ok := s.registry.Lookup(serverName)
467
- if !ok || record == nil || time.Now().After(record.ExpiresAt) || !s.registry.policy.IsIdentityRoutable(record.Key()) || record.stream == nil {
467
+ if !ok || record == nil || time.Now().After(record.ExpiresAt) || record.stream == nil || !s.registry.policy.IsIdentityRoutable(record.Key()) {
468
_ = wrappedConn.Close()
469
return
470
}
@@ -478,7 +478,7 @@ func (s *Server) runSNIListener(ctx context.Context) error {
478
return
479
}
480
481
- s.proxy.Bridge(wrappedConn, session)
481
+ s.proxy.bridge(wrappedConn, session, record.Key(), s.registry.policy.BPSManager())
482
}(conn)
483
case errors.Is(err, net.ErrClosed):
484
return nil
portal/transport/stream_port_relay.go
+8
-36
@@ -20,16 +20,18 @@ type RelayTCPPort struct {
20
port int
21
listener net.Listener
22
stream *RelayStream
23
+ bridge func(net.Conn, net.Conn)
24
25
cancel context.CancelFunc
26
closeOnce sync.Once
27
}
28
28
-func NewRelayTCPPort(identityKey string, port int, stream *RelayStream) *RelayTCPPort {
29
+func NewRelayTCPPort(identityKey string, port int, stream *RelayStream, bridge func(net.Conn, net.Conn)) *RelayTCPPort {
30
return &RelayTCPPort{
31
identityKey: identityKey,
32
port: port,
33
stream: stream,
34
+ bridge: bridge,
35
}
36
}
37
@@ -123,40 +125,10 @@ func (t *RelayTCPPort) handleConn(ctx context.Context, conn net.Conn) {
125
return
126
}
127
126
- bridgeConns(conn, session)
127
-}
128
-
129
-// bridgeConns copies data bidirectionally between two connections.
130
-func bridgeConns(left, right net.Conn) {
131
- defer left.Close()
132
- defer right.Close()
133
-
134
- done := make(chan struct{})
135
- go func() {
136
- defer close(done)
137
- copyAndCloseWrite(right, left)
138
- }()
139
- copyAndCloseWrite(left, right)
140
- <-done
141
-}
142
-
143
-func copyAndCloseWrite(dst, src net.Conn) {
144
- buf := make([]byte, 32*1024)
145
- for {
146
- nr, readErr := src.Read(buf)
147
- if nr > 0 {
148
- if _, writeErr := dst.Write(buf[:nr]); writeErr != nil {
149
- break
150
- }
151
- }
152
- if readErr != nil {
153
- break
154
- }
155
- }
156
- type closeWriter interface {
157
- CloseWrite() error
158
- }
159
- if cw, ok := dst.(closeWriter); ok {
160
- _ = cw.CloseWrite()
128
+ if t.bridge == nil {
129
+ _ = conn.Close()
130
+ _ = session.Close()
131
+ return
132
}
133
+ t.bridge(conn, session)
134
}