From 0985f1e8ff630346835dd1126b97b164a138f283 Mon Sep 17 00:00:00 2001
From: Miek Gieben <miek@miek.nl>
Date: Sun, 28 Feb 2016 16:36:19 +0000
Subject: [PATCH] Fix TCP rtt

TCP wasn't returning rrt info anymore, fix this. Also add
an issue_test.go where fixes for specific issues can be put.

Pull the rtt for udp message up into the function where we now
also set the rrt for tcp (for symmetry).
---
 client.go     |  5 ++---
 issue_test.go | 26 +++++++++-----------------
 2 files changed, 11 insertions(+), 20 deletions(-)

diff --git a/client.go b/client.go
index 881e9057..9afa7506 100644
--- a/client.go
+++ b/client.go
@@ -258,7 +258,7 @@ func (co *Conn) ReadMsgHeader(hdr *Header) ([]byte, error) {
 		}
 		p = make([]byte, l)
 		n, err = tcpRead(r, p)
-
+		co.rtt = time.Since(co.t)
 	default:
 		if co.UDPSize > MinMsgSize {
 			p = make([]byte, co.UDPSize)
@@ -266,6 +266,7 @@ func (co *Conn) ReadMsgHeader(hdr *Header) ([]byte, error) {
 			p = make([]byte, MinMsgSize)
 		}
 		n, err = co.Read(p)
+		co.rtt = time.Since(co.t)
 	}
 
 	if err != nil {
@@ -342,8 +343,6 @@ func (co *Conn) Read(p []byte) (n int, err error) {
 	if err != nil {
 		return n, err
 	}
-
-	co.rtt = time.Since(co.t)
 	return n, err
 }
 
diff --git a/issue_test.go b/issue_test.go
index 878b7f66..45ea8829 100644
--- a/issue_test.go
+++ b/issue_test.go
@@ -10,22 +10,14 @@ func TestTCPRtt(t *testing.T) {
 	m.SetQuestion("example.org.", TypeA)
 
 	c := &Client{}
-	in, rtt, err := c.Exchange(m, "8.8.4.4:53")
-	if err != nil {
-		t.Fatal(err)
+	for _, proto := range []string{"udp", "tcp"} {
+		c.Net = proto
+		_, rtt, err := c.Exchange(m, "8.8.4.4:53")
+		if err != nil {
+			t.Fatal(err)
+		}
+		if rtt == 0 {
+			t.Fatalf("expecting non zero rtt %s, got zero", c.Net)
+		}
 	}
-	if rtt == 0 {
-		t.Fatalf("expecting non zero rtt, got zero")
-	}
-	t.Logf("%s", in)
-
-	c.Net = "tcp"
-	in, rtt, err = c.Exchange(m, "8.8.4.4:53")
-	if err != nil {
-		t.Fatal(err)
-	}
-	if rtt == 0 {
-		t.Fatalf("expecting non zero rtt, got zero")
-	}
-	t.Logf("%s", in)
 }