Merge pull request samuel#48 from macb/add_authenticate_timeouts

Add timeouts around IO calls during authenticate
sherlock-help · Jan 19, 2015 · fa6674a · fa6674a
2 parents bb79e7c + 4f64bfc
commit fa6674a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/zk/conn.go b/zk/conn.go
@@ -365,7 +365,9 @@ func (c *Conn) authenticate() error {
 
 	binary.BigEndian.PutUint32(buf[:4], uint32(n))
 
+	c.conn.SetWriteDeadline(time.Now().Add(c.recvTimeout * 10))
 	_, err = c.conn.Write(buf[:n+4])
+	c.conn.SetWriteDeadline(time.Time{})
 	if err != nil {
 		return err
 	}
@@ -375,7 +377,9 @@ func (c *Conn) authenticate() error {
 	// connect response
 
 	// package length
+	c.conn.SetReadDeadline(time.Now().Add(c.recvTimeout * 10))
 	_, err = io.ReadFull(c.conn, buf[:4])
+	c.conn.SetReadDeadline(time.Time{})
 	if err != nil {
 		return err
 	}

diff --git a/zk/zk_test.go b/zk/zk_test.go
@@ -3,7 +3,6 @@ package zk
 import (
 	"fmt"
 	"io"
-	"log"
 	"net"
 	"strings"
 	"testing"
@@ -491,7 +490,7 @@ func startSlowProxy(t *testing.T, up, down throttle.Rate, upstream string, adj f
 				defer cn.Close()
 				upcn, err := net.Dial("tcp", upstream)
 				if err != nil {
-					log.Print(err)
+					t.Log(err)
 					return
 				}
 				// This will leave hanging goroutines util stopCh is closed