diff --git a/defs/defs.go b/defs/defs.go index 44a8dad..54f2f9b 100644 --- a/defs/defs.go +++ b/defs/defs.go @@ -14,6 +14,14 @@ const ( StackDual ) +type PingType uint8 + +const ( + ICMP PingType = iota + UDP + HTTP +) + var ( BuildDate string ProgName string diff --git a/defs/options.go b/defs/options.go index 1e91c85..b2307e1 100644 --- a/defs/options.go +++ b/defs/options.go @@ -8,7 +8,8 @@ const ( OptionIPv6Alt = "6" OptionNoDownload = "no-download" OptionNoUpload = "no-upload" - OptionNoICMP = "no-icmp" + OptionPingType = "ping" + OptionPingTypeAlt = "p" OptionConcurrent = "concurrent" OptionConcurrentAlt = "n" OptionPingCount = "ping-count" diff --git a/defs/server.go b/defs/server.go index 2b9238b..59916e7 100644 --- a/defs/server.go +++ b/defs/server.go @@ -17,7 +17,6 @@ import ( "github.com/briandowns/spinner" "github.com/prometheus-community/pro-bing" log "github.com/sirupsen/logrus" - "github.com/syndtr/gocapability/capability" ) type ServerType uint8 @@ -47,7 +46,7 @@ type Server struct { UploadURI string `json:"upload"` PingURI string `json:"ping"` Type ServerType `json:"type"` - NoICMP bool `json:"-"` + PingType PingType `json:"-"` } func (s *Server) GetHost() string { @@ -148,28 +147,13 @@ func (s *Server) IsUp() bool { // ICMPPingAndJitter pings the server via ICMP echos and calculate the average ping and jitter func (s *Server) ICMPPingAndJitter(count int, srcIp, network string) (float64, float64, error) { - if s.NoICMP { - log.Debugf("Skipping ICMP for server %s, will use HTTP ping", s.Name) + if s.PingType == HTTP { return s.PingAndJitter(count + 2) } p := probing.New(s.Target) - if os.Getuid() <= 0 { + if s.PingType == ICMP { p.SetPrivileged(true) - } else { - if caps, err := capability.NewPid2(0); err == nil { - if err = caps.Load(); err == nil { - if caps.Get(capability.EFFECTIVE, capability.CAP_NET_RAW) { - p.SetPrivileged(true) - } else { - log.Warnf("ICMP ping requires `cap_net_raw` privilege, will use UDP ping") - } - } else { - log.Debugf("Failed to load capabilities: %s", err) - } - } else { - log.Debugf("Failed to load capabilities: %s", err) - } } p.SetNetwork(network) p.Count = count @@ -204,8 +188,8 @@ func (s *Server) ICMPPingAndJitter(count int, srcIp, network string) (float64, f } if len(stats.Rtts) == 0 { - s.NoICMP = true - log.Debugf("No ICMP pings returned for server %s (%s), trying TCP ping", s.Name, s.ID) + s.PingType = HTTP + log.Debugf("No ICMP/UDP pings returned for server %s (%s), trying TCP ping", s.Name, s.ID) return s.PingAndJitter(count + 2) } diff --git a/main.go b/main.go index d9f80b0..4d030ee 100644 --- a/main.go +++ b/main.go @@ -64,10 +64,15 @@ func main() { Usage: "Do not perform upload test", Hidden: true, }, - &cli.BoolFlag{ - Name: defs.OptionNoICMP, - Usage: "Do not use ICMP ping. ICMP doesn't work well under Linux\n" + - "\tat this moment, so you might want to disable it\n\t", + &cli.StringFlag{ + Name: defs.OptionPingType, + Aliases: []string{defs.OptionPingTypeAlt}, + Usage: "Change ping `TYPE`. Can be `icmp`, `udp` or `http`.\n" + + "\tFor Linux user, icmp ping needs you run as root or give\n" + + "\t`cap_net_raw` capability, unprivileged UDP ping needs you\n" + + "\tset `net.ipv4.ping_group_range` parameter cover all groups.\n" + + "\tFor Windows user, udp ping is not available\n\t", + Value: "icmp", }, &cli.IntFlag{ Name: defs.OptionConcurrent, diff --git a/speedtest/helper.go b/speedtest/helper.go index 5bb0dbf..b9061b5 100644 --- a/speedtest/helper.go +++ b/speedtest/helper.go @@ -317,7 +317,7 @@ func MatchISP(isp string) uint8 { } // doSpeedTest is where the actual speed test happens -func doSpeedTest(c *cli.Context, servers []defs.Server, network string, silent, noICMP bool, ispInfo *defs.IPInfoResponse) error { +func doSpeedTest(c *cli.Context, servers []defs.Server, network string, silent bool, pingType defs.PingType, ispInfo *defs.IPInfoResponse) error { if !silent || c.Bool(defs.OptionSimple) { if serverCount := len(servers); serverCount > 1 { fmt.Printf("Testing against %d servers: [ %s ]\n", serverCount, strings.Join(func() []string { @@ -383,7 +383,7 @@ func doSpeedTest(c *cli.Context, servers []defs.Server, network string, silent, } // skip ICMP if option given - currentServer.NoICMP = noICMP + currentServer.PingType = pingType p, jitter, err := currentServer.ICMPPingAndJitter(c.Int(defs.OptionPingCount), c.String(defs.OptionSource), network) if err != nil { diff --git a/speedtest/speedtest.go b/speedtest/speedtest.go index 9d73bbd..4abe6db 100644 --- a/speedtest/speedtest.go +++ b/speedtest/speedtest.go @@ -7,11 +7,13 @@ import ( "errors" "fmt" "github.com/jedib0t/go-pretty/v6/table" + "github.com/syndtr/gocapability/capability" "math" "math/rand" "net" "net/http" "os" + "runtime" "strconv" "strings" "sync" @@ -110,7 +112,33 @@ func SpeedTest(c *cli.Context) error { forceIPv4 := c.Bool(defs.OptionIPv4) forceIPv6 := c.Bool(defs.OptionIPv6) - noICMP := c.Bool(defs.OptionNoICMP) + var pingType defs.PingType + switch c.String(defs.OptionPingType) { + case "udp": + pingType = defs.UDP + if runtime.GOOS == "windows" { + log.Warn("UDP ping is not supported on Windows, will use ICMP ping") + pingType = defs.ICMP + } + case "http": + pingType = defs.HTTP + default: + pingType = defs.ICMP + if runtime.GOOS == "linux" { + if os.Getuid() > 0 { + if caps, err := capability.NewPid2(0); err != nil { + log.Debugf("Failed to load capabilities: %s, will use UDP ping", err) + pingType = defs.UDP + } else if err = caps.Load(); err != nil { + log.Debugf("Failed to load capabilities: %s, will use UDP ping", err) + pingType = defs.UDP + } else if !caps.Get(capability.EFFECTIVE, capability.CAP_NET_RAW) { + log.Warn("ICMP ping requires `cap_net_raw` privilege, will use UDP ping") + pingType = defs.UDP + } + } + } + } var network string var stack defs.Stack @@ -156,7 +184,10 @@ func SpeedTest(c *cli.Context) error { if iface != "" { defaultDialer = newInterfaceDialer(iface) - noICMP = true + if pingType != defs.HTTP { + log.Warnf("ICMP/UDP ping is disabled when using interface binding") + pingType = defs.HTTP + } } else { defaultDialer = &net.Dialer{ Timeout: 30 * time.Second, @@ -254,7 +285,7 @@ func SpeedTest(c *cli.Context) error { servers = append(servers, serversT...) } else { log.Debugf("Find %d servers", len(serversT)) - if server, ok := selectServer("", serversT, network, c, noICMP); ok { + if server, ok := selectServer("", serversT, network, c, pingType); ok { servers = append(servers, server) } } @@ -367,7 +398,7 @@ func SpeedTest(c *cli.Context) error { logPre := fmt.Sprintf("[%s%s] ", provinceMap[uint8(province)].Short, defs.ISPMap[uint8(isp)].Name) log.Debugf("%sFind %d servers", logPre, len(serversT)) if len(serversT) > 0 { - if server, ok := selectServer(logPre, serversT, network, c, noICMP); ok { + if server, ok := selectServer(logPre, serversT, network, c, pingType); ok { servers = append(servers, server) } } @@ -414,7 +445,7 @@ func SpeedTest(c *cli.Context) error { return nil } - return doSpeedTest(c, servers, network, silent, noICMP, ispInfo) + return doSpeedTest(c, servers, network, silent, pingType, ispInfo) } func initProvinceMap() map[uint8]defs.ProvinceInfo { @@ -427,7 +458,7 @@ func initProvinceMap() map[uint8]defs.ProvinceInfo { return provinceMap } -func selectServer(logPre string, servers []defs.Server, network string, c *cli.Context, noICMP bool) (defs.Server, bool) { +func selectServer(logPre string, servers []defs.Server, network string, c *cli.Context, pingType defs.PingType) (defs.Server, bool) { if len(servers) > 10 { r := rand.New(rand.NewSource(time.Now().Unix())) r.Shuffle(len(servers), func(i int, j int) { @@ -447,7 +478,7 @@ func selectServer(logPre string, servers []defs.Server, network string, c *cli.C // spawn 10 concurrent pingers for i := 0; i < 10; i++ { - go pingWorker(jobs, results, &wg, c.String(defs.OptionSource), network, noICMP) + go pingWorker(jobs, results, &wg, c.String(defs.OptionSource), network, pingType) } // send ping jobs to workers @@ -490,7 +521,7 @@ Loop: return servers[serverIdx], true } -func pingWorker(jobs <-chan PingJob, results chan<- PingResult, wg *sync.WaitGroup, srcIp, network string, noICMP bool) { +func pingWorker(jobs <-chan PingJob, results chan<- PingResult, wg *sync.WaitGroup, srcIp, network string, pingType defs.PingType) { for { job := <-jobs server := job.Server @@ -498,7 +529,7 @@ func pingWorker(jobs <-chan PingJob, results chan<- PingResult, wg *sync.WaitGro // check the server is up by accessing the ping URL and checking its returned value == empty and status code == 200 if server.IsUp() { // skip ICMP if option given - server.NoICMP = noICMP + server.PingType = pingType // if server is up, get ping ping, _, err := server.ICMPPingAndJitter(1, srcIp, network)