From b75a098f99848bf03f07bb63a23c332851ec8e50 Mon Sep 17 00:00:00 2001
From: Gouryella <q6221745@gmail.com>
Date: Sat, 20 Dec 2025 10:25:13 +0800
Subject: [PATCH] feat(tcp): Fix reconnect behavior to keep stable subdomain
 and TCP port

Persist the assigned subdomain after first connect so reconnects reuse it.
Allow reserving a specific TCP port when the subdomain is tcp-<port> to prevent port drift.
---
 internal/client/cli/tunnel_runner.go  |  6 ++++
 internal/server/tcp/connection.go     | 44 ++++++++++++++++++++++-----
 internal/server/tcp/port_allocator.go | 22 ++++++++++++++
 3 files changed, 64 insertions(+), 8 deletions(-)

diff --git a/internal/client/cli/tunnel_runner.go b/internal/client/cli/tunnel_runner.go
index fa104ec..d69d2e2 100644
--- a/internal/client/cli/tunnel_runner.go
+++ b/internal/client/cli/tunnel_runner.go
@@ -58,6 +58,12 @@ func runTunnelWithUI(connConfig *tcp.ConnectorConfig, daemonInfo *DaemonInfo) er
 		}
 
 		reconnectAttempts = 0
+		if assignedSubdomain := connector.GetSubdomain(); assignedSubdomain != "" {
+			connConfig.Subdomain = assignedSubdomain
+			if daemonInfo != nil {
+				daemonInfo.Subdomain = assignedSubdomain
+			}
+		}
 
 		if daemonInfo != nil {
 			daemonInfo.URL = connector.GetURL()
diff --git a/internal/server/tcp/connection.go b/internal/server/tcp/connection.go
index d0ffdd7..bde9558 100644
--- a/internal/server/tcp/connection.go
+++ b/internal/server/tcp/connection.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -131,15 +132,24 @@ func (c *Connection) Handle() error {
 			return fmt.Errorf("port allocator not configured")
 		}
 
-		port, err := c.portAlloc.Allocate()
-		if err != nil {
-			c.sendError("port_allocation_failed", err.Error())
-			return fmt.Errorf("failed to allocate port: %w", err)
-		}
-		c.port = port
+		if requestedPort, ok := parseTCPSubdomainPort(req.CustomSubdomain); ok {
+			port, err := c.portAlloc.AllocateSpecific(requestedPort)
+			if err != nil {
+				c.sendError("port_allocation_failed", err.Error())
+				return fmt.Errorf("failed to allocate requested port %d: %w", requestedPort, err)
+			}
+			c.port = port
+		} else {
+			port, err := c.portAlloc.Allocate()
+			if err != nil {
+				c.sendError("port_allocation_failed", err.Error())
+				return fmt.Errorf("failed to allocate port: %w", err)
+			}
+			c.port = port
 
-		if req.CustomSubdomain == "" {
-			req.CustomSubdomain = fmt.Sprintf("tcp-%d", port)
+			if req.CustomSubdomain == "" {
+				req.CustomSubdomain = fmt.Sprintf("tcp-%d", port)
+			}
 		}
 	}
 
@@ -383,6 +393,24 @@ func min(a, b int) int {
 	return b
 }
 
+func parseTCPSubdomainPort(subdomain string) (int, bool) {
+	if !strings.HasPrefix(subdomain, "tcp-") {
+		return 0, false
+	}
+
+	portStr := strings.TrimPrefix(subdomain, "tcp-")
+	if portStr == "" {
+		return 0, false
+	}
+
+	port, err := strconv.Atoi(portStr)
+	if err != nil || port < 1 || port > 65535 {
+		return 0, false
+	}
+
+	return port, true
+}
+
 func (c *Connection) handleFrames(reader *bufio.Reader) error {
 	for {
 		select {
diff --git a/internal/server/tcp/port_allocator.go b/internal/server/tcp/port_allocator.go
index b140228..69e0e71 100644
--- a/internal/server/tcp/port_allocator.go
+++ b/internal/server/tcp/port_allocator.go
@@ -56,6 +56,28 @@ func (p *PortAllocator) Allocate() (int, error) {
 	return 0, fmt.Errorf("no available port in range %d-%d", p.min, p.max)
 }
 
+// AllocateSpecific reserves a specific port if it is within range and available.
+func (p *PortAllocator) AllocateSpecific(port int) (int, error) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if port < p.min || port > p.max {
+		return 0, fmt.Errorf("requested port %d outside range %d-%d", port, p.min, p.max)
+	}
+	if p.used[port] {
+		return 0, fmt.Errorf("requested port %d already in use", port)
+	}
+
+	ln, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", port))
+	if err != nil {
+		return 0, fmt.Errorf("requested port %d unavailable: %w", port, err)
+	}
+	_ = ln.Close()
+
+	p.used[port] = true
+	return port, nil
+}
+
 // Release frees a previously allocated port.
 func (p *PortAllocator) Release(port int) {
 	p.mu.Lock()