feat: refactor lease management and improve cleanup logic in server and client
Kim committed
Apr 17, 2026 at 19:01 UTC
2de5ae27b959fa5c6bee3db63f8b485e6fb0c5b9
4 files changed
+80
-108
portal/api_server.go
+25
-20
@@ -454,20 +454,28 @@ func (s *Server) handleUnregister(w http.ResponseWriter, r *http.Request) {
454
writeAPIErrorResponse(w, err)
455
return
456
}
457
+ s.cleanupRemovedRecord(context.Background(), record, "delete lease remote state")
458
+
459
+ utils.WriteAPIData(w, http.StatusOK, map[string]any{})
460
+}
461
+
462
+func (s *Server) cleanupRemovedRecord(ctx context.Context, record *leaseRecord, logMessage string) {
463
+ if record == nil {
464
+ return
465
+ }
466
if record.isPublicEntry() && s.acmeManager != nil {
458
- deleteCtx, cancel := context.WithTimeout(context.Background(), defaultClaimTimeout)
459
- if err := s.acmeManager.DeleteENSGaslessHostname(deleteCtx, record.Hostname); err != nil {
467
+ deleteCtx, cancel := context.WithTimeout(ctx, defaultClaimTimeout)
468
+ err := s.acmeManager.DeleteENSGaslessHostname(deleteCtx, record.Hostname)
469
+ cancel()
470
+ if err != nil {
471
log.Warn().
472
Err(err).
473
Str("hostname", record.Hostname).
474
Str("address", record.Address).
464
- Msg("delete lease remote state")
475
+ Msg(logMessage)
476
}
466
- cancel()
477
}
478
record.Close()
469
-
470
- utils.WriteAPIData(w, http.StatusOK, map[string]any{})
479
}
480
481
func (s *Server) handleHop(w http.ResponseWriter, r *http.Request) {
@@ -504,17 +512,7 @@ func (s *Server) handleHop(w http.ResponseWriter, r *http.Request) {
512
}
513
if r.Method == http.MethodDelete {
514
record := s.registry.DeleteHopRoute(&route)
507
- if record != nil && record.isPublicEntry() && s.acmeManager != nil {
508
- deleteCtx, cancel := context.WithTimeout(context.Background(), defaultClaimTimeout)
509
- if err := s.acmeManager.DeleteENSGaslessHostname(deleteCtx, record.Hostname); err != nil {
510
- log.Warn().
511
- Err(err).
512
- Str("hostname", record.Hostname).
513
- Str("address", record.Address).
514
- Msg("delete hop route remote state")
515
- }
516
- cancel()
517
- }
515
+ s.cleanupRemovedRecord(context.Background(), record, "delete hop route remote state")
516
utils.WriteAPIData(w, http.StatusOK, map[string]any{})
517
return
518
}
@@ -551,7 +549,11 @@ func (s *Server) handleHop(w http.ResponseWriter, r *http.Request) {
549
syncCtx, cancel := context.WithTimeout(context.Background(), defaultClaimTimeout)
550
if err := s.acmeManager.SyncENSGaslessHostname(syncCtx, record.Hostname, record.Address); err != nil {
551
cancel()
554
- _ = s.registry.DeleteHopRoute(&route)
552
+ removed := s.registry.DeleteHopRoute(&route)
553
+ if removed == nil {
554
+ removed = record
555
+ }
556
+ s.cleanupRemovedRecord(context.Background(), removed, "delete hop route remote state after sync failure")
557
writeAPIErrorResponse(w, err)
558
return
559
}
@@ -802,8 +804,11 @@ func (s *Server) registerLease(req types.RegisterChallengeRequest, clientIP, rep
804
syncCtx, cancel := context.WithTimeout(context.Background(), defaultClaimTimeout)
805
defer cancel()
806
if err := s.acmeManager.SyncENSGaslessHostname(syncCtx, record.Hostname, record.Address); err != nil {
805
- _, _ = s.registry.Unregister(record.Key())
806
- record.Close()
807
+ removed, _ := s.registry.Unregister(record.Key())
808
+ if removed == nil {
809
+ removed = record
810
+ }
811
+ s.cleanupRemovedRecord(context.Background(), removed, "delete lease remote state after sync failure")
812
return types.RegisterResponse{}, err
813
}
814
}
portal/lease_test.go
+3
-3
@@ -219,12 +219,12 @@ func TestLeaseRegistryCleanupExpiredClosesBroker(t *testing.T) {
219
}
220
}
221
222
-func TestServerRunLeaseJanitorRejectsNonPositiveInterval(t *testing.T) {
222
+func TestServerRunRegistryJanitorRejectsNonPositiveInterval(t *testing.T) {
223
t.Parallel()
224
225
server := &Server{registry: newTestRegistry(t)}
226
- err := server.runLeaseJanitor(context.Background(), 0)
226
+ err := server.runRegistryJanitor(context.Background(), 0)
227
if err == nil {
228
- t.Fatal("runLeaseJanitor() error = nil, want validation error")
228
+ t.Fatal("runRegistryJanitor() error = nil, want validation error")
229
}
230
}
portal/server.go
+44
-75
@@ -253,20 +253,20 @@ func (s *Server) Start(ctx context.Context, apiMux *http.ServeMux) error {
253
started = true
254
255
group.Go(s.runAPIServer)
256
- group.Go(func() error { return s.runSNIListener(groupCtx) })
257
- group.Go(func() error { return s.runLeaseJanitor(groupCtx, 5*time.Second) })
258
- if s.cfg.DiscoveryEnabled {
259
- group.Go(func() error { return s.runRelayDiscoveryLoop(groupCtx) })
260
- }
256
+ group.Go(func() error { return s.runPublicIngress(groupCtx) })
257
if s.overlay != nil {
258
group.Go(s.overlay.Serve)
263
- }
264
- if s.hopMux != nil {
265
- group.Go(func() error { return s.runHopMux(groupCtx) })
259
+ if s.hopMux != nil {
260
+ group.Go(func() error { return s.runOverlayIngress(groupCtx) })
261
+ }
262
}
263
if s.quicTunnel != nil {
264
group.Go(s.runQUICTunnelListener)
265
}
266
+ group.Go(func() error { return s.runRegistryJanitor(groupCtx, 5*time.Second) })
267
+ if s.cfg.DiscoveryEnabled {
268
+ group.Go(func() error { return s.runRelayDiscoveryLoop(groupCtx) })
269
+ }
270
s.acmeManager.Start(serverCtx)
271
group.Go(func() error {
272
<-groupCtx.Done()
@@ -306,6 +306,34 @@ func (s *Server) Wait() error {
306
return err
307
}
308
309
+func (s *Server) PolicyRuntime() *policy.Runtime {
310
+ if s == nil || s.registry == nil {
311
+ return nil
312
+ }
313
+ return s.registry.policy
314
+}
315
+
316
+func (s *Server) PortalURL() string {
317
+ if s == nil {
318
+ return ""
319
+ }
320
+ return s.cfg.PortalURL
321
+}
322
+
323
+func (s *Server) PublicLeases() []types.Lease {
324
+ if s == nil || s.registry == nil {
325
+ return nil
326
+ }
327
+ return s.registry.PublicLeases(time.Now())
328
+}
329
+
330
+func (s *Server) AdminLeases() []types.AdminLease {
331
+ if s == nil || s.registry == nil {
332
+ return nil
333
+ }
334
+ return s.registry.AdminLeases(time.Now())
335
+}
336
+
337
func (s *Server) RelayIdentity() types.RelayIdentity {
338
if s == nil {
339
return types.RelayIdentity{}
@@ -321,20 +349,7 @@ func (s *Server) Shutdown(ctx context.Context) error {
349
}
350
351
for _, lease := range s.registry.CloseAll() {
324
- if lease != nil {
325
- if lease.isPublicEntry() && s.acmeManager != nil {
326
- deleteCtx, cancel := context.WithTimeout(ctx, defaultClaimTimeout)
327
- if err := s.acmeManager.DeleteENSGaslessHostname(deleteCtx, lease.Hostname); err != nil {
328
- log.Warn().
329
- Err(err).
330
- Str("hostname", lease.Hostname).
331
- Str("address", lease.Address).
332
- Msg("delete lease remote state during shutdown")
333
- }
334
- cancel()
335
- }
336
- lease.Close()
337
- }
352
+ s.cleanupRemovedRecord(ctx, lease, "delete lease remote state during shutdown")
353
}
354
355
if s.quicTunnel != nil {
@@ -370,34 +385,6 @@ func (s *Server) Shutdown(ctx context.Context) error {
385
return shutdownErr
386
}
387
373
-func (s *Server) PolicyRuntime() *policy.Runtime {
374
- if s == nil || s.registry == nil {
375
- return nil
376
- }
377
- return s.registry.policy
378
-}
379
-
380
-func (s *Server) PortalURL() string {
381
- if s == nil {
382
- return ""
383
- }
384
- return s.cfg.PortalURL
385
-}
386
-
387
-func (s *Server) PublicLeases() []types.Lease {
388
- if s == nil || s.registry == nil {
389
- return nil
390
- }
391
- return s.registry.PublicLeases(time.Now())
392
-}
393
-
394
-func (s *Server) AdminLeases() []types.AdminLease {
395
- if s == nil || s.registry == nil {
396
- return nil
397
- }
398
- return s.registry.AdminLeases(time.Now())
399
-}
400
-
388
func (s *Server) prepareAPITLS(ctx context.Context) (keyless.TLSMaterialConfig, *acme.Manager, error) {
389
acmeCfg := s.cfg.ACME
390
if baseDomain := utils.NormalizeHostname(acmeCfg.BaseDomain); baseDomain != "" && baseDomain != s.identity.Name {
@@ -435,7 +422,7 @@ func (s *Server) prepareAPITLS(ctx context.Context) (keyless.TLSMaterialConfig,
422
return apiTLS, manager, nil
423
}
424
438
-func (s *Server) runSNIListener(ctx context.Context) error {
425
+func (s *Server) runPublicIngress(ctx context.Context) error {
426
for {
427
conn, err := s.sniListener.Accept()
428
switch {
@@ -478,7 +465,7 @@ func (s *Server) runSNIListener(ctx context.Context) error {
465
return
466
}
467
if err := s.bridgeLeaseConn(ctx, wrappedConn, record); err != nil {
481
- log.Warn().Err(err).Str("server_name", serverName).Msg("bridge lease connection")
468
+ log.Warn().Err(err).Str("server_name", serverName).Msg("bridge public ingress")
469
_ = wrappedConn.Close()
470
return
471
}
@@ -494,12 +481,7 @@ func (s *Server) runSNIListener(ctx context.Context) error {
481
}
482
}
483
497
-func (s *Server) runHopMux(ctx context.Context) error {
498
- if s == nil || s.hopMux == nil {
499
- <-ctx.Done()
500
- return nil
501
- }
502
-
484
+func (s *Server) runOverlayIngress(ctx context.Context) error {
485
group, groupCtx := errgroup.WithContext(ctx)
486
group.Go(func() error { return s.hopMux.Serve(groupCtx) })
487
group.Go(func() error {
@@ -523,6 +505,7 @@ func (s *Server) runHopMux(ctx context.Context) error {
505
hopRole = "middle"
506
}
507
log.Info().Str("remote_addr", stream.RemoteAddr).Str("hop_role", hopRole).Msg("hop stream received")
508
+
509
if err := s.bridgeLeaseConn(groupCtx, stream.Conn, record); err != nil {
510
log.Warn().Err(err).Str("remote_addr", stream.RemoteAddr).Msg("hop stream bridge failed")
511
_ = stream.Conn.Close()
@@ -534,13 +517,11 @@ func (s *Server) runHopMux(ctx context.Context) error {
517
}
518
519
func (s *Server) bridgeLeaseConn(ctx context.Context, conn net.Conn, record *leaseRecord) error {
537
- if s == nil || s.registry == nil || record == nil || time.Now().After(record.ExpiresAt) {
520
+ if record.isExpired(time.Now()) {
521
return errLeaseNotFound
522
}
523
if overlayIPv4, forwardToken, hasNextHop := record.nextHop(); hasNextHop {
524
switch {
542
- case s.hopMux == nil:
543
- return errFeatureUnavailable
525
case overlayIPv4 == "":
526
return errors.New("next hop overlay ipv4 is required")
527
case forwardToken == "":
@@ -584,7 +565,7 @@ func (s *Server) bridgeLeaseConn(ctx context.Context, conn net.Conn, record *lea
565
return nil
566
}
567
587
-func (s *Server) runLeaseJanitor(ctx context.Context, interval time.Duration) error {
568
+func (s *Server) runRegistryJanitor(ctx context.Context, interval time.Duration) error {
569
if interval <= 0 {
570
return errors.New("janitor interval must be positive")
571
}
@@ -598,19 +579,7 @@ func (s *Server) runLeaseJanitor(ctx context.Context, interval time.Duration) er
579
return nil
580
case <-ticker.C:
581
for _, lease := range s.registry.cleanupExpired(time.Now()) {
601
- if lease.isPublicEntry() && s.acmeManager != nil {
602
- deleteCtx, cancel := context.WithTimeout(context.Background(), defaultClaimTimeout)
603
- err := s.acmeManager.DeleteENSGaslessHostname(deleteCtx, lease.Hostname)
604
- cancel()
605
- if err != nil {
606
- log.Warn().
607
- Err(err).
608
- Str("hostname", lease.Hostname).
609
- Str("address", lease.Address).
610
- Msg("delete expired lease remote state")
611
- }
612
- }
613
- lease.Close()
582
+ s.cleanupRemovedRecord(context.Background(), lease, "delete expired lease remote state")
583
}
584
}
585
}
sdk/api_client.go
+8
-10
@@ -115,8 +115,9 @@ func (l *listener) registerLease(ctx context.Context, ttl time.Duration, udpEnab
115
}
116
keylessURL = hopPath[0].APIHTTPSAddr
117
118
- tokens := make([]string, len(hopPath)-1)
119
- for i := range tokens {
118
+ hopRoutes = make([]types.HopRoute, 0, len(hopPath)-1)
119
+ var previousHopToken string
120
+ for i := 0; i < len(hopPath)-1; i++ {
121
token, err := l.identity.DeriveToken(
122
"hop-token",
123
publicHostname,
@@ -127,24 +128,21 @@ func (l *listener) registerLease(ctx context.Context, ttl time.Duration, udpEnab
128
if err != nil {
129
return types.RegisterResponse{}, nil, err
130
}
130
- tokens[i] = "hpt_" + token
131
- }
132
-
133
- hopRoutes = make([]types.HopRoute, 0, len(tokens))
134
- for i := range tokens {
131
+ forwardToken := "hpt_" + token
132
route := types.HopRoute{
133
RelayURL: hopPath[i].APIHTTPSAddr,
134
ForwardRelay: hopPath[i+1],
138
- ForwardToken: tokens[i],
135
+ ForwardToken: forwardToken,
136
}
137
if i == 0 {
138
route.MatchHostname = publicHostname
139
} else {
143
- route.MatchToken = tokens[i-1]
140
+ route.MatchToken = previousHopToken
141
}
142
hopRoutes = append(hopRoutes, route)
143
+ previousHopToken = forwardToken
144
}
147
- exitHopToken = tokens[len(tokens)-1]
145
+ exitHopToken = previousHopToken
146
}
147
148
var challenge types.RegisterChallengeResponse