pilse-5.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/httptest"
	"net/http/httptrace"
	"time"
)

const (
	useragent = "TurboBytes-Pulse/1.2" //Default user agent
)

var (
	tlshandshaketimeout = time.Second * 15 //Timeout for TLS handshake
	dialtimeout         = time.Second * 15 //Timeout for Dial (DNS + TCP connect)
	responsetimeout     = time.Second * 25 //Time out for response header
	keepalive           = time.Second * 30 //Keepalive timeout
)

type conInfo struct {
	DNS     time.Duration
	Connect time.Duration
	SSL     time.Duration
	TTFB    time.Duration
	Total   time.Duration
	//Transfer    time.Duration No Transfer time because we don't consume body
	Addr string
}

func (ci *conInfo) String() string {
	return fmt.Sprintf("Addr: %v\nDNS: %v\nConnect: %v\nSSL: %v\nTTFB: %v\nTotal: %v\n", ci.Addr, sincems(ci.DNS), sincems(ci.Connect), sincems(ci.SSL), sincems(ci.TTFB), sincems(ci.Total))
}

type conTrack struct {
	DNSStart             time.Time
	DNSDone              time.Time
	ConnectStart         map[string]time.Time
	ConnectDone          map[string]time.Time
	Addr                 string
	WroteRequest         time.Time
	GotFirstResponseByte time.Time
}

func sincems(t time.Duration) int64 {
	d := t.Nanoseconds()
	return d / (1000 * 1000)
}

func sincemstime(t time.Time) int64 {
	return sincems(time.Since(t))
}

func (ct *conTrack) getConInfo() *conInfo {
	ci := &conInfo{
		Addr: ct.Addr,
	}
	if ct.GotFirstResponseByte.After(ct.WroteRequest) {
		ci.TTFB = ct.GotFirstResponseByte.Sub(ct.WroteRequest)
	}
	if ct.DNSDone.After(ct.DNSStart) {
		ci.DNS = ct.DNSDone.Sub(ct.DNSStart)
	}
	if ct.Addr == "" && len(ct.ConnectStart) > 0 { //If no addr(cause FAIL) but map has key(s) use any
		for ct.Addr, _ = range ct.ConnectStart {
			//log.Println(ct.Addr)
		}
	}
	cs := ct.ConnectStart[ct.Addr]
	cd, ok := ct.ConnectDone[ct.Addr]
	if !ok {
		cd = time.Now() //If connect was never Done then use now to indicate how long we waited...
	}
	if cd.After(cs) {
		ci.Connect = cd.Sub(cs)
	}
	if ct.WroteRequest.After(cd) {
		ci.SSL = ct.WroteRequest.Sub(cd)
	}
	ci.Total = ci.DNS + ci.Connect + ci.SSL + ci.TTFB
	return ci
}

func dialContext(ctx context.Context, network, address string) (net.Conn, error) {
	con, err := (&net.Dialer{
		Timeout:   dialtimeout, //DNS + Connect
		KeepAlive: keepalive,
	}).DialContext(ctx, network, address)
	return con, err
}

var (
	ssl      bool
	endpoint string
	host     string
	path     string
)

func init() {
	flag.BoolVar(&ssl, "ssl", false, "run on https?")
	flag.StringVar(&endpoint, "endpoint", "", "what endpoint to hit. blank to copy from host")
	flag.StringVar(&host, "host", "www.cloudflare.com", "what hostname to access")
	flag.StringVar(&path, "path", "/", "what path")
	flag.Parse()
}

