ping.go

package ping

import (
	"bytes"
	"crypto/tls"
	"fmt"

	//"net"
	"net/http"
	"net/http/httptrace"
	"net/url"
	"strings"
	"time"
)

type Target struct {
	URL           *url.URL    `json:"url"`
	RequestHeader http.Header `json:"request_header"`
	RequestBody   string      `json:"request_body"`
	RequestMethod string      `json:"request_method"`
}

type Metric struct {
	StatusCode       int    `json:"status-code"`
	Status           string `json:"status"`
	DNSLookup        string `json:"dns_lookup"`
	TCPConnection    string `json:"tcp_connection"`
	TLSHandshake     string `json:"tls_handshake"`
	ServerProcessing string `json:"server_processing"`
	TTFB             string `json:"ttfb"` // Time To First(response) Byte
	Total            string `json:"total"`
}

type Ping struct {
	Timestamp      string      `json:"timestamp"` // UNIX timestamp
	Target         *Target     `json:"target"`
	Metric         Metric      `json:"metric"`
	IPAddr         []string    `json:"ip"`
	ResponseHeader http.Header `json:"response_header"`
	ResponseBody   string      `json:"response_body"`
}

func SanitizeURL(u string) *url.URL {
	if strings.Contains(u, "://") == false {
		u = "https://" + u
	}
	out, _ := url.Parse(u)
	return out
}

func IsValidUrl(u string) bool {
	_, err := url.ParseRequestURI(u)
	if err != nil {
		return false
	} else {
		return true
	}
}

// NewTarget returns a sane Target struct that when can use to determine
// latency.
// Parameters:
//	u - URL of target
//	h - sets custom headers for the request
//	b - sets a body; defaults to ""
//	r - request method; defaults to GET
func NewTarget(u string, h []string, b string, r string) (*Target, error) {
	// URL
	t := new(Target)
	if IsValidUrl(u) == false {
		return nil, fmt.Errorf("URL %s is not valid", u)
	}
	t.URL = SanitizeURL(u)

	// HTTP Headers
	t.RequestHeader = make(http.Header)
	for _, elem := range h {
		split := strings.Split(elem, ":")
		t.RequestHeader.Set(split[0], split[1])
	}

	// Body of request
	t.RequestBody = b

	// HTTP request method
	t.RequestMethod = r

	return t, nil
}

func (t *Target) Latency(printBody bool, printHeader bool) (*Ping, error) {
	var m Metric
	var p Ping

	body := strings.NewReader(t.RequestBody)

	req, err := http.NewRequest(t.RequestMethod, t.URL.String(), body)
	if err != nil {
		return nil, err
	}

	req.Header = t.RequestHeader

	var sDNS, dDNS, sConn, dConn, tGotConn, tGotFB, sTLS, dTLS time.Time
	trace := &httptrace.ClientTrace{
		DNSStart: func(dsi httptrace.DNSStartInfo) {
			sDNS = time.Now()
		},
		DNSDone: func(ddi httptrace.DNSDoneInfo) {
			dDNS = time.Now()
			for _, ip := range ddi.Addrs {
				p.IPAddr = append(p.IPAddr, ip.String())
			}
		},
		ConnectStart: func(network, addr string) {
			sConn = time.Now()
		},
		ConnectDone: func(network, addr string, err error) {
			dConn = time.Now()
		},
		GotConn: func(gci httptrace.GotConnInfo) {
			tGotConn = time.Now()
		},
		GotFirstResponseByte: func() {
			tGotFB = time.Now()
		},
		TLSHandshakeStart: func() {
			sTLS = time.Now()
		},
		TLSHandshakeDone: func(cs tls.ConnectionState, err error) {
			dTLS = time.Now()
		},
	}

	req = req.WithContext(httptrace.WithClientTrace(req.Context(), trace))
	tStart := time.Now()

	tr := &http.Transport{
		MaxIdleConns:          1,
		IdleConnTimeout:       10 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
		DisableKeepAlives:     true,
	}

	client := &http.Client{
		Transport: tr,
	}

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	bodyBuf := new(bytes.Buffer)
	bodyBuf.ReadFrom(resp.Body)

	resp.Body.Close()
	tStop := time.Now()

	m.StatusCode = resp.StatusCode
	m.Status = resp.Status
	m.DNSLookup = dDNS.Sub(sDNS).Round(time.Millisecond).String()
	m.TCPConnection = dConn.Sub(sConn).Round(time.Millisecond).String()
	m.TLSHandshake = dTLS.Sub(sTLS).Round(time.Millisecond).String()
	m.ServerProcessing = tGotFB.Sub(tGotConn).Round(time.Millisecond).String()
	m.TTFB = tGotFB.Sub(tStart).Round(time.Millisecond).String()
	m.Total = tStop.Sub(tStart).Round(time.Millisecond).String()

	p.Timestamp = time.Now().Format(time.UnixDate)
	p.Target = t
	p.Metric = m
	if printHeader {
		p.ResponseHeader = resp.Header
	} else {
		p.ResponseHeader = nil
	}
	if printBody {
		p.ResponseBody = bodyBuf.String()
	} else {
		p.ResponseBody = ""
	}

	return &p, nil
}

// Declutter main()
func (p *Ping) TerminalPrint(printBody bool, printRespHeader bool) {
	fmt.Printf("Target:\t\t\t%s\t\tIP: %s\n", p.Target.URL.Host, stringify(p.IPAddr))
	fmt.Printf("Timestamp:\t\t%s\n\n", p.Timestamp)
	fmt.Printf("Status:\t\t\t%s\n", p.Metric.Status)
	fmt.Printf("DNS Lookup:\t\t%s\n", p.Metric.DNSLookup)
	fmt.Printf("TCP Connection:\t\t%s\n", p.Metric.TCPConnection)
	fmt.Printf("TLS Handshake:\t\t%s\n", p.Metric.TLSHandshake)
	fmt.Printf("Server Processing:\t%s\n", p.Metric.ServerProcessing)
	fmt.Printf("Time To First Byte:\t%s\n", p.Metric.TTFB)
	fmt.Printf("Total:\t\t\t%s\n\n", p.Metric.Total)
	fmt.Printf("Request Headers:\n")
	for k, v := range p.Target.RequestHeader {
		fmt.Printf("\t%s:%s\n", k, stringify(v))
	}
	if printRespHeader == true {
		fmt.Printf("\n")
		fmt.Printf("Response Headers:\n")
		for k, v := range p.ResponseHeader {
			fmt.Printf("\t%s:%s\n", k, stringify(v))
		}
	}
	if printBody == true {
		fmt.Printf("\n")
		fmt.Printf("Body:\n")
		fmt.Printf("\t%s\n", p.ResponseBody)
	}
}

func stringify(s []string) string {
	return strings.Join(s[:], ", ")
}