Skip to content

Commit

Permalink
feat: Update template files
Browse files Browse the repository at this point in the history
  • Loading branch information
erezrokah committed May 23, 2024
1 parent 27ec222 commit 10ff630
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 0 deletions.
55 changes: 55 additions & 0 deletions .generator/templates/client.mustache
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ import (
"regexp"
"strconv"
"strings"
"sync"
"time"
"unicode/utf8"

Expand Down Expand Up @@ -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 {
Expand All @@ -65,6 +72,9 @@ type APIClient struct {
tokenCache *goCache.Cache
freshcache bool
rateLimit *RateLimit
rateLimitLock sync.Mutex
// API Services
{{#apiInfo}}
{{#apis}}
Expand Down Expand Up @@ -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 {
Expand All @@ -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
Expand Down
7 changes: 7 additions & 0 deletions .generator/templates/configuration.mustache
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 10ff630

Please sign in to comment.