From d043f02901007e1951408bfde78e882f9ef2e81f Mon Sep 17 00:00:00 2001
From: Aurora Gaffney <aurora@blinklabs.io>
Date: Fri, 21 Jun 2024 13:34:16 -0500
Subject: [PATCH] feat: CBOR encode support for indefinite-length bytestrings

This also slightly changes the way that indefinite-length lists work, as
well as adds tests for indefinite-length lists

Fixes #659
---
 cbor/encode.go                    | 30 +++++++++++++++++++----
 cbor/encode_test.go               | 40 +++++++++++++++++++++++++++++++
 protocol/txsubmission/messages.go | 12 +++-------
 3 files changed, 68 insertions(+), 14 deletions(-)

diff --git a/cbor/encode.go b/cbor/encode.go
index c6df0729..ca1bedee 100644
--- a/cbor/encode.go
+++ b/cbor/encode.go
@@ -67,16 +67,14 @@ func EncodeGeneric(src interface{}) ([]byte, error) {
 	return cborData, nil
 }
 
-type IndefLengthList struct {
-	Items []any
-}
+type IndefLengthList []any
 
-func (i *IndefLengthList) MarshalCBOR() ([]byte, error) {
+func (i IndefLengthList) MarshalCBOR() ([]byte, error) {
 	ret := []byte{
 		// Start indefinite-length list
 		0x9f,
 	}
-	for _, item := range i.Items {
+	for _, item := range []any(i) {
 		data, err := Encode(&item)
 		if err != nil {
 			return nil, err
@@ -90,3 +88,25 @@ func (i *IndefLengthList) MarshalCBOR() ([]byte, error) {
 	)
 	return ret, nil
 }
+
+type IndefLengthByteString []any
+
+func (i IndefLengthByteString) MarshalCBOR() ([]byte, error) {
+	ret := []byte{
+		// Start indefinite-length bytestring
+		0x5f,
+	}
+	for _, item := range []any(i) {
+		data, err := Encode(&item)
+		if err != nil {
+			return nil, err
+		}
+		ret = append(ret, data...)
+	}
+	ret = append(
+		ret,
+		// End indefinite length bytestring
+		byte(0xff),
+	)
+	return ret, nil
+}
diff --git a/cbor/encode_test.go b/cbor/encode_test.go
index eb774764..eab9c7ea 100644
--- a/cbor/encode_test.go
+++ b/cbor/encode_test.go
@@ -50,3 +50,43 @@ func TestEncode(t *testing.T) {
 		}
 	}
 }
+
+func TestEncodeIndefLengthList(t *testing.T) {
+	expectedCborHex := "9f1904d219162eff"
+	tmpData := cbor.IndefLengthList{
+		1234,
+		5678,
+	}
+	cborData, err := cbor.Encode(tmpData)
+	if err != nil {
+		t.Fatalf("failed to encode object to CBOR: %s", err)
+	}
+	cborHex := hex.EncodeToString(cborData)
+	if cborHex != expectedCborHex {
+		t.Fatalf(
+			"object did not encode to expected CBOR\n  got %s\n  wanted: %s",
+			cborHex,
+			expectedCborHex,
+		)
+	}
+}
+
+func TestEncodeIndefLengthByteString(t *testing.T) {
+	expectedCborHex := "5f440102030443abcdefff"
+	tmpData := cbor.IndefLengthByteString{
+		[]byte{1, 2, 3, 4},
+		[]byte{0xab, 0xcd, 0xef},
+	}
+	cborData, err := cbor.Encode(tmpData)
+	if err != nil {
+		t.Fatalf("failed to encode object to CBOR: %s", err)
+	}
+	cborHex := hex.EncodeToString(cborData)
+	if cborHex != expectedCborHex {
+		t.Fatalf(
+			"object did not encode to expected CBOR\n  got %s\n  wanted: %s",
+			cborHex,
+			expectedCborHex,
+		)
+	}
+}
diff --git a/protocol/txsubmission/messages.go b/protocol/txsubmission/messages.go
index 25346d73..40bb691b 100644
--- a/protocol/txsubmission/messages.go
+++ b/protocol/txsubmission/messages.go
@@ -93,9 +93,7 @@ func (m *MsgReplyTxIds) MarshalCBOR() ([]byte, error) {
 	}
 	tmp := []any{
 		MessageTypeReplyTxIds,
-		cbor.IndefLengthList{
-			Items: items,
-		},
+		cbor.IndefLengthList(items),
 	}
 	return cbor.Encode(tmp)
 }
@@ -122,9 +120,7 @@ func (m *MsgRequestTxs) MarshalCBOR() ([]byte, error) {
 	}
 	tmp := []any{
 		MessageTypeRequestTxs,
-		cbor.IndefLengthList{
-			Items: items,
-		},
+		cbor.IndefLengthList(items),
 	}
 	return cbor.Encode(tmp)
 }
@@ -151,9 +147,7 @@ func (m *MsgReplyTxs) MarshalCBOR() ([]byte, error) {
 	}
 	tmp := []any{
 		MessageTypeReplyTxs,
-		cbor.IndefLengthList{
-			Items: items,
-		},
+		cbor.IndefLengthList(items),
 	}
 	return cbor.Encode(tmp)
 }