[add] http method support

Author: cloverstd

main.go

@@ -20,6 +20,11 @@ var (
 	timeout     string
 	interval    string
 	sigs        chan os.Signal
+
+	httpMode bool
+	httpHead bool
+	httpPost bool
+	httpUA   string
 )
 
 var rootCmd = cobra.Command{
@@ -82,11 +87,16 @@ var rootCmd = cobra.Command{
 				return
 			}
 		}
-		protocol, err := ping.NewProtocol(schema)
-		if err != nil {
-			fmt.Println(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
+			}
 		}
 		target := ping.Target{
 			Timeout:  timeoutDuration,
@@ -99,8 +109,17 @@ var rootCmd = cobra.Command{
 		var pinger ping.Pinger
 		if schema == ping.TCP.String() {
 			pinger = ping.NewTCPing()
-		} else if schema == ping.HTTP.String() {
-			pinger = ping.NewHTTPing()
+		} else if schema == ping.HTTP.String() || schema == ping.HTTPS.String() {
+			var httpMethod string
+			switch {
+			case httpHead:
+				httpMethod = "HEAD"
+			case httpPost:
+				httpMethod = "POST"
+			default:
+				httpMethod = "GET"
+			}
+			pinger = ping.NewHTTPing(httpMethod)
 		} else {
 			fmt.Printf("schema: %s not support\n", schema)
 			cmd.Usage()
@@ -124,6 +143,12 @@ func init() {
 	rootCmd.Flags().IntVarP(&counter, "counter", "c", 4, "ping counter")
 	rootCmd.Flags().StringVarP(&timeout, "timeout", "T", "1s", `connect timeout, units are "ns", "us" (or "µs"), "ms", "s", "m", "h"`)
 	rootCmd.Flags().StringVarP(&interval, "interval", "I", "1s", `ping interval, units are "ns", "us" (or "µs"), "ms", "s", "m", "h"`)
+
+	rootCmd.Flags().BoolVarP(&httpMode, "http", "H", false, `Use "HTTP" mode. will ignore URI Schema, force to http`)
+	rootCmd.Flags().BoolVar(&httpHead, "head", false, `Use POST instead of GET in http mode.`)
+	rootCmd.Flags().BoolVar(&httpPost, "post", false, `Use HEAD instead of GET in http mode.`)
+	rootCmd.Flags().StringVar(&httpUA, "user-agent", "tcping", `Use custom UA in http mode.`)
+
 }
 
 func main() {

ping/http.go

@@ -1,6 +1,7 @@
 package ping
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -19,10 +20,10 @@ type HTTPing struct {
 var _ Pinger = (*HTTPing)(nil)
 
 // NewHTTPing return new HTTPing
-func NewHTTPing() *HTTPing {
+func NewHTTPing(method string) *HTTPing {
 	return &HTTPing{
 		done:   make(chan struct{}),
-		Method: "GET",
+		Method: method,
 	}
 }
 
@@ -53,7 +54,7 @@ func (ping *HTTPing) Start() <-chan struct{} {
 				} else {
 					defer resp.Body.Close()
 					length, _ := io.Copy(ioutil.Discard, resp.Body)
-					fmt.Printf("Ping %s - HTTP is open - time=%s status=%d bytes=%d\n", ping.target, duration, resp.StatusCode, length)
+					fmt.Printf("Ping %s - %s is open - time=%s method=%s status=%d bytes=%d\n", ping.target, ping.target.Protocol, duration, ping.Method, resp.StatusCode, length)
 					if ping.result.MinDuration == 0 {
 						ping.result.MinDuration = duration
 					}
@@ -88,10 +89,16 @@ func (ping *HTTPing) Stop() {
 
 func (ping HTTPing) ping() (time.Duration, *http.Response, error) {
 	var resp *http.Response
-	req, err := http.NewRequest(ping.Method, ping.target.String(), nil)
+	var body io.Reader
+	if ping.Method == "POST" {
+		body = bytes.NewBufferString("{}")
+	}
+	req, err := http.NewRequest(ping.Method, ping.target.String(), body)
+	req.Header.Set(http.CanonicalHeaderKey("User-Agent"), "tcping")
 	if err != nil {
 		return 0, nil, err
 	}
+
 	duration, errIfce := timeIt(func() interface{} {
 		client := http.Client{Timeout: ping.target.Timeout}
 		resp, err = client.Do(req)