Merge branch 'jc/add-tree-modification-helpers' into 'master'

localrepo: Add tree modification helpers

See merge request https://gitlab.com/gitlab-org/gitaly/-/merge_requests/5567

Merged-by: Patrick Steinhardt <psteinhardt@gitlab.com>
Approved-by: Patrick Steinhardt <psteinhardt@gitlab.com>
Reviewed-by: James Fargher <proglottis@gmail.com>
Reviewed-by: Patrick Steinhardt <psteinhardt@gitlab.com>
Reviewed-by: karthik nayak <knayak@gitlab.com>
Co-authored-by: John Cai <jcai@gitlab.com>

5 files changed, 1468 insertions, 66 deletions

diff --git a/internal/featureflag/ff_submodule_with_tree_write_api.go b/internal/featureflag/ff_submodule_with_tree_write_api.go
new file mode 100644
index 000000000..deea347fc
--- /dev/null
+++ b/internal/featureflag/ff_submodule_with_tree_write_api.go
@@ -0,0 +1,10 @@
+package featureflag
+
+// SubmoduleWithTreeAPI will enable the UserUpdateSubmodule RPC to use the
+// localrepo package's TreeEntry APIs to modify an existing entry.
+var SubmoduleWithTreeAPI = NewFeatureFlag(
+	"submodule_with_tree_api",
+	"v16.1.0",
+	"https://gitlab.com/gitlab-org/gitaly/-/issues/5040",
+	false,
+)
diff --git a/internal/git/localrepo/tree.go b/internal/git/localrepo/tree.go
index 884df3801..da82d2b04 100644
--- a/internal/git/localrepo/tree.go
+++ b/internal/git/localrepo/tree.go
@@ -111,8 +111,280 @@ var (
 	// ErrNotTreeish indicates that the requested revision or path does not resolve to a tree
 	// object.
 	ErrNotTreeish = errors.New("not treeish")
+	// ErrEntryNotFound indicates an entry could not be found.
+	ErrEntryNotFound = errors.New("entry not found")
+	// ErrEntryExists indicates an entry already exists.
+	ErrEntryExists = errors.New("entry already exists")
+	// ErrPathTraversal indicates a path contains a traversal.
+	ErrPathTraversal = errors.New("path contains traversal")
+	// ErrAbsolutePath indicates the path is an absolute path
+	ErrAbsolutePath = errors.New("path is absolute")
+	// ErrEmptyPath indicates the path is an absolute path
+	ErrEmptyPath = errors.New("path is empty")
 )
 
