Merge branch '351443_incorrect_pagination_cursor' into 'master'

Fix pagination cursor generation for TreeEntries

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

Merged-by: Patrick Steinhardt <psteinhardt@gitlab.com>
Approved-by: Patrick Steinhardt <psteinhardt@gitlab.com>
Co-authored-by: Vasilii Iakliushin <viakliushin@gitlab.com>

3 files changed, 211 insertions, 90 deletions

diff --git a/internal/gitaly/service/commit/tree_entries.go b/internal/gitaly/service/commit/tree_entries.go
index 2a2489526..5549631a0 100644
--- a/internal/gitaly/service/commit/tree_entries.go
+++ b/internal/gitaly/service/commit/tree_entries.go
@@ -2,6 +2,8 @@ package commit
 
 import (
 	"context"
+	"encoding/base64"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"sort"
@@ -13,6 +15,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git/localrepo"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git/lstree"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/helper/chunk"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/metadata/featureflag"
 	"gitlab.com/gitlab-org/gitaly/v15/proto/go/gitalypb"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -176,7 +179,7 @@ func (s *server) sendTreeEntries(
 
 	cursor := ""
 	if p != nil {
-		entries, cursor, err = paginateTreeEntries(entries, p)
+		entries, cursor, err = paginateTreeEntries(ctx, entries, p)
 		if err != nil {
 			return err
 		}
@@ -300,9 +303,9 @@ func (s *server) GetTreeEntries(in *gitalypb.GetTreeEntriesRequest, stream gital
 	return s.sendTreeEntries(stream, repo, revision, path, in.Recursive, in.SkipFlatPaths, in.GetSort(), in.GetPaginationParams())
 }
 
-func paginateTreeEntries(entries []*gitalypb.TreeEntry, p *gitalypb.PaginationParameter) ([]*gitalypb.TreeEntry, string, error) {
+func paginateTreeEntries(ctx context.Context, entries []*gitalypb.TreeEntry, p *gitalypb.PaginationParameter) ([]*gitalypb.TreeEntry, string, error) {
 	limit := int(p.GetLimit())
-	start := p.GetPageToken()
+	start, tokenType := decodePageToken(p.GetPageToken())
 	index := -1
 
 	// No token means we should start from the top
@@ -310,7 +313,7 @@ func paginateTreeEntries(entries []*gitalypb.TreeEntry, p *gitalypb.PaginationPa
 		index = 0
 	} else {
 		for i, entry := range entries {
-			if entry.GetOid() == start {
+			if buildEntryToken(entry, tokenType) == start {
 				index = i + 1
 				break
 			}
@@ -330,5 +333,68 @@ func paginateTreeEntries(entries []*gitalypb.TreeEntry, p *gitalypb.PaginationPa
 	}
 
 	paginated := entries[index : index+limit]
+
+	if featureflag.TreeEntriesNewPageTokenFormat.IsEnabled(ctx) {
+		newPageToken, err := encodePageToken(paginated[len(paginated)-1])
+		if err != nil {
+			return nil, "", fmt.Errorf("encode page token: %w", err)
+		}
+
+		return paginated, newPageToken, nil
+	}
+
 	return paginated, paginated[len(paginated)-1].GetOid(), nil
 }
+
+func buildEntryToken(entry *gitalypb.TreeEntry, tokenType pageTokenType) string {
+	if tokenType == pageTokenTypeOID {
+		return entry.GetOid()
+	}
+
+	return string(entry.GetPath())
+}
+
+type pageToken struct {
+	// FileName is the name of the tree entry that acts as continuation point.
+	FileName string `json:"file_name"`
+}
+
+type pageTokenType bool
+
+const (
+	// pageTokenTypeOID is an old-style page token that contains the object ID a tree
+	// entry is pointing to. This is ambiguous and thus deprecated.
+	pageTokenTypeOID pageTokenType = false
+	// pageTokenTypeFilename is a page token that contains the tree entry path.
+	pageTokenTypeFilename pageTokenType = true
+)
+
+// decodePageToken decodes the given Base64-encoded page token. It returns the
+// continuation point of the token and its type.
+func decodePageToken(token string) (string, pageTokenType) {
+	var pageToken pageToken
+
+	decodedString, err := base64.StdEncoding.DecodeString(token)
+	if err != nil {
+		return token, pageTokenTypeOID
+	}
+
+	if err := json.Unmarshal(decodedString, &pageToken); err != nil {
+		return token, pageTokenTypeOID
+	}
+
+	return pageToken.FileName, pageTokenTypeFilename
+}
+
+// encodePageToken returns a page token with the TreeEntry's path as the continuation point for
+// the next page. The page token serialized by first JSON marshaling it and then base64 encoding it.
+func encodePageToken(entry *gitalypb.TreeEntry) (string, error) {
+	jsonEncoded, err := json.Marshal(pageToken{FileName: string(entry.GetPath())})
+	if err != nil {
+		return "", err
+	}
+
+	encoded := base64.StdEncoding.EncodeToString(jsonEncoded)
+
+	return encoded, err
+}
diff --git a/internal/gitaly/service/commit/tree_entries_test.go b/internal/gitaly/service/commit/tree_entries_test.go
index 000cbb19c..0b0cc2270 100644
--- a/internal/gitaly/service/commit/tree_entries_test.go
+++ b/internal/gitaly/service/commit/tree_entries_test.go
@@ -3,6 +3,7 @@
 package commit
 
 import (
+	"context"
 	"errors"
 	"io"
 	"strconv"
@@ -12,6 +13,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git/gittest"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/metadata/featureflag"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v15/proto/go/gitalypb"
 	"google.golang.org/grpc/codes"
@@ -342,17 +344,28 @@ func TestGetTreeEntries_successful(t *testing.T) {
 		return strippedEntries
 	}
 
+	getPageToken := func(tb testing.TB, entry *gitalypb.TreeEntry) string {
+		tb.Helper()
+
+		pageToken, err := encodePageToken(entry)
+		require.NoError(tb, err)
+
+		return pageToken
+	}
+
 	testCases := []struct {
-		description   string
-		revision      []byte
-		path          []byte
-		recursive     bool
-		sortBy        gitalypb.GetTreeEntriesRequest_SortBy
-		entries       []*gitalypb.TreeEntry
-		pageToken     string
-		pageLimit     int32
-		cursor        string
-		skipFlatPaths bool
+		description     string
+		revision        []byte
+		path            []byte
+		recursive       bool
+		sortBy          gitalypb.GetTreeEntriesRequest_SortBy
+		entries         []*gitalypb.TreeEntry
+		legacyPageToken string
+		pageToken       string
+		pageLimit       int32
+		cursor          string
+		legacyCursor    string
+		skipFlatPaths   bool
 	}{
 		{
 			description: "with root path",
@@ -427,86 +440,119 @@ func TestGetTreeEntries_successful(t *testing.T) {
 			sortBy:      gitalypb.GetTreeEntriesRequest_TREES_FIRST,
 		},
 		{
-			description: "with root path and sorted by trees with pagination",
-			revision:    []byte(commitID),
-			path:        []byte("."),
-			entries:     sortedAndPaginated,
-			pageLimit:   4,
-			sortBy:      gitalypb.GetTreeEntriesRequest_TREES_FIRST,
-			cursor:      "fd90a3d2d21d6b4f9bec2c33fb7f49780c55f0d2",
-		},
-		{
-			description: "with pagination parameters",
-			revision:    []byte(commitID),
-			path:        []byte("."),
-			entries:     rootEntries[3:6],
-			pageToken:   "fdaada1754989978413d618ee1fb1c0469d6a664",
-			pageLimit:   3,
-			cursor:      rootEntries[5].Oid,
-		},
-		{
-			description: "with pagination parameters larger than length",
-			revision:    []byte(commitID),
-			path:        []byte("."),
-			entries:     rootEntries[12:],
-			pageToken:   "b4a3321157f6e80c42b031ecc9ba79f784c8a557",
-			pageLimit:   20,
-		},
-		{
-			description: "with pagination limit of -1",
-			revision:    []byte(commitID),
-			path:        []byte("."),
-			entries:     rootEntries[2:],
-			pageToken:   "470ad2fcf1e33798f1afc5781d08e60c40f51e7a",
-			pageLimit:   -1,
-		},
-		{
-			description: "with pagination limit of 0",
-			revision:    []byte(commitID),
-			path:        []byte("."),
-			pageToken:   "470ad2fcf1e33798f1afc5781d08e60c40f51e7a",
-			pageLimit:   0,
-		},
-		{
-			description: "with a blank pagination token",
-			revision:    []byte(commitID),
-			path:        []byte("."),
-			pageToken:   "",
-			entries:     rootEntries[0:2],
-			pageLimit:   2,
-			cursor:      rootEntries[1].Oid,
+			description:  "with root path and sorted by trees with pagination",
+			revision:     []byte(commitID),
+			path:         []byte("."),
+			entries:      sortedAndPaginated,
+			pageLimit:    4,
+			sortBy:       gitalypb.GetTreeEntriesRequest_TREES_FIRST,
+			cursor:       "eyJmaWxlX25hbWUiOiIuRFNfU3RvcmUifQ==",
+			legacyCursor: "fd90a3d2d21d6b4f9bec2c33fb7f49780c55f0d2",
+		},
+		{
+			description:     "with pagination parameters",
+			revision:        []byte(commitID),
+			path:            []byte("."),
+			entries:         rootEntries[3:6],
+			legacyPageToken: "fdaada1754989978413d618ee1fb1c0469d6a664",
+			pageToken:       getPageToken(t, rootEntries[2]),
+			pageLimit:       3,
+			cursor:          "eyJmaWxlX25hbWUiOiJMSUNFTlNFIn0=",
+			legacyCursor:    "50b27c6518be44c42c4d87966ae2481ce895624c",
+		},
+		{
+			description:     "with pagination parameters larger than length",
+			revision:        []byte(commitID),
+			path:            []byte("."),
+			entries:         rootEntries[12:],
+			legacyPageToken: "b4a3321157f6e80c42b031ecc9ba79f784c8a557",
+			pageToken:       getPageToken(t, rootEntries[11]),
+			pageLimit:       20,
+		},
+		{
+			description:     "with pagination limit of -1",
+			revision:        []byte(commitID),
+			path:            []byte("."),
+			entries:         rootEntries[2:],
+			legacyPageToken: "470ad2fcf1e33798f1afc5781d08e60c40f51e7a",
+			pageToken:       getPageToken(t, rootEntries[1]),
+			pageLimit:       -1,
+		},
+		{
+			description:     "with pagination limit of 0",
+			revision:        []byte(commitID),
+			path:            []byte("."),
+			legacyPageToken: "470ad2fcf1e33798f1afc5781d08e60c40f51e7a",
+			pageToken:       getPageToken(t, rootEntries[0]),
+			pageLimit:       0,
+		},
+		{
+			description:  "with a blank pagination token",
+			revision:     []byte(commitID),
+			path:         []byte("."),
+			pageToken:    "",
+			entries:      rootEntries[0:2],
+			pageLimit:    2,
+			cursor:       "eyJmaWxlX25hbWUiOiIuZ2l0aWdub3JlIn0=",
+			legacyCursor: "470ad2fcf1e33798f1afc5781d08e60c40f51e7a",
 		},
 	}
 
 	for _, testCase := range testCases {
-		t.Run(testCase.description, func(t *testing.T) {
-			request := &gitalypb.GetTreeEntriesRequest{
-				Repository:    repo,
-				Revision:      testCase.revision,
-				Path:          testCase.path,
-				Recursive:     testCase.recursive,
-				Sort:          testCase.sortBy,
-				SkipFlatPaths: testCase.skipFlatPaths,
-			}
-
-			if testCase.pageToken != "" || testCase.pageLimit > 0 {
-				request.PaginationParams = &gitalypb.PaginationParameter{
-					PageToken: testCase.pageToken,
-					Limit:     testCase.pageLimit,
-				}
-			}
-
-			c, err := client.GetTreeEntries(ctx, request)
-
-			require.NoError(t, err)
-			fetchedEntries, cursor := getTreeEntriesFromTreeEntryClient(t, c, nil)
-			testhelper.ProtoEqual(t, testCase.entries, fetchedEntries)
-
-			if testCase.pageLimit > 0 && len(testCase.entries) < len(rootEntries) {
-				require.NotNil(t, cursor)
-				require.Equal(t, testCase.cursor, cursor.NextCursor)
-			}
-		})
+		testhelper.NewFeatureSets(featureflag.TreeEntriesNewPageTokenFormat).Run(
+			t,
+			func(t *testing.T, ctx context.Context) {
+				t.Run(testCase.description, func(t *testing.T) {
+					for _, tc := range []struct {
+						desc      string
+						pageToken string
+					}{
+						{
+							desc:      "legacy token",
+							pageToken: testCase.legacyPageToken,
+						},
+						{
+							desc:      "structured page token",
+							pageToken: testCase.pageToken,
+						},
+					} {
+						t.Run(tc.desc, func(t *testing.T) {
+							request := &gitalypb.GetTreeEntriesRequest{
+								Repository:    repo,
+								Revision:      testCase.revision,
+								Path:          testCase.path,
+								Recursive:     testCase.recursive,
+								Sort:          testCase.sortBy,
+								SkipFlatPaths: testCase.skipFlatPaths,
+							}
+
+							if tc.pageToken != "" || testCase.pageLimit > 0 {
+								request.PaginationParams = &gitalypb.PaginationParameter{
+									PageToken: tc.pageToken,
+									Limit:     testCase.pageLimit,
+								}
+							}
+
+							c, err := client.GetTreeEntries(ctx, request)
+
+							require.NoError(t, err)
+							fetchedEntries, cursor := getTreeEntriesFromTreeEntryClient(t, c, nil)
+							testhelper.ProtoEqual(t, testCase.entries, fetchedEntries)
+
+							if testCase.pageLimit > 0 && len(testCase.entries) < len(rootEntries) {
+								require.NotNil(t, cursor)
+
+								if featureflag.TreeEntriesNewPageTokenFormat.IsEnabled(ctx) {
+									require.Equal(t, testCase.cursor, cursor.NextCursor)
+								} else {
+									require.Equal(t, testCase.legacyCursor, cursor.NextCursor)
+								}
+							}
+						})
+					}
+				})
+			},
+		)
 	}
 }
 
diff --git a/internal/metadata/featureflag/ff_tree_entries_new_page_token_format.go b/internal/metadata/featureflag/ff_tree_entries_new_page_token_format.go
new file mode 100644
index 000000000..c26194cfe
--- /dev/null
+++ b/internal/metadata/featureflag/ff_tree_entries_new_page_token_format.go
@@ -0,0 +1,9 @@
+package featureflag
+
+// TreeEntriesNewPageTokenFormat enables support for new page token format in TreeEntries RPC
+var TreeEntriesNewPageTokenFormat = NewFeatureFlag(
+	"tree_entries_new_page_token_format",
+	"v15.5.0",
+	"https://gitlab.com/gitlab-org/gitaly/-/issues/4542",
+	false,
+)