Merge branch 'main' into INS-127

pkg/sources/postman/postman.go

@@ -186,6 +186,9 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk, _ .
 	for _, workspaceID := range s.conn.Workspaces {
 		w, err := s.client.GetWorkspace(ctx, workspaceID)
 		if err != nil {
+			if errors.Is(err, errAbortScanDueToAPIRateLimit) {
+				return err
+			}
 			// Log and move on, because sometimes the Postman API seems to give us workspace IDs
 			// that we don't have access to, so we don't want to kill the scan because of it.
 			ctx.Logger().Error(err, "error getting workspace %s", workspaceID)
@@ -206,6 +209,9 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk, _ .
 
 		collection, err := s.client.GetCollection(ctx, collectionID)
 		if err != nil {
+			if errors.Is(err, errAbortScanDueToAPIRateLimit) {
+				return err
+			}
 			// Log and move on, because sometimes the Postman API seems to give us collection IDs
 			// that we don't have access to, so we don't want to kill the scan because of it.
 			ctx.Logger().Error(err, "error getting collection %s", collectionID)
@@ -279,6 +285,9 @@ func (s *Source) scanWorkspace(ctx context.Context, chunksChan chan *sources.Chu
 	for _, envID := range workspace.Environments {
 		envVars, err := s.client.GetEnvironmentVariables(ctx, envID.Uid)
 		if err != nil {
+			if errors.Is(err, errAbortScanDueToAPIRateLimit) {
+				return err
+			}
 			ctx.Logger().Error(err, "could not get env variables", "environment_uuid", envID.Uid)
 			continue
 		}
@@ -316,6 +325,9 @@ func (s *Source) scanWorkspace(ctx context.Context, chunksChan chan *sources.Chu
 		}
 		collection, err := s.client.GetCollection(ctx, collectionID.Uid)
 		if err != nil {
+			if errors.Is(err, errAbortScanDueToAPIRateLimit) {
+				return err
+			}
 			// Log and move on, because sometimes the Postman API seems to give us collection IDs
 			// that we don't have access to, so we don't want to kill the scan because of it.
 			ctx.Logger().Error(err, "error getting collection %s", collectionID)

pkg/sources/postman/postman_client.go

@@ -23,8 +23,11 @@ const (
 	COLLECTIONS_URL               = "https://api.getpostman.com/collections/%s"
 	userAgent                     = "PostmanRuntime/7.26.8"
 	defaultContentType            = "*"
+	abortScanAPIReqLimitThreshold = 0.8 // Abort scan if 80% of monthly api request limit is reached
 )
 
+var errAbortScanDueToAPIRateLimit = fmt.Errorf("aborting scan due to Postman API monthly requests limit being used over %f%%", abortScanAPIReqLimitThreshold*100)
+
 type Workspace struct {
 	Id              string      `json:"id"`
 	Name            string      `json:"name"`
@@ -284,22 +287,9 @@ func (c *Client) getPostmanResponseBodyBytes(ctx trContext.Context, urlString st
 
 	c.Metrics.apiRequests.WithLabelValues(urlString).Inc()
 
-	rateLimitRemainingMonthValue := resp.Header.Get("RateLimit-Remaining-Month")
+	err = c.handleRateLimits(ctx, resp)
-	if rateLimitRemainingMonthValue == "" {
-		rateLimitRemainingMonthValue = resp.Header.Get("X-RateLimit-Remaining-Month")
-	}
-
-	if rateLimitRemainingMonthValue != "" {
-		rateLimitRemainingMonth, err := strconv.Atoi(rateLimitRemainingMonthValue)
 	if err != nil {
-			ctx.Logger().Error(err, "Couldn't convert RateLimit-Remaining-Month to an int",
+		return nil, err
-				"header_value", rateLimitRemainingMonthValue,
-			)
-		} else {
-			c.Metrics.apiMonthlyRequestsRemaining.WithLabelValues().Set(
-				float64(rateLimitRemainingMonth),
-			)
-		}
 	}
 
 	body, err := io.ReadAll(resp.Body)
@@ -315,6 +305,60 @@ func (c *Client) getPostmanResponseBodyBytes(ctx trContext.Context, urlString st
 	return body, nil
 }
 
+// handleRateLimits processes the rate limit headers from the Postman API response
+// and updates the client's metrics accordingly. If the monthly rate limit usage exceeds
+// the set threshold, an error is returned to abort further processing.
+func (c *Client) handleRateLimits(ctx trContext.Context, resp *http.Response) error {
+	rateLimitRemainingMonthValue := resp.Header.Get("RateLimit-Remaining-Month")
+	if rateLimitRemainingMonthValue == "" {
+		rateLimitRemainingMonthValue = resp.Header.Get("X-RateLimit-Remaining-Month")
+	}
+	if rateLimitRemainingMonthValue == "" {
+		ctx.Logger().V(2).Info("RateLimit-Remaining-Month header not found in response")
+		return nil
+	}
+
+	rateLimitRemainingMonth, err := strconv.Atoi(rateLimitRemainingMonthValue)
+	if err != nil {
+		ctx.Logger().Error(err, "Couldn't convert RateLimit-Remaining-Month to an int",
+			"header_value", rateLimitRemainingMonthValue,
+		)
+		return nil
+	}
+	c.Metrics.apiMonthlyRequestsRemaining.WithLabelValues().Set(
+		float64(rateLimitRemainingMonth),
+	)
+
+	rateLimitTotalMonthValue := resp.Header.Get("RateLimit-Limit-Month")
+	if rateLimitTotalMonthValue == "" {
+		rateLimitTotalMonthValue = resp.Header.Get("X-RateLimit-Limit-Month")
+	}
+	if rateLimitTotalMonthValue == "" {
+		ctx.Logger().V(2).Info("RateLimit-Limit-Month header not found in response")
+		return nil
+	}
+	rateLimitTotalMonth, err := strconv.Atoi(rateLimitTotalMonthValue)
+	if err != nil {
+		ctx.Logger().Error(err, "Couldn't convert RateLimit-Limit-Month to an int",
+			"header_value", rateLimitTotalMonthValue,
+		)
+		return nil
+	}
+
+	if rateLimitTotalMonth == 0 {
+		ctx.Logger().V(2).Info("RateLimit-Limit-Month is zero, cannot compute usage percentage")
+		return nil
+	}
+
+	// failsafe to abandon scan if we are over the threshold of monthly API requests used
+	percentageUsed := float64(rateLimitTotalMonth-rateLimitRemainingMonth) / float64(rateLimitTotalMonth)
+	if percentageUsed > abortScanAPIReqLimitThreshold {
+		return errAbortScanDueToAPIRateLimit
+	}
+
+	return nil
+}
+
 // EnumerateWorkspaces returns the workspaces for a given user (both private, public, team and personal).
 // Consider adding additional flags to support filtering.
 func (c *Client) EnumerateWorkspaces(ctx trContext.Context) ([]Workspace, error) {

pkg/sources/postman/postman_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/go-errors/errors"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
@@ -734,3 +735,67 @@ func TestSource_ResponseID(t *testing.T) {
 	}
 
 }
+
+func TestSource_AbortScanRateLimitCrossedThreshold(t *testing.T) {
+	defer gock.Off()
+	// Mock the API response for getting collection
+	gock.New("https://api.getpostman.com").
+		Get("/collections/1234-abc1").
+		Reply(200).
+		AddHeader("X-RateLimit-Remaining-Month", "19").
+		AddHeader("X-RateLimit-Limit-Month", "100").
+		BodyString(`{"collection":{"info":{"_postman_id":"abc1","name":"test-collection-1","schema":"https://schema.postman.com/json/collection/v2.1.0/collection.json",
+		"updatedAt":"2025-03-21T17:39:25.000Z","createdAt":"2025-03-21T17:37:13.000Z","lastUpdatedBy":"1234","uid":"1234-abc1"},
+		"item":[]}}`)
+
+	ctx := context.Background()
+	s, conn := createTestSource(&sourcespb.Postman{
+		Credential: &sourcespb.Postman_Token{
+			Token: "super-secret-token",
+		},
+	})
+
+	err := s.Init(ctx, "test - postman", 0, 1, false, conn, 1)
+	if err != nil {
+		t.Fatalf("init error: %v", err)
+	}
+	gock.InterceptClient(s.client.HTTPClient)
+	defer gock.RestoreClient(s.client.HTTPClient)
+
+	_, err = s.client.GetCollection(ctx, "1234-abc1")
+	if !errors.Is(err, errAbortScanDueToAPIRateLimit) {
+		t.Errorf("expected abort scan error due to rate limit, got: %v", err)
+	}
+}
+
+func TestSource_AbortScanRateLimitBelowThreshold(t *testing.T) {
+	defer gock.Off()
+	// Mock the API response for getting collection
+	gock.New("https://api.getpostman.com").
+		Get("/collections/1234-abc1").
+		Reply(200).
+		AddHeader("X-RateLimit-Remaining-Month", "90").
+		AddHeader("X-RateLimit-Limit-Month", "100").
+		BodyString(`{"collection":{"info":{"_postman_id":"abc1","name":"test-collection-1","schema":"https://schema.postman.com/json/collection/v2.1.0/collection.json",
+		"updatedAt":"2025-03-21T17:39:25.000Z","createdAt":"2025-03-21T17:37:13.000Z","lastUpdatedBy":"1234","uid":"1234-abc1"},
+		"item":[]}}`)
+
+	ctx := context.Background()
+	s, conn := createTestSource(&sourcespb.Postman{
+		Credential: &sourcespb.Postman_Token{
+			Token: "super-secret-token",
+		},
+	})
+
+	err := s.Init(ctx, "test - postman", 0, 1, false, conn, 1)
+	if err != nil {
+		t.Fatalf("init error: %v", err)
+	}
+	gock.InterceptClient(s.client.HTTPClient)
+	defer gock.RestoreClient(s.client.HTTPClient)
+
+	_, err = s.client.GetCollection(ctx, "1234-abc1")
+	if err != nil {
+		t.Errorf("expected no error, got: %v", err)
+	}
+}