From 10ff63008c2709c81ffc62211514bae0989435d8 Mon Sep 17 00:00:00 2001
From: erezrokah <erezrokah@users.noreply.github.com>
Date: Thu, 18 Apr 2024 19:12:05 +0100
Subject: [PATCH] feat: Update template files

---
 .generator/templates/client.mustache        | 55 +++++++++++++++++++++
 .generator/templates/configuration.mustache |  7 +++
 2 files changed, 62 insertions(+)

diff --git a/.generator/templates/client.mustache b/.generator/templates/client.mustache
index 2c8a2b71..4f96810d 100644
--- a/.generator/templates/client.mustache
+++ b/.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
diff --git a/.generator/templates/configuration.mustache b/.generator/templates/configuration.mustache
index 278432db..5ac1a0b6 100644
--- a/.generator/templates/configuration.mustache
+++ b/.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