fix(manifest): resolve node persistence after removal (#4819)

Author: asabya

pkg/manifest/mantaray/marshal.go

@@ -287,6 +287,9 @@ func (n *Node) UnmarshalBinary(data []byte) error {
 		bb.fromBytes(data[offset:])
 		offset += 32 // skip forks
 		return bb.iter(func(b byte) error {
+			if refBytesSize == 0 {
+				return nil
+			}
 			f := &fork{}
 
 			if len(data) < offset+nodeForkTypeBytesSize {
@@ -296,7 +299,6 @@ func (n *Node) UnmarshalBinary(data []byte) error {
 			nodeType := data[offset]
 
 			nodeForkSize := nodeForkPreReferenceSize + refBytesSize
-
 			if nodeTypeIsWithMetadataType(nodeType) {
 				if len(data) < offset+nodeForkPreReferenceSize+refBytesSize+nodeForkMetadataBytesSize {
 					return fmt.Errorf("not enough bytes for node fork: %d (%d) on byte '%x': %w", (len(data) - offset), (nodeForkPreReferenceSize + refBytesSize + nodeForkMetadataBytesSize), []byte{b}, ErrInvalidManifest)

pkg/manifest/mantaray/node.go

@@ -311,8 +311,11 @@ func (n *Node) Remove(ctx context.Context, path []byte, ls LoadSaver) error {
 		return ErrNotFound
 	}
 	rest := path[len(f.prefix):]
+	defer func() {
+		n.ref = nil
+	}()
 	if len(rest) == 0 {
-		// full path matched
+
 		delete(n.forks, path[0])
 		return nil
 	}

pkg/manifest/mantaray/persist.go

@@ -7,7 +7,6 @@ package mantaray
 import (
 	"context"
 	"errors"
-
 	"golang.org/x/sync/errgroup"
 )
 

pkg/manifest/mantaray/persist_test.go

@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"context"
 	"crypto/sha256"
+	"errors"
 	"sync"
 	"testing"
 
@@ -62,6 +63,117 @@ func TestPersistIdempotence(t *testing.T) {
 	}
 }
 
+func TestPersistRemove(t *testing.T) {
+	t.Parallel()
+
+	for _, tc := range []struct {
+		name     string
+		toAdd    []mantaray.NodeEntry
+		toRemove [][]byte
+	}{
+		{
+			name: "simple",
+			toAdd: []mantaray.NodeEntry{
+				{
+					Path: []byte("/"),
+					Metadata: map[string]string{
+						"index-document": "index.html",
+					},
+				},
+				{
+					Path: []byte("index.html"),
+				},
+				{
+					Path: []byte("img/1.png"),
+				},
+				{
+					Path: []byte("img/2.png"),
+				},
+				{
+					Path: []byte("robots.txt"),
+				},
+			},
+			toRemove: [][]byte{
+				[]byte("img/2.png"),
+			},
+		},
+		{
+			name: "nested-prefix-is-not-collapsed",
+			toAdd: []mantaray.NodeEntry{
+				{
+					Path: []byte("index.html"),
+				},
+				{
+					Path: []byte("img/1.png"),
+				},
+				{
+					Path: []byte("img/2/test1.png"),
+				},
+				{
+					Path: []byte("img/2/test2.png"),
+				},
+				{
+					Path: []byte("robots.txt"),
+				},
+			},
+			toRemove: [][]byte{
+				[]byte("img/2/test1.png"),
+			},
+		},
+	} {
+		ctx := context.Background()
+		var ls mantaray.LoadSaver = newMockLoadSaver()
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			// add and persist
+			n := mantaray.New()
+			for i := 0; i < len(tc.toAdd); i++ {
+				c := tc.toAdd[i].Path
+				e := tc.toAdd[i].Entry
+				if len(e) == 0 {
+					e = append(make([]byte, 32-len(c)), c...)
+				}
+				m := tc.toAdd[i].Metadata
+				err := n.Add(ctx, c, e, m, ls)
+				if err != nil {
+					t.Fatalf("expected no error, got %v", err)
+				}
+			}
+			err := n.Save(ctx, ls)
+			if err != nil {
+				t.Fatalf("expected no error, got %v", err)
+			}
+			ref := n.Reference()
+			// reload and remove
+			nn := mantaray.NewNodeRef(ref)
+			for i := 0; i < len(tc.toRemove); i++ {
+				c := tc.toRemove[i]
+				err := nn.Remove(ctx, c, ls)
+				if err != nil {
+					t.Fatalf("expected no error, got %v", err)
+				}
+			}
+
+			err = nn.Save(ctx, ls)
+			if err != nil {
+				t.Fatalf("expected no error, got %v", err)
+			}
+
+			ref = nn.Reference()
+			// reload and lookup removed node
+			nnn := mantaray.NewNodeRef(ref)
+			for i := 0; i < len(tc.toRemove); i++ {
+				c := tc.toRemove[i]
+				_, err = nnn.LookupNode(ctx, c, ls)
+				if !errors.Is(err, mantaray.ErrNotFound) {
+					t.Fatalf("expected not found error, got %v", err)
+				}
+			}
+		})
+	}
+}
+
 type addr [32]byte
 type mockLoadSaver struct {
 	mtx   sync.Mutex