Minor refactoring351443_incorrect_pagination_cursor

2 files changed, 44 insertions, 38 deletions

diff --git a/internal/gitaly/service/commit/tree_entries.go b/internal/gitaly/service/commit/tree_entries.go
index 8a6667f3a..f7510e6f3 100644
--- a/internal/gitaly/service/commit/tree_entries.go
+++ b/internal/gitaly/service/commit/tree_entries.go
@@ -302,28 +302,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())
 }
 
-type PageToken struct {
-	FileName string
-}
-
-func decodedPageToken(token string) (string, string) {
-	var pageToken PageToken
-
-	decodedString, err := base64.StdEncoding.DecodeString(token)
-	if err != nil {
-		return token, "oid"
-	}
-
-	if err := json.Unmarshal([]byte(decodedString), &pageToken); err != nil {
-		return token, "oid"
-	} else {
-		return pageToken.FileName, "filename"
-	}
-}
-
 func paginateTreeEntries(entries []*gitalypb.TreeEntry, p *gitalypb.PaginationParameter) ([]*gitalypb.TreeEntry, string, error) {
 	limit := int(p.GetLimit())
-	start, tokenType := decodedPageToken(p.GetPageToken())
+	start, tokenType := decodePageToken(p.GetPageToken())
 	index := -1
 
 	// No token means we should start from the top
@@ -331,16 +312,9 @@ func paginateTreeEntries(entries []*gitalypb.TreeEntry, p *gitalypb.PaginationPa
 		index = 0
 	} else {
 		for i, entry := range entries {
-			if tokenType == "oid" {
-				if entry.GetOid() == start {
-					index = i + 1
-					break
-				}
-			} else {
-				if string(entry.GetPath()) == start {
-					index = i + 1
-					break
-				}
+			if buildEntryToken(entry, tokenType) == start {
+				index = i + 1
+				break
 			}
 		}
 	}
@@ -359,23 +333,55 @@ func paginateTreeEntries(entries []*gitalypb.TreeEntry, p *gitalypb.PaginationPa
 
 	paginated := entries[index : index+limit]
 
-	newCursor, err := generateCursor(paginated[len(paginated)-1])
+	newPageToken, err := encodePageToken(paginated[len(paginated)-1])
 	if err != nil {
 		return nil, "", fmt.Errorf("cannot marshal JSON token: %s", err)
 	}
 
-	return paginated, newCursor, nil
+	return paginated, newPageToken, nil
+}
+
+func buildEntryToken(entry *gitalypb.TreeEntry, tokenType string) string {
+	if tokenType == "oid" {
+		return entry.GetOid()
+	}
+
+	if tokenType == "filename" {
+		return string(entry.GetPath())
+	}
+
+	return ""
+}
+
+type pageToken struct {
+	FileName string
+}
+
+// Decodes Base64 encoded, json marshalled string
+func decodePageToken(token string) (string, string) {
+	var pageToken pageToken
+
+	decodedString, err := base64.StdEncoding.DecodeString(token)
+	if err != nil {
+		return token, "oid"
+	}
+
+	if err := json.Unmarshal(decodedString, &pageToken); err != nil {
+		return token, "oid"
+	}
+
+	return pageToken.FileName, "filename"
 }
 
 // Oid is not a unique value and cannot be used for the cursor generation.
 // Instead we use the file name that should be unique
-func generateCursor(entry *gitalypb.TreeEntry) (string, error) {
-	jsonEncoded, err := json.Marshal(PageToken{FileName: string(entry.GetPath())})
+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 string(encoded), err
+	return encoded, err
 }
diff --git a/internal/gitaly/service/commit/tree_entries_test.go b/internal/gitaly/service/commit/tree_entries_test.go
index fd84fcfd2..3d7e2cf13 100644
--- a/internal/gitaly/service/commit/tree_entries_test.go
+++ b/internal/gitaly/service/commit/tree_entries_test.go
@@ -495,7 +495,7 @@ func TestGetTreeEntries_successful(t *testing.T) {
 					pageToken: testCase.legacyPageToken,
 				},
 				{
-					desc:      "page token",
+					desc:      "structured page token",
 					pageToken: testCase.pageToken,
 				},
 			} {
@@ -534,10 +534,10 @@ func TestGetTreeEntries_successful(t *testing.T) {
 }
 
 func getPageToken(t *testing.T, entry *gitalypb.TreeEntry) string {
-	newCursor, err := generateCursor(entry)
+	pageToken, err := encodePageToken(entry)
 	require.NoError(t, err)
 
-	return newCursor
+	return pageToken
 }
 
 func TestGetTreeEntries_unsuccessful(t *testing.T) {