+func validateFileCreationPath(path string) (string, error) {
+	path = filepath.Clean(path)
+
+	if filepath.IsAbs(path) {
+		return "", ErrAbsolutePath
+	}
+
+	if strings.HasPrefix(path, "..") {
+		return "", ErrPathTraversal
+	}
+
+	if path == "." {
+		return "", ErrEmptyPath
+	}
+
+	return path, nil
+}
+
+type addTreeEntryConfig struct {
+	overwriteFile bool
+	overwriteDir  bool
+}
+
+// AddTreeEntryOption is a function that sets options on the addTreeEntryConfig
+// struct.
+type AddTreeEntryOption func(*addTreeEntryConfig)
+
+// WithOverwriteFile allows Add to overwrite a file
+func WithOverwriteFile() AddTreeEntryOption {
+	return func(a *addTreeEntryConfig) {
+		a.overwriteFile = true
+	}
+}
+
+// WithOverwriteDirectory allows Add to overwrite a directory
+func WithOverwriteDirectory() AddTreeEntryOption {
+	return func(a *addTreeEntryConfig) {
+		a.overwriteDir = true
+	}
+}
+
+// Add takes an entry and adds it to an existing TreeEntry
+// based on path.
+// It works by walking down the TreeEntry's Entries, layer by layer, based on
+// path. If it reaches the limit when walking down the tree, that means the
+// rest of the directories path will need to be created.
+// If we're able to walk all the way down the tree based on path, then we
+// either add the new entry to the last subtree's entries if it doesn't yet
+// exist, or optionally overwrite the entry if it already exists.
+func (t *TreeEntry) Add(
+	path string,
+	newEntry TreeEntry,
+	options ...AddTreeEntryOption,
+) error {
+	path, err := validateFileCreationPath(path)
+	if err != nil {
+		return err
+	}
+
+	var cfg addTreeEntryConfig
+
+	for _, option := range options {
+		option(&cfg)
+	}
+
+	var firstComponent string
+	var parentDirs string
+
+	secondComponent := path
+
+	// currentTree keeps track of the current tree we are examining.
+	currentTree := t
+
+Loop:
+	for {
+		// loop through each layer of the tree based on the directory
+		// structure specified by path.
+		// firstComponent is the current directory name, and secondComponent is the rest
+		// of the sub directories we have yet to walk down into.
+		firstComponent, secondComponent, _ = strings.Cut(secondComponent, "/")
+
+		// look through the current tree's entries.
+		for i, entry := range currentTree.Entries {
+			// If the entry's Path does not match firstComponent, we don't
+			// want to do anything with it.
+			if entry.Path != firstComponent {
+				continue
+			}
+
+			// If the entry's Path does match firstComponent, then either
+			// we found the next directory to walk into, or we've
+			// found the entry we want to replace.
+
+			// We found an entry in the tree that matches the entry
+			// we want to add. Replace the entry or throw an error,
+			// depending on the options passed in via options
+			if filepath.Join(parentDirs, entry.Path) == path {
+				if (entry.Type == Tree && !cfg.overwriteDir) ||
+					(entry.Type == Blob && !cfg.overwriteFile) {
+					return ErrEntryExists
+				}
+
+				currentTree.OID = ""
+				currentTree.Entries[i] = &newEntry
+
+				return nil
+			}
+
+			// We found the next directory to walk into.
+			currentTree.OID = ""
+			currentTree = entry
+			parentDirs = filepath.Join(parentDirs, entry.Path)
+			continue Loop
+		}
+
+		// If we get here, that means we didn't find any directories to
+		// recurse into, which means we need to create a brand new
+		// tree
+		if firstComponent == filepath.Base(path) {
+			currentTree.OID = ""
+			currentTree.Entries = append(
+				currentTree.Entries,
+				&newEntry,
+			)
+
+			return nil
+		}
+
+		currentTree.Entries = append(currentTree.Entries, &TreeEntry{
+			Mode: "040000",
+			Type: Tree,
+			Path: firstComponent,
+		})
+
+		currentTree.OID = ""
+		currentTree = currentTree.Entries[len(currentTree.Entries)-1]
+
+		if secondComponent == "" {
+			return nil
+		}
+	}
+}
+
+// recurse is a common function to recurse down into a TreeEntry based on path,
+// and take some action once we've found the entry in question.
+// entryFn is called on the entry specified by path.
+// treeFn is called on each tree visited during the walk.
+func (t *TreeEntry) recurse(
+	path string,
+	entryFn func(currentTree, entry *TreeEntry, i int) error,
+	treeFn func(entry *TreeEntry) error,
+) error {
+	var firstComponent string
+	var parentDirs string
+	secondComponent := path
+
+	currentTree := t
+	for {
+		// Loop through each layer of the tree based on the directory
+		// structure specified by path.
+		// firstComponent is the current directory name, and secondComponent is the rest
+		// of the sub directories we have yet to walk down into.
+		firstComponent, secondComponent, _ = strings.Cut(secondComponent, "/")
+
+		// Look through the current tree's entries.
+		for i, entry := range currentTree.Entries {
+			if firstComponent != entry.Path {
+				continue
+			}
+
+			if filepath.Join(parentDirs, entry.Path) == path {
+				currentTree.OID = ""
+
+				// Once we find the entry in question, apply the
+				// function to modify the current tree and/or
+				// entry.
+				return entryFn(currentTree, entry, i)
+			}
+
+			if err := treeFn(currentTree); err != nil {
+				return err
+			}
+			currentTree = entry
+			parentDirs = filepath.Join(parentDirs, entry.Path)
+		}
+
+		if secondComponent == "" {
+			return ErrEntryNotFound
+		}
+	}
+}
+
+// Delete deletes the entry of a current tree based on the path.
+func (t *TreeEntry) Delete(
+	path string,
+) error {
+	path, err := validateFileCreationPath(path)
+	if err != nil {
+		return err
+	}
+
+	return t.recurse(
+		path,
+		func(currentTree, entry *TreeEntry, i int) error {
+			currentTree.Entries = append(
+				currentTree.Entries[:i],
+				currentTree.Entries[i+1:]...)
+
+			return nil
+		},
+		func(entry *TreeEntry) error {
+			entry.OID = ""
+			return nil
+		},
+	)
+}
+
+// Get gets an entry of a current tree based on the path.
+func (t *TreeEntry) Get(
+	path string,
+) (*TreeEntry, error) {
+	if path == "" || path == "." {
+		return t, nil
+	}
+
+	var result *TreeEntry
+
+	if err := t.recurse(path, func(currentTree, entry *TreeEntry, i int) error {
+		result = entry
+		return nil
+	},
+		func(entry *TreeEntry) error {
+			return nil
+		},
+	); err != nil {
+		return nil, err
+	}
+
+	return result, nil
+}
+
+// Modify modifies an existing TreeEntry based on a path and a function to
+// modify an entry.
+func (t *TreeEntry) Modify(
+	path string,
+	modifyEntry func(*TreeEntry) error,
+) error {
+	path, err := validateFileCreationPath(path)
+	if err != nil {
+		return err
+	}
+
+	return t.recurse(path, func(currentTree, entry *TreeEntry, _ int) error {
+		return modifyEntry(entry)
+	},
+		func(entry *TreeEntry) error {
+			entry.OID = ""
+			return nil
+		},
+	)
+}
+
 // listEntries lists tree entries for the given treeish.
 func (repo *Repo) listEntries(
 	ctx context.Context,
diff --git a/internal/git/localrepo/tree_test.go b/internal/git/localrepo/tree_test.go
index e4794f996..828538267 100644
--- a/internal/git/localrepo/tree_test.go
+++ b/internal/git/localrepo/tree_test.go
@@ -957,3 +957,1047 @@ func TestReadTree_WithRecursive(t *testing.T) {
 		})
 	}
 }
