From 42aa774f6aecdb0309cadf04aa0b76e76913c9ea Mon Sep 17 00:00:00 2001
From: Damien Degois <damien@degois.info>
Date: Tue, 31 Oct 2023 01:11:24 +0100
Subject: [PATCH] Add ip family reslution hint in tcp probe and ping

Introduce ip:// notation
Introduce proto<ip family>:// hint
---
 README.md      | 11 ++++++++++-
 main.go        | 13 +++++++++++++
 pingwrapper.go | 42 ++++++++++++++++++++++++++++++++----------
 3 files changed, 55 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 981240f..21632dc 100644
--- a/README.md
+++ b/README.md
@@ -27,6 +27,11 @@ Available probing means are:
 
 Pure go is the default option but for unprivileged users ([see linux notes](#linux-notes-on-pure-go-ping)), OS/system's ping command (usually available on OS with specific cap or setuid) can be used with a background spawn model with `-s` flag. Privileged mode (default when user is root or on windows) can be forcefully enabled with `-privileged`.
 
+Hint ca be given about address family resolution using `ip<family>://`, `ip://` is the default, `ip4://` to force IPv4 and `ip6://` to force IPv6, example:
+ - `google.com` is equivalent to `ip://google.com`
+ - `ip4://google.com` forces resolution of google.com as ipv4
+ - `ip6://google.com` forces resolution of google.com as ipv6
+
 ### TCP probing
 
 For tcp probing, on linux, freebsd and openbsd, S/SA/R pattern is used. This allows to probe tcp ports without really triggering an accept on the listening app. Issue is if a device in between perform syn poxing, the result might not reflect reality.
@@ -37,6 +42,10 @@ tcp probing example syntax:
 - `tcp://192.168.0.1:443`
 - `tcp://[::1]:22`
 
+As for `ip://`, `tcp://` can also have hint of address family:
+- `tcp4://google.com:80` forces resolution of google.com as ipv4
+- `tcp6://google.com:80` forces resolution of google.com as ipv6
+
 ### Transition logging
 
 Transition logging can be enabled using `-log filename`.
@@ -71,6 +80,6 @@ Github repository: https://github.com/babs/multiping
 
 ### libs used
 
-* https://github.com/pterm/pterm 
+* https://github.com/pterm/pterm
 * https://github.com/prometheus-community/pro-bing
 * https://github.com/tevino/tcp-shaker
diff --git a/main.go b/main.go
index 2d74f1a..b90aa19 100644
--- a/main.go
+++ b/main.go
@@ -92,4 +92,17 @@ func usage() {
 	fmt.Print(VersionStringLong())
 	fmt.Fprintf(flag.CommandLine.Output(), "Usage of %s:\n", os.Args[0])
 	flag.PrintDefaults()
+	fmt.Println(`  host [hosts...]
+
+Hosts can have the following form:
+- hostname or ip or ip://hostname => ping (implementation used depends on '-s' flag)
+- tcp://hostname:port or tcp://[ipv6]:port => tcp probing
+    While using ip addresses, tcp:// can take IPv4 or IPv6 (w/ brackets), tcp4:// can only take IPv4 and tcp6:// only IPv6 (w/ brackets)
+
+Hint on address family can be provided with the following form:
+- ip://hostname and tcp://hostname resolves as default
+- ip4://hostname and tcp4://hostname resolves as IPv4
+- ip6://hostname and tcp6://hostname resolves as IPv6
+
+Notes about implementation: tcp implementation between probing (S/SA/R) and full handshake depends on the platform`)
 }
diff --git a/pingwrapper.go b/pingwrapper.go
index 18ccacf..b9f7b83 100644
--- a/pingwrapper.go
+++ b/pingwrapper.go
@@ -15,41 +15,63 @@ type PingWrapperInterface interface {
 	CalcStats(int64) PWStats
 }
 
-var re_host_w_proto = regexp.MustCompile(`^(tcp)://(\[?.+?\]?):(\d+)$`)
+var re_host_w_proto = regexp.MustCompile(`^(tcp|ip)([46])?://(\[?.+?\]?)(?::(\d+))?$`)
 
 func NewPingWrapper(host string, options Options, transition_writer *TransitionWriter) PingWrapperInterface {
 
 	host_findings := re_host_w_proto.FindAllStringSubmatch(host, -1)
+
+	var found_proto, found_ip_family, found_host, found_port string
+	var found_port_int int
+
 	if len(host_findings) > 0 {
-		port, err := strconv.Atoi(host_findings[0][3])
+		found_proto = host_findings[0][1]
+		found_ip_family = host_findings[0][2]
+		found_host = host_findings[0][3]
+		found_port = host_findings[0][4]
+	} else {
+		found_host = host
+	}
+
+	if found_proto == "tcp" {
+
+		if found_port == "" {
+			log.Fatalf("%v: tcp probing requested but no port given\n", host)
+		}
+		port, err := strconv.Atoi(found_port)
 		if err != nil {
-			log.Fatal(err)
+			log.Fatalf("%v: %v\n", host, err)
 		}
+		if port <= 0 || port > 65535 {
+			log.Fatalf("%v: tcp probing port invalid: %v\n", host, port)
+		}
+		found_port_int = port
+
 		return &TCPPingWrapper{
-			host:  host_findings[0][2],
-			ip:    mustResolve(host_findings[0][2]),
-			port:  port,
+			host:  found_host,
+			ip:    mustResolve(found_host, found_ip_family),
+			port:  found_port_int,
 			stats: &PWStats{transition_writer: transition_writer},
 		}
 	} else if *options.system {
 		return &SystemPingWrapper{
 			host:  host,
-			ip:    mustResolve(host),
+			ip:    mustResolve(found_host, found_ip_family),
 			stats: &PWStats{transition_writer: transition_writer},
 		}
 	} else {
 		return &ProbingWrapper{
 			host:       host,
-			ip:         mustResolve(host),
+			ip:         mustResolve(found_host, found_ip_family),
 			privileged: *options.privileged,
 			stats:      &PWStats{transition_writer: transition_writer},
 		}
 	}
 }
 
-func mustResolve(host string) *net.IPAddr {
+func mustResolve(host string, ip_family string) *net.IPAddr {
 	host = strings.Trim(host, "[]")
-	ipaddr, err := net.ResolveIPAddr("ip", host)
+	ipaddr, err := net.ResolveIPAddr("ip"+ip_family, host)
 	if err != nil {
 		log.Fatal(err)
 	}