Skip to content

Commit 6649880

Browse files
cx-anurag-dalkecx-anurag-dalke
authored
Merge pull request #1320 from Checkmarx/other/AST-114776
Fix GitHub API rate limit exceeds for contributor count command(AST-114776)
2 parents 3bd04e3 + b838e49 commit 6649880

File tree

7 files changed

+279
-6
lines changed

7 files changed

+279
-6
lines changed

‎internal/wrappers/azure-http.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,12 @@ func (g *AzureHTTPWrapper) get(
108108
queryParams map[string]string,
109109
authFormat string,
110110
) (bool, error) {
111-
resp, err := GetWithQueryParams(g.client, url, token, authFormat, queryParams)
111+
resp, err := WithSCMRateLimitRetry(
112+
AzureRateLimitConfig,
113+
func() (*http.Response, error) {
114+
return GetWithQueryParams(g.client, url, token, authFormat, queryParams)
115+
},
116+
)
112117
if err != nil {
113118
return false, err
114119
}

‎internal/wrappers/bitbucket-http.go‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,12 @@ func (g *BitBucketHTTPWrapper) getFromBitBucket(
150150

151151
logger.PrintIfVerbose(fmt.Sprintf("Request to %s", url))
152152

153-
resp, err := GetWithQueryParams(g.client, url, token, basicFormat, queryParams)
153+
resp, err := WithSCMRateLimitRetry(
154+
BitbucketRateLimitConfig,
155+
func() (*http.Response, error) {
156+
return GetWithQueryParams(g.client, url, token, basicFormat, queryParams)
157+
},
158+
)
154159
if err != nil {
155160
return err
156161
}
@@ -264,7 +269,12 @@ func collectPageBitBucket(
264269
}
265270

266271
func getBitBucket(client *http.Client, token, url string, target interface{}, queryParams map[string]string) error {
267-
resp, err := GetWithQueryParams(client, url, token, basicFormat, queryParams)
272+
resp, err := WithSCMRateLimitRetry(
273+
BitbucketRateLimitConfig,
274+
func() (*http.Response, error) {
275+
return GetWithQueryParams(client, url, token, basicFormat, queryParams)
276+
},
277+
)
268278
if err != nil {
269279
return err
270280
}

‎internal/wrappers/bitbucketserver/bitbucket-server-http.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,12 @@ func getBitBucketServer(
162162
}
163163

164164
req.URL.RawQuery = q.Encode()
165-
resp, err := client.Do(req)
165+
resp, err := wrappers.WithSCMRateLimitRetry(
166+
wrappers.BitbucketRateLimitConfig,
167+
func() (*http.Response, error) {
168+
return client.Do(req)
169+
},
170+
)
166171
if err != nil {
167172
return err
168173
}

‎internal/wrappers/github-http.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -244,7 +244,12 @@ func get(client *http.Client, url string, target interface{}, queryParams map[st
244244
req.Header.Add(acceptHeader, apiVersion)
245245
token := viper.GetString(params.SCMTokenFlag)
246246
logger.PrintRequest(req)
247-
resp, err := GetWithQueryParamsAndCustomRequest(client, req, url, token, tokenFormat, queryParams)
247+
resp, err := WithSCMRateLimitRetry(
248+
GitHubRateLimitConfig,
249+
func() (*http.Response, error) {
250+
return GetWithQueryParamsAndCustomRequest(client, req, url, token, tokenFormat, queryParams)
251+
},
252+
)
248253
if err != nil {
249254
return nil, err
250255
}

‎internal/wrappers/gitlab-http.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,12 @@ func getFromGitLab(
137137

138138
logger.PrintRequest(req)
139139

140-
resp, err := GetWithQueryParamsAndCustomRequest(client, req, requestURL, token, bearerFormat, queryParams)
140+
resp, err := WithSCMRateLimitRetry(
141+
GitLabRateLimitConfig,
142+
func() (*http.Response, error) {
143+
return GetWithQueryParamsAndCustomRequest(client, req, requestURL, token, bearerFormat, queryParams)
144+
},
145+
)
141146
if err != nil {
142147
return nil, err
143148
}

‎internal/wrappers/rate-limit.go‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
package wrappers
2+
3+
import (
4+
"log"
5+
"net/http"
6+
"strconv"
7+
"strings"
8+
"time"
9+
10+
"github.com/pkg/errors"
11+
)
12+
13+
const defaultRateLimitWaitSeconds = 60
14+
15+
// SCMRateLimitConfig holds rate limit configuration for different SCM providers
16+
type SCMRateLimitConfig struct {
17+
Provider string
18+
ResetHeaderName string
19+
RemainingHeaderName string
20+
LimitHeaderName string
21+
RateLimitStatusCodes []int
22+
DefaultWaitTime time.Duration
23+
}
24+
25+
// Common SCM rate limit configurations
26+
var (
27+
GitHubRateLimitConfig = &SCMRateLimitConfig{
28+
Provider: "GitHub",
29+
ResetHeaderName: "X-RateLimit-Reset",
30+
RemainingHeaderName: "X-RateLimit-Remaining",
31+
LimitHeaderName: "X-RateLimit-Limit",
32+
RateLimitStatusCodes: []int{403, 429},
33+
DefaultWaitTime: defaultRateLimitWaitSeconds * time.Second,
34+
}
35+
36+
GitLabRateLimitConfig = &SCMRateLimitConfig{
37+
Provider: "GitLab",
38+
ResetHeaderName: "RateLimit-Reset",
39+
RemainingHeaderName: "RateLimit-Remaining",
40+
LimitHeaderName: "RateLimit-Limit",
41+
RateLimitStatusCodes: []int{429},
42+
DefaultWaitTime: defaultRateLimitWaitSeconds * time.Second,
43+
}
44+
45+
BitbucketRateLimitConfig = &SCMRateLimitConfig{
46+
Provider: "Bitbucket",
47+
ResetHeaderName: "X-RateLimit-Reset",
48+
RemainingHeaderName: "X-RateLimit-Remaining",
49+
LimitHeaderName: "X-RateLimit-Limit",
50+
RateLimitStatusCodes: []int{429},
51+
DefaultWaitTime: defaultRateLimitWaitSeconds * time.Second,
52+
}
53+
54+
AzureRateLimitConfig = &SCMRateLimitConfig{
55+
Provider: "Azure",
56+
ResetHeaderName: "X-Ratelimit-Reset",
57+
RemainingHeaderName: "X-Ratelimit-Remaining",
58+
LimitHeaderName: "X-Ratelimit-Limit",
59+
RateLimitStatusCodes: []int{429},
60+
DefaultWaitTime: defaultRateLimitWaitSeconds * time.Second,
61+
}
62+
)
63+
64+
// SCMRateLimitError represents a rate limit error from any SCM provider
65+
type SCMRateLimitError struct {
66+
Provider string
67+
ResetTime int64
68+
Message string
69+
}
70+
71+
func (e *SCMRateLimitError) Error() string {
72+
if e.Message != "" {
73+
return e.Message
74+
}
75+
return e.Provider + " API rate limit exceeded"
76+
}
77+
78+
func (e *SCMRateLimitError) RetryAfter() time.Duration {
79+
if e.ResetTime > 0 {
80+
reset := time.Unix(e.ResetTime, 0)
81+
now := time.Now()
82+
if reset.After(now) {
83+
return reset.Sub(now) + (defaultRateLimitWaitSeconds * time.Second) // add buffer for 60 seconds
84+
}
85+
}
86+
return defaultRateLimitWaitSeconds * time.Second
87+
}
88+
89+
// WithSCMRateLimitRetry wraps any SCM API call with rate limit retry logic
90+
func WithSCMRateLimitRetry(config *SCMRateLimitConfig, apiCall func() (*http.Response, error)) (*http.Response, error) {
91+
maxRetries := 3
92+
retryCount := 0
93+
94+
for {
95+
resp, err := apiCall()
96+
if err != nil {
97+
return nil, err
98+
}
99+
100+
// Check if it's a rate limit error
101+
if isRateLimitStatusCode(resp.StatusCode, config) {
102+
rateLimitErr := ParseRateLimitHeaders(resp.Header, config)
103+
wait := config.DefaultWaitTime
104+
if rateLimitErr != nil {
105+
wait = rateLimitErr.RetryAfter()
106+
}
107+
if retryCount >= maxRetries {
108+
return nil, errors.Errorf("%s API rate limit exceeded after %d retries", config.Provider, maxRetries)
109+
}
110+
log.Printf("%s API rate limit exceeded (status %d). Waiting %v until %v before retrying... (attempt %d/%d)",
111+
config.Provider, resp.StatusCode, wait, time.Now().Add(wait), retryCount+1, maxRetries)
112+
time.Sleep(wait)
113+
// Reset Authorization header before retry
114+
if resp.Request != nil {
115+
resetAuthorizationHeader(resp.Request)
116+
}
117+
retryCount++
118+
continue
119+
}
120+
return resp, err
121+
}
122+
}
123+
124+
// ParseRateLimitHeaders extracts rate limit information from HTTP response headers
125+
func ParseRateLimitHeaders(headers map[string][]string, config *SCMRateLimitConfig) *SCMRateLimitError {
126+
resetHeader := getHeaderValue(headers, config.ResetHeaderName)
127+
if resetHeader == "" {
128+
return nil
129+
}
130+
131+
resetTime, err := strconv.ParseInt(resetHeader, 10, 64)
132+
if err != nil {
133+
return nil
134+
}
135+
136+
return &SCMRateLimitError{
137+
Provider: config.Provider,
138+
ResetTime: resetTime,
139+
}
140+
}
141+
142+
// getHeaderValue retrieves a header value in a case-insensitive manner
143+
func getHeaderValue(headers map[string][]string, headerName string) string {
144+
for name, values := range headers {
145+
if strings.EqualFold(name, headerName) && len(values) > 0 {
146+
return values[0]
147+
}
148+
}
149+
return ""
150+
}
151+
152+
// isRateLimitStatusCode checks if the status code indicates a rate limit error
153+
func isRateLimitStatusCode(statusCode int, config *SCMRateLimitConfig) bool {
154+
for _, code := range config.RateLimitStatusCodes {
155+
if statusCode == code {
156+
return true
157+
}
158+
}
159+
return false
160+
}
161+
162+
// resetAuthorizationHeader removes the Authorization header from the request
163+
func resetAuthorizationHeader(req *http.Request) {
164+
req.Header.Del("Authorization")
165+
}

‎test/integration/rate-limit_test.go‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
package integration
2+
3+
import (
4+
"io"
5+
"net/http"
6+
"net/http/httptest"
7+
"strconv"
8+
"strings"
9+
"testing"
10+
"time"
11+
12+
"github.com/checkmarx/ast-cli/internal/wrappers"
13+
"github.com/stretchr/testify/assert"
14+
)
15+
16+
func mockAPI(repeatCode, repeatCount int, headerName, headerValue string) func() (*http.Response, error) {
17+
attempt := 0
18+
return func() (*http.Response, error) {
19+
rec := httptest.NewRecorder()
20+
if attempt < repeatCount {
21+
rec.Code = repeatCode
22+
if headerName != "" {
23+
rec.Header().Set(headerName, headerValue)
24+
}
25+
} else {
26+
rec.Code = http.StatusOK
27+
}
28+
attempt++
29+
resp := rec.Result()
30+
resp.Body = io.NopCloser(strings.NewReader(""))
31+
return resp, nil
32+
}
33+
}
34+
35+
func runRateLimitTest(t *testing.T, config *wrappers.SCMRateLimitConfig, repeatCode, repeatCount int, headerName string) {
36+
reset := strconv.FormatInt(time.Now().Unix(), 10) // simulate immediate retry
37+
38+
//nolint:bodyclose // safe in test, body closed later
39+
api := mockAPI(repeatCode, repeatCount, headerName, reset)
40+
41+
start := time.Now()
42+
resp, err := wrappers.WithSCMRateLimitRetry(config, api)
43+
if resp != nil {
44+
defer resp.Body.Close()
45+
}
46+
47+
assert := assert.New(t)
48+
assert.NoError(err)
49+
assert.NotNil(resp)
50+
assert.Equal(http.StatusOK, resp.StatusCode)
51+
52+
elapsed := time.Since(start)
53+
assert.GreaterOrEqual(elapsed, config.DefaultWaitTime)
54+
}
55+
56+
func TestGitHubRateLimit_SuccessAfterRetryOne(t *testing.T) {
57+
runRateLimitTest(t, wrappers.GitHubRateLimitConfig, 429, 1, "X-RateLimit-Reset")
58+
}
59+
60+
func TestGitHubRateLimit_SuccessAfterRetryTwo(t *testing.T) {
61+
runRateLimitTest(t, wrappers.GitHubRateLimitConfig, 429, 2, "X-RateLimit-Reset")
62+
}
63+
64+
func TestGitHubRateLimit_SuccessAfterRetryThree(t *testing.T) {
65+
runRateLimitTest(t, wrappers.GitHubRateLimitConfig, 403, 3, "X-RateLimit-Reset")
66+
}
67+
68+
func TestGitLabRateLimit_SuccessAfterRetryOne(t *testing.T) {
69+
runRateLimitTest(t, wrappers.GitLabRateLimitConfig, 429, 1, "RateLimit-Reset")
70+
}
71+
72+
func TestBitBucketRateLimit_SuccessAfterRetryOne(t *testing.T) {
73+
runRateLimitTest(t, wrappers.BitbucketRateLimitConfig, 429, 1, "X-RateLimit-Reset")
74+
}
75+
76+
func TestAzureRateLimit_SuccessAfterRetryOne(t *testing.T) {
77+
runRateLimitTest(t, wrappers.AzureRateLimitConfig, 429, 1, "X-Ratelimit-Reset")
78+
}

0 commit comments

Comments
 (0)