func main() {
	if endpoint == "" {
		endpoint = host
	}
	var url string
	if ssl {
		url = fmt.Sprintf("https://%s%s", endpoint, path)
	} else {
		url = fmt.Sprintf("http://%s%s", endpoint, path)
	}
	//Create a request object
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Fatal(err)
	}
	req.Header.Set("User-Agent", useragent)
	//Override Host header if needed
	tlshost := endpoint //Validate with endpoint if no host given
	if host != "" {
		req.Host = host
		tlshost = host //Validate with Host hdr if present
	}

	// Currently the transport leaks FD because currently http2
	// does not respect IdleConnTimeout
	// https://github.com/golang/go/issues/16808

	//Configure our transport, new one for each request
	transport := &http.Transport{
		Proxy:                 http.ProxyFromEnvironment,
		DialContext:           dialContext,
		MaxIdleConns:          100,              //Irrelevant
		IdleConnTimeout:       90 * time.Second, //Irrelevant
		TLSHandshakeTimeout:   tlshandshaketimeout,
		ExpectContinueTimeout: 1 * time.Second,
		ResponseHeaderTimeout: responsetimeout,
	}

	// Due to #16808, transport going out of scope does not cleanup
	// idle connections. We must do it by hand using CloseIdleConnections()
	defer func() {
		// There is something racey going on, noticed an issue on my dev machine
		// but not on prod. Does not hurt to sleep for a sec.
		time.Sleep(time.Second)
		transport.CloseIdleConnections()
	}()

	//Initialize our client
	client := http.Client{
		Transport: transport,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		}, //Since we now use high-level client we must stop redirects.
	}

	// A dilema... Configuring our own TLSClientConfig causes http
	// package to not kick in http2. Doing http2.ConfigureTransport
	// manually messes up httptrace. As a workaround we ho not configure
	// TLSClientConfig at first, then, if needed, we do a mock request
	// to fire off transport.onceSetNextProtoDefaults() and then sneak
	// in the transport.TLSClientConfig.ServerName that we want to configure.

	if ssl {
		if tlshost != endpoint {
			//Begin hacky workaround...
			//Make mock req to kick in onceSetNextProtoDefaults()
			server := httptest.NewServer(http.HandlerFunc(http.NotFound))
			reqtmp, _ := http.NewRequest("GET", server.URL, nil)
			client.Do(reqtmp) //Don't care about response..
			//Now mess with TLSClientConfig
			transport.TLSClientConfig.ServerName = tlshost
			//transport.TLSClientConfig = &tls.Config{ServerName: tlshost}
			//Not closing server was causing FD leak in prod but not in dev.. weird
			server.Close()
		}
	}

	//Initialize connection tracker
	ct := &conTrack{
		ConnectStart: make(map[string]time.Time),
		ConnectDone:  make(map[string]time.Time),
	}
	debugst := time.Now()
	//Initialize httptrace
	trace := &httptrace.ClientTrace{
		GotConn: func(connInfo httptrace.GotConnInfo) {
			ct.Addr = connInfo.Conn.RemoteAddr().String()
			log.Println(sincemstime(debugst), "GotConn: ", connInfo.Conn.RemoteAddr().String())
		},
		DNSStart: func(ds httptrace.DNSStartInfo) {
			ct.DNSStart = time.Now()
			log.Println(sincemstime(debugst), "DNSStart: ", ds)
		},
		DNSDone: func(dd httptrace.DNSDoneInfo) {
			ct.DNSDone = time.Now()
			log.Println(sincemstime(debugst), "DNSDone: ", dd)
		},
		ConnectStart: func(network, addr string) {
			ct.ConnectStart[addr] = time.Now()
			log.Println(sincemstime(debugst), "ConnectStart: ", network, addr)
		},
		ConnectDone: func(network, addr string, err error) {
			ct.ConnectDone[addr] = time.Now()
			log.Println(sincemstime(debugst), "ConnectDone: ", network, addr, err)
		},
		GotFirstResponseByte: func() {
			ct.GotFirstResponseByte = time.Now()
			log.Println(sincemstime(debugst), "GotFirstResponseByte")
		},
		WroteRequest: func(wr httptrace.WroteRequestInfo) {
			ct.WroteRequest = time.Now()
			log.Println(sincemstime(debugst), "WroteRequest: ", wr)
		},
	}
	//Wrap trace into req
	req = req.WithContext(httptrace.WithClientTrace(req.Context(), trace))
	//Make the request
	debugst = time.Now()
	resp, err := client.Do(req)
	ti := ct.getConInfo()
	//log.Println(ct, ti)
	//populate the result with timing info regardless of failure
	fmt.Println(ti)
	//On error stamp err and return
	if err != nil {
		log.Fatal(err)
	}
	resp.Body.Close()
	//Not a fail, extract more info
	log.Println(resp.Proto)
}