+
+func TestTreeEntry_Modify(t *testing.T) {
+	testCases := []struct {
+		desc            string
+		tree            *TreeEntry
+		pathToModify    string
+		modifyEntryFunc func(*TreeEntry) error
+		expectedTree    TreeEntry
+		expectedErr     error
+	}{
+		{
+			desc: "flat tree",
+			tree: &TreeEntry{
+				OID:  "abcd",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+					{
+						OID:  "cba",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file2",
+					},
+				},
+			},
+			pathToModify: "file2",
+			modifyEntryFunc: func(t *TreeEntry) error {
+				t.OID = "deadbeef"
+				t.Mode = "100755"
+				t.Path = "new-file-name"
+				return nil
+			},
+			expectedTree: TreeEntry{
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+					{
+						OID:  "deadbeef",
+						Type: Blob,
+						Mode: "100755",
+						Path: "new-file-name",
+					},
+				},
+			},
+		},
+		{
+			desc: "nested tree",
+			tree: &TreeEntry{
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+						Entries: []*TreeEntry{
+							{
+								OID:  "def",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirB",
+								Entries: []*TreeEntry{
+									{
+										OID:  "",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file1",
+									},
+								},
+							},
+							{
+								OID:  "aa123",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirC",
+								Entries: []*TreeEntry{
+									{
+										OID:  "abcd",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file2",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			pathToModify: "dirA/dirB/file1",
+			modifyEntryFunc: func(t *TreeEntry) error {
+				t.OID = "deadbeef"
+				t.Mode = "100755"
+				t.Path = "new-file-name"
+				return nil
+			},
+			expectedTree: TreeEntry{
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+						Entries: []*TreeEntry{
+							{
+								OID:  "",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirB",
+								Entries: []*TreeEntry{
+									{
+										OID:  "deadbeef",
+										Type: Blob,
+										Mode: "100755",
+										Path: "new-file-name",
+									},
+								},
+							},
+							{
+								OID:  "aa123",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirC",
+								Entries: []*TreeEntry{
+									{
+										OID:  "abcd",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file2",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			desc: "nested tree with duplicate file",
+			tree: &TreeEntry{
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+						Entries: []*TreeEntry{
+							{
+								OID:  "def",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirB",
+								Entries: []*TreeEntry{
+									{
+										OID:  "",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file1",
+									},
+								},
+							},
+							{
+								OID:  "aa123",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirC",
+								Entries: []*TreeEntry{
+									{
+										OID:  "abcd",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file2",
+									},
+								},
+							},
+							{
+								OID:  "feedface",
+								Type: Blob,
+								Mode: "100644",
+								Path: "file1",
+							},
+						},
+					},
+				},
+			},
+			pathToModify: "dirA/dirB/file1",
+			modifyEntryFunc: func(t *TreeEntry) error {
+				t.OID = "deadbeef"
+				t.Mode = "100755"
+				t.Path = "new-file-name"
+				return nil
+			},
+			expectedTree: TreeEntry{
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+						Entries: []*TreeEntry{
+							{
+								OID:  "",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirB",
+								Entries: []*TreeEntry{
+									{
+										OID:  "deadbeef",
+										Type: Blob,
+										Mode: "100755",
+										Path: "new-file-name",
+									},
+								},
+							},
+							{
+								OID:  "aa123",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirC",
+								Entries: []*TreeEntry{
+									{
+										OID:  "abcd",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file2",
+									},
+								},
+							},
+							{
+								OID:  "feedface",
+								Type: Blob,
+								Mode: "100644",
+								Path: "file1",
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			desc: "entry does not exist",
+			tree: &TreeEntry{
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+					{
+						OID:  "cba",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file2",
+					},
+				},
+			},
+			pathToModify: "file3",
+			modifyEntryFunc: func(t *TreeEntry) error {
+				t.OID = "deadbeef"
+				t.Mode = "100755"
+				t.Path = "new-file-name"
+				return nil
+			},
+			expectedErr: ErrEntryNotFound,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			err := tc.tree.Modify(
+				tc.pathToModify,
+				tc.modifyEntryFunc,
+			)
+			require.Equal(t, tc.expectedErr, err)
+			if tc.expectedErr != nil {
+				return
+			}
+
+			require.Equal(t, tc.expectedTree, *tc.tree)
+		})
+	}
+}
+
+func TestTreeEntry_Delete(t *testing.T) {
+	testCases := []struct {
+		desc         string
+		tree         *TreeEntry
+		pathToDelete string
+		expectedTree TreeEntry
+		expectedErr  error
+	}{
+		{
+			desc: "flat tree",
+			tree: &TreeEntry{
+				OID:  "def",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+					{
+						OID:  "cba",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file2",
+					},
+				},
+			},
+			pathToDelete: "file2",
+			expectedTree: TreeEntry{
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+				},
+			},
+		},
+		{
+			desc: "nested tree",
+			tree: &TreeEntry{
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+						Entries: []*TreeEntry{
+							{
+								OID:  "",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirB",
+								Entries: []*TreeEntry{
+									{
+										OID:  "defg",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file1",
+									},
+								},
+							},
+							{
+								OID:  "aa123",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirC",
+								Entries: []*TreeEntry{
+									{
+										OID:  "abcd",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file2",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			pathToDelete: "dirA/dirB/file1",
+			expectedTree: TreeEntry{
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+						Entries: []*TreeEntry{
+							{
+								OID:     "",
+								Type:    Tree,
+								Mode:    "040000",
+								Path:    "dirB",
+								Entries: []*TreeEntry{},
+							},
+							{
+								OID:  "aa123",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirC",
+								Entries: []*TreeEntry{
+									{
+										OID:  "abcd",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file2",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			desc: "entry does not exist",
+			tree: &TreeEntry{
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+					{
+						OID:  "cba",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file2",
+					},
+				},
+			},
+			pathToDelete: "file3",
+			expectedErr:  ErrEntryNotFound,
+		},
+		{
+			desc: "path contains traversal",
+			tree: &TreeEntry{
+				OID:  "def",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+					{
+						OID:  "cba",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file2",
+					},
+				},
+			},
+			pathToDelete: "a/../../something",
+			expectedErr:  ErrPathTraversal,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			err := tc.tree.Delete(
+				tc.pathToDelete,
+			)
+
+			require.Equal(t, tc.expectedErr, err)
+			if tc.expectedErr != nil {
+				return
+			}
+			require.Equal(t, tc.expectedTree, *tc.tree)
+		})
+	}
+}
+
+func TestTreeEntry_Add(t *testing.T) {
+	testCases := []struct {
+		desc           string
+		tree           TreeEntry
+		pathToAdd      string
+		entryToAdd     TreeEntry
+		expectedTree   TreeEntry
+		addTreeOptions []AddTreeEntryOption
+		expectedErr    error
+	}{
+		{
+			desc: "empty tree",
+			tree: TreeEntry{
+				OID:  "abc",
+				Type: Tree,
+				Mode: "040000",
+			},
+			pathToAdd: "dirA/dirB/file1",
+			entryToAdd: TreeEntry{
+				OID:  "d1",
+				Type: Blob,
+				Mode: "100644",
+				Path: "file1",
+			},
+			expectedTree: TreeEntry{
+				OID:  "",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+						Entries: []*TreeEntry{
+							{
+								OID:  "",
+								Mode: "040000",
+								Type: Tree,
+								Path: "dirB",
+								Entries: []*TreeEntry{
+									{
+										OID:  "d1",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file1",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			desc: "cannot add entry into existing tree",
+			tree: TreeEntry{
+				OID:  "abc",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "def",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+						Entries: []*TreeEntry{
+							{
+								OID:  "gab",
+								Mode: "040000",
+								Type: Tree,
+								Path: "dirB",
+								Entries: []*TreeEntry{
+									{
+										OID:  "d1",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file1",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			pathToAdd: "dirA/dirB/file1",
+			entryToAdd: TreeEntry{
+				OID:  "e1",
+				Type: Blob,
+				Mode: "100644",
+				Path: "file2",
+			},
+			expectedErr: ErrEntryExists,
+		},
+		{
+			desc: "add entry into existing tree",
+			tree: TreeEntry{
+				OID:  "abc",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "def",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+						Entries: []*TreeEntry{
+							{
+								OID:  "gab",
+								Mode: "040000",
+								Type: Tree,
+								Path: "dirB",
+								Entries: []*TreeEntry{
+									{
+										OID:  "d1",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file1",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			pathToAdd: "dirA/dirB/file2",
+			entryToAdd: TreeEntry{
+				OID:  "e1",
+				Type: Blob,
+				Mode: "100644",
+				Path: "file2",
+			},
+			expectedTree: TreeEntry{
+				OID:  "",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+						Entries: []*TreeEntry{
+							{
+								OID:  "",
+								Mode: "040000",
+								Type: Tree,
+								Path: "dirB",
+								Entries: []*TreeEntry{
+									{
+										OID:  "d1",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file1",
+									},
+									{
+										OID:  "e1",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file2",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			desc: "overwrite file",
+			tree: TreeEntry{
+				OID:  "abc",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "def",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+				},
+			},
+			pathToAdd: "file1",
+			entryToAdd: TreeEntry{
+				OID:  "e1",
+				Type: Blob,
+				Mode: "100644",
+				Path: "file2",
+			},
+			expectedTree: TreeEntry{
+				OID:  "",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "e1",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file2",
+					},
+				},
+			},
+			addTreeOptions: []AddTreeEntryOption{
+				WithOverwriteFile(),
+			},
+		},
+		{
+			desc: "cannot overwrite file",
+			tree: TreeEntry{
+				OID:  "abc",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "def",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+				},
+			},
+			pathToAdd: "file1",
+			entryToAdd: TreeEntry{
+				OID:  "e1",
+				Type: Blob,
+				Mode: "100644",
+				Path: "file2",
+			},
+			expectedErr: ErrEntryExists,
+		},
+		{
+			desc: "overwrite directory",
+			tree: TreeEntry{
+				OID:  "abc",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "def",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+					},
+				},
+			},
+			pathToAdd: "dirA",
+			entryToAdd: TreeEntry{
+				OID:  "e1",
+				Type: Blob,
+				Mode: "100644",
+				Path: "file2",
+			},
+			expectedTree: TreeEntry{
+				OID:  "",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "e1",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file2",
+					},
+				},
+			},
+			addTreeOptions: []AddTreeEntryOption{
+				WithOverwriteDirectory(),
+			},
+		},
+		{
+			desc: "cannot overwrite directory",
+			tree: TreeEntry{
+				OID:  "abc",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "def",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+					},
+				},
+			},
+			pathToAdd: "dirA",
+			entryToAdd: TreeEntry{
+				OID:  "e1",
+				Type: Blob,
+				Mode: "100644",
+				Path: "file2",
+			},
+			expectedErr: ErrEntryExists,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			err := tc.tree.Add(
+				tc.pathToAdd,
+				tc.entryToAdd,
+				tc.addTreeOptions...,
+			)
+
+			require.Equal(t, tc.expectedErr, err)
+			if tc.expectedErr != nil {
+				return
+			}
+			require.Equal(t, tc.expectedTree, tc.tree)
+		})
+	}
+}
+
+func TestTreeEntry_Get(t *testing.T) {
+	testCases := []struct {
+		desc          string
+		tree          *TreeEntry
+		path          string
+		expectedEntry TreeEntry
+		expectedErr   error
+	}{
+		{
+			desc: "flat tree",
+			tree: &TreeEntry{
+				OID:  "def",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+					{
+						OID:  "cba",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file2",
+					},
+				},
+			},
+			path: "file2",
+			expectedEntry: TreeEntry{
+				OID:  "cba",
+				Type: Blob,
+				Mode: "100644",
+				Path: "file2",
+			},
+		},
+		{
+			desc: "get root",
+			tree: &TreeEntry{
+				OID:  "def",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+					{
+						OID:  "cba",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file2",
+					},
+				},
+			},
+			path: "",
+			expectedEntry: TreeEntry{
+				OID:  "def",
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+					{
+						OID:  "cba",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file2",
+					},
+				},
+			},
+		},
+		{
+			desc: "nested tree",
+			tree: &TreeEntry{
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "",
+						Type: Tree,
+						Mode: "040000",
+						Path: "dirA",
+						Entries: []*TreeEntry{
+							{
+								OID:  "",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirB",
+								Entries: []*TreeEntry{
+									{
+										OID:  "defg",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file1",
+									},
+								},
+							},
+							{
+								OID:  "aa123",
+								Type: Tree,
+								Mode: "040000",
+								Path: "dirC",
+								Entries: []*TreeEntry{
+									{
+										OID:  "abcd",
+										Type: Blob,
+										Mode: "100644",
+										Path: "file2",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			path: "dirA/dirB/file1",
+			expectedEntry: TreeEntry{
+				OID:  "defg",
+				Type: Blob,
+				Mode: "100644",
+				Path: "file1",
+			},
+		},
+		{
+			desc: "entry does not exist",
+			tree: &TreeEntry{
+				Type: Tree,
+				Mode: "040000",
+				Entries: []*TreeEntry{
+					{
+						OID:  "abc",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file1",
+					},
+					{
+						OID:  "cba",
+						Type: Blob,
+						Mode: "100644",
+						Path: "file2",
+					},
+				},
+			},
+			path:        "file3",
+			expectedErr: ErrEntryNotFound,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			entry, err := tc.tree.Get(
+				tc.path,
+			)
+
+			require.Equal(t, tc.expectedErr, err)
+			if tc.expectedErr != nil {
+				return
+			}
+			require.Equal(t, tc.expectedEntry, *entry)
+		})
+	}
+}
+
+func TestValidatePath(t *testing.T) {
+	testCases := []struct {
+		desc         string
+		path         string
+		expectedPath string
+		expectedErr  error
+	}{
+		{
+			desc:         "normal path",
+			path:         "a/b/c/d/e/f/g",
+			expectedPath: "a/b/c/d/e/f/g",
+			expectedErr:  nil,
+		},
+		{
+			desc:         "weird path",
+			path:         "a/b..c/d..e/f/g",
+			expectedPath: "a/b..c/d..e/f/g",
+			expectedErr:  nil,
+		},
+		{
+			desc:        "starts with traversal",
+			path:        "../../a",
+			expectedErr: ErrPathTraversal,
+		},
+		{
+			desc:        "contains traversal",
+			path:        "a/../../../b",
+			expectedErr: ErrPathTraversal,
+		},
+		{
+			desc:         "does not contain traversal",
+			path:         "a/b/../c",
+			expectedPath: "a/c",
+			expectedErr:  nil,
+		},
+		{
+			desc:        "absolute path",
+			path:        "/a/b/c",
+			expectedErr: ErrAbsolutePath,
+		},
+		{
+			desc:        "empty path",
+			path:        "",
+			expectedErr: ErrEmptyPath,
+		},
+		{
+			desc:        "empty path",
+			path:        ".",
+			expectedErr: ErrEmptyPath,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			path, err := validateFileCreationPath(tc.path)
+			require.Equal(
+				t,
+				tc.expectedPath,
+				path,
+			)
+			require.Equal(
+				t,
+				tc.expectedErr,
+				err,
+			)
+		})
+	}
+}
diff --git a/internal/gitaly/service/operations/submodules.go b/internal/gitaly/service/operations/submodules.go
index b40ea1b4c..ff409a214 100644
--- a/internal/gitaly/service/operations/submodules.go
+++ b/internal/gitaly/service/operations/submodules.go
@@ -9,6 +9,7 @@ import (
 	"strings"
 
 	"github.com/grpc-ecosystem/go-grpc-middleware/logging/logrus/ctxlogrus"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/featureflag"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/localrepo"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/updateref"
@@ -59,22 +60,27 @@ func validateUserUpdateSubmoduleRequest(req *gitalypb.UserUpdateSubmoduleRequest
 	return nil
 }
 
+// legacyGit2GoSubmoduleAlreadyAtShaErr is used to maintain backwards
+// compatibility with the git2go error.
+type legacyGit2GoSubmoduleAlreadyAtShaError struct {
+	submodulePath string
+	commitSha     string
+}
+
+func (l *legacyGit2GoSubmoduleAlreadyAtShaError) Error() string {
+	return fmt.Sprintf("The submodule %s is already at %s", l.submodulePath, l.commitSha)
+}
+
 func (s *Server) updateSubmodule(ctx context.Context, quarantineRepo *localrepo.Repo, req *gitalypb.UserUpdateSubmoduleRequest) (string, error) {
-	path := filepath.Dir(string(req.GetSubmodule()))
-	base := filepath.Base(string(req.GetSubmodule()))
-	replaceWith := git.ObjectID(req.GetCommitSha())
+	var treeID git.ObjectID
 
-	var submoduleFound bool
+	path := filepath.Dir(string(req.GetSubmodule()))
 
-	// First find the tree containing the submodule to be updated.
-	// Write a new tree object abc with the updated submodule. Then, write a new
-	// tree with the new tree abcabc. Continue iterating up the tree,
-	// writing a new tree object each time.
-	for {
-		tree, err := quarantineRepo.ReadTree(
+	if featureflag.SubmoduleWithTreeAPI.IsEnabled(ctx) {
+		fullTree, err := quarantineRepo.ReadTree(
 			ctx,
-			git.Revision("refs/heads/"+string(req.GetBranch())),
-			localrepo.WithRelativePath(path),
+			git.NewReferenceNameFromBranchName(string(req.GetBranch())).Revision(),
+			localrepo.WithRecursive(),
 		)
 		if err != nil {
 			if errors.Is(err, git.ErrReferenceNotFound) {
@@ -84,71 +90,136 @@ func (s *Server) updateSubmodule(ctx context.Context, quarantineRepo *localrepo.
 			return "", fmt.Errorf("error reading tree: %w", err)
 		}
 
-		var newEntries []*localrepo.TreeEntry
-		if err := tree.Walk(func(_ string, entry *localrepo.TreeEntry) error {
-			if entry.OID == tree.OID {
+		if err := fullTree.Modify(
+			string(req.GetSubmodule()),
+			func(t *localrepo.TreeEntry) error {
+				replaceWith := git.ObjectID(req.GetCommitSha())
+
+				if t.Type != localrepo.Submodule {
+					return fmt.Errorf("submodule: %s", git2go.LegacyErrPrefixInvalidSubmodulePath)
+				}
+
+				if replaceWith == t.OID {
+					return &legacyGit2GoSubmoduleAlreadyAtShaError{
+						submodulePath: string(req.GetSubmodule()),
+						commitSha:     string(replaceWith),
+					}
+				}
+
+				t.OID = replaceWith
+
 				return nil
+			},
+		); err != nil {
+			if err == localrepo.ErrEntryNotFound {
+				return "", fmt.Errorf("submodule: %s", git2go.LegacyErrPrefixInvalidSubmodulePath)
 			}
 
-			// If the entry's path does not match, then we simply
-			// want to retain this tree entry.
-			if entry.Path != base {
-				newEntries = append(newEntries, entry)
-				return nil
+			var git2GoErr *legacyGit2GoSubmoduleAlreadyAtShaError
+			if errors.As(err, &git2GoErr) {
+				return "", err
 			}
 
-			// If we are at the submodule we want to replace, check
-			// if it's already at the value we want to replace, or
-			// if it's not a submodule.
-			if filepath.Join(path, entry.Path) == string(req.GetSubmodule()) {
-				if string(entry.OID) == req.GetCommitSha() {
-					//nolint:stylecheck
-					return fmt.Errorf(
-						"The submodule %s is already at %s",
-						req.GetSubmodule(),
-						replaceWith,
-					)
-				}
+			return "", fmt.Errorf("modifying tree: %w", err)
+		}
 
-				if entry.Type != localrepo.Submodule {
-					return fmt.Errorf("submodule: %s", git2go.LegacyErrPrefixInvalidSubmodulePath)
+		if err := fullTree.Write(ctx, quarantineRepo); err != nil {
+			return "", fmt.Errorf("writing tree: %w", err)
+		}
+
+		treeID = fullTree.OID
+	} else {
+		base := filepath.Base(string(req.GetSubmodule()))
+		replaceWith := git.ObjectID(req.GetCommitSha())
+
+		var submoduleFound bool
+
+		// First find the tree containing the submodule to be updated.
+		// Write a new tree object abc with the updated submodule. Then, write a new
+		// tree with the new tree abcabc. Continue iterating up the tree,
+		// writing a new tree object each time.
+		for {
+			tree, err := quarantineRepo.ReadTree(
+				ctx,
+				git.Revision("refs/heads/"+string(req.GetBranch())),
+				localrepo.WithRelativePath(path),
+			)
+			if err != nil {
+				if errors.Is(err, git.ErrReferenceNotFound) {
+					return "", fmt.Errorf("submodule: %s", git2go.LegacyErrPrefixInvalidSubmodulePath)
 				}
+
+				return "", fmt.Errorf("error reading tree: %w", err)
 			}
 
-			// Otherwise, create a new tree entry
-			submoduleFound = true
+			var newEntries []*localrepo.TreeEntry
+			if err := tree.Walk(func(_ string, entry *localrepo.TreeEntry) error {
+				if entry.OID == tree.OID {
+					return nil
+				}
 
-			newEntries = append(newEntries, &localrepo.TreeEntry{
-				Mode: entry.Mode,
-				Path: entry.Path,
-				OID:  replaceWith,
-			})
+				// If the entry's path does not match, then we simply
+				// want to retain this tree entry.
+				if entry.Path != base {
+					newEntries = append(newEntries, entry)
+					return nil
+				}
 
-			return nil
-		}); err != nil {
-			return "", err
-		}
+				// If we are at the submodule we want to replace, check
+				// if it's already at the value we want to replace, or
+				// if it's not a submodule.
+				if filepath.Join(path, entry.Path) == string(req.GetSubmodule()) {
+					if string(entry.OID) == req.GetCommitSha() {
+						//nolint:stylecheck
+						return fmt.Errorf(
+							"The submodule %s is already at %s",
+							req.GetSubmodule(),
+							replaceWith,
+						)
+					}
+
+					if entry.Type != localrepo.Submodule {
+						return fmt.Errorf("submodule: %s", git2go.LegacyErrPrefixInvalidSubmodulePath)
+					}
+				}
 
-		newTree := &localrepo.TreeEntry{
-			Type:    localrepo.Tree,
-			Mode:    "040000",
-			Entries: newEntries,
-		}
-		if err := newTree.Write(ctx, quarantineRepo); err != nil {
-			return "", fmt.Errorf("write tree: %w", err)
-		}
-		replaceWith = newTree.OID
+				// Otherwise, create a new tree entry
+				submoduleFound = true
+
+				newEntries = append(newEntries, &localrepo.TreeEntry{
+					Mode: entry.Mode,
+					Path: entry.Path,
+					OID:  replaceWith,
+				})
+
+				return nil
+			}); err != nil {
+				return "", err
+			}
+
+			newTree := &localrepo.TreeEntry{
+				Type:    localrepo.Tree,
+				Mode:    "040000",
+				Entries: newEntries,
+			}
+			if err := newTree.Write(ctx, quarantineRepo); err != nil {
+				return "", fmt.Errorf("write tree: %w", err)
+			}
+			replaceWith = newTree.OID
 
-		if path == "." {
-			break
+			if path == "." {
+				break
+			}
+
+			base = filepath.Base(path)
+			path = filepath.Dir(path)
 		}
 
-		base = filepath.Base(path)
-		path = filepath.Dir(path)
-	}
+		if !submoduleFound {
+			return "", fmt.Errorf("submodule: %s", git2go.LegacyErrPrefixInvalidSubmodulePath)
+		}
 
-	if !submoduleFound {
-		return "", fmt.Errorf("submodule: %s", git2go.LegacyErrPrefixInvalidSubmodulePath)
+		treeID = replaceWith
 	}
 
 	currentBranchCommit, err := quarantineRepo.ResolveRevision(ctx, git.Revision(req.GetBranch()))
@@ -170,7 +241,7 @@ func (s *Server) updateSubmodule(ctx context.Context, quarantineRepo *localrepo.
 		CommitterEmail: string(req.GetUser().GetEmail()),
 		CommitterDate:  authorDate,
 		Message:        string(req.GetCommitMessage()),
-		TreeID:         replaceWith,
+		TreeID:         treeID,
 	})
 	if err != nil {
 		return "", fmt.Errorf("creating commit %w", err)
@@ -226,8 +297,7 @@ func (s *Server) userUpdateSubmodule(ctx context.Context, req *gitalypb.UserUpda
 
 	commitID, err := s.updateSubmodule(ctx, quarantineRepo, req)
 	if err != nil {
-		errStr := strings.TrimPrefix(err.Error(), "submodule: ")
-		errStr = strings.TrimSpace(errStr)
+		errStr := strings.TrimSpace(err.Error())
 
 		var resp *gitalypb.UserUpdateSubmoduleResponse
 		for _, legacyErr := range []string{
@@ -235,7 +305,7 @@ func (s *Server) userUpdateSubmodule(ctx context.Context, req *gitalypb.UserUpda
 			git2go.LegacyErrPrefixInvalidSubmodulePath,
 			git2go.LegacyErrPrefixFailedCommit,
 		} {
-			if strings.HasPrefix(errStr, legacyErr) {
+			if strings.Contains(errStr, legacyErr) {
 				resp = &gitalypb.UserUpdateSubmoduleResponse{
 					CommitError: legacyErr,
 				}
diff --git a/internal/gitaly/service/operations/submodules_test.go b/internal/gitaly/service/operations/submodules_test.go
index bb689e78f..9dbe897e7 100644
--- a/internal/gitaly/service/operations/submodules_test.go
+++ b/internal/gitaly/service/operations/submodules_test.go
@@ -4,11 +4,13 @@ package operations
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"path/filepath"
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/featureflag"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/gittest"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/git/localrepo"
@@ -21,7 +23,11 @@ import (
 func TestUserUpdateSubmodule(t *testing.T) {
 	t.Parallel()
 
-	ctx := testhelper.Context(t)
+	testhelper.NewFeatureSets(featureflag.SubmoduleWithTreeAPI).
+		Run(t, testUserUpdateSubmodule)
+}
+
+func testUserUpdateSubmodule(t *testing.T, ctx context.Context) {
 	ctx, cfg, client := setupOperationsServiceWithoutRepo(t, ctx)
 
 	type setupData struct {