protocol: reset timeout after processing header

Fixes #75

Signed-off-by: Pires <[email protected]>

protocol.go

@@ -54,6 +54,8 @@ func (p *Listener) Accept() (net.Conn, error) {
 	}
 
 	if d := p.ReadHeaderTimeout; d != 0 {
+		// The deadline will be reset after parsing the header.
+		// Otherwise, future p.conn.Read() will timeout.
 		conn.SetReadDeadline(time.Now().Add(d))
 	}
 
@@ -106,10 +108,12 @@ func NewConn(conn net.Conn, opts ...func(*Conn)) *Conn {
 func (p *Conn) Read(b []byte) (int, error) {
 	p.once.Do(func() {
 		p.readErr = p.readHeader()
+		p.conn.SetReadDeadline(time.Time{})
 	})
 	if p.readErr != nil {
 		return 0, p.readErr
 	}
+
 	return p.bufReader.Read(b)
 }
 

protocol_test.go

@@ -31,7 +31,7 @@ func TestPassthrough(t *testing.T) {
 		conn, err := net.Dial("tcp", pl.Addr().String())
 		if err != nil {
 			t.Fatalf("err: %v", err)
-			}
+		}
 		defer conn.Close()
 
 		conn.Write([]byte("ping"))
@@ -99,7 +99,7 @@ func TestReadHeaderTimeout(t *testing.T) {
 	recv := make([]byte, 4)
 	_, err = conn.Read(recv)
 
-	if err != nil && !errors.Is(err, os.ErrDeadlineExceeded){
+	if err != nil && !errors.Is(err, os.ErrDeadlineExceeded) {
 		t.Fatal("should timeout")
 	}
 }
@@ -113,7 +113,7 @@ func TestReadHeaderTimeoutIsReset(t *testing.T) {
 	}
 
 	pl := &Listener{
-		Listener: l,
+		Listener:          l,
 		ReadHeaderTimeout: timeout,
 	}