Revert "refactor(lease, relay): implement event loops for LeaseManager and RelayServer to centralize state and remove explicit locks"
This reverts commit 50c5584555292efd9372aef761d3ca7cd44f4aa2.
cognitive-glitch committed
Dec 9, 2025 at 13:33 UTC
7c19f8650f9fa699206f0eb96dd6217636e3e1f8
4 files changed
+236
-544
portal/handlers.go
+41
-39
@@ -66,11 +66,11 @@ func (g *RelayServer) handleLeaseUpdateRequest(ctx *StreamContext, packet *rdver
66
if g.leaseManager.UpdateLease(req.Lease, ctx.ConnectionID) {
67
resp.Code = rdverb.ResponseCode_RESPONSE_CODE_ACCEPTED
68
69
- // Register lease connection via event loop
69
+ // Register lease connection
70
leaseID := string(req.Lease.Identity.Id)
71
- done := make(chan struct{}, 1)
72
- g.cmdCh <- &cmdRegisterLeaseConn{leaseID: leaseID, conn: ctx.Connection, done: done}
73
- <-done
71
+ g.leaseConnectionsLock.Lock()
72
+ g.leaseConnections[leaseID] = ctx.Connection
73
+ g.leaseConnectionsLock.Unlock()
74
75
// Log lease update completion
76
log.Debug().
@@ -123,11 +123,11 @@ func (g *RelayServer) handleLeaseDeleteRequest(ctx *StreamContext, packet *rdver
123
if g.leaseManager.DeleteLease(req.Identity) {
124
resp.Code = rdverb.ResponseCode_RESPONSE_CODE_ACCEPTED
125
126
- // Remove lease connection via event loop
126
+ // Remove lease connection
127
leaseID := string(req.Identity.Id)
128
- done := make(chan struct{}, 1)
129
- g.cmdCh <- &cmdUnregisterLeaseConn{leaseID: leaseID, done: done}
130
- <-done
128
+ g.leaseConnectionsLock.Lock()
129
+ delete(g.leaseConnections, leaseID)
130
+ g.leaseConnectionsLock.Unlock()
131
132
// Log lease deletion completion
133
log.Debug().
@@ -168,12 +168,12 @@ func (g *RelayServer) handleConnectionRequest(ctx *StreamContext, packet *rdverb
168
return g.sendConnectionResponse(ctx.Stream, rdverb.ResponseCode_RESPONSE_CODE_INVALID_IDENTITY)
169
}
170
171
- // Get the lease connection via event loop
172
- reply := make(chan *Connection, 1)
173
- g.cmdCh <- &cmdGetConnByLeaseEntry{connID: leaseEntry.ConnectionID, reply: reply}
174
- leaseConn := <-reply
171
+ // Get the lease connection
172
+ g.connectionsLock.RLock()
173
+ leaseConn, leaseExists := g.connections[leaseEntry.ConnectionID]
174
+ g.connectionsLock.RUnlock()
175
176
- if leaseConn == nil {
176
+ if !leaseExists {
177
return g.sendConnectionResponse(ctx.Stream, rdverb.ResponseCode_RESPONSE_CODE_INVALID_IDENTITY)
178
}
179
@@ -190,12 +190,11 @@ func (g *RelayServer) handleConnectionRequest(ctx *StreamContext, packet *rdverb
190
// Enforce relayed connection limits
191
if respCode == rdverb.ResponseCode_RESPONSE_CODE_ACCEPTED {
192
leaseID := string(leaseEntry.Lease.Identity.Id)
193
- // Atomically check and increment limit via event loop
194
- reply := make(chan bool, 1)
195
- g.cmdCh <- &cmdCheckAndIncLimit{leaseID: leaseID, reply: reply}
196
- underLimit := <-reply
193
+ g.limitsLock.Lock()
194
+ overPerLease := g.maxRelayedPerLease > 0 && g.relayedPerLeaseCount[leaseID] >= g.maxRelayedPerLease
195
+ g.limitsLock.Unlock()
196
198
- if !underLimit {
197
+ if overPerLease {
198
log.Warn().Str("lease_id", leaseID).Msg("[RelayServer] Relayed connection per-lease limit reached")
199
respCode = rdverb.ResponseCode_RESPONSE_CODE_REJECTED
200
leaseStream.Close()
@@ -204,13 +203,6 @@ func (g *RelayServer) handleConnectionRequest(ctx *StreamContext, packet *rdverb
203
204
// Send response to client
205
if err := g.sendConnectionResponse(ctx.Stream, respCode); err != nil {
207
- // If we already incremented the limit, decrement it
208
- if respCode == rdverb.ResponseCode_RESPONSE_CODE_ACCEPTED {
209
- leaseID := string(leaseEntry.Lease.Identity.Id)
210
- done := make(chan struct{}, 1)
211
- g.cmdCh <- &cmdDecLimit{leaseID: leaseID, done: done}
212
- <-done
213
- }
206
leaseStream.Close()
207
return err
208
}
@@ -281,23 +273,33 @@ func (g *RelayServer) sendConnectionResponse(stream *yamux.Stream, code rdverb.R
273
}
274
275
func (g *RelayServer) establishRelayedConnection(clientStream, leaseStream *yamux.Stream, leaseID string) {
284
- // Register stream for tracking via event loop
285
- // Note: limit was already incremented in handleConnectionRequest
286
- done := make(chan struct{}, 1)
287
- g.cmdCh <- &cmdAddRelayed{leaseID: leaseID, stream: clientStream, done: done}
288
- <-done
276
+ // Register connection for tracking
277
+ g.limitsLock.Lock()
278
+ g.relayedPerLeaseCount[leaseID]++
279
+ g.limitsLock.Unlock()
280
+
281
+ g.relayedConnectionsLock.Lock()
282
+ g.relayedConnections[leaseID] = append(g.relayedConnections[leaseID], clientStream)
283
+ g.relayedConnectionsLock.Unlock()
284
285
// Cleanup function
286
defer func() {
292
- // Decrement limit via event loop
293
- done := make(chan struct{}, 1)
294
- g.cmdCh <- &cmdDecLimit{leaseID: leaseID, done: done}
295
- <-done
296
-
297
- // Remove stream from tracking via event loop
298
- done2 := make(chan struct{}, 1)
299
- g.cmdCh <- &cmdRemoveRelayed{leaseID: leaseID, stream: clientStream, done: done2}
300
- <-done2
287
+ g.limitsLock.Lock()
288
+ if g.relayedPerLeaseCount[leaseID] > 0 {
289
+ g.relayedPerLeaseCount[leaseID]--
290
+ }
291
+ g.limitsLock.Unlock()
292
+
293
+ g.relayedConnectionsLock.Lock()
294
+ if streams, ok := g.relayedConnections[leaseID]; ok {
295
+ for i, s := range streams {
296
+ if s == clientStream {
297
+ g.relayedConnections[leaseID] = append(streams[:i], streams[i+1:]...)
298
+ break
299
+ }
300
+ }
301
+ }
302
+ g.relayedConnectionsLock.Unlock()
303
}()
304
305
// Use callback for actual relay (handles BPS limiting in relay-server)
portal/lease.go
+118
-296
@@ -25,114 +25,16 @@ type LeaseEntry struct {
25
Lease *rdverb.Lease
26
Expires time.Time
27
LastSeen time.Time
28
- ConnectionID int64 // Store the connection ID
28
+ ConnectionID int64
29
ParsedMetadata *ParsedMetadata // Cached parsed metadata
30
}
31
32
-// leaseCmd is the command interface for LeaseManager event loop.
33
-type leaseCmd interface{ leaseCmd() }
34
-
35
-type cmdUpdateLease struct {
36
- lease *rdverb.Lease
37
- connID int64
38
- reply chan<- bool
39
-}
40
-
41
-type cmdDeleteLease struct {
42
- identity *rdsec.Identity
43
- reply chan<- bool
44
-}
45
-
46
-type cmdGetLease struct {
47
- identity *rdsec.Identity
48
- reply chan<- leaseResult
49
-}
50
-
51
-type cmdGetLeaseByID struct {
52
- leaseID string
53
- reply chan<- leaseResult
54
-}
55
-
56
-type cmdGetAllLeases struct {
57
- reply chan<- []*rdverb.Lease
58
-}
59
-
60
-type cmdCleanupByConnID struct {
61
- connID int64
62
- reply chan<- []string
63
-}
64
-
65
-type cmdBan struct {
66
- leaseID string
67
- done chan<- struct{}
68
-}
69
-
70
-type cmdUnban struct {
71
- leaseID string
72
- done chan<- struct{}
73
-}
74
-
75
-type cmdGetBanned struct {
76
- reply chan<- [][]byte
77
-}
78
-
79
-type cmdSetPattern struct {
80
- pattern string
81
- reply chan<- error
82
-}
83
-
84
-type cmdSetTTL struct {
85
- min time.Duration
86
- max time.Duration
87
- done chan<- struct{}
88
-}
89
-
90
-type cmdCleanExpired struct {
91
- done chan<- struct{}
92
-}
93
-
94
-type cmdGetAllEntries struct {
95
- reply chan<- []*LeaseEntry
96
-}
97
-
98
-type cmdGetLeaseALPNs struct {
99
- leaseID string
100
- reply chan<- []string
101
-}
102
-
103
-// leaseResult carries a LeaseEntry lookup result.
104
-type leaseResult struct {
105
- entry *LeaseEntry
106
- exists bool
107
-}
108
-
109
-// Command interface markers
110
-func (cmdUpdateLease) leaseCmd() {}
111
-func (cmdDeleteLease) leaseCmd() {}
112
-func (cmdGetLease) leaseCmd() {}
113
-func (cmdGetLeaseByID) leaseCmd() {}
114
-func (cmdGetAllLeases) leaseCmd() {}
115
-func (cmdCleanupByConnID) leaseCmd() {}
116
-func (cmdBan) leaseCmd() {}
117
-func (cmdUnban) leaseCmd() {}
118
-func (cmdGetBanned) leaseCmd() {}
119
-func (cmdSetPattern) leaseCmd() {}
120
-func (cmdSetTTL) leaseCmd() {}
121
-func (cmdCleanExpired) leaseCmd() {}
122
-func (cmdGetAllEntries) leaseCmd() {}
123
-func (cmdGetLeaseALPNs) leaseCmd() {}
124
-
125
-// LeaseManager manages lease registrations using a single-threaded event loop.
126
-// All state mutations are processed sequentially via the command channel.
32
type LeaseManager struct {
33
leases map[string]*LeaseEntry // Key: identity ID
34
+ leasesLock sync.RWMutex
35
stopCh chan struct{}
36
ttlInterval time.Duration
37
132
- // Event loop
133
- cmdCh chan leaseCmd
134
- runWg sync.WaitGroup
135
-
38
// policy controls
39
bannedLeases map[string]struct{}
40
namePattern *regexp.Regexp
@@ -145,156 +47,48 @@ func NewLeaseManager(ttlInterval time.Duration) *LeaseManager {
47
leases: make(map[string]*LeaseEntry),
48
stopCh: make(chan struct{}),
49
ttlInterval: ttlInterval,
148
- cmdCh: make(chan leaseCmd, 256),
50
bannedLeases: make(map[string]struct{}),
51
}
52
}
53
54
func (lm *LeaseManager) Start() {
154
- lm.runWg.Add(1)
155
- go lm.run()
55
go lm.ttlWorker()
56
}
57
58
func (lm *LeaseManager) Stop() {
59
close(lm.stopCh)
161
- lm.runWg.Wait()
60
}
61
164
-// run is the main event loop that processes all commands sequentially.
165
-func (lm *LeaseManager) run() {
166
- defer lm.runWg.Done()
62
+func (lm *LeaseManager) ttlWorker() {
63
+ ticker := time.NewTicker(lm.ttlInterval)
64
+ defer ticker.Stop()
65
+
66
for {
67
select {
169
- case cmd := <-lm.cmdCh:
170
- lm.handleCmd(cmd)
68
+ case <-ticker.C:
69
+ lm.cleanupExpiredLeases()
70
case <-lm.stopCh:
71
return
72
}
73
}
74
}
75
177
-// handleCmd dispatches commands to their handlers.
178
-func (lm *LeaseManager) handleCmd(cmd leaseCmd) {
179
- switch c := cmd.(type) {
180
- case *cmdUpdateLease:
181
- c.reply <- lm.updateLeaseInternal(c.lease, c.connID)
182
-
183
- case *cmdDeleteLease:
184
- identityID := string(c.identity.Id)
185
- if _, exists := lm.leases[identityID]; exists {
186
- delete(lm.leases, identityID)
187
- c.reply <- true
188
- } else {
189
- c.reply <- false
190
- }
191
-
192
- case *cmdGetLease:
193
- identityID := string(c.identity.Id)
194
- entry, exists := lm.leases[identityID]
195
- if !exists || time.Now().After(entry.Expires) {
196
- c.reply <- leaseResult{nil, false}
197
- } else {
198
- c.reply <- leaseResult{entry, true}
199
- }
200
-
201
- case *cmdGetLeaseByID:
202
- if _, banned := lm.bannedLeases[c.leaseID]; banned {
203
- c.reply <- leaseResult{nil, false}
204
- return
205
- }
206
- entry, exists := lm.leases[c.leaseID]
207
- if !exists || time.Now().After(entry.Expires) {
208
- c.reply <- leaseResult{nil, false}
209
- } else {
210
- c.reply <- leaseResult{entry, true}
211
- }
212
-
213
- case *cmdGetAllLeases:
214
- now := time.Now()
215
- var valid []*rdverb.Lease
216
- for _, entry := range lm.leases {
217
- if now.Before(entry.Expires) {
218
- valid = append(valid, entry.Lease)
219
- }
220
- }
221
- c.reply <- valid
222
-
223
- case *cmdCleanupByConnID:
224
- var cleaned []string
225
- for leaseID, entry := range lm.leases {
226
- if entry.ConnectionID == c.connID {
227
- delete(lm.leases, leaseID)
228
- cleaned = append(cleaned, leaseID)
229
- }
230
- }
231
- c.reply <- cleaned
232
-
233
- case *cmdBan:
234
- lm.bannedLeases[c.leaseID] = struct{}{}
235
- c.done <- struct{}{}
236
-
237
- case *cmdUnban:
238
- delete(lm.bannedLeases, c.leaseID)
239
- c.done <- struct{}{}
240
-
241
- case *cmdGetBanned:
242
- banned := make([][]byte, 0, len(lm.bannedLeases))
243
- for id := range lm.bannedLeases {
244
- banned = append(banned, []byte(id))
245
- }
246
- c.reply <- banned
76
+func (lm *LeaseManager) cleanupExpiredLeases() {
77
+ lm.leasesLock.Lock()
78
+ defer lm.leasesLock.Unlock()
79
248
- case *cmdSetPattern:
249
- if c.pattern == "" {
250
- lm.namePattern = nil
251
- c.reply <- nil
252
- return
253
- }
254
- re, err := regexp.Compile(c.pattern)
255
- if err != nil {
256
- c.reply <- err
257
- return
258
- }
259
- lm.namePattern = re
260
- c.reply <- nil
261
-
262
- case *cmdSetTTL:
263
- lm.minTTL = c.min
264
- lm.maxTTL = c.max
265
- c.done <- struct{}{}
266
-
267
- case *cmdCleanExpired:
268
- now := time.Now()
269
- for id, entry := range lm.leases {
270
- if now.After(entry.Expires) {
271
- delete(lm.leases, id)
272
- }
273
- }
274
- c.done <- struct{}{}
275
-
276
- case *cmdGetAllEntries:
277
- now := time.Now()
278
- var entries []*LeaseEntry
279
- for _, entry := range lm.leases {
280
- if now.Before(entry.Expires) {
281
- entries = append(entries, entry)
282
- }
283
- }
284
- c.reply <- entries
285
-
286
- case *cmdGetLeaseALPNs:
287
- entry, exists := lm.leases[c.leaseID]
288
- if !exists || time.Now().After(entry.Expires) {
289
- c.reply <- nil
290
- } else {
291
- c.reply <- entry.Lease.Alpn
80
+ now := time.Now()
81
+ for id, lease := range lm.leases {
82
+ if now.After(lease.Expires) {
83
+ delete(lm.leases, id)
84
}
85
}
86
}
87
296
-// updateLeaseInternal contains the lease update logic, called within the event loop.
297
-func (lm *LeaseManager) updateLeaseInternal(lease *rdverb.Lease, connectionID int64) bool {
88
+func (lm *LeaseManager) UpdateLease(lease *rdverb.Lease, connectionID int64) bool {
89
+ lm.leasesLock.Lock()
90
+ defer lm.leasesLock.Unlock()
91
+
92
identityID := string(lease.Identity.Id)
93
expires := time.Unix(lease.Expires, 0)
94
@@ -310,6 +104,7 @@ func (lm *LeaseManager) updateLeaseInternal(lease *rdverb.Lease, connectionID in
104
if lm.namePattern != nil && lease.Name != "" && !lm.namePattern.MatchString(lease.Name) {
105
return false
106
}
107
+ // reserved prefix check removed
108
if lm.minTTL > 0 || lm.maxTTL > 0 {
109
ttl := time.Until(expires)
110
if lm.minTTL > 0 && ttl < lm.minTTL {
@@ -323,10 +118,13 @@ func (lm *LeaseManager) updateLeaseInternal(lease *rdverb.Lease, connectionID in
118
// Check for name conflicts (only if name is not empty)
119
if lease.Name != "" && lease.Name != "(unnamed)" {
120
for existingID, existingEntry := range lm.leases {
121
+ // Skip if it's the same identity (updating own lease)
122
if existingID == identityID {
123
continue
124
}
125
+ // Check if another identity is using the same name
126
if existingEntry.Lease.Name == lease.Name {
127
+ // Name conflict with a different identity
128
return false
129
}
130
}
@@ -364,109 +162,133 @@ func (lm *LeaseManager) updateLeaseInternal(lease *rdverb.Lease, connectionID in
162
return true
163
}
164
367
-func (lm *LeaseManager) ttlWorker() {
368
- ticker := time.NewTicker(lm.ttlInterval)
369
- defer ticker.Stop()
165
+func (lm *LeaseManager) DeleteLease(identity *rdsec.Identity) bool {
166
+ lm.leasesLock.Lock()
167
+ defer lm.leasesLock.Unlock()
168
371
- for {
372
- select {
373
- case <-ticker.C:
374
- done := make(chan struct{}, 1)
375
- select {
376
- case lm.cmdCh <- &cmdCleanExpired{done: done}:
377
- <-done
378
- case <-lm.stopCh:
379
- return
380
- }
381
- case <-lm.stopCh:
382
- return
383
- }
169
+ identityID := string(identity.Id)
170
+ if _, exists := lm.leases[identityID]; exists {
171
+ delete(lm.leases, identityID)
172
+ return true
173
}
174
+ return false
175
}
176
387
-func (lm *LeaseManager) UpdateLease(lease *rdverb.Lease, connectionID int64) bool {
388
- reply := make(chan bool, 1)
389
- lm.cmdCh <- &cmdUpdateLease{lease: lease, connID: connectionID, reply: reply}
390
- return <-reply
391
-}
177
+func (lm *LeaseManager) GetLease(identity *rdsec.Identity) (*LeaseEntry, bool) {
178
+ lm.leasesLock.RLock()
179
+ defer lm.leasesLock.RUnlock()
180
393
-func (lm *LeaseManager) DeleteLease(identity *rdsec.Identity) bool {
394
- reply := make(chan bool, 1)
395
- lm.cmdCh <- &cmdDeleteLease{identity: identity, reply: reply}
396
- return <-reply
397
-}
181
+ identityID := string(identity.Id)
182
399
-func (lm *LeaseManager) GetLease(identity *rdsec.Identity) (*LeaseEntry, bool) {
400
- reply := make(chan leaseResult, 1)
401
- lm.cmdCh <- &cmdGetLease{identity: identity, reply: reply}
402
- result := <-reply
403
- return result.entry, result.exists
183
+ lease, exists := lm.leases[identityID]
184
+ if !exists {
185
+ return nil, false
186
+ }
187
+
188
+ // Check if lease is expired
189
+ if time.Now().After(lease.Expires) {
190
+ return nil, false
191
+ }
192
+
193
+ return lease, true
194
}
195
196
func (lm *LeaseManager) GetLeaseByID(leaseID string) (*LeaseEntry, bool) {
407
- reply := make(chan leaseResult, 1)
408
- lm.cmdCh <- &cmdGetLeaseByID{leaseID: leaseID, reply: reply}
409
- result := <-reply
410
- return result.entry, result.exists
197
+ lm.leasesLock.RLock()
198
+ defer lm.leasesLock.RUnlock()
199
+
200
+ // Check if banned
201
+ if _, banned := lm.bannedLeases[leaseID]; banned {
202
+ return nil, false
203
+ }
204
+
205
+ lease, exists := lm.leases[leaseID]
206
+ if !exists {
207
+ return nil, false
208
+ }
209
+
210
+ // Check if lease is expired
211
+ if time.Now().After(lease.Expires) {
212
+ return nil, false
213
+ }
214
+
215
+ return lease, true
216
}
217
218
func (lm *LeaseManager) GetAllLeases() []*rdverb.Lease {
414
- reply := make(chan []*rdverb.Lease, 1)
415
- lm.cmdCh <- &cmdGetAllLeases{reply: reply}
416
- return <-reply
219
+ lm.leasesLock.RLock()
220
+ defer lm.leasesLock.RUnlock()
221
+
222
+ now := time.Now()
223
+ var validLeases []*rdverb.Lease
224
+
225
+ for _, lease := range lm.leases {
226
+ if now.Before(lease.Expires) {
227
+ validLeases = append(validLeases, lease.Lease)
228
+ }
229
+ }
230
+
231
+ return validLeases
232
}
233
234
// Lease policy configuration helpers
235
func (lm *LeaseManager) BanLease(leaseID string) {
421
- done := make(chan struct{}, 1)
422
- lm.cmdCh <- &cmdBan{leaseID: leaseID, done: done}
423
- <-done
236
+ lm.leasesLock.Lock()
237
+ lm.bannedLeases[leaseID] = struct{}{}
238
+ lm.leasesLock.Unlock()
239
}
240
241
func (lm *LeaseManager) UnbanLease(leaseID string) {
427
- done := make(chan struct{}, 1)
428
- lm.cmdCh <- &cmdUnban{leaseID: leaseID, done: done}
429
- <-done
242
+ lm.leasesLock.Lock()
243
+ delete(lm.bannedLeases, leaseID)
244
+ lm.leasesLock.Unlock()
245
}
246
247
func (lm *LeaseManager) GetBannedLeases() [][]byte {
433
- reply := make(chan [][]byte, 1)
434
- lm.cmdCh <- &cmdGetBanned{reply: reply}
435
- return <-reply
248
+ lm.leasesLock.RLock()
249
+ defer lm.leasesLock.RUnlock()
250
+ banned := make([][]byte, 0, len(lm.bannedLeases))
251
+ for id := range lm.bannedLeases {
252
+ banned = append(banned, []byte(id))
253
+ }
254
+ return banned
255
}
256
257
func (lm *LeaseManager) SetNamePattern(pattern string) error {
439
- reply := make(chan error, 1)
440
- lm.cmdCh <- &cmdSetPattern{pattern: pattern, reply: reply}
441
- return <-reply
258
+ lm.leasesLock.Lock()
259
+ defer lm.leasesLock.Unlock()
260
+ if pattern == "" {
261
+ lm.namePattern = nil
262
+ return nil
263
+ }
264
+ re, err := regexp.Compile(pattern)
265
+ if err != nil {
266
+ return err
267
+ }
268
+ lm.namePattern = re
269
+ return nil
270
}
271
272
// SetReservedPrefixes removed: reserved prefix policy no longer supported
273
274
func (lm *LeaseManager) SetTTLBounds(min, max time.Duration) {
447
- done := make(chan struct{}, 1)
448
- lm.cmdCh <- &cmdSetTTL{min: min, max: max, done: done}
449
- <-done
275
+ lm.leasesLock.Lock()
276
+ lm.minTTL = min
277
+ lm.maxTTL = max
278
+ lm.leasesLock.Unlock()
279
}
280
281
func (lm *LeaseManager) CleanupLeasesByConnectionID(connectionID int64) []string {
453
- reply := make(chan []string, 1)
454
- lm.cmdCh <- &cmdCleanupByConnID{connID: connectionID, reply: reply}
455
- return <-reply
456
-}
457
-
458
-// GetAllEntries returns all valid (non-expired) lease entries.
459
-// This method is used by RelayServer.GetAllLeaseEntries.
460
-func (lm *LeaseManager) GetAllEntries() []*LeaseEntry {
461
- reply := make(chan []*LeaseEntry, 1)
462
- lm.cmdCh <- &cmdGetAllEntries{reply: reply}
463
- return <-reply
464
-}
282
+ lm.leasesLock.Lock()
283
+ defer lm.leasesLock.Unlock()
284
+
285
+ var cleanedLeaseIDs []string
286
+ for leaseID, lease := range lm.leases {
287
+ if lease.ConnectionID == connectionID {
288
+ delete(lm.leases, leaseID)
289
+ cleanedLeaseIDs = append(cleanedLeaseIDs, leaseID)
290
+ }
291
+ }
292
466
-// GetLeaseALPNs returns the ALPN identifiers for a given lease ID.
467
-// This method is used by RelayServer.GetLeaseALPNs.
468
-func (lm *LeaseManager) GetLeaseALPNs(leaseID string) []string {
469
- reply := make(chan []string, 1)
470
- lm.cmdCh <- &cmdGetLeaseALPNs{leaseID: leaseID, reply: reply}
471
- return <-reply
293
+ return cleanedLeaseIDs
294
}
portal/lease_test.go
-5
@@ -10,7 +10,6 @@ import (
10
11
func TestLeaseManager_NameConflict(t *testing.T) {
12
lm := NewLeaseManager(30 * time.Second)
13
- lm.Start()
13
defer lm.Stop()
14
15
// Create two different identities
@@ -68,7 +67,6 @@ func TestLeaseManager_NameConflict(t *testing.T) {
67
68
func TestLeaseManager_SameIdentityUpdate(t *testing.T) {
69
lm := NewLeaseManager(30 * time.Second)
71
- lm.Start()
70
defer lm.Stop()
71
72
identity := &rdsec.Identity{
@@ -114,7 +112,6 @@ func TestLeaseManager_SameIdentityUpdate(t *testing.T) {
112
113
func TestLeaseManager_EmptyNameAllowed(t *testing.T) {
114
lm := NewLeaseManager(30 * time.Second)
117
- lm.Start()
115
defer lm.Stop()
116
117
identity1 := &rdsec.Identity{
@@ -153,7 +150,6 @@ func TestLeaseManager_EmptyNameAllowed(t *testing.T) {
150
151
func TestLeaseManager_UnnamedAllowed(t *testing.T) {
152
lm := NewLeaseManager(30 * time.Second)
156
- lm.Start()
153
defer lm.Stop()
154
155
identity1 := &rdsec.Identity{
@@ -192,7 +188,6 @@ func TestLeaseManager_UnnamedAllowed(t *testing.T) {
188
189
func TestLeaseManager_UnicodeNameConflict(t *testing.T) {
190
lm := NewLeaseManager(30 * time.Second)
195
- lm.Start()
191
defer lm.Stop()
192
193
identity1 := &rdsec.Identity{
portal/relay.go
+77
-204
@@ -3,7 +3,6 @@ package portal
3
import (
4
"io"
5
"sync"
6
- "sync/atomic"
6
"time"
7
8
"github.com/hashicorp/yamux"
@@ -21,102 +20,22 @@ type Connection struct {
20
streamsLock sync.Mutex
21
}
22
24
-// relayCmd is the command interface for RelayServer event loop.
25
-type relayCmd interface{ relayCmd() }
26
-
27
-// Connection management commands
28
-type cmdRegisterConn struct {
29
- id int64
30
- conn *Connection
31
- done chan<- struct{}
32
-}
33
-
34
-type cmdRemoveConn struct {
35
- id int64
36
- done chan<- []string // returns cleaned lease IDs
37
-}
38
-
39
-type cmdIsConnActive struct {
40
- id int64
41
- reply chan<- bool
42
-}
43
-
44
-type cmdGetConnByLeaseEntry struct {
45
- connID int64
46
- reply chan<- *Connection
47
-}
48
-
49
-// Lease connection management commands
50
-type cmdRegisterLeaseConn struct {
51
- leaseID string
52
- conn *Connection
53
- done chan<- struct{}
54
-}
55
-
56
-type cmdUnregisterLeaseConn struct {
57
- leaseID string
58
- done chan<- struct{}
59
-}
60
-
61
-// Relayed connection tracking commands
62
-type cmdAddRelayed struct {
63
- leaseID string
64
- stream *yamux.Stream
65
- done chan<- struct{}
66
-}
67
-
68
-type cmdRemoveRelayed struct {
69
- leaseID string
70
- stream *yamux.Stream
71
- done chan<- struct{}
72
-}
73
-
74
-// Limit management commands
75
-type cmdCheckAndIncLimit struct {
76
- leaseID string
77
- reply chan<- bool // true if under limit and incremented
78
-}
79
-
80
-type cmdDecLimit struct {
81
- leaseID string
82
- done chan<- struct{}
83
-}
84
-
85
-type cmdSetMaxRelayed struct {
86
- max int
87
- done chan<- struct{}
88
-}
89
-
90
-// Command interface markers
91
-func (cmdRegisterConn) relayCmd() {}
92
-func (cmdRemoveConn) relayCmd() {}
93
-func (cmdIsConnActive) relayCmd() {}
94
-func (cmdGetConnByLeaseEntry) relayCmd() {}
95
-func (cmdRegisterLeaseConn) relayCmd() {}
96
-func (cmdUnregisterLeaseConn) relayCmd() {}
97
-func (cmdAddRelayed) relayCmd() {}
98
-func (cmdRemoveRelayed) relayCmd() {}
99
-func (cmdCheckAndIncLimit) relayCmd() {}
100
-func (cmdDecLimit) relayCmd() {}
101
-func (cmdSetMaxRelayed) relayCmd() {}
102
-
103
-// RelayServer handles relay connections using a single-threaded event loop.
104
-// All state mutations are processed sequentially via the command channel.
23
type RelayServer struct {
24
credential *cryptoops.Credential
25
identity *rdsec.Identity
26
address []string
27
110
- connidCounter int64
111
- connections map[int64]*Connection
112
- leaseConnections map[string]*Connection // Key: lease ID, Value: Connection
113
- relayedConnections map[string][]*yamux.Stream // Key: lease ID, Value: slice of relayed streams
28
+ connidCounter int64
29
+ connections map[int64]*Connection
30
+ connectionsLock sync.RWMutex
31
115
- leaseManager *LeaseManager
32
+ leaseConnections map[string]*Connection // Key: lease ID, Value: Connection
33
+ leaseConnectionsLock sync.RWMutex
34
+
35
+ relayedConnections map[string][]*yamux.Stream // Key: lease ID, Value: slice of relayed streams
36
+ relayedConnectionsLock sync.RWMutex
37
117
- // Event loop
118
- cmdCh chan relayCmd
119
- runWg sync.WaitGroup
38
+ leaseManager *LeaseManager
39
40
stopch chan struct{}
41
waitgroup sync.WaitGroup
@@ -124,6 +43,7 @@ type RelayServer struct {
43
// Traffic control limits and counters
44
maxRelayedPerLease int
45
relayedPerLeaseCount map[string]int
46
+ limitsLock sync.Mutex
47
48
// Callback for relay connection establishment (set by relay-server for BPS handling)
49
onEstablishRelay func(clientStream, leaseStream *yamux.Stream, leaseID string)
@@ -142,7 +62,6 @@ func NewRelayServer(credential *cryptoops.Credential, address []string) *RelaySe
62
leaseConnections: make(map[string]*Connection),
63
relayedConnections: make(map[string][]*yamux.Stream),
64
leaseManager: NewLeaseManager(30 * time.Second), // TTL check every 30 seconds
145
- cmdCh: make(chan relayCmd, 256),
65
stopch: make(chan struct{}),
66
relayedPerLeaseCount: make(map[string]int),
67
}
@@ -162,10 +81,8 @@ func (g *RelayServer) handleConn(id int64, connection *Connection) {
81
defer func() {
82
log.Debug().Int64("conn_id", id).Msg("[RelayServer] Connection closing, cleaning up")
83
165
- // Remove connection and clean up all associated state via event loop
166
- done := make(chan []string, 1)
167
- g.cmdCh <- &cmdRemoveConn{id: id, done: done}
168
- cleanedLeaseIDs := <-done
84
+ // Clean up leases associated with this connection when it closes
85
+ cleanedLeaseIDs := g.leaseManager.CleanupLeasesByConnectionID(id)
86
87
if len(cleanedLeaseIDs) > 0 {
88
log.Debug().
@@ -174,6 +91,31 @@ func (g *RelayServer) handleConn(id int64, connection *Connection) {
91
Msg("[RelayServer] Cleaned up leases for connection")
92
}
93
94
+ // Also clean up lease connections mapping
95
+ g.leaseConnectionsLock.Lock()
96
+ for _, leaseID := range cleanedLeaseIDs {
97
+ delete(g.leaseConnections, leaseID)
98
+ }
99
+ g.leaseConnectionsLock.Unlock()
100
+
101
+ // Clean up relayed connections for these leases
102
+ g.relayedConnectionsLock.Lock()
103
+ for _, leaseID := range cleanedLeaseIDs {
104
+ if streams, exists := g.relayedConnections[leaseID]; exists {
105
+ // Close all relayed streams
106
+ for _, stream := range streams {
107
+ stream.Close()
108
+ }
109
+ delete(g.relayedConnections, leaseID)
110
+ }
111
+ }
112
+ g.relayedConnectionsLock.Unlock()
113
+
114
+ // Remove the connection itself
115
+ g.connectionsLock.Lock()
116
+ delete(g.connections, id)
117
+ g.connectionsLock.Unlock()
118
+
119
// Close the underlying connection
120
connection.conn.Close()
121
@@ -297,17 +239,16 @@ func (g *RelayServer) HandleConnection(conn io.ReadWriteCloser) error {
239
return err
240
}
241
300
- connID := atomic.AddInt64(&g.connidCounter, 1)
242
+ g.connectionsLock.Lock()
243
+ g.connidCounter++
244
+ connID := g.connidCounter
245
connection := &Connection{
246
conn: conn,
247
sess: sess,
248
streams: make(map[uint32]*yamux.Stream),
249
}
306
-
307
- // Register connection via event loop
308
- done := make(chan struct{}, 1)
309
- g.cmdCh <- &cmdRegisterConn{id: connID, conn: connection, done: done}
310
- <-done
250
+ g.connections[connID] = connection
251
+ g.connectionsLock.Unlock()
252
253
log.Debug().Int64("conn_id", connID).Msg("[RelayServer] Connection registered, starting handler")
254
go g.handleConn(connID, connection)
@@ -330,131 +271,63 @@ func (g *RelayServer) GetLeaseManager() *LeaseManager {
271
272
// IsConnectionActive checks if a connection with the given ID is still active
273
func (g *RelayServer) IsConnectionActive(connectionID int64) bool {
333
- reply := make(chan bool, 1)
334
- g.cmdCh <- &cmdIsConnActive{id: connectionID, reply: reply}
335
- return <-reply
274
+ g.connectionsLock.RLock()
275
+ defer g.connectionsLock.RUnlock()
276
+
277
+ _, exists := g.connections[connectionID]
278
+ return exists
279
}
280
281
// GetAllLeaseEntries returns all lease entries from the lease manager
282
func (g *RelayServer) GetAllLeaseEntries() []*LeaseEntry {
340
- return g.leaseManager.GetAllEntries()
283
+ g.leaseManager.leasesLock.RLock()
284
+ defer g.leaseManager.leasesLock.RUnlock()
285
+
286
+ var entries []*LeaseEntry
287
+ now := time.Now()
288
+
289
+ for _, entry := range g.leaseManager.leases {
290
+ if now.Before(entry.Expires) {
291
+ entries = append(entries, entry)
292
+ }
293
+ }
294
+
295
+ return entries
296
}
297
298
// GetLeaseALPNs returns the ALPN identifiers for a given lease ID
299
func (g *RelayServer) GetLeaseALPNs(leaseID string) []string {
345
- return g.leaseManager.GetLeaseALPNs(leaseID)
300
+ g.leaseManager.leasesLock.RLock()
301
+ defer g.leaseManager.leasesLock.RUnlock()
302
+
303
+ entry, exists := g.leaseManager.leases[leaseID]
304
+ if !exists {
305
+ return nil
306
+ }
307
+
308
+ now := time.Now()
309
+ if now.After(entry.Expires) {
310
+ return nil
311
+ }
312
+
313
+ return entry.Lease.Alpn
314
}
315
316
func (g *RelayServer) Start() {
349
- g.runWg.Add(1)
350
- go g.run()
317
g.leaseManager.Start()
318
}
319
320
func (g *RelayServer) Stop() {
321
close(g.stopch)
356
- g.runWg.Wait()
322
g.leaseManager.Stop()
323
g.waitgroup.Wait()
324
}
325
361
-// run is the main event loop that processes all commands sequentially.
362
-func (g *RelayServer) run() {
363
- defer g.runWg.Done()
364
- for {
365
- select {
366
- case cmd := <-g.cmdCh:
367
- g.handleCmd(cmd)
368
- case <-g.stopch:
369
- return
370
- }
371
- }
372
-}
373
-
374
-// handleCmd dispatches commands to their handlers.
375
-func (g *RelayServer) handleCmd(cmd relayCmd) {
376
- switch c := cmd.(type) {
377
- case *cmdRegisterConn:
378
- g.connections[c.id] = c.conn
379
- c.done <- struct{}{}
380
-
381
- case *cmdRemoveConn:
382
- delete(g.connections, c.id)
383
- // Clean up leases associated with this connection
384
- cleanedLeaseIDs := g.leaseManager.CleanupLeasesByConnectionID(c.id)
385
- // Clean up lease connections and relayed connections
386
- for _, leaseID := range cleanedLeaseIDs {
387
- delete(g.leaseConnections, leaseID)
388
- if streams, exists := g.relayedConnections[leaseID]; exists {
389
- for _, stream := range streams {
390
- stream.Close()
391
- }
392
- delete(g.relayedConnections, leaseID)
393
- }
394
- delete(g.relayedPerLeaseCount, leaseID)
395
- }
396
- c.done <- cleanedLeaseIDs
397
-
398
- case *cmdIsConnActive:
399
- _, exists := g.connections[c.id]
400
- c.reply <- exists
401
-
402
- case *cmdGetConnByLeaseEntry:
403
- conn, exists := g.connections[c.connID]
404
- if exists {
405
- c.reply <- conn
406
- } else {
407
- c.reply <- nil
408
- }
409
-
410
- case *cmdRegisterLeaseConn:
411
- g.leaseConnections[c.leaseID] = c.conn
412
- c.done <- struct{}{}
413
-
414
- case *cmdUnregisterLeaseConn:
415
- delete(g.leaseConnections, c.leaseID)
416
- c.done <- struct{}{}
417
-
418
- case *cmdAddRelayed:
419
- g.relayedConnections[c.leaseID] = append(g.relayedConnections[c.leaseID], c.stream)
420
- c.done <- struct{}{}
421
-
422
- case *cmdRemoveRelayed:
423
- if streams, ok := g.relayedConnections[c.leaseID]; ok {
424
- for i, s := range streams {
425
- if s == c.stream {
426
- g.relayedConnections[c.leaseID] = append(streams[:i], streams[i+1:]...)
427
- break
428
- }
429
- }
430
- }
431
- c.done <- struct{}{}
432
-
433
- case *cmdCheckAndIncLimit:
434
- if g.maxRelayedPerLease > 0 && g.relayedPerLeaseCount[c.leaseID] >= g.maxRelayedPerLease {
435
- c.reply <- false
436
- } else {
437
- g.relayedPerLeaseCount[c.leaseID]++
438
- c.reply <- true
439
- }
440
-
441
- case *cmdDecLimit:
442
- if g.relayedPerLeaseCount[c.leaseID] > 0 {
443
- g.relayedPerLeaseCount[c.leaseID]--
444
- }
445
- c.done <- struct{}{}
446
-
447
- case *cmdSetMaxRelayed:
448
- g.maxRelayedPerLease = c.max
449
- c.done <- struct{}{}
450
- }
451
-}
452
-
326
// Traffic control setters
327
func (g *RelayServer) SetMaxRelayedPerLease(n int) {
455
- done := make(chan struct{}, 1)
456
- g.cmdCh <- &cmdSetMaxRelayed{max: n, done: done}
457
- <-done
328
+ g.limitsLock.Lock()
329
+ g.maxRelayedPerLease = n
330
+ g.limitsLock.Unlock()
331
}
332
333
// SetEstablishRelayCallback sets the callback for relay connection establishment