Gitlab Source: Backoff from Scan2 which is experimental to legacy pagination API call (#4608)

* Backoff from Scan2 which is experimental to legacy pagination API call

This commit rewrite simplifiedGitlabEnumeration to use legacy pagination API call with keyset pagination instead of Scan2 which is currently in experimental state. Note
that this doesn't promise to fix this problem it's just a test to check. It also adds a retry logic in case any 500 error occurs. I added some logs as well to keep track
of no of projects being enumerated.

* implemented builtin retry mechanism for gitlab and proper handling of next page

* fixed basic auth

* Some enhancements

Reversed the gitlab cloud logic to add membership flag, so that we use the default false for non gitlab.com instances.
This can help if the issue really was membership flag as mentioned in some gitlab issues.
Also added simple flag in list projects to get only minimal fields in response instead of big json response for each project.

Added test case as well.

* enhance the test case

Author: kashifkhan0771

pkg/sources/gitlab/gitlab.go

@@ -8,6 +8,7 @@ import (
 	"slices"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
@@ -471,14 +472,25 @@ func (s *Source) newClient() (*gitlab.Client, error) {
 	// Initialize a new api instance.
 	switch s.authMethod {
 	case "OAUTH":
-		apiClient, err := gitlab.NewOAuthClient(s.token, gitlab.WithBaseURL(s.url))
+		apiClient, err := gitlab.NewOAuthClient(
+			s.token,
+			gitlab.WithBaseURL(s.url),
+			gitlab.WithCustomRetryWaitMinMax(time.Second, 5*time.Second),
+			gitlab.WithCustomRetryMax(3),
+		)
 		if err != nil {
 			return nil, fmt.Errorf("could not create Gitlab OAUTH client for %q: %w", s.url, err)
 		}
 		return apiClient, nil
 
 	case "BASIC_AUTH":
-		apiClient, err := gitlab.NewBasicAuthClient(s.user, s.password, gitlab.WithBaseURL(s.url))
+		apiClient, err := gitlab.NewBasicAuthClient(
+			s.user,
+			s.password,
+			gitlab.WithBaseURL(s.url),
+			gitlab.WithCustomRetryWaitMinMax(time.Second, 5*time.Second),
+			gitlab.WithCustomRetryMax(3),
+		)
 		if err != nil {
 			return nil, fmt.Errorf("could not create Gitlab BASICAUTH client for %q: %w", s.url, err)
 		}
@@ -491,7 +503,12 @@ func (s *Source) newClient() (*gitlab.Client, error) {
 		}
 		fallthrough
 	case "TOKEN":
-		apiClient, err := gitlab.NewOAuthClient(s.token, gitlab.WithBaseURL(s.url))
+		apiClient, err := gitlab.NewOAuthClient(
+			s.token,
+			gitlab.WithBaseURL(s.url),
+			gitlab.WithCustomRetryWaitMinMax(time.Second, 5*time.Second),
+			gitlab.WithCustomRetryMax(3),
+		)
 		if err != nil {
 			return nil, fmt.Errorf("could not create Gitlab TOKEN client for %q: %w", s.url, err)
 		}
@@ -699,74 +716,91 @@ func (s *Source) getAllProjectReposV2(
 
 	projectQueryOptions := &gitlab.ListProjectsOptions{
 		ListOptions: listOpts,
-		Membership:  gitlab.Ptr(true),
+		// Return only limited fields for each project
+		Simple: gitlab.Ptr(true),
 	}
 
-	// for non gitlab.com instances, include all available projects (public + membership).
-	if s.url != gitlabBaseURL {
-		projectQueryOptions.Membership = gitlab.Ptr(false)
+	// for gitlab.com instance, include only projects where the user is a member.
+	if s.url == gitlabBaseURL {
+		projectQueryOptions.Membership = gitlab.Ptr(true)
 	}
 
 	ctx.Logger().Info("starting projects enumeration",
-		"list_options", listOpts,
-		"all_available", *projectQueryOptions.Membership)
-
-	// https://pkg.go.dev/gitlab.com/gitlab-org/api/client-go#Scan2
-	projectsIter := gitlab.Scan2(func(p gitlab.PaginationOptionFunc) ([]*gitlab.Project, *gitlab.Response, error) {
-		return apiClient.Projects.ListProjects(projectQueryOptions, p, gitlab.WithContext(ctx))
-	})
+		"list_options", listOpts)
 
+	// totalCount tracks the total number of projects processed by this enumeration.
+	// It includes all projects fetched from the API, even those later skipped by ignore rules.
 	totalCount := 0
 
-	// process each project
-	for project, projectErr := range projectsIter {
-		if projectErr != nil {
-			err := fmt.Errorf("error during project enumeration: %w", projectErr)
+	requestOptions := []gitlab.RequestOptionFunc{gitlab.WithContext(ctx)}
 
-			if reportErr := reporter.UnitErr(ctx, err); reportErr != nil {
-				return reportErr
+	// Pagination loop: Continue fetching pages until the API indicates there are no more.
+	for {
+		// Fetch a page of projects from the GitLab API using the current query options.
+		projects, resp, err := apiClient.Projects.ListProjects(projectQueryOptions, requestOptions...)
+		if err != nil {
+			err = fmt.Errorf("received error on listing projects, you might not have permissions to do that: %w", err)
+			if err := reporter.UnitErr(ctx, err); err != nil {
+				return err
 			}
-
-			continue
+			// break on error as with error we will not have any response and no next page
+			break
 		}
 
-		totalCount++
+		// Log the batch size for debugging and monitoring.
+		ctx.Logger().V(3).Info("listed projects batch", "batch_size", len(projects), "running_total", totalCount)
+		// Process each project in the current page.
+		for _, project := range projects {
+			projCtx := context.WithValues(ctx,
+				"project_id", project.ID,
+				"project_name", project.NameWithNamespace)
 
-		projCtx := context.WithValues(ctx,
-			"project_id", project.ID,
-			"project_name", project.NameWithNamespace)
+			totalCount++
 
-		// skip projects configured to be ignored.
-		if ignoreRepo(project.PathWithNamespace) {
-			projCtx.Logger().V(3).Info("skipping project", "reason", "ignored in config")
+			// skip projects configured to be ignored.
+			if ignoreRepo(project.PathWithNamespace) {
+				projCtx.Logger().V(3).Info("skipping project", "reason", "ignored in config")
 
-			continue
-		}
+				continue
+			}
 
-		// report an error if we could not convert the project into a URL.
-		if _, err := url.Parse(project.HTTPURLToRepo); err != nil {
-			projCtx.Logger().V(3).Info("skipping project",
-				"reason", "URL parse failure",
-				"url", project.HTTPURLToRepo,
-				"parse_error", err)
+			// report an error if we could not convert the project into a URL.
+			if _, err := url.Parse(project.HTTPURLToRepo); err != nil {
+				projCtx.Logger().V(3).Info("skipping project",
+					"reason", "URL parse failure",
+					"url", project.HTTPURLToRepo,
+					"parse_error", err)
 
-			err = fmt.Errorf("could not parse url %q given by project: %w", project.HTTPURLToRepo, err)
-			if err := reporter.UnitErr(ctx, err); err != nil {
-				return err
+				err = fmt.Errorf("could not parse url %q given by project: %w", project.HTTPURLToRepo, err)
+				if err := reporter.UnitErr(ctx, err); err != nil {
+					return err
+				}
+
+				continue
 			}
 
-			continue
-		}
+			// report the unit.
+			projCtx.Logger().V(3).Info("accepting project")
 
-		// report the unit.
-		projCtx.Logger().V(3).Info("accepting project")
+			s.cacheGitlabProject(project)
+			unit := git.SourceUnit{Kind: git.UnitRepo, ID: project.HTTPURLToRepo}
+			gitlabReposEnumerated.WithLabelValues(s.name).Inc()
 
-		s.cacheGitlabProject(project)
-		unit := git.SourceUnit{Kind: git.UnitRepo, ID: project.HTTPURLToRepo}
-		gitlabReposEnumerated.WithLabelValues(s.name).Inc()
+			if err := reporter.UnitOk(ctx, unit); err != nil {
+				return err
+			}
+		}
 
-		if err := reporter.UnitOk(ctx, unit); err != nil {
-			return err
+		// if next page is empty, break the loop
+		if resp == nil || resp.NextLink == "" {
+			// No more pages to fetch. This is the normal loop exit condition.
+			// It also acts as a safety stop if the current request failed.
+			break
+		}
+		// Only update the token for the next page if we have a valid, non-empty link.
+		requestOptions = []gitlab.RequestOptionFunc{
+			gitlab.WithContext(ctx),
+			gitlab.WithKeysetPaginationParameters(resp.NextLink),
 		}
 	}
 

pkg/sources/gitlab/gitlab_integration_test.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/feature"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/credentialspb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/source_metadatapb"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/sourcespb"
@@ -813,3 +814,77 @@ func TestSource_Enumerate_ProjectDetailsInChunkMetadata(t *testing.T) {
 		t.Errorf("0 chunks scanned.")
 	}
 }
+
+// TestSource_Chunks_SimplifiedGitlabEnumeration enumerates GitLab projects
+// using a stored GitLab secret in GCP with the `UseSimplifiedGitlabEnumeration`
+// feature flag enabled. When enabled, the enumeration path is redirected to
+// `getAllProjectReposV2`, validating project listing via keyset pagination.
+func TestSource_Chunks_SimplifiedGitlabEnumeration(t *testing.T) {
+	// Preserve and restore the feature flag to avoid cross-test contamination
+	prev := feature.UseSimplifiedGitlabEnumeration.Load()
+	// enable the simplified gitlab enumeration flag
+	feature.UseSimplifiedGitlabEnumeration.Store(true)
+	defer feature.UseSimplifiedGitlabEnumeration.Store(prev)
+
+	// Create a bounded context for the entire test
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+
+	// Retrieve test secret containing the GitLab token
+	secret, err := common.GetTestSecret(ctx)
+	require.NoError(t, err, "failed to access test secret")
+
+	token := secret.MustGetField("GITLAB_TOKEN")
+
+	// Initialize the GitLab source with token-based authentication
+	s := Source{}
+	conn, err := anypb.New(&sourcespb.GitLab{
+		Credential: &sourcespb.GitLab_Token{
+			Token: token,
+		},
+	})
+	require.NoError(t, err)
+
+	err = s.Init(ctx, "enumerate gitlab projects with V2", 0, 0, false, conn, 10)
+	require.NoError(t, err, "failed during Source.Init")
+
+	// Enumerate GitLab projects
+	testReporter := sourcestest.TestReporter{}
+	err = s.Enumerate(ctx, &testReporter)
+	require.NoError(t, err, "enumeration should not fail")
+
+	// Ensure enumeration actually produced units
+	require.NotEmpty(t, testReporter.Units, "enumeration returned no units")
+
+	// Clear project cache to force project-detail lookups during chunking
+	clear(s.repoToProjCache.cache)
+
+	// Channel-based reporter to capture emitted chunks
+	chunksCh := make(chan *sources.Chunk, 1)
+	chanReporter := sources.ChanReporter{Ch: chunksCh}
+
+	// Chunk all enumerated units asynchronously
+	go func() {
+		defer close(chunksCh)
+		for _, unit := range testReporter.Units {
+			if err := s.ChunkUnit(ctx, unit, chanReporter); err != nil {
+				t.Errorf("Source.ChunkUnit() error = %v", err)
+			}
+		}
+	}()
+
+	// Validate produced chunks and their GitLab metadata
+	gotChunks := false
+	for chunk := range chunksCh {
+		gotChunks = true
+
+		meta, ok := chunk.SourceMetadata.Data.(*source_metadatapb.MetaData_Gitlab)
+		require.True(t, ok, "unexpected metadata type")
+
+		assert.NotZero(t, meta.Gitlab.ProjectId, "missing project ID in chunk metadata")
+		assert.NotEmpty(t, meta.Gitlab.ProjectName, "missing project name in chunk metadata")
+	}
+
+	// Ensure at least one chunk was produced
+	assert.True(t, gotChunks, "expected at least one chunk, got zero")
+}