fix(codecs): error when encoding maps whose lengths don't match entry…

… count Ref: #332
ipld · Jan 18, 2022 · de5750b · de5750b
1 parent a16dd99
commit de5750b
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/codec/dagcbor/marshal.go b/codec/dagcbor/marshal.go
@@ -162,7 +162,8 @@ func marshal(n datamodel.Node, tk *tok.Token, sink shared.TokenSink, options Enc
 func marshalMap(n datamodel.Node, tk *tok.Token, sink shared.TokenSink, options EncodeOptions) error {
 	// Emit start of map.
 	tk.Type = tok.TMapOpen
-	tk.Length = int(n.Length()) // TODO: overflow check
+	expectedLength := int(n.Length())
+	tk.Length = expectedLength // TODO: overflow check
 	if _, err := sink.Step(tk); err != nil {
 		return err
 	}
@@ -184,6 +185,9 @@ func marshalMap(n datamodel.Node, tk *tok.Token, sink shared.TokenSink, options
 			}
 			entries = append(entries, entry{keyStr, v})
 		}
+		if len(entries) != expectedLength {
+			return fmt.Errorf("map Length() does not match number of MapIterator() entries")
+		}
 		// Apply the desired sort function.
 		switch options.MapSortMode {
 		case codec.MapSortMode_Lexical:
@@ -213,11 +217,13 @@ func marshalMap(n datamodel.Node, tk *tok.Token, sink shared.TokenSink, options
 		}
 	} else { // no sorting
 		// Emit map contents (and recurse).
+		var entryCount int
 		for itr := n.MapIterator(); !itr.Done(); {
 			k, v, err := itr.Next()
 			if err != nil {
 				return err
 			}
+			entryCount++
 			tk.Type = tok.TString
 			tk.Str, err = k.AsString()
 			if err != nil {
@@ -230,6 +236,9 @@ func marshalMap(n datamodel.Node, tk *tok.Token, sink shared.TokenSink, options
 				return err
 			}
 		}
+		if entryCount != expectedLength {
+			return fmt.Errorf("map Length() does not match number of MapIterator() entries")
+		}
 	}
 	// Emit map close.
 	tk.Type = tok.TMapClose

diff --git a/codec/dagjson/marshal.go b/codec/dagjson/marshal.go
@@ -61,7 +61,8 @@ func Marshal(n datamodel.Node, sink shared.TokenSink, options EncodeOptions) err
 	case datamodel.Kind_Map:
 		// Emit start of map.
 		tk.Type = tok.TMapOpen
-		tk.Length = int(n.Length()) // TODO: overflow check
+		expectedLength := int(n.Length())
+		tk.Length = expectedLength // TODO: overflow check
 		if _, err := sink.Step(&tk); err != nil {
 			return err
 		}
@@ -83,6 +84,9 @@ func Marshal(n datamodel.Node, sink shared.TokenSink, options EncodeOptions) err
 				}
 				entries = append(entries, entry{keyStr, v})
 			}
+			if len(entries) != expectedLength {
+				return fmt.Errorf("map Length() does not match number of MapIterator() entries")
+			}
 			// Apply the desired sort function.
 			switch options.MapSortMode {
 			case codec.MapSortMode_Lexical:
@@ -100,16 +104,21 @@ func Marshal(n datamodel.Node, sink shared.TokenSink, options EncodeOptions) err
 				})
 			}
 			// Emit map contents (and recurse).
+			var entryCount int
 			for _, e := range entries {
 				tk.Type = tok.TString
 				tk.Str = e.key
+				entryCount++
 				if _, err := sink.Step(&tk); err != nil {
 					return err
 				}
 				if err := Marshal(e.value, sink, options); err != nil {
 					return err
 				}
 			}
+			if entryCount != expectedLength {
+				return fmt.Errorf("map Length() does not match number of MapIterator() entries")
+			}
 		} else {
 			// Don't sort map, emit map contents (and recurse).
 			for itr := n.MapIterator(); !itr.Done(); {