From fa0cc9bb7a341c42b32221a4dc63343fff92b1c2 Mon Sep 17 00:00:00 2001
From: naiba <hi@nai.ba>
Date: Sun, 28 Jul 2024 18:48:44 +0800
Subject: [PATCH] fix: tcp-ping

---
 cmd/agent/main.go      | 12 ++++++++++--
 cmd/agent/main_test.go | 15 +++++++++++++++
 pkg/monitor/myip.go    |  9 ++++++---
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/cmd/agent/main.go b/cmd/agent/main.go
index ec317c1..960dc4d 100644
--- a/cmd/agent/main.go
+++ b/cmd/agent/main.go
@@ -474,13 +474,21 @@ func handleUpgradeTask(*pb.Task, *pb.TaskResult) {
 }
 
 func handleTcpPingTask(task *pb.Task, result *pb.TaskResult) {
-	ipAddr, err := lookupIP(task.GetData())
+	host, port, err := net.SplitHostPort(task.GetData())
 	if err != nil {
 		result.Data = err.Error()
 		return
 	}
+	ipAddr, err := lookupIP(host)
+	if err != nil {
+		result.Data = err.Error()
+		return
+	}
+	if strings.Contains(ipAddr, ":") {
+		ipAddr = fmt.Sprintf("[%s]", ipAddr)
+	}
 	start := time.Now()
-	conn, err := net.DialTimeout("tcp", ipAddr, time.Second*10)
+	conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%s", ipAddr, port), time.Second*10)
 	if err == nil {
 		conn.Write([]byte("ping\n"))
 		conn.Close()
diff --git a/cmd/agent/main_test.go b/cmd/agent/main_test.go
index 3dd5aac..717cce9 100644
--- a/cmd/agent/main_test.go
+++ b/cmd/agent/main_test.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"net"
 	"reflect"
 	"testing"
 )
@@ -31,4 +32,18 @@ func TestLookupIP(t *testing.T) {
 	if err != nil {
 		t.Errorf("lookupIP failed: %v", err)
 	}
+	_, err = net.ResolveIPAddr("ip", "www.google.com")
+	if err != nil {
+		t.Errorf("ResolveIPAddr failed: %v", err)
+	}
+
+	ip, err = lookupIP("ipv6.google.com")
+	fmt.Printf("ip: %v, err: %v\n", ip, err)
+	if err != nil {
+		t.Errorf("lookupIP failed: %v", err)
+	}
+	_, err = net.ResolveIPAddr("ip", "ipv6.google.com")
+	if err != nil {
+		t.Errorf("ResolveIPAddr failed: %v", err)
+	}
 }
diff --git a/pkg/monitor/myip.go b/pkg/monitor/myip.go
index 9dd2766..e0acc34 100644
--- a/pkg/monitor/myip.go
+++ b/pkg/monitor/myip.go
@@ -39,14 +39,17 @@ func UpdateIP(useIPv6CountryCode bool, period uint32) {
 			continue
 		}
 		if ipv4 == "" || ipv6 == "" {
+			if ipv4 == "" {
+				GeoQueryIP = ipv6
+			} else {
+				GeoQueryIP = ipv4
+			}
 			CachedIP = fmt.Sprintf("%s%s", ipv4, ipv6)
 		} else {
 			CachedIP = fmt.Sprintf("%s/%s", ipv4, ipv6)
 		}
 
-		if !useIPv6CountryCode {
-			GeoQueryIP = ipv4
-		} else {
+		if useIPv6CountryCode {
 			GeoQueryIP = ipv6
 		}