feat: implement traffic control limits and per-lease byte rate limiting
Kim committed
Nov 13, 2025 at 16:46 UTC
25126a737c2c0e61d2cd5fa970066b12f269c275
5 files changed
+303
-14
cmd/relay-server/main.go
+11
@@ -23,6 +23,8 @@ var (
23
flagPort int
24
flagStaticDir string
25
flagPortalHost string
26
+ flagMaxLease int
27
+ flagLeaseBPS int
28
)
29
30
func main() {
@@ -56,6 +58,8 @@ func main() {
58
flag.IntVar(&flagPort, "port", 4017, "admin UI and HTTP proxy port")
59
flag.StringVar(&flagStaticDir, "static-dir", defaultStaticDir, "static files directory for portal frontend (env: STATIC_DIR)")
60
flag.StringVar(&flagPortalHost, "portal-host", defaultPortalHost, "portal host for frontend serving (env: PORTAL_HOST)")
61
+ flag.IntVar(&flagMaxLease, "max-lease", 0, "maximum active relayed connections per lease (0 = unlimited)")
62
+ flag.IntVar(&flagLeaseBPS, "lease-bps", 0, "default bytes-per-second limit per lease (0 = unlimited)")
63
64
flag.Parse()
65
@@ -120,6 +124,13 @@ func runServer() error {
124
cred := sdk.NewCredential()
125
126
serv := portal.NewRelayServer(cred, flagBootstraps)
127
+ // Apply traffic controls if configured
128
+ if flagMaxLease > 0 {
129
+ serv.SetMaxRelayedPerLease(flagMaxLease)
130
+ }
131
+ if flagLeaseBPS > 0 {
132
+ serv.GetLeaseManager().SetDefaultBPS(int64(flagLeaseBPS))
133
+ }
134
serv.Start()
135
defer serv.Stop()
136
portal/handlers.go
+33
-4
@@ -11,6 +11,7 @@ import (
11
"gosuda.org/portal/portal/core/cryptoops"
12
"gosuda.org/portal/portal/core/proto/rdsec"
13
"gosuda.org/portal/portal/core/proto/rdverb"
14
+ "gosuda.org/portal/portal/utils/ratelimit"
15
)
16
17
type StreamContext struct {
@@ -326,6 +327,21 @@ func (g *RelayServer) handleConnectionRequest(ctx *StreamContext, packet *rdverb
327
Str("response_code", resp.Code.String()).
328
Msg("[RelayServer] Received response from lease holder, sending to client")
329
330
+ // Enforce relayed connection limits if currently accepted
331
+ if resp.Code == rdverb.ResponseCode_RESPONSE_CODE_ACCEPTED {
332
+ leaseID := string(leaseEntry.Lease.Identity.Id)
333
+ g.limitsLock.Lock()
334
+ overPerLease := g.maxRelayedPerLease > 0 && g.relayedPerLeaseCount[leaseID] >= g.maxRelayedPerLease
335
+ if overPerLease {
336
+ log.Warn().
337
+ Str("lease_id", leaseID).
338
+ Bool("over_per_lease", overPerLease).
339
+ Msg("[RelayServer] Relayed connection per-lease limit reached, rejecting")
340
+ resp.Code = rdverb.ResponseCode_RESPONSE_CODE_REJECTED
341
+ }
342
+ g.limitsLock.Unlock()
343
+ }
344
+
345
// Send response to client
346
response, err := resp.MarshalVT()
347
if err != nil {
@@ -350,6 +366,10 @@ func (g *RelayServer) handleConnectionRequest(ctx *StreamContext, packet *rdverb
366
ctx.Hijack()
367
368
leaseID := string(leaseEntry.Lease.Identity.Id)
369
+ // Increment counters for active relayed connections
370
+ g.limitsLock.Lock()
371
+ g.relayedPerLeaseCount[leaseID] = g.relayedPerLeaseCount[leaseID] + 1
372
+ g.limitsLock.Unlock()
373
g.relayedConnectionsLock.Lock()
374
g.relayedConnections[leaseID] = append(g.relayedConnections[leaseID], ctx.Stream)
375
g.relayedConnectionsLock.Unlock()
@@ -358,10 +378,10 @@ func (g *RelayServer) handleConnectionRequest(ctx *StreamContext, packet *rdverb
378
var wg sync.WaitGroup
379
wg.Add(2)
380
361
- // Copy from client to lease holder
381
+ // Copy from client to lease holder (with optional per-lease BPS limit)
382
go func() {
383
defer wg.Done()
364
- n, err := io.Copy(leaseStream, ctx.Stream)
384
+ n, err := ratelimit.Copy(leaseStream, ctx.Stream, g.getLeaseBPSBucket(leaseID))
385
log.Debug().
386
Str("lease_id", leaseID).
387
Int64("bytes", n).
@@ -370,10 +390,10 @@ func (g *RelayServer) handleConnectionRequest(ctx *StreamContext, packet *rdverb
390
leaseStream.Close()
391
}()
392
373
- // Copy from lease holder to client
393
+ // Copy from lease holder to client (with optional per-lease BPS limit)
394
go func() {
395
defer wg.Done()
376
- n, err := io.Copy(ctx.Stream, leaseStream)
396
+ n, err := ratelimit.Copy(ctx.Stream, leaseStream, g.getLeaseBPSBucket(leaseID))
397
log.Debug().
398
Str("lease_id", leaseID).
399
Int64("bytes", n).
@@ -385,6 +405,15 @@ func (g *RelayServer) handleConnectionRequest(ctx *StreamContext, packet *rdverb
405
wg.Wait()
406
log.Debug().Str("lease_id", leaseID).Msg("[RelayServer] Connection forwarding completed successfully")
407
408
+ // Decrement counters after forwarding completes
409
+ g.limitsLock.Lock()
410
+ if v := g.relayedPerLeaseCount[leaseID]; v > 1 {
411
+ g.relayedPerLeaseCount[leaseID] = v - 1
412
+ } else {
413
+ delete(g.relayedPerLeaseCount, leaseID)
414
+ }
415
+ g.limitsLock.Unlock()
416
+
417
// Clean up relayed connection tracking
418
g.relayedConnectionsLock.Lock()
419
if streams, exists := g.relayedConnections[leaseID]; exists {
portal/lease.go
+112
-3
@@ -1,6 +1,7 @@
1
package portal
2
3
import (
4
+ "regexp"
5
"sync"
6
"time"
7
@@ -20,13 +21,25 @@ type LeaseManager struct {
21
leasesLock sync.RWMutex
22
stopCh chan struct{}
23
ttlInterval time.Duration
24
+
25
+ // policy controls
26
+ bannedLeases map[string]struct{}
27
+ namePattern *regexp.Regexp
28
+ minTTL time.Duration // 0 = no bound
29
+ maxTTL time.Duration // 0 = no bound
30
+ // per-lease byte limit
31
+ bpsLimits map[string]int64 // leaseID -> bytes-per-second (0 = unlimited)
32
+ defaultBPS int64 // default bytes-per-second for new/updated leases (0 = none)
33
}
34
35
func NewLeaseManager(ttlInterval time.Duration) *LeaseManager {
36
return &LeaseManager{
27
- leases: make(map[string]*LeaseEntry),
28
- stopCh: make(chan struct{}),
29
- ttlInterval: ttlInterval,
37
+ leases: make(map[string]*LeaseEntry),
38
+ stopCh: make(chan struct{}),
39
+ ttlInterval: ttlInterval,
40
+ bannedLeases: make(map[string]struct{}),
41
+ bpsLimits: make(map[string]int64),
42
+ defaultBPS: 0,
43
}
44
}
45
@@ -60,6 +73,7 @@ func (lm *LeaseManager) cleanupExpiredLeases() {
73
for id, lease := range lm.leases {
74
if now.After(lease.Expires) {
75
delete(lm.leases, id)
76
+ delete(lm.bpsLimits, id)
77
}
78
}
79
}
@@ -76,6 +90,24 @@ func (lm *LeaseManager) UpdateLease(lease *rdverb.Lease, connectionID int64) boo
90
return false
91
}
92
93
+ // policy checks
94
+ if _, banned := lm.bannedLeases[identityID]; banned {
95
+ return false
96
+ }
97
+ if lm.namePattern != nil && lease.Name != "" && !lm.namePattern.MatchString(lease.Name) {
98
+ return false
99
+ }
100
+ // reserved prefix check removed
101
+ if lm.minTTL > 0 || lm.maxTTL > 0 {
102
+ ttl := time.Until(expires)
103
+ if lm.minTTL > 0 && ttl < lm.minTTL {
104
+ return false
105
+ }
106
+ if lm.maxTTL > 0 && ttl > lm.maxTTL {
107
+ return false
108
+ }
109
+ }
110
+
111
// Check for name conflicts (only if name is not empty)
112
if lease.Name != "" && lease.Name != "(unnamed)" {
113
for existingID, existingEntry := range lm.leases {
@@ -98,6 +130,13 @@ func (lm *LeaseManager) UpdateLease(lease *rdverb.Lease, connectionID int64) boo
130
ConnectionID: connectionID,
131
}
132
133
+ // Apply default BPS limit for this lease if configured and no explicit limit set
134
+ if lm.defaultBPS > 0 {
135
+ if _, exists := lm.bpsLimits[identityID]; !exists {
136
+ lm.bpsLimits[identityID] = lm.defaultBPS
137
+ }
138
+ }
139
+
140
return true
141
}
142
@@ -108,6 +147,7 @@ func (lm *LeaseManager) DeleteLease(identity *rdsec.Identity) bool {
147
identityID := string(identity.Id)
148
if _, exists := lm.leases[identityID]; exists {
149
delete(lm.leases, identityID)
150
+ delete(lm.bpsLimits, identityID)
151
return true
152
}
153
return false
@@ -164,6 +204,43 @@ func (lm *LeaseManager) GetAllLeases() []*rdverb.Lease {
204
return validLeases
205
}
206
207
+// Lease policy configuration helpers
208
+func (lm *LeaseManager) BanLease(leaseID string) {
209
+ lm.leasesLock.Lock()
210
+ lm.bannedLeases[leaseID] = struct{}{}
211
+ lm.leasesLock.Unlock()
212
+}
213
+
214
+func (lm *LeaseManager) UnbanLease(leaseID string) {
215
+ lm.leasesLock.Lock()
216
+ delete(lm.bannedLeases, leaseID)
217
+ lm.leasesLock.Unlock()
218
+}
219
+
220
+func (lm *LeaseManager) SetNamePattern(pattern string) error {
221
+ lm.leasesLock.Lock()
222
+ defer lm.leasesLock.Unlock()
223
+ if pattern == "" {
224
+ lm.namePattern = nil
225
+ return nil
226
+ }
227
+ re, err := regexp.Compile(pattern)
228
+ if err != nil {
229
+ return err
230
+ }
231
+ lm.namePattern = re
232
+ return nil
233
+}
234
+
235
+// SetReservedPrefixes removed: reserved prefix policy no longer supported
236
+
237
+func (lm *LeaseManager) SetTTLBounds(min, max time.Duration) {
238
+ lm.leasesLock.Lock()
239
+ lm.minTTL = min
240
+ lm.maxTTL = max
241
+ lm.leasesLock.Unlock()
242
+}
243
+
244
func (lm *LeaseManager) CleanupLeasesByConnectionID(connectionID int64) []string {
245
lm.leasesLock.Lock()
246
defer lm.leasesLock.Unlock()
@@ -172,9 +249,41 @@ func (lm *LeaseManager) CleanupLeasesByConnectionID(connectionID int64) []string
249
for leaseID, lease := range lm.leases {
250
if lease.ConnectionID == connectionID {
251
delete(lm.leases, leaseID)
252
+ delete(lm.bpsLimits, leaseID)
253
cleanedLeaseIDs = append(cleanedLeaseIDs, leaseID)
254
}
255
}
256
257
return cleanedLeaseIDs
258
}
259
+
260
+// Per-lease BPS limit configuration
261
+func (lm *LeaseManager) SetBPSLimit(leaseID string, bps int64) {
262
+ lm.leasesLock.Lock()
263
+ defer lm.leasesLock.Unlock()
264
+ if bps <= 0 {
265
+ delete(lm.bpsLimits, leaseID)
266
+ return
267
+ }
268
+ lm.bpsLimits[leaseID] = bps
269
+}
270
+
271
+func (lm *LeaseManager) GetBPSLimit(leaseID string) int64 {
272
+ lm.leasesLock.RLock()
273
+ defer lm.leasesLock.RUnlock()
274
+ if v, ok := lm.bpsLimits[leaseID]; ok {
275
+ return v
276
+ }
277
+ return 0
278
+}
279
+
280
+// SetDefaultBPS sets a default bytes-per-second limit applied to leases on update/registration
281
+// If set to 0, no default is applied. Existing explicit per-lease limits are not overwritten.
282
+func (lm *LeaseManager) SetDefaultBPS(bps int64) {
283
+ lm.leasesLock.Lock()
284
+ defer lm.leasesLock.Unlock()
285
+ if bps < 0 {
286
+ bps = 0
287
+ }
288
+ lm.defaultBPS = bps
289
+}
portal/relay.go
+53
-7
@@ -10,6 +10,7 @@ import (
10
"gosuda.org/portal/portal/core/cryptoops"
11
"gosuda.org/portal/portal/core/proto/rdsec"
12
"gosuda.org/portal/portal/core/proto/rdverb"
13
+ "gosuda.org/portal/portal/utils/ratelimit"
14
)
15
16
type Connection struct {
@@ -39,6 +40,15 @@ type RelayServer struct {
40
41
stopch chan struct{}
42
waitgroup sync.WaitGroup
43
+
44
+ // Traffic control limits and counters
45
+ maxRelayedPerLease int
46
+ relayedPerLeaseCount map[string]int
47
+ limitsLock sync.Mutex
48
+
49
+ // Per-lease byte rate limit (BPS for relay throughput)
50
+ leaseBPS map[string]*ratelimit.Bucket
51
+ leaseBPSRate map[string]int64
52
}
53
54
func NewRelayServer(credential *cryptoops.Credential, address []string) *RelayServer {
@@ -48,18 +58,46 @@ func NewRelayServer(credential *cryptoops.Credential, address []string) *RelaySe
58
Id: credential.ID(),
59
PublicKey: credential.PublicKey(),
60
},
51
- address: address,
52
- connidCounter: 0,
53
- connections: make(map[int64]*Connection),
54
- leaseConnections: make(map[string]*Connection),
55
- relayedConnections: make(map[string][]*yamux.Stream),
56
- leaseManager: NewLeaseManager(30 * time.Second), // TTL check every 30 seconds
57
- stopch: make(chan struct{}),
61
+ address: address,
62
+ connidCounter: 0,
63
+ connections: make(map[int64]*Connection),
64
+ leaseConnections: make(map[string]*Connection),
65
+ relayedConnections: make(map[string][]*yamux.Stream),
66
+ leaseManager: NewLeaseManager(30 * time.Second), // TTL check every 30 seconds
67
+ stopch: make(chan struct{}),
68
+ relayedPerLeaseCount: make(map[string]int),
69
+ leaseBPS: make(map[string]*ratelimit.Bucket),
70
+ leaseBPSRate: make(map[string]int64),
71
}
72
}
73
74
var _yamux_config = yamux.DefaultConfig()
75
76
+func (g *RelayServer) getLeaseBPSBucket(leaseID string) *ratelimit.Bucket {
77
+ // Lookup per-lease BPS from LeaseManager (0 = unlimited)
78
+ bps := g.leaseManager.GetBPSLimit(leaseID)
79
+ if bps <= 0 {
80
+ return nil
81
+ }
82
+ desired := bps
83
+ g.limitsLock.Lock()
84
+ defer g.limitsLock.Unlock()
85
+ if b, ok := g.leaseBPS[leaseID]; ok {
86
+ if g.leaseBPSRate[leaseID] == desired {
87
+ return b
88
+ }
89
+ // replace with new rate
90
+ nb := ratelimit.NewBucket(desired, desired)
91
+ g.leaseBPS[leaseID] = nb
92
+ g.leaseBPSRate[leaseID] = desired
93
+ return nb
94
+ }
95
+ b := ratelimit.NewBucket(desired, desired)
96
+ g.leaseBPS[leaseID] = b
97
+ g.leaseBPSRate[leaseID] = desired
98
+ return b
99
+}
100
+
101
func (g *RelayServer) handleConn(id int64, connection *Connection) {
102
log.Debug().Int64("conn_id", id).Msg("[RelayServer] Handling new connection")
103
@@ -103,6 +141,7 @@ func (g *RelayServer) handleConn(id int64, connection *Connection) {
141
142
// Close the underlying connection
143
connection.conn.Close()
144
+
145
log.Debug().Int64("conn_id", id).Msg("[RelayServer] Connection cleanup complete")
146
}()
147
@@ -306,3 +345,10 @@ func (g *RelayServer) Stop() {
345
g.leaseManager.Stop()
346
g.waitgroup.Wait()
347
}
348
+
349
+// Traffic control setters
350
+func (g *RelayServer) SetMaxRelayedPerLease(n int) {
351
+ g.limitsLock.Lock()
352
+ g.maxRelayedPerLease = n
353
+ g.limitsLock.Unlock()
354
+}
portal/utils/ratelimit/bucket.go
new
+94
@@ -0,0 +1,94 @@
1
+package ratelimit
2
+
3
+import (
4
+ "io"
5
+ "sync"
6
+ "time"
7
+)
8
+
9
+// Bucket is a simple, precise byte-rate limiter (thread-safe).
10
+// All fields are integer-based (bytes, ns) to avoid float overhead.
11
+type Bucket struct {
12
+ mu sync.Mutex
13
+ rateBps int64 // bytes per second
14
+ capacity int64 // max tokens (burst), typically = rateBps
15
+ tokens int64 // current tokens in bytes
16
+ last time.Time // last refill time
17
+}
18
+
19
+// NewBucket creates a bucket with the given rate and burst (bytes).
20
+// If burst <= 0, it defaults to rateBps.
21
+func NewBucket(rateBps int64, burst int64) *Bucket {
22
+ if burst <= 0 {
23
+ burst = rateBps
24
+ }
25
+ return &Bucket{rateBps: rateBps, capacity: burst, tokens: burst, last: time.Now()}
26
+}
27
+
28
+// Take blocks until n bytes worth of tokens are available, then consumes them.
29
+func (b *Bucket) Take(n int64) {
30
+ for {
31
+ var sleep time.Duration
32
+ b.mu.Lock()
33
+ now := time.Now()
34
+ elapsed := now.Sub(b.last)
35
+ if elapsed > 0 {
36
+ refill := (elapsed.Nanoseconds() * b.rateBps) / int64(time.Second)
37
+ if refill > 0 {
38
+ b.tokens += refill
39
+ if b.tokens > b.capacity {
40
+ b.tokens = b.capacity
41
+ }
42
+ b.last = now
43
+ }
44
+ }
45
+ if b.tokens >= n {
46
+ b.tokens -= n
47
+ b.mu.Unlock()
48
+ return
49
+ }
50
+ deficit := n - b.tokens
51
+ nsNeeded := max((deficit*int64(time.Second))/b.rateBps, int64(time.Millisecond))
52
+ sleep = time.Duration(nsNeeded)
53
+ b.mu.Unlock()
54
+ time.Sleep(sleep)
55
+ }
56
+}
57
+
58
+// internal buffer pool for Copy
59
+var bufPool = sync.Pool{New: func() any { return make([]byte, 64*1024) }}
60
+
61
+// Copy copies from src to dst, enforcing the provided byte-rate bucket if not nil.
62
+// Returns bytes written and any copy error encountered.
63
+func Copy(dst io.Writer, src io.Reader, b *Bucket) (int64, error) {
64
+ if b == nil {
65
+ return io.Copy(dst, src)
66
+ }
67
+ buf := bufPool.Get().([]byte)
68
+ defer bufPool.Put(buf)
69
+
70
+ var total int64
71
+ for {
72
+ nr, er := src.Read(buf)
73
+ if nr > 0 {
74
+ b.Take(int64(nr))
75
+ nw, ew := dst.Write(buf[:nr])
76
+ if nw > 0 {
77
+ total += int64(nw)
78
+ }
79
+ if ew != nil {
80
+ return total, ew
81
+ }
82
+ if nr != nw {
83
+ return total, io.ErrShortWrite
84
+ }
85
+ }
86
+ if er != nil {
87
+ if er == io.EOF {
88
+ break
89
+ }
90
+ return total, er
91
+ }
92
+ }
93
+ return total, nil
94
+}