Repository: pires/go-proxyproto
Pull request #38 (2020-10-02 17:22:26)

Complete PR36
Fixes #36 

Thank you for your contributions @emersion @mschneider82 🙇 


Patch:
diff --git a/header.go b/header.go
index 5c5f869..1d27805 100644
--- a/header.go
+++ b/header.go
@@ -175,6 +175,17 @@ func (header *Header) TLVs() ([]TLV, error) {
 	return SplitTLVs(header.rawTLVs)
 }
 
+// SetTLVs sets the TLVs stored in this header. This method replaces any
+// previous TLV.
+func (header *Header) SetTLVs(tlvs []TLV) error {
+	raw, err := JoinTLVs(tlvs)
+	if err != nil {
+		return err
+	}
+	header.rawTLVs = raw
+	return nil
+}
+
 // Read identifies the proxy protocol version and reads the remaining of
 // the header, accordingly.
 //
diff --git a/tlv.go b/tlv.go
index e6a2149..6873aa0 100644
--- a/tlv.go
+++ b/tlv.go
@@ -6,6 +6,8 @@ package proxyproto
 import (
 	"encoding/binary"
 	"errors"
+	"fmt"
+	"math"
 )
 
 const (
@@ -73,6 +75,25 @@ func SplitTLVs(raw []byte) ([]TLV, error) {
 	return tlvs, nil
 }
 
+// JoinTLVs joins multiple Type-Length-Value records.
+func JoinTLVs(tlvs []TLV) ([]byte, error) {
+	var raw []byte
+	for _, tlv := range tlvs {
+		if tlv.Length > math.MaxUint16 {
+			return nil, fmt.Errorf("proxyproto: cannot format TLV %v with length %d", tlv.Type, tlv.Length)
+		}
+		if tlv.Length != len(tlv.Value) {
+			return nil, fmt.Errorf("proxyproto: TLV %v length doesn't match value length", tlv.Type)
+		}
+		var length [2]byte
+		binary.BigEndian.PutUint16(length[:], uint16(tlv.Length))
+		raw = append(raw, byte(tlv.Type))
+		raw = append(raw, length[:]...)
+		raw = append(raw, tlv.Value...)
+	}
+	return raw, nil
+}
+
 // Registered is true if the type is registered in the spec, see section 2.2
 func (p PP2Type) Registered() bool {
 	switch p {


Test patch:
diff --git a/header_test.go b/header_test.go
index a9671e8..1c2bacc 100644
--- a/header_test.go
+++ b/header_test.go
@@ -184,6 +184,57 @@ func TestLocalAddr(t *testing.T) {
 	}
 }
 
+func TestSetTLVs(t *testing.T) {
+	tests := []struct {
+		header    *Header
+		name      string
+		tlvs      []TLV
+		expectErr bool
+	}{
+		{
+			name: "add authority TLV",
+			header: &Header{
+				Version:            1,
+				Command:            PROXY,
+				TransportProtocol:  TCPv4,
+				SourceAddress:      net.ParseIP("10.1.1.1"),
+				SourcePort:         1000,
+				DestinationAddress: net.ParseIP("20.2.2.2"),
+				DestinationPort:    2000,
+			},
+			tlvs: []TLV{{
+				Type:   PP2_TYPE_AUTHORITY,
+				Length: 11,
+				Value:  []byte("example.org"),
+			}},
+		},
+		{
+			name: "add wrong length",
+			header: &Header{
+				Version:            1,
+				Command:            PROXY,
+				TransportProtocol:  TCPv4,
+				SourceAddress:      net.ParseIP("10.1.1.1"),
+				SourcePort:         1000,
+				DestinationAddress: net.ParseIP("20.2.2.2"),
+				DestinationPort:    2000,
+			},
+			tlvs: []TLV{{
+				Type:   PP2_TYPE_AUTHORITY,
+				Length: 1,
+				Value:  []byte("example.org"),
+			}},
+			expectErr: true,
+		},
+	}
+	for _, tt := range tests {
+		err := tt.header.SetTLVs(tt.tlvs)
+		if err != nil && !tt.expectErr {
+			t.Fatalf("shouldn't have thrown error %q", err.Error())
+		}
+	}
+}
+
 func TestRemoteAddr(t *testing.T) {
 	var headers = []struct {
 		header       *Header
diff --git a/tlv_test.go b/tlv_test.go
index e699cdd..8968473 100644
--- a/tlv_test.go
+++ b/tlv_test.go
@@ -3,7 +3,6 @@ package proxyproto
 import (
 	"bufio"
 	"bytes"
-	"encoding/binary"
 	"testing"
 )
 
@@ -38,13 +37,6 @@ func checkTLVs(t *testing.T, name string, raw []byte, expected []PP2Type) []TLV
 	return tlvs
 }
 
-func formatTLV(tlv TLV) []byte {
-	out := make([]byte, 3) // 1 = type + 2 = uint16 length
-	out[0] = byte(tlv.Type)
-	binary.BigEndian.PutUint16(out[1:3], uint16(tlv.Length))
-	return append(out, tlv.Value...)
-}
-
 var invalidTLVTests = []struct {
 	name          string
 	reader        *bufio.Reader
@@ -128,3 +120,39 @@ func TestV2TLVPP2Registered(t *testing.T) {
 		t.Fatalf("TestV2TLVPP2Registered: type %x unexpectedly registered", lastType)
 	}
 }
+
+func TestJoinTLVs(t *testing.T) {
+	tests := []struct {
+		name string
+		raw  []byte
+		tlvs []TLV
+	}{
+		{
+			name: "authority TLV",
+			raw:  append([]byte{byte(PP2_TYPE_AUTHORITY), 0x00, 0x0B}, []byte("example.org")...),
+			tlvs: []TLV{{
+				Type:   PP2_TYPE_AUTHORITY,
+				Length: 11,
+				Value:  []byte("example.org"),
+			}},
+		},
+		{
+			name: "empty TLV",
+			raw:  []byte{byte(PP2_TYPE_NOOP), 0x00, 0x00},
+			tlvs: []TLV{{
+				Type:   PP2_TYPE_NOOP,
+				Length: 0,
+				Value:  nil,
+			}},
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			if raw, err := JoinTLVs(tc.tlvs); err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			} else if !bytes.Equal(raw, tc.raw) {
+				t.Errorf("expected %#v, got %#v", tc.raw, raw)
+			}
+		})
+	}
+}