Merge pull request #209 from michaelharo/client

Cleanup Client.exchange

client.go

@@ -128,31 +128,39 @@ func (c *Client) Exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err erro
 	return r, rtt, nil
 }
 
-func (c *Client) exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
-	timeout := dnsTimeout
-	var co *Conn
+func (c *Client) dialTimeout() time.Duration {
 	if c.DialTimeout != 0 {
-		timeout = c.DialTimeout
+		return c.DialTimeout
 	}
+	return dnsTimeout
+}
+
+func (c *Client) readTimeout() time.Duration {
+	if c.ReadTimeout != 0 {
+		return c.ReadTimeout
+	}
+	return dnsTimeout
+}
+
+func (c *Client) writeTimeout() time.Duration {
+	if c.WriteTimeout != 0 {
+		return c.WriteTimeout
+	}
+	return dnsTimeout
+}
+
+func (c *Client) exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
+	var co *Conn
 	if c.Net == "" {
-		co, err = DialTimeout("udp", a, timeout)
+		co, err = DialTimeout("udp", a, c.dialTimeout())
 	} else {
-		co, err = DialTimeout(c.Net, a, timeout)
+		co, err = DialTimeout(c.Net, a, c.dialTimeout())
 	}
 	if err != nil {
 		return nil, 0, err
 	}
-	timeout = dnsTimeout
-	if c.ReadTimeout != 0 {
-		timeout = c.ReadTimeout
-	}
-	co.SetReadDeadline(time.Now().Add(timeout))
-	timeout = dnsTimeout
-	if c.WriteTimeout != 0 {
-		timeout = c.WriteTimeout
-	}
-	co.SetWriteDeadline(time.Now().Add(timeout))
 	defer co.Close()
+
 	opt := m.IsEdns0()
 	// If EDNS0 is used use that for size.
 	if opt != nil && opt.UDPSize() >= MinMsgSize {
@@ -162,6 +170,10 @@ func (c *Client) exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err erro
 	if opt == nil && c.UDPSize >= MinMsgSize {
 		co.UDPSize = c.UDPSize
 	}
+
+	co.SetReadDeadline(time.Now().Add(c.readTimeout()))
+	co.SetWriteDeadline(time.Now().Add(c.writeTimeout()))
+
 	co.TsigSecret = c.TsigSecret
 	if err = co.WriteMsg(m); err != nil {
 		return nil, 0, err