pktline: Add helper function to retrieve and verify data

While the pktline package already provides the `Data()` function to
retrieve a pktline's payload, this function doesn't perform any sanity
checks on retrieved data at all.

Add a new function `Payload()` that verifies that the pktline header's
claimed length matches the actual data's length. While it would be
preferable to just change `Data()`, chances are high that this would
cause unintended side effects.

2 files changed, 79 insertions, 0 deletions

diff --git a/internal/git/pktline/pkt_line_test.go b/internal/git/pktline/pkt_line_test.go
index e3192a011..ef9daaf78 100644
--- a/internal/git/pktline/pkt_line_test.go
+++ b/internal/git/pktline/pkt_line_test.go
@@ -332,3 +332,58 @@ func TestEachSidebandPacket(t *testing.T) {
 		})
 	}
 }
+
+func TestPayload(t *testing.T) {
+	for _, tc := range []struct {
+		desc            string
+		input           string
+		expectedErr     string
+		expectedPayload []byte
+	}{
+		{
+			desc:        "packet too small",
+			input:       "123",
+			expectedErr: "packet too small",
+		},
+		{
+			desc:        "invalid length prefix",
+			input:       "something",
+			expectedErr: "parsing length header \"some\": strconv.ParseUint: parsing \"some\": invalid syntax",
+		},
+		{
+			desc:        "flush packet",
+			input:       "0000",
+			expectedErr: "flush packets do not have a payload",
+		},
+		{
+			desc:        "packet with trailing bytes",
+			input:       "0005atrailing",
+			expectedErr: "packet length 13 does not match header length 5",
+		},
+		{
+			desc:        "packet with missing bytes",
+			input:       "0006a",
+			expectedErr: "packet length 5 does not match header length 6",
+		},
+		{
+			desc:            "empty packet",
+			input:           "0004",
+			expectedPayload: []byte{},
+		},
+		{
+			desc:            "packet with data",
+			input:           "0068" + strings.Repeat("x", 100),
+			expectedPayload: bytes.Repeat([]byte("x"), 100),
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			payload, err := Payload([]byte(tc.input))
+			if tc.expectedErr == "" {
+				require.NoError(t, err)
+			} else {
+				require.EqualError(t, err, tc.expectedErr)
+			}
+			require.Equal(t, tc.expectedPayload, payload)
+		})
+	}
+}
diff --git a/internal/git/pktline/pktline.go b/internal/git/pktline/pktline.go
index e97531fa3..3cdeaf4e9 100644
--- a/internal/git/pktline/pktline.go
+++ b/internal/git/pktline/pktline.go
@@ -39,6 +39,30 @@ func Data(pkt []byte) []byte {
 	return pkt[4:]
 }
 
+// Payload returns the pktline's data. It verifies that the length header matches what we expect as
+// data.
+func Payload(pkt []byte) ([]byte, error) {
+	if len(pkt) < 4 {
+		return nil, fmt.Errorf("packet too small")
+	}
+
+	if IsFlush(pkt) {
+		return nil, fmt.Errorf("flush packets do not have a payload")
+	}
+
+	lengthHeader := string(pkt[:4])
+	length, err := strconv.ParseUint(lengthHeader, 16, 16)
+	if err != nil {
+		return nil, fmt.Errorf("parsing length header %q: %w", lengthHeader, err)
+	}
+
+	if uint64(len(pkt)) != length {
+		return nil, fmt.Errorf("packet length %d does not match header length %d", len(pkt), length)
+	}
+
+	return pkt[4:], nil
+}
+
 // IsFlush detects the special flush packet '0000'
 func IsFlush(pkt []byte) bool {
 	return bytes.Equal(pkt, PktFlush())