From 6bc4abeaf810cc8d58a09399b1ae5cc1d7f19095 Mon Sep 17 00:00:00 2001
From: sjlleo <sjlleo001@gmail.com>
Date: Sat, 4 Jun 2022 21:42:44 +0800
Subject: [PATCH] =?UTF-8?q?update:=20TCP/UDP=E8=B7=AF=E7=94=B1=E8=B7=9F?=
 =?UTF-8?q?=E8=B8=AA=E6=97=B6=EF=BC=8C=E5=B0=86=E4=B8=8D=E4=BC=9A=E5=86=8D?=
 =?UTF-8?q?=E6=98=BE=E7=A4=BAIPv6=20IP?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 main.go                 | 10 +++++++++-
 printer/printer_test.go |  9 +++++----
 util/util.go            | 21 ++++++++++++---------
 3 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/main.go b/main.go
index 76f4381..150481c 100644
--- a/main.go
+++ b/main.go
@@ -4,6 +4,7 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"net"
 	"os"
 	"strings"
 	"time"
@@ -95,7 +96,14 @@ func main() {
 		*dataOrigin = configData.DataOrigin
 	}
 
-	ip := util.DomainLookUp(domain)
+	var ip net.IP
+
+	if *tcpSYNFlag || *udpPackageFlag {
+		ip = util.DomainLookUp(domain, true)
+	} else {
+		ip = util.DomainLookUp(domain, false)
+	}
+
 	printer.PrintTraceRouteNav(ip, domain, *dataOrigin)
 
 	var m trace.Method = ""
diff --git a/printer/printer_test.go b/printer/printer_test.go
index d285211..25460f4 100644
--- a/printer/printer_test.go
+++ b/printer/printer_test.go
@@ -2,16 +2,17 @@ package printer
 
 import (
 	"errors"
-	"github.com/xgadget-lab/nexttrace/ipgeo"
-	"github.com/xgadget-lab/nexttrace/trace"
-	"github.com/xgadget-lab/nexttrace/util"
 	"net"
 	"testing"
 	"time"
+
+	"github.com/xgadget-lab/nexttrace/ipgeo"
+	"github.com/xgadget-lab/nexttrace/trace"
+	"github.com/xgadget-lab/nexttrace/util"
 )
 
 func TestPrintTraceRouteNav(t *testing.T) {
-	PrintTraceRouteNav(util.DomainLookUp("1.1.1.1"), "1.1.1.1", "dataOrigin")
+	PrintTraceRouteNav(util.DomainLookUp("1.1.1.1", false), "1.1.1.1", "dataOrigin")
 }
 
 var testGeo = &ipgeo.IPGeoData{
diff --git a/util/util.go b/util/util.go
index e1f2bcf..b981b38 100644
--- a/util/util.go
+++ b/util/util.go
@@ -25,7 +25,7 @@ func LocalIPPort(dstip net.IP) (net.IP, int) {
 	return nil, -1
 }
 
-func DomainLookUp(host string) net.IP {
+func DomainLookUp(host string, ipv4Only bool) net.IP {
 	ips, err := net.LookupIP(host)
 	if err != nil {
 		fmt.Println("Domain " + host + " Lookup Fail.")
@@ -36,17 +36,20 @@ func DomainLookUp(host string) net.IP {
 	var ipv6Flag = false
 
 	for _, ip := range ips {
-		ipSlice = append(ipSlice, ip)
-		// 仅返回ipv4的ip
-		// if ip.To4() != nil {
-		// 	ipSlice = append(ipSlice, ip)
-		// } else {
-		// 	ipv6Flag = true
-		// }
+		if ipv4Only {
+			// 仅返回ipv4的ip
+			if ip.To4() != nil {
+				ipSlice = append(ipSlice, ip)
+			} else {
+				ipv6Flag = true
+			}
+		} else {
+			ipSlice = append(ipSlice, ip)
+		}
 	}
 
 	if ipv6Flag {
-		fmt.Println("[Info] IPv6 Traceroute is not supported right now.")
+		fmt.Println("[Info] IPv6 TCP/UDP Traceroute is not supported right now.")
 		if len(ipSlice) == 0 {
 			os.Exit(0)
 		}