fix: use correct header length for indef-length map on decode (#125)

This also adds support for encode/decode of indef-length maps

Signed-off-by: Aurora Gaffney <[email protected]>

Author: agaffney

data/data.go

@@ -65,7 +65,8 @@ func NewConstrDefIndef(useIndef bool, tag uint, fields ...PlutusData) PlutusData
 // Map
 
 type Map struct {
-	Pairs [][2]PlutusData // Each pair is [key, value]
+	Pairs    [][2]PlutusData // Each pair is [key, value]
+	useIndef *bool
 }
 
 func (Map) isPlutusData() {}
@@ -79,6 +80,11 @@ func NewMap(pairs [][2]PlutusData) PlutusData {
 	return &Map{Pairs: pairs}
 }
 
+// NewMapDefIndef creates a new Map with the ability to specify whether it should use definite- or indefinite-length encoding
+func NewMapDefIndef(useIndef bool, pairs [][2]PlutusData) PlutusData {
+	return &Map{Pairs: pairs, useIndef: &useIndef}
+}
+
 // Integer
 
 type Integer struct {

data/data_test.go

@@ -52,7 +52,8 @@ var testDefs = []struct {
 		CborHex: "9f0102ff",
 	},
 	{
-		Data: NewMap(
+		Data: NewMapDefIndef(
+			false,
 			[][2]PlutusData{
 				{
 					NewInteger(big.NewInt(1)),
@@ -63,7 +64,8 @@ var testDefs = []struct {
 		CborHex: "a10102",
 	},
 	{
-		Data: NewMap(
+		Data: NewMapDefIndef(
+			false,
 			[][2]PlutusData{
 				{
 					NewConstrDefIndef(
@@ -86,7 +88,8 @@ var testDefs = []struct {
 		Data: NewConstrDefIndef(
 			true,
 			0,
-			NewMap(
+			NewMapDefIndef(
+				false,
 				[][2]PlutusData{
 					{
 						NewConstrDefIndef(
@@ -117,15 +120,17 @@ var testDefs = []struct {
 		CborHex: "d866821903e79f0607ff",
 	},
 	{
-		Data: NewMap(
+		Data: NewMapDefIndef(
+			false,
 			[][2]PlutusData{
 				{
 					NewListDefIndef(
 						true,
 						NewInteger(big.NewInt(1)),
 						NewInteger(big.NewInt(2)),
 					),
-					NewMap(
+					NewMapDefIndef(
+						false,
 						[][2]PlutusData{
 							{
 								NewListDefIndef(
@@ -148,10 +153,12 @@ var testDefs = []struct {
 		CborHex: "a19f0102ffa19f40ffd8799f0201ff",
 	},
 	{
-		Data: NewMap(
+		Data: NewMapDefIndef(
+			false,
 			[][2]PlutusData{
 				{
-					NewMap(
+					NewMapDefIndef(
+						false,
 						[][2]PlutusData{
 							{
 								NewInteger(big.NewInt(1)),
@@ -177,8 +184,10 @@ var testDefs = []struct {
 		),
 		CborHex: "d87a81d87980",
 	},
+	// Map with specific-order keys
 	{
-		Data: NewMap(
+		Data: NewMapDefIndef(
+			false,
 			[][2]PlutusData{
 				{
 					NewInteger(big.NewInt(2)),
@@ -197,6 +206,24 @@ var testDefs = []struct {
 		// {2:2,3:3,1:1}
 		CborHex: "a3020203030101",
 	},
+	// Indef-length map
+	{
+		Data: NewMapDefIndef(
+			true,
+			[][2]PlutusData{
+				{
+					NewByteString([]byte{0x01}),
+					NewInteger(big.NewInt(1)),
+				},
+				{
+					NewByteString([]byte{0x02}),
+					NewInteger(big.NewInt(2)),
+				},
+			},
+		),
+		// {_ h'01': 1, h'02': 2}
+		CborHex: "bf410101410202ff",
+	},
 }
 
 func TestPlutusDataEncode(t *testing.T) {

data/decode.go

@@ -103,22 +103,29 @@ func decodeCborRawMap(data []byte) (any, error) {
 	// order when decoding a map. We decode our map to determine its length, create a dummy
 	// list the same length as our map to determine the header size, and then decode each
 	// key/value pair individually
+	useIndef := (data[0] & CborIndefFlag) == CborIndefFlag
 	var tmpData map[RawMessageStr]RawMessageStr
 	if err := cborUnmarshal(data, &tmpData); err != nil {
 		return nil, err
 	}
-	// Create dummy list of same length to determine map header length
-	tmpList := make([]bool, len(tmpData))
-	tmpListRaw, err := cborMarshal(tmpList)
-	if err != nil {
-		return nil, err
+	if useIndef {
+		// Strip off indef-length map header byte
+		data = data[1:]
+	} else {
+		// Create dummy list of same length to determine map header length
+		tmpList := make([]bool, len(tmpData))
+		tmpListRaw, err := cborMarshal(tmpList)
+		if err != nil {
+			return nil, err
+		}
+		tmpListHeader := tmpListRaw[0 : len(tmpListRaw)-len(tmpData)]
+		// Strip off map header bytes
+		data = data[len(tmpListHeader):]
 	}
-	tmpListHeader := tmpListRaw[0 : len(tmpListRaw)-len(tmpData)]
-	// Strip off map header bytes
-	data = data[len(tmpListHeader):]
 	pairs := make([][2]PlutusData, 0, len(tmpData))
 	var rawKey, rawVal cbor.RawMessage
 	// Read key/value pairs until we have no data left
+	var err error
 	for len(data) > 0 {
 		// Check for "break" at end of indefinite-length map
 		if data[0] == 0xFF {
@@ -159,6 +166,7 @@ func decodeCborRawMap(data []byte) (any, error) {
 		)
 	}
 	ret := NewMap(pairs)
+	ret.(*Map).useIndef = &useIndef
 	return ret, nil
 }
 

data/encode.go

@@ -129,6 +129,11 @@ func encodeMap(m *Map) (any, error) {
 	// steal and modify its header, and build our own output from pieces. This avoids
 	// needing to support 6 different possible encodings of a map's header byte depending
 	// on length
+	useIndef := false
+	if m.useIndef != nil {
+		useIndef = *m.useIndef
+	}
+	// Build encoded pairs
 	tmpPairs := make([][]byte, 0, len(m.Pairs))
 	for _, pair := range m.Pairs {
 		key, err := encodeToRaw(pair[0])
@@ -152,19 +157,27 @@ func encodeMap(m *Map) (any, error) {
 			slices.Concat(keyRaw, valueRaw),
 		)
 	}
-	// Create dummy list with simple (one-byte) values so we can easily extract the header
-	tmpList := make([]bool, len(tmpPairs))
-	tmpListRaw, err := cborMarshal(tmpList)
-	if err != nil {
-		return nil, err
-	}
-	tmpListHeader := tmpListRaw[0 : len(tmpListRaw)-len(tmpPairs)]
-	// Modify header byte to switch type from array to map
-	tmpListHeader[0] |= 0x20
 	// Build return value
 	ret := bytes.NewBuffer(nil)
-	_, _ = ret.Write(tmpListHeader)
+	if useIndef {
+		ret.WriteByte(CborTypeMap | CborIndefFlag)
+	} else {
+		// Create dummy list with simple (one-byte) values so we can easily extract the header
+		tmpList := make([]bool, len(tmpPairs))
+		tmpListRaw, err := cborMarshal(tmpList)
+		if err != nil {
+			return nil, err
+		}
+		tmpListHeader := tmpListRaw[0 : len(tmpListRaw)-len(tmpPairs)]
+		// Modify header byte to switch type from array to map
+		tmpListHeader[0] |= 0x20
+		_, _ = ret.Write(tmpListHeader)
+	}
 	_, _ = ret.Write(slices.Concat(tmpPairs...))
+	if useIndef {
+		// Indef-length "break" byte
+		ret.WriteByte(0xff)
+	}
 	return cbor.RawMessage(ret.Bytes()), nil
 }