fix: partial dag, serialization, verification and tests

Author: f7f376a1fcd0d0e11a10ed1b6577c9

dag/dag.go

@@ -15,7 +15,7 @@ func TestFull(t *testing.T) {
 
 	defer os.RemoveAll(tmpDir)
 
-	GenerateDummyDirectory(filepath.Join(tmpDir, "input"), 6, 6)
+	GenerateDummyDirectory(filepath.Join(tmpDir, "input"), 3, 6, 1, 3)
 	if err != nil {
 		t.Fatalf("Could not generate dummy directory: %s", err)
 	}

dag/dag_test.go

@@ -53,7 +53,7 @@ func TestOutOfRangeLeafRequests(t *testing.T) {
 				return // Expected for invalid ranges
 			}
 			// If we got a partial DAG, verify it's valid
-			if err := partial.VerifyPartial(); err != nil {
+			if err := partial.Verify(); err != nil {
 				t.Errorf("Invalid partial DAG returned for range %d-%d: %v", tt.start, tt.end, err)
 			}
 		})
@@ -210,37 +210,40 @@ func TestInvalidPaths(t *testing.T) {
 }
 
 func TestBrokenDags(t *testing.T) {
-	tmpDir, err := ioutil.TempDir("", "test")
-	if err != nil {
-		t.Fatalf("Could not create temp directory: %s", err)
-	}
-	defer os.RemoveAll(tmpDir)
+	// Create a valid DAG with known structure
+	dagBuilder := CreateDagBuilder()
 
-	// Create a valid DAG first
-	err = ioutil.WriteFile(filepath.Join(tmpDir, "test.txt"), []byte("test content"), 0644)
+	// Create a file leaf
+	fileBuilder := CreateDagLeafBuilder("test.txt")
+	fileBuilder.SetType(FileLeafType)
+	fileBuilder.SetData([]byte("test content"))
+	fileLeaf, err := fileBuilder.BuildLeaf(nil)
 	if err != nil {
-		t.Fatalf("Failed to create test file: %v", err)
+		t.Fatalf("Failed to build file leaf: %v", err)
 	}
+	fileLeaf.SetLabel("1")
+	dagBuilder.AddLeaf(fileLeaf, nil)
 
-	dag, err := CreateDag(tmpDir, false)
+	// Create a directory with the file
+	dirBuilder := CreateDagLeafBuilder("testdir")
+	dirBuilder.SetType(DirectoryLeafType)
+	dirBuilder.AddLink("1", fileLeaf.Hash)
+	dirLeaf, err := dirBuilder.BuildRootLeaf(dagBuilder, nil)
 	if err != nil {
-		t.Fatalf("Failed to create DAG: %v", err)
+		t.Fatalf("Failed to build directory leaf: %v", err)
 	}
+	dagBuilder.AddLeaf(dirLeaf, nil)
+
+	dag := dagBuilder.BuildDag(dirLeaf.Hash)
 
 	t.Run("missing_leaf", func(t *testing.T) {
 		brokenDag := &Dag{
 			Root:  dag.Root,
 			Leafs: make(map[string]*DagLeaf),
 		}
-		// Copy all leaves except one
-		var skippedOne bool
-		for hash, leaf := range dag.Leafs {
-			if !skippedOne {
-				skippedOne = true
-				continue
-			}
-			brokenDag.Leafs[hash] = leaf
-		}
+		// Only copy the root leaf
+		brokenDag.Leafs[dag.Root] = dag.Leafs[dag.Root].Clone()
+
 		if err := brokenDag.Verify(); err == nil {
 			t.Error("Expected verification to fail for DAG with missing leaf")
 		}
@@ -251,11 +254,18 @@ func TestBrokenDags(t *testing.T) {
 			Root:  dag.Root,
 			Leafs: make(map[string]*DagLeaf),
 		}
-		// Copy all leaves but corrupt one's content
+		// Copy all leaves but corrupt file content
 		for hash, leaf := range dag.Leafs {
 			leafCopy := leaf.Clone()
-			if leaf.Type == FileLeafType || leaf.Type == ChunkLeafType {
-				leafCopy.Content = append(leafCopy.Content, []byte("corrupted")...)
+			if leaf.Type == FileLeafType {
+				// Create a new leaf with corrupted content
+				builder := CreateDagLeafBuilder(leaf.ItemName)
+				builder.SetType(leaf.Type)
+				builder.SetData(append(leaf.Content, []byte("corrupted")...))
+				corruptedLeaf, _ := builder.BuildLeaf(nil)
+				// Keep original hash but use corrupted content and hash
+				leafCopy.Content = corruptedLeaf.Content
+				leafCopy.ContentHash = corruptedLeaf.ContentHash
 			}
 			brokenDag.Leafs[hash] = leafCopy
 		}
@@ -273,7 +283,12 @@ func TestBrokenDags(t *testing.T) {
 		for hash, leaf := range dag.Leafs {
 			leafCopy := leaf.Clone()
 			if len(leafCopy.ClassicMerkleRoot) > 0 {
-				leafCopy.ClassicMerkleRoot = append(leafCopy.ClassicMerkleRoot, []byte("corrupted")...)
+				// Create a different merkle root by changing the content
+				builder := CreateDagLeafBuilder(leaf.ItemName)
+				builder.SetType(leaf.Type)
+				builder.AddLink("invalid", "invalid:hash")
+				corruptedLeaf, _ := builder.BuildLeaf(nil)
+				leafCopy.ClassicMerkleRoot = corruptedLeaf.ClassicMerkleRoot
 			}
 			brokenDag.Leafs[hash] = leafCopy
 		}
@@ -291,8 +306,13 @@ func TestBrokenDags(t *testing.T) {
 		for hash, leaf := range dag.Leafs {
 			leafCopy := leaf.Clone()
 			if len(leafCopy.Links) > 0 {
-				// Add invalid link
-				leafCopy.Links["invalid"] = "invalid:hash"
+				// Add invalid link while preserving CurrentLinkCount
+				builder := CreateDagLeafBuilder(leaf.ItemName)
+				builder.SetType(leaf.Type)
+				builder.AddLink("invalid", "invalid:hash")
+				corruptedLeaf, _ := builder.BuildLeaf(nil)
+				leafCopy.Links = corruptedLeaf.Links
+				// CurrentLinkCount stays the same as it's part of the hash
 			}
 			brokenDag.Leafs[hash] = leafCopy
 		}

dag/edge_test.go

@@ -518,6 +518,7 @@ func (leaf *DagLeaf) Clone() *DagLeaf {
 		LeafCount:         leaf.LeafCount,
 		Links:             make(map[string]string),
 		AdditionalData:    make(map[string]string),
+		Proofs:            make(map[string]*ClassicTreeBranch),
 	}
 
 	// Deep copy maps
@@ -527,26 +528,42 @@ func (leaf *DagLeaf) Clone() *DagLeaf {
 	for k, v := range leaf.AdditionalData {
 		cloned.AdditionalData[k] = v
 	}
-
-	// If the original leaf has a merkle tree, preserve it
-	if leaf.MerkleTree != nil && leaf.LeafMap != nil {
-		// Create new leaf map with same data
-		leafMap := make(map[string]merkletree.DataBlock)
-		for k, v := range leaf.LeafMap {
-			if tl, ok := v.(*testLeaf); ok {
-				leafMap[k] = &testLeaf{data: tl.data}
-			} else {
-				// For other types, create a new leaf with the same data
-				data, _ := v.Serialize()
-				leafMap[k] = &testLeaf{data: string(data)}
-			}
+	if leaf.Proofs != nil {
+		for k, v := range leaf.Proofs {
+			cloned.Proofs[k] = v
 		}
+	}
+
+	// Copy root-specific fields if this is the root leaf
+	if leaf.Hash == leaf.ParentHash || leaf.Hash == GetHash(leaf.Hash) {
+		cloned.LatestLabel = leaf.LatestLabel
+		cloned.LeafCount = leaf.LeafCount
+		cloned.ParentHash = cloned.Hash // Root is its own parent
+	} else {
+		cloned.ParentHash = leaf.ParentHash
+	}
 
-		// Build new merkle tree with same data
-		merkleTree, err := merkletree.New(nil, leafMap)
-		if err == nil {
-			cloned.MerkleTree = merkleTree
-			cloned.LeafMap = leafMap
+	// If leaf has multiple children according to CurrentLinkCount,
+	// we need to handle its merkle tree state
+	if leaf.CurrentLinkCount > 1 {
+		if len(leaf.Links) > 1 {
+			// Build merkle tree with current links
+			builder := merkle_tree.CreateTree()
+			for l, h := range leaf.Links {
+				builder.AddLeaf(l, h)
+			}
+			merkleTree, leafMap, err := builder.Build()
+			if err == nil {
+				cloned.MerkleTree = merkleTree
+				cloned.LeafMap = leafMap
+				cloned.ClassicMerkleRoot = merkleTree.Root
+			}
+		} else {
+			// Clear merkle tree if we don't have enough links to rebuild it
+			cloned.MerkleTree = nil
+			cloned.LeafMap = nil
+			// But keep ClassicMerkleRoot as it's part of the leaf's identity
+			cloned.ClassicMerkleRoot = leaf.ClassicMerkleRoot
 		}
 	}
 
@@ -586,7 +603,11 @@ func GetLabel(hash string) string {
 
 func sortMapByKeys(inputMap map[string]string) map[string]string {
 	if inputMap == nil {
-		return inputMap
+		return map[string]string{}
+	}
+
+	if len(inputMap) <= 0 {
+		return map[string]string{}
 	}
 
 	keys := make([]string, 0, len(inputMap))