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)
+		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 net.ParseIP(host) == nil && host != "localhost" && host != "127.0.0.1" {
+	if _, ok := conn.(*socks5PacketConn); ok && network == "udp" {
+		if net.ParseIP(host) == nil {
 			return &SocksAddr{Host: host, Port: port}, nil
 		}
 	}
@@ -103,15 +121,17 @@ func performTest(conn net.PacketConn, serverAddr string, network string, timeout
 		return nil, nil, err
 	}
 	txID := [12]byte{}
-	rand.Read(txID[:])
+	if _, err := rand.Read(txID[:]); err != nil {
+		return nil, nil, err
+	}
 	req := encodeSTUNRequest(txID, changeIP, changePort)
 	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() {
@@ -147,6 +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 {
+			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
@@ -158,7 +192,10 @@ func DetectNAT(conn net.PacketConn, primarySTUN, secondarySTUN, network string,
 			host, port, _ := net.SplitHostPort(primarySTUN)
 			ips, err := net.LookupIP(host)
 			if err == nil && len(ips) > 1 {
-				primaryIP, _ := net.ResolveIPAddr("ip", host)
+				primaryIP, err := net.ResolveIPAddr("ip", host)
+				if err != nil {
+					return res
+				}
 				wantIPv4 := network == "udp4"
 				wantIPv6 := network == "udp6"
 				if network == "udp" {
@@ -182,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)
-		if err == nil && resp3 != nil {
-			mappedAddr3 = resp3.GetMappedAddress()
+		resp2, _, err := performTest(conn, targetSTUN2, network, timeout, false, false)
+		if err == nil && resp2 != nil {
+			mappedAddr2 = resp2.GetMappedAddress()
 		}
 	}
-	if mappedAddr3 != nil {
-		if mappedAddr1.String() == mappedAddr3.String() {
+	if mappedAddr2 == nil {
+		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 resp2 != nil {
+			if primaryIP != nil && changedAddr.Port != 0 {
+				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)
-		if resp4 != nil {
+		respF2, addrF2, _ := performTest(conn, primarySTUN, network, timeout, false, true)
+		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
@@ -214,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连接读取数据
@@ -58,16 +57,13 @@ func (c *socks5PacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error)
 		return 0, nil, err
 	}
 	if n < 10 {
-		return 0, nil, nil
+		return 0, nil, errors.New("packet too short")
 	}
 	atyp := buf[3]
 	var rAddr *net.UDPAddr
 	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)}
@@ -93,9 +89,11 @@ func (c *socks5PacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error)
 		port := binary.BigEndian.Uint16(buf[20:22])
 		rAddr = &net.UDPAddr{IP: ip, Port: int(port)}
 		dataOffset = 22
+	default:
+		return 0, nil, fmt.Errorf("unknown address type: 0x%x", atyp)
 	}
-	copy(p, buf[dataOffset:n])
-	return n - dataOffset, rAddr, nil
+	copied := copy(p, buf[dataOffset:n])
+	return copied, rAddr, nil
 }
 
 // WriteTo 写入数据到目标地址
@@ -142,8 +140,12 @@ func (c *socks5PacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 // Close 关闭连接
 // 返回: err 关闭错误
 func (c *socks5PacketConn) Close() error {
-	c.tcpConn.Close()
-	return c.udpConn.Close()
+	err1 := c.tcpConn.Close()
+	err2 := c.udpConn.Close()
+	if err1 != nil {
+		return err1
+	}
+	return err2
 }
 
 // LocalAddr 获取本地地址
@@ -174,9 +176,9 @@ func (c *socks5PacketConn) SetWriteDeadline(t time.Time) error {
 }
 
 // DialSocks5UDP 建立SOCKS5 UDP关联
-// 入参: proxyAddr 代理服务器地址
+// 入参: proxyAddr 代理服务器地址, timeout 超时时间
 // 返回: conn 数据包连接, err 连接错误
-func DialSocks5UDP(proxyAddr 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)
@@ -187,10 +189,11 @@ func DialSocks5UDP(proxyAddr string) (net.PacketConn, error) {
 	} else {
 		host = proxyAddr
 	}
-	conn, err := net.DialTimeout("tcp", host, 5*time.Second)
+	conn, err := net.DialTimeout("tcp", host, timeout)
 	if err != nil {
 		return nil, err
 	}
+	conn.SetDeadline(time.Now().Add(timeout))
 	_, err = conn.Write([]byte{0x05, 0x01, 0x00})
 	if err != nil {
 		conn.Close()
@@ -253,6 +256,9 @@ func DialSocks5UDP(proxyAddr 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 {
@@ -265,8 +271,15 @@ func DialSocks5UDP(proxyAddr string) (net.PacketConn, error) {
 			relayIP = remoteAddr.IP
 		}
 	}
+	conn.SetDeadline(time.Time{})
 	relayAddr := &net.UDPAddr{IP: relayIP, Port: relayPort}
-	lConn, err := net.ListenUDP("udp", 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

probe_stun.go

@@ -136,9 +136,9 @@ func decodeSTUNResponse(data []byte, txID [12]byte) (*stunMessage, error) {
 }
 
 // parseAddress 解析STUN属性中的地址信息
-// 入参: attrType 属性类型, data 属性值数据
+// 入参: attrType 属性类型, data 属性值数据, txID 事务ID用于XOR解码
 // 返回: addr 解析出的UDP地址, err 解析错误信息
-func parseAddress(attrType uint16, data []byte) (*net.UDPAddr, error) {
+func parseAddress(attrType uint16, data []byte, txID [12]byte) (*net.UDPAddr, error) {
 	if len(data) < 4 {
 		return nil, errors.New("attribute too short")
 	}
@@ -157,12 +157,15 @@ func parseAddress(attrType uint16, data []byte) (*net.UDPAddr, error) {
 	copy(ip, data[4:4+ipLen])
 	if attrType == attrXorMappedAddress {
 		port ^= uint16(stunMagicCookie >> 16)
-		if ipLen == 4 {
 		mc := make([]byte, 4)
 		binary.BigEndian.PutUint32(mc, stunMagicCookie)
 		for i := 0; i < 4; i++ {
 			ip[i] ^= mc[i]
 		}
+		if ipLen == 16 {
+			for i := 4; i < 16; i++ {
+				ip[i] ^= txID[i-4]
+			}
 		}
 	}
 	return &net.UDPAddr{IP: ip, Port: int(port)}, nil
@@ -173,14 +176,14 @@ func parseAddress(attrType uint16, data []byte) (*net.UDPAddr, error) {
 func (m *stunMessage) GetMappedAddress() *net.UDPAddr {
 	for _, attr := range m.Attributes {
 		if attr.Type == attrXorMappedAddress {
-			if addr, err := parseAddress(attr.Type, attr.Value); err == nil {
+			if addr, err := parseAddress(attr.Type, attr.Value, m.Header.ID); err == nil {
 				return addr
 			}
 		}
 	}
 	for _, attr := range m.Attributes {
 		if attr.Type == attrMappedAddress {
-			if addr, err := parseAddress(attr.Type, attr.Value); err == nil {
+			if addr, err := parseAddress(attr.Type, attr.Value, m.Header.ID); err == nil {
 				return addr
 			}
 		}
@@ -193,14 +196,14 @@ func (m *stunMessage) GetMappedAddress() *net.UDPAddr {
 func (m *stunMessage) GetChangedAddress() *net.UDPAddr {
 	for _, attr := range m.Attributes {
 		if attr.Type == attrChangedAddress {
-			if addr, err := parseAddress(attr.Type, attr.Value); err == nil {
+			if addr, err := parseAddress(attr.Type, attr.Value, m.Header.ID); err == nil {
 				return addr
 			}
 		}
 	}
 	for _, attr := range m.Attributes {
 		if attr.Type == attrOtherAddress {
-			if addr, err := parseAddress(attr.Type, attr.Value); err == nil {
+			if addr, err := parseAddress(attr.Type, attr.Value, m.Header.ID); err == nil {
 				return addr
 			}
 		}