feat: support ping with tcp (#156)

* feat: support ping under tcp

* chore: lint fix

* chore: update README.md

* chore: typo fix

README.md

@@ -54,7 +54,7 @@ Flags:
       --search=SEARCH          Fuzzy search servers by a keyword.
       --no-download            Disable download test.
       --no-upload              Disable upload test.
-      --force-http-ping        Force ping using http.
+      --ping-mode              Select a method for Ping. (support icmp/tcp/http)
   -d  --debug                  Enable debug mode.
       --version                Show application version.
 ```

speedtest.go

@@ -12,23 +12,23 @@ import (
 )
 
 var (
-	showList      = kingpin.Flag("list", "Show available speedtest.net servers.").Short('l').Bool()
-	serverIds     = kingpin.Flag("server", "Select server id to run speedtest.").Short('s').Ints()
-	customURL     = kingpin.Flag("custom-url", "Specify the url of the server instead of getting a list from speedtest.net.").String()
-	savingMode    = kingpin.Flag("saving-mode", "Test with few resources, though low accuracy (especially > 30Mbps).").Bool()
-	jsonOutput    = kingpin.Flag("json", "Output results in json format.").Bool()
-	location      = kingpin.Flag("location", "Change the location with a precise coordinate. Format: lat,lon").String()
-	city          = kingpin.Flag("city", "Change the location with a predefined city label.").String()
-	showCityList  = kingpin.Flag("city-list", "List all predefined city labels.").Bool()
-	proxy         = kingpin.Flag("proxy", "Set a proxy(http[s] or socks) for the speedtest.").String()
-	source        = kingpin.Flag("source", "Bind a source interface for the speedtest.").String()
-	multi         = kingpin.Flag("multi", "Enable multi-server mode.").Short('m').Bool()
-	thread        = kingpin.Flag("thread", "Set the number of concurrent connections.").Short('t').Int()
-	search        = kingpin.Flag("search", "Fuzzy search servers by a keyword.").String()
-	noDownload    = kingpin.Flag("no-download", "Disable download test.").Bool()
-	noUpload      = kingpin.Flag("no-upload", "Disable upload test.").Bool()
-	forceHTTPPing = kingpin.Flag("force-http-ping", "Force ping using http.").Bool()
-	debug         = kingpin.Flag("debug", "Enable debug mode.").Short('d').Bool()
+	showList     = kingpin.Flag("list", "Show available speedtest.net servers.").Short('l').Bool()
+	serverIds    = kingpin.Flag("server", "Select server id to run speedtest.").Short('s').Ints()
+	customURL    = kingpin.Flag("custom-url", "Specify the url of the server instead of getting a list from speedtest.net.").String()
+	savingMode   = kingpin.Flag("saving-mode", "Test with few resources, though low accuracy (especially > 30Mbps).").Bool()
+	jsonOutput   = kingpin.Flag("json", "Output results in json format.").Bool()
+	location     = kingpin.Flag("location", "Change the location with a precise coordinate. Format: lat,lon").String()
+	city         = kingpin.Flag("city", "Change the location with a predefined city label.").String()
+	showCityList = kingpin.Flag("city-list", "List all predefined city labels.").Bool()
+	proxy        = kingpin.Flag("proxy", "Set a proxy(http[s] or socks) for the speedtest.").String()
+	source       = kingpin.Flag("source", "Bind a source interface for the speedtest.").String()
+	multi        = kingpin.Flag("multi", "Enable multi-server mode.").Short('m').Bool()
+	thread       = kingpin.Flag("thread", "Set the number of concurrent connections.").Short('t').Int()
+	search       = kingpin.Flag("search", "Fuzzy search servers by a keyword.").String()
+	noDownload   = kingpin.Flag("no-download", "Disable download test.").Bool()
+	noUpload     = kingpin.Flag("no-upload", "Disable upload test.").Bool()
+	pingMode     = kingpin.Flag("ping-mode", "Select a method for Ping. (support icmp/tcp/http)").Default("http").String()
+	debug        = kingpin.Flag("debug", "Enable debug mode.").Short('d').Bool()
 )
 
 func main() {
@@ -44,7 +44,7 @@ func main() {
 			Proxy:        *proxy,
 			Source:       *source,
 			Debug:        *debug,
-			ICMP:         (os.Geteuid() == 0 || os.Geteuid() == -1) && len(*proxy) == 0 && !*forceHTTPPing, // proxy may not support ICMP
+			PingMode:     parseProto(*pingMode), // TCP as default
 			SavingMode:   *savingMode,
 			CityFlag:     *city,
 			LocationFlag: *location,
@@ -176,6 +176,16 @@ func showServerList(servers speedtest.Servers) {
 	}
 }
 
+func parseProto(str string) speedtest.Proto {
+	if str == "icmp" {
+		return speedtest.ICMP
+	} else if str == "tcp" {
+		return speedtest.TCP
+	} else {
+		return speedtest.HTTP
+	}
+}
+
 func AppInfo() {
 	if !*jsonOutput {
 		fmt.Println()

speedtest/request.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/showwin/speedtest-go/speedtest/tcp"
 	"math"
 	"net/http"
 	"net/url"
@@ -22,6 +23,10 @@ var (
 	ulSizes = [...]int{100, 300, 500, 800, 1000, 1500, 2500, 3000, 3500, 4000} // kB
 )
 
+var (
+	ErrConnectTimeout = errors.New("server connect timeout")
+)
+
 func (s *Server) MultiDownloadTestContext(ctx context.Context, servers Servers) error {
 	if s.Context.config.NoDownload {
 		dbg.Println("Download test disabled")
@@ -175,10 +180,12 @@ func (s *Server) PingTest(callback func(latency time.Duration)) error {
 func (s *Server) PingTestContext(ctx context.Context, callback func(latency time.Duration)) (err error) {
 	start := time.Now()
 	var vectorPingResult []int64
-	if s.Context.config.ICMP {
+	if s.Context.config.PingMode == TCP {
+		vectorPingResult, err = s.TCPPing(ctx, 10, time.Millisecond*200, callback)
+	} else if s.Context.config.PingMode == ICMP {
 		vectorPingResult, err = s.ICMPPing(ctx, time.Second*4, 10, time.Millisecond*200, callback)
 	} else {
-		vectorPingResult, err = s.HTTPPing(ctx, 10, time.Millisecond*200, nil)
+		vectorPingResult, err = s.HTTPPing(ctx, 10, time.Millisecond*200, callback)
 	}
 	if err != nil || len(vectorPingResult) == 0 {
 		return err
@@ -208,6 +215,46 @@ func (s *Server) TestAll() error {
 	return s.UploadTest()
 }
 
+func (s *Server) TCPPing(
+	ctx context.Context,
+	echoTimes int,
+	echoFreq time.Duration,
+	callback func(latency time.Duration),
+) (latencies []int64, err error) {
+	var pingDst string
+	if len(s.Host) == 0 {
+		u, err := url.Parse(s.URL)
+		if err != nil || len(u.Host) == 0 {
+			return nil, err
+		}
+		pingDst = u.Host
+	} else {
+		pingDst = s.Host
+	}
+	failTimes := 0
+	client := tcp.NewClient(s.Context.tcpDialer, pingDst)
+	err = client.Connect()
+	if err != nil {
+		return nil, err
+	}
+	for i := 0; i < echoTimes; i++ {
+		latency, err := client.PingContext(ctx)
+		if err != nil {
+			failTimes++
+			continue
+		}
+		latencies = append(latencies, latency)
+		if callback != nil {
+			callback(time.Duration(latency))
+		}
+		time.Sleep(echoFreq)
+	}
+	if failTimes == echoTimes {
+		return nil, ErrConnectTimeout
+	}
+	return
+}
+
 func (s *Server) HTTPPing(
 	ctx context.Context,
 	echoTimes int,
@@ -241,7 +288,7 @@ func (s *Server) HTTPPing(
 		time.Sleep(echoFreq)
 	}
 	if failTimes == echoTimes {
-		return nil, errors.New("server connect timeout")
+		return nil, ErrConnectTimeout
 	}
 	return
 }
@@ -318,7 +365,7 @@ func (s *Server) ICMPPing(
 		time.Sleep(echoFreq)
 	}
 	if failTimes == echoTimes {
-		return nil, errors.New("server connect timeout")
+		return nil, ErrConnectTimeout
 	}
 	return
 }

speedtest/server.go

@@ -93,8 +93,10 @@ func (s *Speedtest) CustomServer(host string) (*Server, error) {
 }
 
 // ServerList list of Server
+// Users(Client) also exists with @param speedTestServersAdvanced
 type ServerList struct {
 	Servers []*Server `xml:"servers>server"`
+	Users   []User    `xml:"client"`
 }
 
 // Servers for sorting servers.
@@ -174,6 +176,13 @@ func (s *Speedtest) FetchServerByIDContext(ctx context.Context, serverID string)
 	for i := range list.Servers {
 		if list.Servers[i].ID == serverID {
 			list.Servers[i].Context = s
+			if len(list.Users) > 0 {
+				sLat, _ := strconv.ParseFloat(list.Servers[i].Lat, 64)
+				sLon, _ := strconv.ParseFloat(list.Servers[i].Lon, 64)
+				uLat, _ := strconv.ParseFloat(list.Users[0].Lat, 64)
+				uLon, _ := strconv.ParseFloat(list.Users[0].Lon, 64)
+				list.Servers[i].Distance = distance(sLat, sLon, uLat, uLon)
+			}
 			return list.Servers[i], err
 		}
 	}
@@ -274,7 +283,9 @@ func (s *Speedtest) FetchServerListContext(ctx context.Context) (Servers, error)
 		wg.Add(1)
 		go func(gs *Server) {
 			var latency []int64
-			if s.config.ICMP {
+			if s.config.PingMode == TCP {
+				latency, err = gs.TCPPing(pCtx, 1, time.Millisecond, nil)
+			} else if s.config.PingMode == ICMP {
 				latency, err = gs.ICMPPing(pCtx, 4*time.Second, 1, time.Millisecond, nil)
 			} else {
 				latency, err = gs.HTTPPing(pCtx, 1, time.Millisecond, nil)
@@ -390,6 +401,9 @@ func (s *Server) String() string {
 	if s.Sponsor == "?" {
 		return fmt.Sprintf("[%4s] %s", s.ID, s.Name)
 	}
+	if len(s.Country) == 0 {
+		return fmt.Sprintf("[%4s] %.2fkm %s by %s", s.ID, s.Distance, s.Name, s.Sponsor)
+	}
 	return fmt.Sprintf("[%4s] %.2fkm %s (%s) by %s", s.ID, s.Distance, s.Name, s.Country, s.Sponsor)
 }
 

speedtest/speedtest.go

@@ -15,6 +15,14 @@ var (
 	DefaultUserAgent = fmt.Sprintf("showwin/speedtest-go %s", version)
 )
 
+type Proto int
+
+const (
+	ICMP Proto = iota
+	TCP
+	HTTP
+)
+
 // Speedtest is a speedtest client.
 type Speedtest struct {
 	User *User
@@ -32,7 +40,7 @@ type UserConfig struct {
 	Proxy     string
 	Source    string
 	Debug     bool
-	ICMP      bool
+	PingMode  Proto
 
 	SavingMode bool
 
@@ -159,7 +167,7 @@ func WithUserConfig(userConfig *UserConfig) Option {
 		dbg.Printf("Proxy: %s\n", s.config.Proxy)
 		dbg.Printf("SavingMode: %v\n", s.config.SavingMode)
 		dbg.Printf("Keyword: %v\n", s.config.Keyword)
-		dbg.Printf("ICMP: %v\n", s.config.ICMP)
+		dbg.Printf("PingType: %v\n", s.config.PingMode)
 		dbg.Printf("OS: %s, ARCH: %s, NumCPU: %d\n", runtime.GOOS, runtime.GOARCH, runtime.NumCPU())
 	}
 }