net/http/httptrace: add ClientTrace.TLSHandshakeStart & TLSHandshakeDone

Fixes #16965

Change-Id: I3638fe280a5b1063ff589e6e1ff8a97c74b77c66
Reviewed-on: https://go-review.googlesource.com/30359
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

4 files changed, 90 insertions, 2 deletions

diff --git a/src/go/build/deps_test.go b/src/go/build/deps_test.go
index 6da1e68fde..cbdcca4ac8 100644
--- a/src/go/build/deps_test.go
+++ b/src/go/build/deps_test.go
@@ -396,7 +396,7 @@ var pkgDeps = map[string][]string{
 		"runtime/debug",
 	},
 	"net/http/internal":  {"L4"},
-	"net/http/httptrace": {"context", "internal/nettrace", "net", "reflect", "time"},
+	"net/http/httptrace": {"context", "crypto/tls", "internal/nettrace", "net", "reflect", "time"},
 
 	// HTTP-using packages.
 	"expvar":             {"L4", "OS", "encoding/json", "net/http"},
diff --git a/src/net/http/httptrace/trace.go b/src/net/http/httptrace/trace.go
index 8c29c4aa6f..5b042c097f 100644
--- a/src/net/http/httptrace/trace.go
+++ b/src/net/http/httptrace/trace.go
@@ -8,6 +8,7 @@ package httptrace
 
 import (
 	"context"
+	"crypto/tls"
 	"internal/nettrace"
 	"net"
 	"reflect"
@@ -119,6 +120,16 @@ type ClientTrace struct {
 	// enabled, this may be called multiple times.
 	ConnectDone func(network, addr string, err error)
 
+	// TLSHandshakeStart is called when the TLS handshake is started. When
+	// connecting to a HTTPS site via a HTTP proxy, the handshake happens after
+	// the CONNECT request is processed by the proxy.
+	TLSHandshakeStart func()
+
+	// TLSHandshakeDone is called after the TLS handshake with either the
+	// successful handshake's connection state, or a non-nil error on handshake
+	// failure.
+	TLSHandshakeDone func(tls.ConnectionState, error)
+
 	// WroteHeaders is called after the Transport has written
 	// the request headers.
 	WroteHeaders func()
diff --git a/src/net/http/transport.go b/src/net/http/transport.go
index 5594c948cd..429f667c14 100644
--- a/src/net/http/transport.go
+++ b/src/net/http/transport.go
@@ -955,6 +955,7 @@ func (t *Transport) dialConn(ctx context.Context, cm connectMethod) (*persistCon
 		writeErrCh:    make(chan error, 1),
 		writeLoopDone: make(chan struct{}),
 	}
+	trace := httptrace.ContextClientTrace(ctx)
 	tlsDial := t.DialTLS != nil && cm.targetScheme == "https" && cm.proxyURL == nil
 	if tlsDial {
 		var err error
@@ -968,11 +969,20 @@ func (t *Transport) dialConn(ctx context.Context, cm connectMethod) (*persistCon
 		if tc, ok := pconn.conn.(*tls.Conn); ok {
 			// Handshake here, in case DialTLS didn't. TLSNextProto below
 			// depends on it for knowing the connection state.
+			if trace != nil && trace.TLSHandshakeStart != nil {
+				trace.TLSHandshakeStart()
+			}
 			if err := tc.Handshake(); err != nil {
 				go pconn.conn.Close()
+				if trace != nil && trace.TLSHandshakeDone != nil {
+					trace.TLSHandshakeDone(tls.ConnectionState{}, err)
+				}
 				return nil, err
 			}
 			cs := tc.ConnectionState()
+			if trace != nil && trace.TLSHandshakeDone != nil {
+				trace.TLSHandshakeDone(cs, nil)
+			}
 			pconn.tlsState = &cs
 		}
 	} else {
@@ -1042,6 +1052,9 @@ func (t *Transport) dialConn(ctx context.Context, cm connectMethod) (*persistCon
 			})
 		}
 		go func() {
+			if trace != nil && trace.TLSHandshakeStart != nil {
+				trace.TLSHandshakeStart()
+			}
 			err := tlsConn.Handshake()
 			if timer != nil {
 				timer.Stop()
@@ -1050,6 +1063,9 @@ func (t *Transport) dialConn(ctx context.Context, cm connectMethod) (*persistCon
 		}()
 		if err := <-errc; err != nil {
 			plainConn.Close()
+			if trace != nil && trace.TLSHandshakeDone != nil {
+				trace.TLSHandshakeDone(tls.ConnectionState{}, err)
+			}
 			return nil, err
 		}
 		if !cfg.InsecureSkipVerify {
@@ -1059,6 +1075,9 @@ func (t *Transport) dialConn(ctx context.Context, cm connectMethod) (*persistCon
 			}
 		}
 		cs := tlsConn.ConnectionState()
+		if trace != nil && trace.TLSHandshakeDone != nil {
+			trace.TLSHandshakeDone(cs, nil)
+		}
 		pconn.tlsState = &cs
 		pconn.conn = tlsConn
 	}
diff --git a/src/net/http/transport_test.go b/src/net/http/transport_test.go
index cef2acc456..147b468e78 100644
--- a/src/net/http/transport_test.go
+++ b/src/net/http/transport_test.go
@@ -3288,6 +3288,12 @@ func testTransportEventTrace(t *testing.T, h2 bool, noHooks bool) {
 			close(gotWroteReqEvent)
 		},
 	}
+	if h2 {
+		trace.TLSHandshakeStart = func() { logf("tls handshake start") }
+		trace.TLSHandshakeDone = func(s tls.ConnectionState, err error) {
+			logf("tls handshake done. ConnectionState = %v \n err = %v", s, err)
+		}
+	}
 	if noHooks {
 		// zero out all func pointers, trying to get some path to crash
 		*trace = httptrace.ClientTrace{}
@@ -3339,7 +3345,10 @@ func testTransportEventTrace(t *testing.T, h2 bool, noHooks bool) {
 	wantOnceOrMore("connected to tcp " + addrStr + " = <nil>")
 	wantOnce("Reused:false WasIdle:false IdleTime:0s")
 	wantOnce("first response byte")
-	if !h2 {
+	if h2 {
+		wantOnce("tls handshake start")
+		wantOnce("tls handshake done")
+	} else {
 		wantOnce("PutIdleConn = <nil>")
 	}
 	wantOnce("Wait100Continue")
@@ -3411,6 +3420,55 @@ func TestTransportEventTraceRealDNS(t *testing.T) {
 	}
 }
 
+// Test the httptrace.TLSHandshake{Start,Done} hooks with a https http1
+// connections. The http2 test is done in TestTransportEventTrace_h2
+func TestTLSHandshakeTrace(t *testing.T) {
+	defer afterTest(t)
+	s := httptest.NewTLSServer(HandlerFunc(func(w ResponseWriter, r *Request) {}))
+	defer s.Close()
+
+	var mu sync.Mutex
+	var start, done bool
+	trace := &httptrace.ClientTrace{
+		TLSHandshakeStart: func() {
+			mu.Lock()
+			defer mu.Unlock()
+			start = true
+		},
+		TLSHandshakeDone: func(s tls.ConnectionState, err error) {
+			mu.Lock()
+			defer mu.Unlock()
+			done = true
+			if err != nil {
+				t.Fatal("Expected error to be nil but was:", err)
+			}
+		},
+	}
+
+	tr := &Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}
+	defer tr.CloseIdleConnections()
+	c := &Client{Transport: tr}
+	req, err := NewRequest("GET", s.URL, nil)
+	if err != nil {
+		t.Fatal("Unable to construct test request:", err)
+	}
+	req = req.WithContext(httptrace.WithClientTrace(req.Context(), trace))
+
+	r, err := c.Do(req)
+	if err != nil {
+		t.Fatal("Unexpected error making request:", err)
+	}
+	r.Body.Close()
+	mu.Lock()
+	defer mu.Unlock()
+	if !start {
+		t.Fatal("Expected TLSHandshakeStart to be called, but wasn't")
+	}
+	if !done {
+		t.Fatal("Expected TLSHandshakeDone to be called, but wasnt't")
+	}
+}
+
 func TestTransportMaxIdleConns(t *testing.T) {
 	defer afterTest(t)
 	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {