jwt transport: fix retry on unauthorized from CAPI(#3006)

pkg/apiclient/auth_jwt.go

@@ -26,6 +26,7 @@ type JWTTransport struct {
  URL *url.URL
  VersionPrefix string
  UserAgent string
+ RetryConfig *RetryConfig
  // Transport is the underlying HTTP transport to use when making requests.
  // It will default to http.DefaultTransport if nil.
  Transport http.RoundTripper
@@ -165,36 +166,67 @@ func (t *JWTTransport) prepareRequest(req *http.Request) (*http.Request, error)
 
 // RoundTrip implements the RoundTripper interface.
 func (t *JWTTransport) RoundTrip(req *http.Request) (*http.Response, error) {
- req, err := t.prepareRequest(req)
- if err != nil {
- return nil, err
- }
 
- if log.GetLevel() >= log.TraceLevel {
- // requestToDump := cloneRequest(req)
- dump, _ := httputil.DumpRequest(req, true)
- log.Tracef("req-jwt: %s", string(dump))
- }
+ var resp *http.Response
+ attemptsCount := make(map[int]int)
 
- // Make the HTTP request.
- resp, err := t.transport().RoundTrip(req)
- if log.GetLevel() >= log.TraceLevel {
- dump, _ := httputil.DumpResponse(resp, true)
- log.Tracef("resp-jwt: %s (err:%v)", string(dump), err)
- }
+ for {
+ if log.GetLevel() >= log.TraceLevel {
+ // requestToDump := cloneRequest(req)
+ dump, _ := httputil.DumpRequest(req, true)
+ log.Tracef("req-jwt: %s", string(dump))
+ }
+ // Make the HTTP request.
+ clonedReq := cloneRequest(req)
 
- if err != nil {
- // we had an error (network error for example, or 401 because token is refused), reset the token?
- t.ResetToken()
+ clonedReq, err := t.prepareRequest(clonedReq)
+ if err != nil {
+ return nil, err
+ }
 
- return resp, fmt.Errorf("performing jwt auth: %w", err)
- }
+ resp, err = t.transport().RoundTrip(clonedReq)
+ if log.GetLevel() >= log.TraceLevel {
+ dump, _ := httputil.DumpResponse(resp, true)
+ log.Tracef("resp-jwt: %s (err:%v)", string(dump), err)
+ }
 
- if resp != nil {
- log.Debugf("resp-jwt: %d", resp.StatusCode)
- }
+ if err != nil {
+ // we had an error (network error for example), reset the token?
+ t.ResetToken()
+ return resp, fmt.Errorf("performing jwt auth: %w", err)
+ }
+
+ if resp != nil {
+ log.Debugf("resp-jwt: %d", resp.StatusCode)
+ }
 
+ config, shouldRetry := t.RetryConfig.StatusCodeConfig[resp.StatusCode]
+ if !shouldRetry {
+ break
+ }
+
+ if attemptsCount[resp.StatusCode] >= config.MaxAttempts {
+ log.Infof("max attempts reached for status code %d", resp.StatusCode)
+ break
+ }
+
+ if config.InvalidateToken {
+ log.Debugf("invalidating token for status code %d", resp.StatusCode)
+ t.ResetToken()
+ }
+
+ log.Debugf("retrying request to %s", req.URL.String())
+ attemptsCount[resp.StatusCode]++
+ log.Infof("attempt %d out of %d", attemptsCount[resp.StatusCode], config.MaxAttempts)
+
+ if config.Backoff {
+ backoff := 2*attemptsCount[resp.StatusCode] + 5
+ log.Infof("retrying in %d seconds (attempt %d of %d)", backoff, attemptsCount[resp.StatusCode], config.MaxAttempts)
+ time.Sleep(time.Duration(backoff) * time.Second)
+ }
+ }
  return resp, nil
+
 }
 
 func (t *JWTTransport) Client() *http.Client {
@@ -211,27 +243,8 @@ func (t *JWTTransport) ResetToken() {
 // transport() returns a round tripper that retries once when the status is unauthorized,
 // and 5 times when the infrastructure is overloaded.
 func (t *JWTTransport) transport() http.RoundTripper {
- transport := t.Transport
- if transport == nil {
- transport = http.DefaultTransport
- }
-
- return &retryRoundTripper{
- next: &retryRoundTripper{
- next: transport,
- maxAttempts: 5,
- withBackOff: true,
- retryStatusCodes: []int{http.StatusTooManyRequests, http.StatusServiceUnavailable, http.StatusGatewayTimeout},
- },
- maxAttempts: 2,
- withBackOff: false,
- retryStatusCodes: []int{http.StatusUnauthorized, http.StatusForbidden},
- onBeforeRequest: func(attempt int) {
- // reset the token only in the second attempt as this is when we know we had a 401 or 403
- // the second attempt is supposed to refresh the token
- if attempt > 0 {
- t.ResetToken()
- }
- },
+ if t.Transport != nil {
+ return t.Transport
  }
+ return http.DefaultTransport
 }

pkg/apiclient/client.go

@@ -72,6 +72,13 @@ func NewClient(config *Config) (*ApiClient, error) {
  UserAgent: config.UserAgent,
  VersionPrefix: config.VersionPrefix,
  UpdateScenario: config.UpdateScenario,
+ RetryConfig: NewRetryConfig(
+ WithStatusCodeConfig(http.StatusUnauthorized, 2, false, true),
+ WithStatusCodeConfig(http.StatusForbidden, 2, false, true),
+ WithStatusCodeConfig(http.StatusTooManyRequests, 5, true, false),
+ WithStatusCodeConfig(http.StatusServiceUnavailable, 5, true, false),
+ WithStatusCodeConfig(http.StatusGatewayTimeout, 5, true, false),
+ ),
  }
 
  transport, baseURL := createTransport(config.URL)

pkg/apiclient/retry_config.go

@@ -0,0 +1,33 @@
+package apiclient
+
+type StatusCodeConfig struct {
+ MaxAttempts int
+ Backoff bool
+ InvalidateToken bool
+}
+
+type RetryConfig struct {
+ StatusCodeConfig map[int]StatusCodeConfig
+}
+
+type RetryConfigOption func(*RetryConfig)
+
+func NewRetryConfig(options ...RetryConfigOption) *RetryConfig {
+ rc := &RetryConfig{
+ StatusCodeConfig: make(map[int]StatusCodeConfig),
+ }
+ for _, opt := range options {
+ opt(rc)
+ }
+ return rc
+}
+
+func WithStatusCodeConfig(statusCode int, maxAttempts int, backOff bool, invalidateToken bool) RetryConfigOption {
+ return func(rc *RetryConfig) {
+ rc.StatusCodeConfig[statusCode] = StatusCodeConfig{
+ MaxAttempts: maxAttempts,
+ Backoff: backOff,
+ InvalidateToken: invalidateToken,
+ }
+ }
+}