fix bugs

rabbitprincess committed Mar 21, 2026 at 00:03 UTC 0aed33f693ddad556d3ed012c3d2c4095eca13d2
3 files changed +128 -21
sdk/expose.go
+8
@@ -519,6 +519,14 @@ func (e *Exposure) runListenerAcceptLoop(listener *Listener) {
519 }
520
521 relayURL := listener.api.baseURL.String()
522 + defer func() {
523 + e.mu.Lock()
524 + if current, ok := e.listeners[relayURL]; ok && current == listener {
525 + delete(e.listeners, relayURL)
526 + }
527 + e.mu.Unlock()
528 + }()
529 +
530 for {
531 conn, err := listener.Accept()
532 if err != nil {
sdk/listener.go
+20 -18
@@ -52,10 +52,11 @@ type Listener struct {
52 cancel context.CancelFunc
53 doneCh <-chan struct{}
54
55 - retryCount int
56 - retryWait time.Duration
57 - leaseTTL time.Duration
58 - renewBefore time.Duration
55 + retryCount int
56 + retryWait time.Duration
57 + leaseTTL time.Duration
58 + renewBefore time.Duration
59 + registerBootstraps []string
60
61 stream *transport.ClientStream
62 datagram *transport.ClientDatagram
@@ -101,16 +102,17 @@ func NewListener(ctx context.Context, relayURL string, cfg ListenerConfig) (*Lis
102 }
103
104 l := &Listener{
104 - doneCh: listenerCtx.Done(),
105 - cancel: cancel,
106 - api: api,
107 - registered: make(chan struct{}),
108 - startupStatus: listenerStatusInactive,
109 - retryCount: cfg.RetryCount,
110 - retryWait: retryWait,
111 - leaseTTL: leaseTTL,
112 - renewBefore: renewBefore,
113 - metadata: cfg.Metadata.Copy(),
105 + doneCh: listenerCtx.Done(),
106 + cancel: cancel,
107 + api: api,
108 + registered: make(chan struct{}),
109 + startupStatus: listenerStatusInactive,
110 + retryCount: cfg.RetryCount,
111 + retryWait: retryWait,
112 + leaseTTL: leaseTTL,
113 + renewBefore: renewBefore,
114 + registerBootstraps: append([]string(nil), initialBootstraps...),
115 + metadata: cfg.Metadata.Copy(),
116 }
117 l.stream = transport.NewClientStream(readyTarget, handshakeTimeout)
118 if cfg.UDPEnabled {
@@ -129,15 +131,15 @@ func NewListener(ctx context.Context, relayURL string, cfg ListenerConfig) (*Lis
131 })
132 }
133
132 - go l.runStartup(listenerCtx, initialBootstraps, readyTarget)
134 + go l.runStartup(listenerCtx, readyTarget)
135 return l, nil
136 }
137
136 -func (l *Listener) runStartup(ctx context.Context, initialBootstraps []string, readyTarget int) {
138 +func (l *Listener) runStartup(ctx context.Context, readyTarget int) {
139 var retries int
140
141 for {
140 - err := l.registerAndConfigure(ctx, initialBootstraps)
142 + err := l.registerAndConfigure(ctx, l.registerBootstraps)
143 switch {
144 case err == nil:
145 for range readyTarget {
@@ -451,7 +453,7 @@ func (l *Listener) renewLease(ctx context.Context) error {
453
454 requestCtx, cancel = context.WithTimeout(ctx, 10*time.Second)
455 defer cancel()
454 - if err := l.registerAndConfigure(requestCtx, nil); err != nil {
456 + if err := l.registerAndConfigure(requestCtx, l.registerBootstraps); err != nil {
457 return err
458 }
459 return nil
sdk/sdk_test.go
+100 -3
@@ -179,6 +179,7 @@ func TestNewListenerRegistersLeaseWithMainContract(t *testing.T) {
179
180 func TestNewListenerReregistersOnLeaseNotFound(t *testing.T) {
181 var registerCount atomic.Int32
182 + registerReqCh := make(chan types.RegisterRequest, 2)
183 server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
184 switch r.URL.Path {
185 case types.PathSDKDomain:
@@ -189,6 +190,14 @@ func TestNewListenerReregistersOnLeaseNotFound(t *testing.T) {
190 },
191 })
192 case types.PathSDKRegister:
193 + var registerReq types.RegisterRequest
194 + if err := json.NewDecoder(r.Body).Decode(&registerReq); err != nil {
195 + t.Fatalf("decode register request: %v", err)
196 + }
197 + select {
198 + case registerReqCh <- registerReq:
199 + default:
200 + }
201 count := registerCount.Add(1)
202 leaseID := "lease-1"
203 if count > 1 {
@@ -231,9 +240,10 @@ func TestNewListenerReregistersOnLeaseNotFound(t *testing.T) {
240 defer server.Close()
241
242 listener, err := NewListener(context.Background(), server.URL, ListenerConfig{
234 - Name: "demo",
235 - LeaseTTL: 80 * time.Millisecond,
236 - RenewBefore: 40 * time.Millisecond,
243 + Name: "demo",
244 + LeaseTTL: 80 * time.Millisecond,
245 + RenewBefore: 40 * time.Millisecond,
246 + RegisterBootstraps: []string{"https://relay-a.example.com", "https://relay-b.example.com"},
247 })
248 if err != nil {
249 t.Fatalf("NewListener() error = %v", err)
@@ -243,6 +253,25 @@ func TestNewListenerReregistersOnLeaseNotFound(t *testing.T) {
253 waitForSDKTest(t, func() bool {
254 return listener.LeaseID() == "lease-2"
255 })
256 +
257 + var requests []types.RegisterRequest
258 + waitForSDKTest(t, func() bool {
259 + for len(requests) < 2 {
260 + select {
261 + case req := <-registerReqCh:
262 + requests = append(requests, req)
263 + default:
264 + return false
265 + }
266 + }
267 + return true
268 + })
269 +
270 + for _, req := range requests {
271 + if len(req.Bootstraps) != 2 || req.Bootstraps[0] != "https://relay-a.example.com" || req.Bootstraps[1] != "https://relay-b.example.com" {
272 + t.Fatalf("register request Bootstraps = %v, want [%q %q]", req.Bootstraps, "https://relay-a.example.com", "https://relay-b.example.com")
273 + }
274 + }
275 }
276
277 func TestNewListenerClosesAfterReverseSessionRetryBudgetExhausted(t *testing.T) {
@@ -447,6 +476,74 @@ func TestExposeRegistersKnownRelayURLs(t *testing.T) {
476 }
477 }
478
479 +func TestExposeRemovesClosedListenersSoRelaysCanRestart(t *testing.T) {
480 + var registerCount atomic.Int32
481 + server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
482 + switch r.URL.Path {
483 + case types.PathSDKDomain:
484 + writeSDKTestEnvelope(w, http.StatusOK, types.APIEnvelope[types.DomainResponse]{
485 + OK: true,
486 + Data: types.DomainResponse{
487 + Version: types.SDKProtocolVersion,
488 + },
489 + })
490 + case types.PathSDKRegister:
491 + if registerCount.Add(1) == 1 {
492 + writeSDKTestEnvelope(w, http.StatusServiceUnavailable, types.APIEnvelope[any]{
493 + OK: false,
494 + Error: &types.APIError{Code: types.APIErrorCodeFeatureUnavailable, Message: "relay unavailable"},
495 + })
496 + return
497 + }
498 + writeSDKTestEnvelope(w, http.StatusCreated, types.APIEnvelope[types.RegisterResponse]{
499 + OK: true,
500 + Data: types.RegisterResponse{
501 + LeaseID: "lease-1",
502 + Hostname: "127.0.0.1",
503 + },
504 + })
505 + case types.PathSDKConnect:
506 + writeSDKTestEnvelope(w, http.StatusForbidden, types.APIEnvelope[any]{
507 + OK: false,
508 + Error: &types.APIError{Code: types.APIErrorCodeUnauthorized, Message: "not used in test"},
509 + })
510 + case types.PathSDKRenew:
511 + writeSDKTestEnvelope(w, http.StatusOK, types.APIEnvelope[types.RenewResponse]{
512 + OK: true,
513 + Data: types.RenewResponse{LeaseID: "lease-1"},
514 + })
515 + case types.PathSDKUnregister:
516 + writeSDKTestEnvelope(w, http.StatusOK, types.APIEnvelope[any]{OK: true})
517 + default:
518 + http.NotFound(w, r)
519 + }
520 + }))
521 + defer server.Close()
522 +
523 + exposure, err := Expose(context.Background(), ExposeConfig{
524 + RelayURLs: []string{server.URL},
525 + Name: "demo",
526 + })
527 + if err != nil {
528 + t.Fatalf("Expose() error = %v", err)
529 + }
530 + defer exposure.Close()
531 +
532 + waitForSDKTest(t, func() bool {
533 + exposure.mu.RLock()
534 + defer exposure.mu.RUnlock()
535 + return registerCount.Load() >= 1 && len(exposure.listeners) == 0
536 + })
537 +
538 + if err := exposure.applyRelayURLs(exposure.RelayURLs(), false); err != nil {
539 + t.Fatalf("applyRelayURLs() error = %v", err)
540 + }
541 +
542 + waitForSDKTest(t, func() bool {
543 + return registerCount.Load() >= 2 && len(exposure.PublicURLs()) == 1
544 + })
545 +}
546 +
547 func TestExposeResolvesOwnerPrivateKey(t *testing.T) {
548 ownerPrivateKey := strings.Repeat("11", 32)
549 identity, err := discovery.ResolveIdentity(ownerPrivateKey)