From 5db6c0142a58b9afdac6413710ad976b5594b2f6 Mon Sep 17 00:00:00 2001
From: Yukinari Toyota <xxseyxx@gmail.com>
Date: Fri, 10 May 2019 02:19:53 +0900
Subject: [PATCH] fix nilp error by #102 for tls conn (again) (#109)

* fix nilp error by #102 for tls conn

* extend recv time for GOMAXPROCS=1 env
---
 xray/client_test.go | 43 +++++++++++++++++++++++++++++++++++++++++++
 xray/httptrace.go   |  2 +-
 xray/util_test.go   |  2 +-
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/xray/client_test.go b/xray/client_test.go
index 58be9d28..00fc1eb7 100644
--- a/xray/client_test.go
+++ b/xray/client_test.go
@@ -279,6 +279,49 @@ func TestRoundTripReuseDatarace(t *testing.T) {
 	wg.Wait()
 }
 
+func TestRoundTripReuseTLSDatarace(t *testing.T) {
+	ts := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		b := []byte(`200 - Nothing to see`)
+		w.WriteHeader(http.StatusOK)
+		w.Write(b)
+	}))
+	defer ts.Close()
+
+	certpool := x509.NewCertPool()
+	certpool.AddCert(ts.Certificate())
+	tr := &http.Transport{
+		TLSClientConfig: &tls.Config{
+			RootCAs: certpool,
+		},
+	}
+	rt := &roundtripper{
+		Base: tr,
+	}
+
+	wg := sync.WaitGroup{}
+	n := 30
+	wg.Add(n)
+	for i := 0; i < n; i++ {
+		go func() {
+			defer wg.Done()
+			reader := strings.NewReader("")
+			ctx, root := BeginSegment(context.Background(), "Test")
+			req, _ := http.NewRequest("GET", ts.URL, reader)
+			req = req.WithContext(ctx)
+			res, err := rt.RoundTrip(req)
+			assert.NoError(t, err)
+			ioutil.ReadAll(res.Body)
+			res.Body.Close() // make net/http/transport.go connection reuse
+			root.Close(nil)
+		}()
+	}
+	for i := 0; i < n; i++ {
+		_, e := TestDaemon.Recv()
+		assert.NoError(t, e)
+	}
+	wg.Wait()
+}
+
 func TestRoundTripHttp2Datarace(t *testing.T) {
 	ts := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		b := []byte(`200 - Nothing to see`)
diff --git a/xray/httptrace.go b/xray/httptrace.go
index ac66a851..c56c40fa 100644
--- a/xray/httptrace.go
+++ b/xray/httptrace.go
@@ -97,7 +97,7 @@ func (xt *HTTPSubsegments) ConnectDone(network, addr string, err error) {
 // TLSHandshakeStart begins a tls subsegment if the HTTP operation
 // subsegment is still in progress.
 func (xt *HTTPSubsegments) TLSHandshakeStart() {
-	if GetSegment(xt.opCtx).safeInProgress() {
+	if GetSegment(xt.opCtx).safeInProgress() && xt.connCtx != nil {
 		xt.tlsCtx, _ = BeginSubsegment(xt.connCtx, "tls")
 	}
 }
diff --git a/xray/util_test.go b/xray/util_test.go
index 5b47ab3d..61ea8b79 100644
--- a/xray/util_test.go
+++ b/xray/util_test.go
@@ -74,7 +74,7 @@ func (td *Testdaemon) Run() {
 }
 
 func (td *Testdaemon) Recv() (*Segment, error) {
-	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
 	defer cancel()
 	select {
 	case r := <-td.Channel: