Add support for GCP extensions

tlvparse/gcp.go

@@ -0,0 +1,47 @@
+package tlvparse
+
+import (
+	"encoding/binary"
+
+	"github.com/pires/go-proxyproto"
+)
+
+const (
+	// PP2_TYPE_GCP indicates a Google Cloud Platform header
+	PP2_TYPE_GCP proxyproto.PP2Type = 0xE0
+)
+
+// ExtractPSCConnectionID returns the first PSC Connection ID in the TLV if it exists and is well-formed and
+// a bool indicating one was found.
+func ExtractPSCConnectionID(tlvs []proxyproto.TLV) (uint64, bool) {
+	for _, tlv := range tlvs {
+		if linkID, err := pscConnectionID(tlv); err == nil {
+			return linkID, true
+		}
+	}
+	return 0, false
+}
+
+// pscConnectionID returns the ID of a GCP PSC extension TLV or errors with ErrIncompatibleTLV or
+// ErrMalformedTLV if it's the wrong TLV type or is malformed.
+//
+//	Field	Length (bytes)	Description
+//	Type	1	PP2_TYPE_GCP (0xE0)
+//	Length	2	Length of value (always 0x0008)
+//	Value	8	The 8-byte PSC Connection ID (decode to uint64; big endian)
+//
+// For example proxyproto.TLV{Type:0xea, Length:8, Value:[]byte{0xff, 0xff, 0xff, 0xff, 0xc0, 0xa8, 0x64, 0x02}}
+// will be decoded as 18446744072646845442.
+//
+// See https://cloud.google.com/vpc/docs/configure-private-service-connect-producer
+func pscConnectionID(t proxyproto.TLV) (uint64, error) {
+	if !isPSCConnectionID(t) {
+		return 0, proxyproto.ErrIncompatibleTLV
+	}
+	linkID := binary.BigEndian.Uint64(t.Value)
+	return linkID, nil
+}
+
+func isPSCConnectionID(t proxyproto.TLV) bool {
+	return t.Type == PP2_TYPE_GCP && len(t.Value) == 8
+}

tlvparse/gcp_test.go

@@ -0,0 +1,82 @@
+package tlvparse
+
+import (
+	"testing"
+
+	"github.com/pires/go-proxyproto"
+)
+
+func TestExtractPSCConnectionID(t *testing.T) {
+	tests := []struct {
+		name                string
+		tlvs                []proxyproto.TLV
+		wantPSCConnectionID uint64
+		wantFound           bool
+	}{
+		{
+			name:      "nil TLVs",
+			tlvs:      nil,
+			wantFound: false,
+		},
+		{
+			name:      "empty TLVs",
+			tlvs:      []proxyproto.TLV{},
+			wantFound: false,
+		},
+		{
+			name: "AWS VPC endpoint ID",
+			tlvs: []proxyproto.TLV{
+				{
+					Type:  0xEA,
+					Value: []byte{0x01, 0x76, 0x70, 0x63, 0x65, 0x2d, 0x61, 0x62, 0x63, 0x31, 0x32, 0x33},
+				},
+			},
+			wantFound: false,
+		},
+		{
+			name: "GCP link ID",
+			tlvs: []proxyproto.TLV{
+				{
+					Type:  PP2_TYPE_GCP,
+					Value: []byte{'\xff', '\xff', '\xff', '\xff', '\xc0', '\xa8', '\x64', '\x02'},
+				},
+			},
+			wantPSCConnectionID: 18446744072646845442,
+			wantFound:           true,
+		},
+		{
+			name: "Multiple TLVs",
+			tlvs: []proxyproto.TLV{
+				{ // AWS
+					Type:  0xEA,
+					Value: []byte{0x01, 0x76, 0x70, 0x63, 0x65, 0x2d, 0x61, 0x62, 0x63, 0x31, 0x32, 0x33},
+				},
+				{ // Azure
+					Type:  0xEE,
+					Value: []byte{0x02, 0x01, 0x01, 0x01, 0x01},
+				},
+				{ // GCP but wrong length
+					Type:  0xE0,
+					Value: []byte{0xff, 0xff, 0xff},
+				},
+				{ // Correct
+					Type:  0xE0,
+					Value: []byte{'\xff', '\xff', '\xff', '\xff', '\xc0', '\xa8', '\x64', '\x02'},
+				},
+			},
+			wantPSCConnectionID: 18446744072646845442,
+			wantFound:           true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			linkID, hasLinkID := ExtractPSCConnectionID(tt.tlvs)
+			if hasLinkID != tt.wantFound {
+				t.Errorf("ExtractPSCConnectionID() got1 = %v, want %v", hasLinkID, tt.wantFound)
+			}
+			if linkID != tt.wantPSCConnectionID {
+				t.Errorf("ExtractPSCConnectionID() got = %v, want %v", linkID, tt.wantPSCConnectionID)
+			}
+		})
+	}
+}