Merge pull request #262 from nxtrace/main

SYNC
nxtrace · Oct 6, 2024 · fa35005 · fa35005
2 parents 6a3ea6a + b0c0f8d
commit fa35005
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 60 deletions.
diff --git a/cmd/cmd.go b/cmd/cmd.go
@@ -3,7 +3,6 @@ package cmd
 import (
 	"encoding/json"
 	"fmt"
-	"github.com/fatih/color"
 	"log"
 	"net"
 	"os"
@@ -12,6 +11,8 @@ import (
 	"strings"
 	"time"
 
+	"github.com/fatih/color"
+
 	"github.com/akamensky/argparse"
 	"github.com/nxtrace/NTrace-core/config"
 	fastTrace "github.com/nxtrace/NTrace-core/fast_trace"
@@ -114,7 +115,7 @@ func Excute() {
 			Timeout:        time.Duration(*timeout) * time.Millisecond,
 			File:           *file,
 			DontFragment:   *dontFragment,
-			Dot:			*dot,
+			Dot:            *dot,
 		}
 
 		fastTrace.FastTest(*tcp, *output, paramsFastTrace)
@@ -186,7 +187,9 @@ func Excute() {
 			w.Interrupt = make(chan os.Signal, 1)
 			signal.Notify(w.Interrupt, os.Interrupt)
 			defer func() {
-				w.Conn.Close()
+				if w.Conn != nil {
+					w.Conn.Close()
+				}
 			}()
 		}
 	}

diff --git a/trace/icmp_ipv4.go b/trace/icmp_ipv4.go
@@ -69,7 +69,7 @@ func (t *ICMPTracer) Execute() (*Result, error) {
 
 	var err error
 
-	t.icmpListen, err = internal.ListenICMP("ip4:1", t.SrcAddr)
+	t.icmpListen, err = internal.ListenICMP("ip4:icmp", t.SrcAddr)
 	if err != nil {
 		return &t.res, err
 	}

diff --git a/trace/internal/icmp_darwin.go b/trace/internal/icmp_darwin.go
@@ -6,7 +6,6 @@ import (
 	"context"
 	"errors"
 	"net"
-	"os"
 	"syscall"
 	"unsafe"
 )
@@ -30,68 +29,69 @@ var (
 )
 
 func ListenICMP(network string, laddr string) (net.PacketConn, error) {
-	if os.Getuid() == 0 { // root
-		return net.ListenPacket(network, laddr)
-	} else {
-		if nw, ok := networkMap[network]; ok {
-			proto := syscall.IPPROTO_ICMP
-			if nw == "udp6" {
-				proto = syscall.IPPROTO_ICMPV6
-			}
+	// 为兼容NE，需要注释掉
+	//if os.Getuid() == 0 { // root
+	//	return net.ListenPacket(network, laddr)
+	//} else {
+	if nw, ok := networkMap[network]; ok {
+		proto := syscall.IPPROTO_ICMP
+		if nw == "udp6" {
+			proto = syscall.IPPROTO_ICMPV6
+		}
 
-			var ifIndex = -1
-			if laddr != "" {
-				la := net.ParseIP(laddr)
-				if ifaces, err := net.Interfaces(); err == nil {
-					for _, iface := range ifaces {
-						addrs, err := iface.Addrs()
-						if err != nil {
-							continue
-						}
-						for _, addr := range addrs {
-							if ipnet, ok := addr.(*net.IPNet); ok {
-								if ipnet.IP.Equal(la) {
-									ifIndex = iface.Index
-									break
-								}
+		var ifIndex = -1
+		if laddr != "" {
+			la := net.ParseIP(laddr)
+			if ifaces, err := net.Interfaces(); err == nil {
+				for _, iface := range ifaces {
+					addrs, err := iface.Addrs()
+					if err != nil {
+						continue
+					}
+					for _, addr := range addrs {
+						if ipnet, ok := addr.(*net.IPNet); ok {
+							if ipnet.IP.Equal(la) {
+								ifIndex = iface.Index
+								break
 							}
 						}
 					}
-					if ifIndex == -1 {
-						return nil, errUnknownIface
-					}
-				} else {
-					return nil, err
 				}
+				if ifIndex == -1 {
+					return nil, errUnknownIface
+				}
+			} else {
+				return nil, err
 			}
+		}
 
-			isock, err := internetSocket(context.Background(), nw, nil, nil, syscall.SOCK_DGRAM, proto, "listen",
-				func(ctx context.Context, network, address string, c syscall.RawConn) error {
-					if ifIndex != -1 {
-						if proto == syscall.IPPROTO_ICMP {
-							return c.Control(func(fd uintptr) {
-								err := syscall.SetsockoptInt(int(fd), syscall.IPPROTO_IP, syscall.IP_BOUND_IF, ifIndex)
-								if err != nil {
-									return
-								}
-							})
-						} else {
-							return c.Control(func(fd uintptr) {
-								err := syscall.SetsockoptInt(int(fd), syscall.IPPROTO_IPV6, syscall.IPV6_BOUND_IF, ifIndex)
-								if err != nil {
-									return
-								}
-							})
-						}
+		isock, err := internetSocket(context.Background(), nw, nil, nil, syscall.SOCK_DGRAM, proto, "listen",
+			func(ctx context.Context, network, address string, c syscall.RawConn) error {
+				if ifIndex != -1 {
+					if proto == syscall.IPPROTO_ICMP {
+						return c.Control(func(fd uintptr) {
+							err := syscall.SetsockoptInt(int(fd), syscall.IPPROTO_IP, syscall.IP_BOUND_IF, ifIndex)
+							if err != nil {
+								return
+							}
+						})
+					} else {
+						return c.Control(func(fd uintptr) {
+							err := syscall.SetsockoptInt(int(fd), syscall.IPPROTO_IPV6, syscall.IPV6_BOUND_IF, ifIndex)
+							if err != nil {
+								return
+							}
+						})
 					}
-					return nil
-				})
-			if err != nil {
-				panic(err)
-			}
-			return newIPConn(isock), nil
-		} else {
-			return nil, errUnknownNetwork
+				}
+				return nil
+			})
+		if err != nil {
+			panic(err)
 		}
+		return newIPConn(isock), nil
+	} else {
+		return nil, errUnknownNetwork
 	}
+	//}
 }
diff --git a/trace/tcp_ipv6.go b/trace/tcp_ipv6.go
@@ -47,7 +47,7 @@ func (t *TCPTracerv6) Execute() (*Result, error) {
 	if err != nil {
 		return nil, err
 	}
-	t.icmp, err = icmp.ListenPacket("ip6:58", "::")
+	t.icmp, err = icmp.ListenPacket("ip6:ipv6-icmp", "::")
 	if err != nil {
 		return &t.res, err
 	}