From efc48d41bcc34a9f8ae22084fa4ef55f07aa44d7 Mon Sep 17 00:00:00 2001
From: Hector Rivas Gandara <keymon@gmail.com>
Date: Wed, 13 Jul 2016 15:20:54 +0100
Subject: [PATCH] Fix Close connection for invalid header

RemoteAddr() must close the connection and clear the buffer if we
receive a invalid PROXY protocol header.

Change the client test as it can get an EOF or a ECONNRESET.
---
 protocol.go      | 2 ++
 protocol_test.go | 3 +--
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/protocol.go b/protocol.go
index e01aeb5..dfa6eb0 100644
--- a/protocol.go
+++ b/protocol.go
@@ -110,6 +110,8 @@ func (p *Conn) RemoteAddr() net.Addr {
 	p.once.Do(func() {
 		if err := p.checkPrefix(); err != nil && err != io.EOF {
 			log.Printf("[ERR] Failed to read proxy prefix: %v", err)
+			p.Close()
+			p.bufReader = bufio.NewReader(p.conn)
 		}
 	})
 	if p.srcAddr != nil {
diff --git a/protocol_test.go b/protocol_test.go
index 6cc7b4c..ab29731 100644
--- a/protocol_test.go
+++ b/protocol_test.go
@@ -2,7 +2,6 @@ package proxyproto
 
 import (
 	"bytes"
-	"io"
 	"net"
 	"testing"
 	"time"
@@ -259,7 +258,7 @@ func TestParse_BadHeader(t *testing.T) {
 
 		recv := make([]byte, 4)
 		_, err = conn.Read(recv)
-		if err != io.EOF {
+		if err == nil {
 			t.Fatalf("err: %v", err)
 		}
 	}()