internal/port/port.go

@@ -5,10 +5,12 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"sync"
 	"tunnel_pls/utils"
 )
 
 type PortManager struct {
+	mu          sync.RWMutex
 	ports       map[uint16]bool
 	sortedPorts []uint16
 }
@@ -37,6 +39,9 @@ func init() {
 }
 
 func (pm *PortManager) AddPortRange(startPort, endPort uint16) error {
+	pm.mu.Lock()
+	defer pm.mu.Unlock()
+
 	if startPort > endPort {
 		return fmt.Errorf("start port cannot be greater than end port")
 	}
@@ -53,8 +58,12 @@ func (pm *PortManager) AddPortRange(startPort, endPort uint16) error {
 }
 
 func (pm *PortManager) GetUnassignedPort() (uint16, bool) {
+	pm.mu.Lock()
+	defer pm.mu.Unlock()
+
 	for _, port := range pm.sortedPorts {
 		if !pm.ports[port] {
+			pm.ports[port] = true
 			return port, true
 		}
 	}
@@ -62,6 +71,9 @@ func (pm *PortManager) GetUnassignedPort() (uint16, bool) {
 }
 
 func (pm *PortManager) SetPortStatus(port uint16, assigned bool) error {
+	pm.mu.Lock()
+	defer pm.mu.Unlock()
+
 	if _, exists := pm.ports[port]; !exists {
 		return fmt.Errorf("port %d is not in the allowed range", port)
 	}
@@ -70,6 +82,9 @@ func (pm *PortManager) SetPortStatus(port uint16, assigned bool) error {
 }
 
 func (pm *PortManager) GetPortStatus(port uint16) (bool, bool) {
+	pm.mu.RLock()
+	defer pm.mu.RUnlock()
+
 	status, exists := pm.ports[port]
 	return status, exists
 }