Refactor parse args

Author: cloverstd

go.mod

@@ -1,6 +1,6 @@
 module github.com/cloverstd/tcping
 
-go 1.17
+go 1.18
 
 require (
 	github.com/smartystreets/goconvey v1.7.2

main.go

@@ -4,10 +4,8 @@ import (
 	"fmt"
 	"os"
 	"os/signal"
-	"syscall"
-	"time"
-
 	"strconv"
+	"syscall"
 
 	"github.com/cloverstd/tcping/ping"
 	"github.com/spf13/cobra"
@@ -43,78 +41,70 @@ var rootCmd = cobra.Command{
   3. ping over http
   	> tcping -H google.com
   4. ping with URI schema
-  	> tcping http://hui.lu
+  	> tcping https://hui.lu
 	`,
 	Run: func(cmd *cobra.Command, args []string) {
 		if showVersion {
 			fmt.Printf("version: %s\n", version)
 			fmt.Printf("git: %s\n", gitCommit)
 			return
 		}
-		if len(args) != 2 && len(args) != 1 {
+		if len(args) == 0 {
 			cmd.Usage()
 			return
 		}
+		if len(args) > 2 {
+			cmd.Println("invalid command arguments")
+			return
+		}
 		host := args[0]
 
+		url, err := ping.ParseAddress(host)
+		if err != nil {
+			fmt.Printf("%s is an invalid target.\n", host)
+			return
+		}
+		defaultPort := "80"
+		if len(args) > 1 {
+			defaultPort = args[1]
+		}
+		port, err := strconv.Atoi(defaultPort)
+		if err != nil {
+			cmd.Printf("%s is invalid port.\n", defaultPort)
+			return
+		}
+		url.Host = fmt.Sprintf("%s:%d", url.Hostname(), port)
+
 		var (
-			err    error
-			port   int
 			schema string
 		)
-		if len(args) == 2 {
-			port, err = strconv.Atoi(args[1])
-			if err != nil {
-				fmt.Println("port should be integer")
-				cmd.Usage()
-				return
-			}
-			schema = ping.TCP.String()
-		} else {
-			var matched bool
-			schema, host, port, matched = ping.CheckURI(host)
-			if !matched {
-				fmt.Println("not a valid uri")
-				cmd.Usage()
-				return
-			}
-		}
-		var timeoutDuration time.Duration
-		if res, err := strconv.Atoi(timeout); err == nil {
-			timeoutDuration = time.Duration(res) * time.Millisecond
-		} else {
-			timeoutDuration, err = time.ParseDuration(timeout)
-			if err != nil {
-				fmt.Println("parse timeout failed", err)
-				cmd.Usage()
-				return
-			}
+
+		timeoutDuration, err := ping.ParseDuration(timeout)
+		if err != nil {
+			cmd.Println("parse timeout failed", err)
+			cmd.Usage()
+			return
 		}
 
-		var intervalDuration time.Duration
-		if res, err := strconv.Atoi(interval); err == nil {
-			intervalDuration = time.Duration(res) * time.Millisecond
-		} else {
-			intervalDuration, err = time.ParseDuration(interval)
-			if err != nil {
-				fmt.Println("parse interval failed", err)
-				cmd.Usage()
-				return
-			}
+		intervalDuration, err := ping.ParseDuration(interval)
+		if err != nil {
+			cmd.Println("parse interval failed", err)
+			cmd.Usage()
+			return
 		}
-		var protocol ping.Protocol
+
 		if httpMode {
-			protocol = ping.HTTP
-		} else {
-			protocol, err = ping.NewProtocol(schema)
-			if err != nil {
-				fmt.Println(err)
-				cmd.Usage()
-				return
-			}
+			url.Scheme = ping.HTTP.String()
+		}
+		protocol, err := ping.NewProtocol(url.Scheme)
+		if err != nil {
+			cmd.Println("invalid protocol", err)
+			cmd.Usage()
+			return
 		}
+
 		if len(dnsServer) != 0 {
-			ping.UseCustomeDNS(dnsServer)
+			ping.UseCustomDNS(dnsServer)
 		}
 
 		parseHost, _ := ping.FormatIP(host)

ping/address.go

@@ -0,0 +1,15 @@
+package ping
+
+import (
+	pkgurl "net/url"
+	"strings"
+)
+
+// ParseAddress will try to parse addr as url.URL.
+func ParseAddress(addr string) (*pkgurl.URL, error) {
+	if strings.Contains(addr, "://") {
+		// it maybe with scheme, try url.Parse
+		return pkgurl.Parse(addr)
+	}
+	return pkgurl.Parse("tcp://" + addr)
+}

ping/ping.go

@@ -22,7 +22,7 @@ func (protocol Protocol) String() string {
 	case HTTPS:
 		return "https"
 	}
-	return "unkown"
+	return "unknown"
 }
 
 const (
@@ -34,7 +34,7 @@ const (
 	HTTPS
 )
 
-// NewProtocol convert protocol stirng to Protocol
+// NewProtocol convert protocol string to Protocol
 func NewProtocol(protocol string) (Protocol, error) {
 	switch strings.ToLower(protocol) {
 	case TCP.String():

ping/utils.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"strconv"
 	"strings"
 	"time"
 )
@@ -15,8 +16,8 @@ func timeIt(f func() interface{}) (int64, interface{}) {
 	return endAt.UnixNano() - startAt.UnixNano(), res
 }
 
-// UseCustomeDNS will set the dns to default DNS resolver for global
-func UseCustomeDNS(dns []string) {
+// UseCustomDNS will set the dns to default DNS resolver for global
+func UseCustomDNS(dns []string) {
 	resolver := net.Resolver{
 		PreferGo: true,
 		Dial: func(ctx context.Context, network, address string) (conn net.Conn, err error) {
@@ -53,3 +54,11 @@ func FormatIP(IP string) (string, error) {
 	}
 	return "", fmt.Errorf("Error IP format")
 }
+
+// ParseDuration parse the t as time.Duration, it will parse t as mills when missing unit.
+func ParseDuration(t string) (time.Duration, error) {
+	if timeout, err := strconv.ParseInt(t, 10, 64); err == nil {
+		return time.Duration(timeout) * time.Millisecond, nil
+	}
+	return time.ParseDuration(t)
+}