@samitouri / QOSAMI-X402Tunnel / commits / cdcb59ee

fix: resolve concurrency race, leak, and hang issues

fatheradvisor committed Mar 6, 2026 at 02:48 UTC cdcb59eeb8c8325a706ee3131c4ff46a6137fa54

6 files changed +190 -17

portal/policy/rate_limiter.go

+118 -3

@@ -97,6 +97,12 @@ func (m *RateLimiter) Copy(dst io.Writer, src io.Reader, leaseID string) (int64,

97 return Copy(dst, src, bucket)

98 }

100 +// CopyInterruptible copies data with rate limiting, interruptible via done channel.

101 +func (m *RateLimiter) CopyInterruptible(dst io.Writer, src io.Reader, leaseID string, done <-chan struct{}) (int64, error) {

102 + bucket := m.GetBucket(leaseID)

103 + return CopyInterruptible(dst, src, bucket, done)

104 +}

105 +

106 // EstablishRelayWithBPS sets up bidirectional relay with BPS limiting.

107 // In the new TLS passthrough architecture, this uses net.Conn.

108 func EstablishRelayWithBPS(clientConn, leaseConn net.Conn, leaseID string, bpsManager *RateLimiter) {

@@ -113,17 +119,24 @@ func EstablishRelayWithBPS(clientConn, leaseConn net.Conn, leaseID string, bpsMa

119 var wg sync.WaitGroup

120 wg.Add(2)

121

122 + // done is closed when either copy direction finishes, unblocking

123 + // any rate-limit sleep in the other direction.

124 + done := make(chan struct{})

125 + var doneOnce sync.Once

126 + closeDone := func() { doneOnce.Do(func() { close(done) }) }

127 +

128 if bpsManager == nil {

117 - // No BPS manager - direct copy without rate limiting

129 go func() {

130 defer wg.Done()

131 _, _ = io.Copy(leaseConn, clientConn)

132 + closeDone()

133 _ = leaseConn.Close()

134 }()

135

136 go func() {

137 defer wg.Done()

138 _, _ = io.Copy(clientConn, leaseConn)

139 + closeDone()

140 _ = clientConn.Close()

141 }()

142 } else {

@@ -136,7 +149,8 @@ func EstablishRelayWithBPS(clientConn, leaseConn net.Conn, leaseID string, bpsMa

149 // Client -> Lease

150 go func() {

151 defer wg.Done()

139 - _, _ = bpsManager.Copy(leaseConn, clientConn, leaseID)

152 + _, _ = bpsManager.CopyInterruptible(leaseConn, clientConn, leaseID, done)

153 + closeDone()

154 if err := leaseConn.Close(); err != nil {

155 log.Debug().Err(err).Str("lease_id", leaseID).Msg("[Relay] failed to close lease connection")

156 }

@@ -145,7 +159,8 @@ func EstablishRelayWithBPS(clientConn, leaseConn net.Conn, leaseID string, bpsMa

159 // Lease -> Client

160 go func() {

161 defer wg.Done()

148 - _, _ = bpsManager.Copy(clientConn, leaseConn, leaseID)

162 + _, _ = bpsManager.CopyInterruptible(clientConn, leaseConn, leaseID, done)

163 + closeDone()

164 if err := clientConn.Close(); err != nil {

165 log.Debug().Err(err).Str("lease_id", leaseID).Msg("[Relay] failed to close client connection")

166 }

@@ -242,6 +257,62 @@ func (b *Bucket) Take(n int64) {

257 }

258 }

259

260 +// TakeInterruptible requests n bytes from the bucket, like Take, but the wait

261 +// can be interrupted via the done channel. Returns false if interrupted.

262 +func (b *Bucket) TakeInterruptible(n int64, done <-chan struct{}) bool {

263 + if b == nil || n <= 0 {

264 + return true

265 + }

266 +

267 + needed := float64(n)

268 +

269 + for {

270 + b.mu.Lock()

271 +

272 + now := time.Now()

273 + elapsed := now.Sub(b.lastRefill).Seconds()

274 + b.tokens += elapsed * float64(b.rateBps)

275 + if b.tokens > b.maxTokens {

276 + b.tokens = b.maxTokens

277 + }

278 + b.lastRefill = now

279 +

280 + if b.tokens >= needed {

281 + b.tokens -= needed

282 + b.mu.Unlock()

283 + atomic.AddInt64(&b.totalBytes, n)

284 + return true

285 + }

286 +

287 + deficit := needed - b.tokens

288 + waitTime := time.Duration(deficit / float64(b.rateBps) * float64(time.Second))

289 +

290 + if b.tokens > 0 {

291 + needed -= b.tokens

292 + b.tokens = 0

293 + }

294 +

295 + b.mu.Unlock()

296 +

297 + if waitTime > 0 {

298 + atomic.AddInt64(&b.throttleHits, 1)

299 + atomic.AddInt64(&b.totalWaited, int64(waitTime))

300 + log.Debug().

301 + Int64("bytes_requested", n).

302 + Int64("rate_bps", b.rateBps).

303 + Dur("wait_time", waitTime).

304 + Msg("[RateLimit] Throttling - waiting for bandwidth")

305 + timer := time.NewTimer(waitTime)

306 + select {

307 + case <-done:

308 + timer.Stop()

309 + return false

310 + case <-timer.C:

311 + }

312 + }

313 + }

314 +}

315 +

316 // TakeWithTimeout requests n bytes but returns false if it would take longer

317 // than maxWait to acquire them. Returns true if tokens were acquired.

318 func (b *Bucket) TakeWithTimeout(n int64, maxWait time.Duration) bool {

@@ -381,6 +452,50 @@ func Copy(dst io.Writer, src io.Reader, b *Bucket) (int64, error) {

452 return total, nil

453 }

454

455 +// CopyInterruptible copies from src to dst with rate limiting, interruptible via done channel.

456 +// When done is closed, the current rate-limit wait is interrupted and the copy returns.

457 +func CopyInterruptible(dst io.Writer, src io.Reader, b *Bucket, done <-chan struct{}) (int64, error) {

458 + if b == nil {

459 + return io.Copy(dst, src)

460 + }

461 + buf := *bufPool.Get().(*[]byte)

462 + defer bufPool.Put(&buf)

463 +

464 + var total int64

465 + startTime := time.Now()

466 +

467 + for {

468 + nr, er := src.Read(buf)

469 + if nr > 0 {

470 + if !b.TakeInterruptible(int64(nr), done) {

471 + logCopyStats(b, total, startTime)

472 + return total, nil

473 + }

474 + nw, ew := dst.Write(buf[:nr])

475 + if nw > 0 {

476 + total += int64(nw)

477 + }

478 + if ew != nil {

479 + logCopyStats(b, total, startTime)

480 + return total, ew

481 + }

482 + if nr != nw {

483 + logCopyStats(b, total, startTime)

484 + return total, io.ErrShortWrite

485 + }

486 + }

487 + if er != nil {

488 + if er == io.EOF {

489 + break

490 + }

491 + logCopyStats(b, total, startTime)

492 + return total, er

493 + }

494 + }

495 + logCopyStats(b, total, startTime)

496 + return total, nil

497 +}

498 +

499 // logCopyStats logs summary statistics when copy completes.

500 func logCopyStats(b *Bucket, totalBytes int64, startTime time.Time) {

501 if b == nil || totalBytes == 0 {

portal/registry.go

@@ -139,10 +139,15 @@ func (g *RelayServer) UnregisterLease(leaseID string) {

139 }

140

141 if g.leaseManager != nil && g.leaseManager.DeleteLease(leaseID) {

142 + // DeleteLease callback (handleLeaseDeleted) already called

143 + // DropLease + UnregisterRouteByLeaseID, so we're done.

144 log.Info().

145 Str("lease_id", leaseID).

146 Msg("[Registry] Lease unregistered")

147 + return

148 }

149 +

150 + // Lease was already removed (e.g. TTL expiry) — defensive cleanup.

151 if g.sniRouter != nil {

152 g.sniRouter.UnregisterRouteByLeaseID(leaseID)

153 }

portal/registry_test.go

+3 -1

@@ -9,12 +9,14 @@ import (

9 )

11 func newTestRegistryRelay(baseHost string) *RelayServer {

12 - return &RelayServer{

12 + s := &RelayServer{

13 BaseHost: baseHost,

14 leaseManager: NewLeaseManager(DefaultLeaseTTL),

15 reverseHub: NewReverseHub(),

16 sniRouter: sni.NewRouter(":0"),

17 }

18 + s.bindLeaseLifecycleHooks()

19 + return s

20 }

22 func newTestLease(id, name, token string) *types.Lease {

portal/reverse_hub.go

+16

@@ -120,6 +120,12 @@ func NewReverseHub() *ReverseHub {

120 }

121

122 func (h *ReverseHub) getOrCreatePool(leaseID string) chan *ReverseConn {

123 + select {

124 + case <-h.stopCh:

125 + return nil

126 + default:

127 + }

128 +

129 h.mu.Lock()

130 defer h.mu.Unlock()

131

@@ -208,6 +214,16 @@ func (h *ReverseHub) Offer(leaseID string, conn *ReverseConn) bool {

214 for range QueueSize + 1 {

215 select {

216 case pool <- conn:

217 + // Re-check: DropLease may have run between getOrCreatePool and send.

218 + h.mu.RLock()

219 + _, dropped := h.dropped[leaseID]

220 + h.mu.RUnlock()

221 + if dropped {

222 + // Pool was drained by DropLease; conn may still be in channel.

223 + // Close it defensively — DropLease drain will also close if it gets it.

224 + conn.Close()

225 + return false

226 + }

227 return true

228 default:

229 }

sdk/listener.go

+7 -9

@@ -82,6 +82,8 @@ type Listener struct {

82 tlsConfig *tls.Config

83 lease *types.Lease

84 httpClient *http.Client

85 + baseCtx context.Context

86 + baseCancel context.CancelFunc

87 stopCh chan struct{}

88 acceptCh chan net.Conn

89 relayAddr string

@@ -135,6 +137,7 @@ func NewListener(relayAddr string, lease *types.Lease, tlsConfig *tls.Config, re

137 }

138 lease.TLS = true

139

140 + baseCtx, baseCancel := context.WithCancel(context.Background())

141 return &Listener{

142 relayAddr: apiURL,

143 lease: lease,

@@ -144,6 +147,8 @@ func NewListener(relayAddr string, lease *types.Lease, tlsConfig *tls.Config, re

147 },

148 tlsConfig: tlsConfig,

149 closeFns: closeFns,

150 + baseCtx: baseCtx,

151 + baseCancel: baseCancel,

152 stopCh: make(chan struct{}),

153 acceptCh: make(chan net.Conn, 128),

154 reverseWorkers: reverseWorkers,

@@ -207,6 +212,7 @@ func (l *Listener) Accept() (net.Conn, error) {

212 func (l *Listener) Close() error {

213 var retErr error

214 l.closeOnce.Do(func() {

215 + l.baseCancel()

216 close(l.stopCh)

217

218 l.mu.Lock()

@@ -573,15 +579,7 @@ func (l *Listener) newStopAwareContext(timeout time.Duration) (context.Context,

579 if timeout <= 0 {

580 timeout = defaultReverseDialTimeout

581 }

576 - ctx, cancel := context.WithTimeout(context.Background(), timeout)

577 - go func() {

578 - select {

579 - case <-l.stopCh:

580 - cancel()

581 - case <-ctx.Done():

582 - }

583 - }()

584 - return ctx, cancel

582 + return context.WithTimeout(l.baseCtx, timeout)

583 }

584

585 func (l *Listener) closeConnOnStop(conn net.Conn) func() {

sdk/listener_test.go

+41 -4

@@ -1,6 +1,7 @@

1 package sdk

3 import (

4 + "context"

5 "crypto/tls"

6 "errors"

7 "fmt"

@@ -187,10 +188,14 @@ func TestBuildReverseConnectRequest(t *testing.T) {

188 func TestOpenReverseConnection_RejectsNonHTTPSRelay(t *testing.T) {

189 t.Parallel()

190

191 + baseCtx, baseCancel := context.WithCancel(context.Background())

192 + defer baseCancel()

193 l := &Listener{

194 relayAddr: "http://localhost:4017",

195 lease: &types.Lease{ID: "lease-1", ReverseToken: "token-1"},

196 reverseDialTimeout: 2 * time.Second,

197 + baseCtx: baseCtx,

198 + baseCancel: baseCancel,

199 stopCh: make(chan struct{}),

200 }

201

@@ -226,10 +231,14 @@ func TestOpenReverseConnection_StopUnblocksTLSHandshake(t *testing.T) {

231 _, _ = conn.Read(buf)

232 }()

233

234 + baseCtx, baseCancel := context.WithCancel(context.Background())

235 + defer baseCancel()

236 l := &Listener{

237 relayAddr: "https://" + ln.Addr().String(),

238 lease: &types.Lease{ID: "lease-1", ReverseToken: "token-1"},

239 reverseDialTimeout: 5 * time.Second,

240 + baseCtx: baseCtx,

241 + baseCancel: baseCancel,

242 stopCh: make(chan struct{}),

243 }

244

@@ -272,9 +281,13 @@ func TestWriteReverseConnectRequest_RespectsWriteDeadline(t *testing.T) {

281 t.Fatalf("parse request URL: %v", err)

282 }

283

284 + baseCtx, baseCancel := context.WithCancel(context.Background())

285 + defer baseCancel()

286 l := &Listener{

287 lease: &types.Lease{ReverseToken: "token-1"},

288 reverseDialTimeout: 25 * time.Millisecond,

289 + baseCtx: baseCtx,

290 + baseCancel: baseCancel,

291 stopCh: make(chan struct{}),

292 }

293

@@ -304,8 +317,12 @@ func TestReadReverseConnectResponse_RespectsReadDeadline(t *testing.T) {

317 defer local.Close()

318 defer peer.Close()

319

320 + baseCtx, baseCancel := context.WithCancel(context.Background())

321 + defer baseCancel()

322 l := &Listener{

323 reverseDialTimeout: 25 * time.Millisecond,

324 + baseCtx: baseCtx,

325 + baseCancel: baseCancel,

326 stopCh: make(chan struct{}),

327 }

328

@@ -385,8 +402,12 @@ func TestReadReverseConnectResponseParsesEnvelopeError(t *testing.T) {

402 defer local.Close()

403 defer peer.Close()

404

405 + baseCtx, baseCancel := context.WithCancel(context.Background())

406 + defer baseCancel()

407 l := &Listener{

408 reverseDialTimeout: 500 * time.Millisecond,

409 + baseCtx: baseCtx,

410 + baseCancel: baseCancel,

411 stopCh: make(chan struct{}),

412 }

413

@@ -493,7 +514,9 @@ func TestReverseConnectRejectionErrorIsFatal(t *testing.T) {

514 func TestWaitForReverseStart_HTTPMode(t *testing.T) {

515 t.Parallel()

516

496 - l := &Listener{stopCh: make(chan struct{})}

517 + baseCtx, baseCancel := context.WithCancel(context.Background())

518 + defer baseCancel()

519 + l := &Listener{baseCtx: baseCtx, baseCancel: baseCancel, stopCh: make(chan struct{})}

520 local, peer := net.Pipe()

521 defer local.Close()

522 defer peer.Close()

@@ -521,7 +544,11 @@ func TestWaitForReverseStart_HTTPMode(t *testing.T) {

544 func TestWaitForReverseStart_TLSMode(t *testing.T) {

545 t.Parallel()

546

547 + baseCtx, baseCancel := context.WithCancel(context.Background())

548 + defer baseCancel()

549 l := &Listener{

550 + baseCtx: baseCtx,

551 + baseCancel: baseCancel,

552 stopCh: make(chan struct{}),

553 tlsConfig: &tls.Config{MinVersion: tls.VersionTLS12},

554 }

@@ -552,7 +579,9 @@ func TestWaitForReverseStart_TLSMode(t *testing.T) {

579 func TestWaitForReverseStart_IgnoresKeepaliveMarker(t *testing.T) {

580 t.Parallel()

581

555 - l := &Listener{stopCh: make(chan struct{})}

582 + baseCtx, baseCancel := context.WithCancel(context.Background())

583 + defer baseCancel()

584 + l := &Listener{baseCtx: baseCtx, baseCancel: baseCancel, stopCh: make(chan struct{})}

585 local, peer := net.Pipe()

586 defer local.Close()

587 defer peer.Close()

@@ -584,7 +613,11 @@ func TestWaitForReverseStart_IgnoresKeepaliveMarker(t *testing.T) {

613 func TestWaitForReverseStart_TLSRejectsHTTPMarker(t *testing.T) {

614 t.Parallel()

615

616 + baseCtx, baseCancel := context.WithCancel(context.Background())

617 + defer baseCancel()

618 l := &Listener{

619 + baseCtx: baseCtx,

620 + baseCancel: baseCancel,

621 stopCh: make(chan struct{}),

622 tlsConfig: &tls.Config{MinVersion: tls.VersionTLS12},

623 }

@@ -615,7 +648,9 @@ func TestWaitForReverseStart_TLSRejectsHTTPMarker(t *testing.T) {

648 func TestWaitForReverseStart_HTTPRejectsTLSMarker(t *testing.T) {

649 t.Parallel()

650

618 - l := &Listener{stopCh: make(chan struct{})}

651 + baseCtx, baseCancel := context.WithCancel(context.Background())

652 + defer baseCancel()

653 + l := &Listener{baseCtx: baseCtx, baseCancel: baseCancel, stopCh: make(chan struct{})}

654 local, peer := net.Pipe()

655 defer local.Close()

656 defer peer.Close()

@@ -643,7 +678,9 @@ func TestWaitForReverseStart_HTTPRejectsTLSMarker(t *testing.T) {

678 func TestWaitForReverseStart_StopCancelsWait(t *testing.T) {

679 t.Parallel()

680

646 - l := &Listener{stopCh: make(chan struct{})}

681 + baseCtx, baseCancel := context.WithCancel(context.Background())

682 + defer baseCancel()

683 + l := &Listener{baseCtx: baseCtx, baseCancel: baseCancel, stopCh: make(chan struct{})}

684 local, peer := net.Pipe()

685 defer local.Close()

686