feat: Add option to prevent hitting the rate limit

.generator/templates/client.mustache

@@ -28,6 +28,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 	"unicode/utf8"
 
@@ -56,6 +57,12 @@ const (
 	DpopAccessTokenPrivateKey = "DPOP_OKTA_ACCESS_TOKEN_PRIVATE_KEY"
 )
 
+type RateLimit struct {
+	Limit     int
+	Remaining int
+	Reset     int64
+}
+
 // APIClient manages communication with the {{appName}} API v{{version}}
 // In most cases there should be only one, shared, APIClient.
 type APIClient struct {
@@ -65,6 +72,9 @@ type APIClient struct {
 	tokenCache *goCache.Cache
 	freshcache bool
 
+	rateLimit     *RateLimit
+	rateLimitLock sync.Mutex
+
 	// API Services
 {{#apiInfo}}
 {{#apis}}
@@ -926,6 +936,26 @@ func (c *APIClient) RefreshNext() *APIClient {
 	return c
 }
 
+func parseRateLimit(resp *http.Response) (*RateLimit, error) {
+	limit, err := strconv.Atoi(resp.Header.Get("X-Rate-Limit-Limit"))
+	if err != nil {
+		return nil, err
+	}
+	remaining, err := strconv.Atoi(resp.Header.Get("X-Rate-Limit-Remaining"))
+	if err != nil {
+		return nil, err
+	}
+	reset, err := Get429BackoffTime(resp)
+	if err != nil {
+		return nil, err
+	}
+	return &RateLimit{
+		Limit:     limit,
+		Remaining: remaining,
+		Reset:     reset,
+	}, nil
+}
+
 func (c *APIClient) do(ctx context.Context, req *http.Request)(*http.Response, error){
 	cacheKey := CreateCacheKey(req)
 	if req.Method != http.MethodGet {
@@ -938,11 +968,36 @@ func (c *APIClient) do(ctx context.Context, req *http.Request)(*http.Response, e
 		c.freshcache = false
 	}
 	if !inCache {
+		if c.cfg.Okta.Client.RateLimit.Prevent {
+			c.rateLimitLock.Lock()
+			limit := c.rateLimit
+			c.rateLimitLock.Unlock()
+			if limit != nil && limit.Remaining <= 0 {
+				timer := time.NewTimer(time.Second * time.Duration(limit.Reset))
+				select {
+				case <-ctx.Done():
+					if !timer.Stop() {
+						<-timer.C
+					}
+					return nil, ctx.Err()
+				case <-timer.C:
+				}
+			}
+		}
+
 		resp, err := c.doWithRetries(ctx, req)
 		if err != nil {
 			return nil, err
 		}
 		if resp.StatusCode >= 200 && resp.StatusCode <= 299 && req.Method == http.MethodGet {
+			if c.cfg.Okta.Client.RateLimit.Prevent {
+				c.rateLimitLock.Lock()
+				newLimit, err := parseRateLimit(resp)
+				if err == nil {
+					c.rateLimit = newLimit
+				}
+				c.rateLimitLock.Unlock()
+			}
 			c.cache.Set(cacheKey, resp)
 		}
 		return resp, err

.generator/templates/configuration.mustache

@@ -144,6 +144,7 @@ type Configuration struct {
 			RateLimit         struct {
 				MaxRetries int32 `yaml:"maxRetries" envconfig:"OKTA_CLIENT_RATE_LIMIT_MAX_RETRIES"`
 				MaxBackoff int64 `yaml:"maxBackoff" envconfig:"OKTA_CLIENT_RATE_LIMIT_MAX_BACKOFF"`
+				Prevent bool `yaml:"prevent" envconfig:"OKTA_CLIENT_RATE_LIMIT_PREVENT"`
 			} `yaml:"rateLimit"`
 			OrgUrl            string   `yaml:"orgUrl" envconfig:"OKTA_CLIENT_ORGURL"`
 			Token             string   `yaml:"token" envconfig:"OKTA_CLIENT_TOKEN"`
@@ -542,6 +543,12 @@ func WithRateLimitMaxRetries(maxRetries int32) ConfigSetter {
 	}
 }
 
+func WithRateLimitPrevent(prevent bool) ConfigSetter {
+	return func(c *Configuration) {
+		c.Okta.Client.RateLimit.Prevent = prevent
+	}
+}
+
 func WithRateLimitMaxBackOff(maxBackoff int64) ConfigSetter {
 	return func(c *Configuration) {
 		c.Okta.Client.RateLimit.MaxBackoff = maxBackoff

okta/retry_logic_test.go

@@ -35,6 +35,42 @@ func Test_429_Will_Automatically_Retry(t *testing.T) {
 	require.Equal(t, 2, info["GET /api/v1/users"], "did not make exactly 2 call to /api/v1/users")
 }
 
+func Test_RateLimitPrevent_Blocks(t *testing.T) {
+	httpmock.Activate()
+	defer httpmock.DeactivateAndReset()
+	configuration, err := NewConfiguration(WithRateLimitPrevent(true))
+	require.NoError(t, err, "Creating a new config should not error")
+	configuration.Debug = true
+	proxyClient := NewAPIClient(configuration)
+	resetTime := time.Now().Add(time.Second * 5)
+	count := 0
+	httpmock.RegisterResponder("GET", "/api/v1/users",
+		func(*http.Request) (*http.Response, error) {
+			if count > 2 {
+				return nil, fmt.Errorf("should not have made more than 2 calls")
+			}
+			count++
+			return MockRateLimitedResponse(count, resetTime), nil
+		},
+	)
+	// First request withing the limit
+	_, resp, err := proxyClient.UserAPI.ListUsers(apiClient.cfg.Context).Execute()
+	require.Nil(t, err, "Error should have been nil")
+	require.NotNil(t, resp, "Response was nil")
+
+	// Second request should block until the reset time
+	_, resp, err = proxyClient.UserAPI.ListUsers(apiClient.cfg.Context).Execute()
+	require.Nil(t, err, "Error should have been nil")
+	require.NotNil(t, resp, "Response was nil")
+
+	secondRequestDate, err := time.Parse("Mon, 02 Jan 2006 15:04:05 GMT", resp.Header.Get("Date"))
+	require.NoError(t, err, "should have been able to parse the date")
+	require.True(t, secondRequestDate.After(resetTime), "second request should have been after the reset time")
+
+	info := httpmock.GetCallCountInfo()
+	require.Equal(t, 2, info["GET /api/v1/users"], "did not make exactly 2 call to /api/v1/users")
+}
+
 func MockResponse(responses ...*http.Response) httpmock.Responder {
 	return func(req *http.Request) (*http.Response, error) {
 		httpmock.GetTotalCallCount()
@@ -82,3 +118,23 @@ func MockValidResponse() *http.Response {
 		ContentLength: -1,
 	}
 }
+
+func MockRateLimitedResponse(requestId int, reset time.Time) *http.Response {
+	loc, _ := time.LoadLocation("UTC")
+	zulu := time.Now().In(loc)
+	header := http.Header{}
+	header.Add("X-Rate-Limit-Limit", strconv.Itoa(50))
+	header.Add("X-Rate-Limit-Remaining", strconv.Itoa(0))
+	header.Add("X-Rate-Limit-Reset", strconv.FormatInt(reset.Unix(), 10))
+	header.Add("X-Okta-Request-id", strconv.Itoa(requestId))
+	header.Add("Content-Type", "application/json")
+	header.Add("Date", zulu.Format("Mon, 02 Jan 2006 15:04:05 GMT"))
+
+	return &http.Response{
+		Status:        strconv.Itoa(200),
+		StatusCode:    200,
+		Body:          httpmock.NewRespBodyFromString("[]"),
+		Header:        header,
+		ContentLength: -1,
+	}
+}