Merge pull request #21 from igorwwwwwwwwwwwwwwwwwwww/patch-1

header: return err from reader.Peek()

Author: pires

header.go

@@ -123,10 +123,34 @@ func (header *Header) TLVs() ([]TLV, error) {
 // Also, this operation will block until enough bytes are available for peeking.
 func Read(reader *bufio.Reader) (*Header, error) {
 	// In order to improve speed for small non-PROXYed packets, take a peek at the first byte alone.
-	if b1, err := reader.Peek(1); err == nil && (bytes.Equal(b1[:1], SIGV1[:1]) || bytes.Equal(b1[:1], SIGV2[:1])) {
-		if signature, err := reader.Peek(5); err == nil && bytes.Equal(signature[:5], SIGV1) {
+	b1, err := reader.Peek(1)
+	if err != nil {
+		if err == io.EOF {
+			return nil, ErrNoProxyProtocol
+		}
+		return nil, err
+	}
+
+	if bytes.Equal(b1[:1], SIGV1[:1]) || bytes.Equal(b1[:1], SIGV2[:1]) {
+		signature, err := reader.Peek(5)
+		if err != nil {
+			if err == io.EOF {
+				return nil, ErrNoProxyProtocol
+			}
+			return nil, err
+		}
+		if bytes.Equal(signature[:5], SIGV1) {
 			return parseVersion1(reader)
-		} else if signature, err := reader.Peek(12); err == nil && bytes.Equal(signature[:12], SIGV2) {
+		}
+
+		signature, err = reader.Peek(12)
+		if err != nil {
+			if err == io.EOF {
+				return nil, ErrNoProxyProtocol
+			}
+			return nil, err
+		}
+		if bytes.Equal(signature[:12], SIGV2) {
 			return parseVersion2(reader)
 		}
 	}

header_test.go

@@ -3,6 +3,7 @@ package proxyproto
 import (
 	"bufio"
 	"bytes"
+	"errors"
 	"net"
 	"testing"
 	"time"
@@ -21,6 +22,8 @@ const (
 var (
 	v4addr = net.ParseIP(IP4_ADDR).To4()
 	v6addr = net.ParseIP(IP6_ADDR).To16()
+
+	errReadIntentionallyBroken = errors.New("read is intentionally broken")
 )
 
 type timeoutReader []byte
@@ -30,6 +33,12 @@ func (t *timeoutReader) Read([]byte) (int, error) {
 	return 0, nil
 }
 
+type errorReader []byte
+
+func (e *errorReader) Read([]byte) (int, error) {
+	return 0, errReadIntentionallyBroken
+}
+
 func TestReadTimeoutV1Invalid(t *testing.T) {
 	var b timeoutReader
 	reader := bufio.NewReader(&b)
@@ -41,6 +50,18 @@ func TestReadTimeoutV1Invalid(t *testing.T) {
 	}
 }
 
+func TestReadTimeoutPropagatesReadError(t *testing.T) {
+	var e errorReader
+	reader := bufio.NewReader(&e)
+	_, err := ReadTimeout(reader, 50*time.Millisecond)
+
+	if err == nil {
+		t.Fatalf("expected error %s", errReadIntentionallyBroken)
+	} else if err != errReadIntentionallyBroken {
+		t.Fatalf("expected error %s, actual %s", errReadIntentionallyBroken, err)
+	}
+}
+
 func TestEqualsTo(t *testing.T) {
 	var headersEqual = []struct {
 		this, that *Header