Jelajahi Sumber

Merge pull request #3435 from roylee17/3149-fix_port_mapping_of_tcp_plus_udp

Fix #3149: port mapping of TCP + UDP
Michael Crosby 11 tahun lalu
induk
melakukan
b8b18a2b42
1 mengubah file dengan 10 tambahan dan 5 penghapusan
  1. 10 5
      network.go

+ 10 - 5
network.go

@@ -248,12 +248,12 @@ type PortMapper struct {
 }
 
 func (mapper *PortMapper) Map(ip net.IP, port int, backendAddr net.Addr) error {
-	mapKey := (&net.TCPAddr{Port: port, IP: ip}).String()
-	if _, exists := mapper.tcpProxies[mapKey]; exists {
-		return fmt.Errorf("Port %s is already in use", mapKey)
-	}
 
 	if _, isTCP := backendAddr.(*net.TCPAddr); isTCP {
+		mapKey := (&net.TCPAddr{Port: port, IP: ip}).String()
+		if _, exists := mapper.tcpProxies[mapKey]; exists {
+			return fmt.Errorf("TCP Port %s is already in use", mapKey)
+		}
 		backendPort := backendAddr.(*net.TCPAddr).Port
 		backendIP := backendAddr.(*net.TCPAddr).IP
 		if mapper.iptables != nil {
@@ -270,6 +270,10 @@ func (mapper *PortMapper) Map(ip net.IP, port int, backendAddr net.Addr) error {
 		mapper.tcpProxies[mapKey] = proxy
 		go proxy.Run()
 	} else {
+		mapKey := (&net.UDPAddr{Port: port, IP: ip}).String()
+		if _, exists := mapper.udpProxies[mapKey]; exists {
+			return fmt.Errorf("UDP: Port %s is already in use", mapKey)
+		}
 		backendPort := backendAddr.(*net.UDPAddr).Port
 		backendIP := backendAddr.(*net.UDPAddr).IP
 		if mapper.iptables != nil {
@@ -290,8 +294,8 @@ func (mapper *PortMapper) Map(ip net.IP, port int, backendAddr net.Addr) error {
 }
 
 func (mapper *PortMapper) Unmap(ip net.IP, port int, proto string) error {
-	mapKey := (&net.TCPAddr{Port: port, IP: ip}).String()
 	if proto == "tcp" {
+		mapKey := (&net.TCPAddr{Port: port, IP: ip}).String()
 		backendAddr, ok := mapper.tcpMapping[mapKey]
 		if !ok {
 			return fmt.Errorf("Port tcp/%s is not mapped", mapKey)
@@ -307,6 +311,7 @@ func (mapper *PortMapper) Unmap(ip net.IP, port int, proto string) error {
 		}
 		delete(mapper.tcpMapping, mapKey)
 	} else {
+		mapKey := (&net.UDPAddr{Port: port, IP: ip}).String()
 		backendAddr, ok := mapper.udpMapping[mapKey]
 		if !ok {
 			return fmt.Errorf("Port udp/%s is not mapped", mapKey)