feat(优化代码): 优化代码

README.md

@@ -27,7 +27,7 @@ probe -ip 6
 probe -s5 127.0.0.1:54321
 
 # 使用自定义服务器（STUN协议）
-probe -s1 stun.miwifi.com:3478
+probe -s1 stun.hot-chilli.net:3478
 ```
 
 # 授权协议

probe.go

@@ -31,7 +31,7 @@ func main() {
 		to time.Duration
 	)
 	flag.StringVar(&ip, "ip", "", "协议版本")
-	flag.StringVar(&s1, "s1", "stun.l.google.com:19302", "主服务器")
+	flag.StringVar(&s1, "s1", "stun.nextcloud.com:3478", "主服务器")
 	flag.StringVar(&s2, "s2", "", "次服务器")
 	flag.StringVar(&s5, "s5", "", "代理地址")
 	flag.DurationVar(&to, "to", 10*time.Second, "超时时间")
@@ -52,7 +52,7 @@ func main() {
 	var err error
 	if s5 != "" {
 		fmt.Printf("通过代理探测: %s\n", s5)
-		conn, err = DialSocks5UDP(s5, network)
+		conn, err = DialSocks5UDP(s5, to)
 		if err != nil {
 			fmt.Printf("连接代理失败: %v\n", err)
 			return

probe_nat.go

@@ -74,6 +74,24 @@ type NATResult struct {
 	MappedIP  string
 }
 
+// isLocalIP 判断IP是否为本地IP
+// 入参: ip 待判断IP
+// 返回: bool 是否为本地IP
+func isLocalIP(ip net.IP) bool {
+	addrs, err := net.InterfaceAddrs()
+	if err != nil {
+		return false
+	}
+	for _, addr := range addrs {
+		if ipNet, ok := addr.(*net.IPNet); ok {
+			if ipNet.IP.Equal(ip) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // resolveAddr 解析探测目标地址
 // 入参: conn 当前连接, addrStr 目标地址字符串, network 网络协议
 // 返回: addr 解析后的网络地址, err 解析错误
@@ -86,8 +104,8 @@ func resolveAddr(conn net.PacketConn, addrStr, network string) (net.Addr, error)
 	if err != nil {
 		return nil, err
 	}
-	if _, ok := conn.(*socks5PacketConn); ok {
+	if _, ok := conn.(*socks5PacketConn); ok && network == "udp" {
-		if net.ParseIP(host) == nil && host != "localhost" && host != "127.0.0.1" {
+		if net.ParseIP(host) == nil {
 			return &SocksAddr{Host: host, Port: port}, nil
 		}
 	}
@@ -110,10 +128,10 @@ func performTest(conn net.PacketConn, serverAddr string, network string, timeout
 	if _, err := conn.WriteTo(req, dst); err != nil {
 		return nil, nil, err
 	}
-	conn.SetReadDeadline(time.Now().Add(timeout))
 	defer conn.SetReadDeadline(time.Time{})
 	buf := make([]byte, 2048)
 	for i := 0; i < 3; i++ {
+		conn.SetReadDeadline(time.Now().Add(timeout))
 		n, addr, err := conn.ReadFrom(buf)
 		if err != nil {
 			if nerr, ok := err.(net.Error); ok && nerr.Timeout() {
@@ -149,14 +167,20 @@ func DetectNAT(conn net.PacketConn, primarySTUN, secondarySTUN, network string,
 		return res
 	}
 	res.MappedIP = mappedAddr1.String()
+	if _, isSocks5 := conn.(*socks5PacketConn); !isSocks5 {
 		if localAddr, ok := conn.LocalAddr().(*net.UDPAddr); ok {
-		if localAddr.IP.Equal(mappedAddr1.IP) && localAddr.Port == mappedAddr1.Port {
+			isLocal := localAddr.IP.Equal(mappedAddr1.IP)
+			if !isLocal && localAddr.IP.IsUnspecified() {
+				isLocal = isLocalIP(mappedAddr1.IP)
+			}
+			if isLocal && localAddr.Port == mappedAddr1.Port {
 				res.Type = NATOpen
 				res.Mapping = MappingEndpointIndependent
 				res.Filtering = FilteringEndpointIndependent
 				return res
 			}
 		}
+	}
 	var targetSTUN2 string
 	if secondarySTUN != "" {
 		targetSTUN2 = secondarySTUN
@@ -195,28 +219,80 @@ func DetectNAT(conn net.PacketConn, primarySTUN, secondarySTUN, network string,
 			}
 		}
 	}
-	var mappedAddr3 *net.UDPAddr
+	var mappedAddr2 *net.UDPAddr
 	if targetSTUN2 != "" {
-		resp3, _, err := performTest(conn, targetSTUN2, network, timeout, false, false)
+		resp2, _, err := performTest(conn, targetSTUN2, network, timeout, false, false)
-		if err == nil && resp3 != nil {
+		if err == nil && resp2 != nil {
-			mappedAddr3 = resp3.GetMappedAddress()
+			mappedAddr2 = resp2.GetMappedAddress()
 		}
 	}
-	if mappedAddr3 != nil {
+	if mappedAddr2 == nil {
-		if mappedAddr1.String() == mappedAddr3.String() {
+		res.Mapping = MappingUnknown
+	} else if mappedAddr1.IP.Equal(mappedAddr2.IP) && mappedAddr1.Port == mappedAddr2.Port {
 		res.Mapping = MappingEndpointIndependent
 	} else {
 		res.Mapping = MappingAddressPortDependent
+		changedAddr := resp1.GetChangedAddress()
+		if changedAddr != nil {
+			host, _, _ := net.SplitHostPort(primarySTUN)
+			primaryIP := net.ParseIP(host)
+			if primaryIP == nil {
+				if ipAddr, err := net.ResolveIPAddr("ip", host); err == nil {
+					primaryIP = ipAddr.IP
 				}
-	} else {
-		res.Mapping = MappingUnknown
 			}
-	resp2, _, _ := performTest(conn, primarySTUN, network, timeout, true, true)
+			if primaryIP != nil && changedAddr.Port != 0 {
-	if resp2 != nil {
+				altPortSTUN := net.JoinHostPort(primaryIP.String(), strconv.Itoa(changedAddr.Port))
+				if altPortSTUN != primarySTUN {
+					resp3, _, err := performTest(conn, altPortSTUN, network, timeout, false, false)
+					if err == nil && resp3 != nil {
+						if mappedAddr3 := resp3.GetMappedAddress(); mappedAddr3 != nil {
+							if mappedAddr1.IP.Equal(mappedAddr3.IP) && mappedAddr1.Port == mappedAddr3.Port {
+								res.Mapping = MappingAddressDependent
+							}
+						}
+					}
+				}
+			}
+		}
+	}
+	respF1, addrF1, _ := performTest(conn, primarySTUN, network, timeout, true, true)
+	if respF1 != nil {
+		if dst, err := resolveAddr(conn, primarySTUN, network); err == nil {
+			var checkIP net.IP
+			var checkPort int
+			if sAddr, ok := dst.(*net.UDPAddr); ok {
+				checkIP = sAddr.IP
+				checkPort = sAddr.Port
+			} else if sAddr, ok := dst.(*SocksAddr); ok {
+				if ipAddr, err := net.ResolveIPAddr("ip", sAddr.Host); err == nil {
+					checkIP = ipAddr.IP
+					checkPort = sAddr.Port
+				}
+			}
+			if checkIP != nil && addrF1 != nil && checkIP.Equal(addrF1.IP) && addrF1.Port == checkPort {
+				respF1 = nil
+			}
+		}
+	}
+	if respF1 != nil {
 		res.Filtering = FilteringEndpointIndependent
 	} else {
-		resp4, _, _ := performTest(conn, primarySTUN, network, timeout, false, true)
+		respF2, addrF2, _ := performTest(conn, primarySTUN, network, timeout, false, true)
-		if resp4 != nil {
+		if respF2 != nil {
+			if dst, err := resolveAddr(conn, primarySTUN, network); err == nil {
+				var checkPort int
+				if sAddr, ok := dst.(*net.UDPAddr); ok {
+					checkPort = sAddr.Port
+				} else if sAddr, ok := dst.(*SocksAddr); ok {
+					checkPort = sAddr.Port
+				}
+				if checkPort != 0 && addrF2 != nil && checkPort == addrF2.Port {
+					respF2 = nil
+				}
+			}
+		}
+		if respF2 != nil {
 			res.Filtering = FilteringAddressDependent
 		} else {
 			res.Filtering = FilteringAddressPortDependent
@@ -227,12 +303,16 @@ func DetectNAT(conn net.PacketConn, primarySTUN, secondarySTUN, network string,
 		if res.Mapping == MappingUnknown {
 			res.Mapping = MappingEndpointIndependent
 		}
+	} else if res.Mapping == MappingUnknown {
+		res.Type = NATUnknown
 	} else if res.Mapping == MappingEndpointIndependent {
 		switch res.Filtering {
 		case FilteringAddressDependent:
 			res.Type = NATRestricted
 		case FilteringAddressPortDependent:
 			res.Type = NATPortRestricted
+		default:
+			res.Type = NATUnknown
 		}
 	} else {
 		res.Type = NATSymmetric

probe_socks5.go

@@ -45,7 +45,6 @@ type socks5PacketConn struct {
 	tcpConn   net.Conn
 	udpConn   *net.UDPConn
 	relayAddr *net.UDPAddr
-	targetAddr net.Addr
 }
 
 // ReadFrom 从UDP连接读取数据
@@ -65,9 +64,6 @@ func (c *socks5PacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error)
 	var dataOffset int
 	switch atyp {
 	case 0x01:
-		if n < 10 {
-			return 0, nil, errors.New("short packet")
-		}
 		ip := net.IP(buf[4:8])
 		port := binary.BigEndian.Uint16(buf[8:10])
 		rAddr = &net.UDPAddr{IP: ip, Port: int(port)}
@@ -96,8 +92,8 @@ func (c *socks5PacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error)
 	default:
 		return 0, nil, fmt.Errorf("unknown address type: 0x%x", atyp)
 	}
-	copy(p, buf[dataOffset:n])
+	copied := copy(p, buf[dataOffset:n])
-	return n - dataOffset, rAddr, nil
+	return copied, rAddr, nil
 }
 
 // WriteTo 写入数据到目标地址
@@ -180,9 +176,9 @@ func (c *socks5PacketConn) SetWriteDeadline(t time.Time) error {
 }
 
 // DialSocks5UDP 建立SOCKS5 UDP关联
-// 入参: proxyAddr 代理服务器地址, network 网络协议(udp/udp4/udp6)
+// 入参: proxyAddr 代理服务器地址, timeout 超时时间
 // 返回: conn 数据包连接, err 连接错误
-func DialSocks5UDP(proxyAddr, network string) (net.PacketConn, error) {
+func DialSocks5UDP(proxyAddr string, timeout time.Duration) (net.PacketConn, error) {
 	var host string
 	if strings.Contains(proxyAddr, "://") {
 		u, err := url.Parse(proxyAddr)
@@ -193,17 +189,11 @@ func DialSocks5UDP(proxyAddr, network string) (net.PacketConn, error) {
 	} else {
 		host = proxyAddr
 	}
-	tcpNetwork := "tcp"
+	conn, err := net.DialTimeout("tcp", host, timeout)
-	switch network {
-	case "udp4":
-		tcpNetwork = "tcp4"
-	case "udp6":
-		tcpNetwork = "tcp6"
-	}
-	conn, err := net.DialTimeout(tcpNetwork, host, 5*time.Second)
 	if err != nil {
 		return nil, err
 	}
+	conn.SetDeadline(time.Now().Add(timeout))
 	_, err = conn.Write([]byte{0x05, 0x01, 0x00})
 	if err != nil {
 		conn.Close()
@@ -266,6 +256,9 @@ func DialSocks5UDP(proxyAddr, network string) (net.PacketConn, error) {
 			return nil, err
 		}
 		relayIP = net.IP(b)
+	default:
+		conn.Close()
+		return nil, fmt.Errorf("unsupported address type: 0x%x", header[3])
 	}
 	pb := make([]byte, 2)
 	if _, err := io.ReadFull(conn, pb); err != nil {
@@ -278,8 +271,15 @@ func DialSocks5UDP(proxyAddr, network string) (net.PacketConn, error) {
 			relayIP = remoteAddr.IP
 		}
 	}
+	conn.SetDeadline(time.Time{})
 	relayAddr := &net.UDPAddr{IP: relayIP, Port: relayPort}
-	lConn, err := net.ListenUDP(network, nil)
+	listenNetwork := "udp"
+	if relayIP.To4() != nil {
+		listenNetwork = "udp4"
+	} else {
+		listenNetwork = "udp6"
+	}
+	lConn, err := net.ListenUDP(listenNetwork, nil)
 	if err != nil {
 		conn.Close()
 		return nil, err