From 9c2564a94c2c0fd48386d323db8823f9b5f32f14 Mon Sep 17 00:00:00 2001
From: Sergiu Ghitea <28300158+sergiught@users.noreply.github.com>
Date: Thu, 24 Aug 2023 10:10:37 +0300
Subject: [PATCH] Use a custom client with retries for auth0 api client config

---
 go.mod                    |   2 +-
 internal/config/config.go | 101 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 101 insertions(+), 2 deletions(-)

diff --git a/go.mod b/go.mod
index 97413c45f..2ac88abb3 100644
--- a/go.mod
+++ b/go.mod
@@ -3,6 +3,7 @@ module github.com/auth0/terraform-provider-auth0
 go 1.20
 
 require (
+	github.com/PuerkitoBio/rehttp v1.2.0
 	github.com/auth0/go-auth0 v1.0.0
 	github.com/google/go-cmp v0.5.9
 	github.com/hashicorp/go-cty v1.4.1-0.20200414143053-d3edf31b6320
@@ -19,7 +20,6 @@ require (
 	github.com/Masterminds/semver/v3 v3.1.1 // indirect
 	github.com/Masterminds/sprig/v3 v3.2.2 // indirect
 	github.com/ProtonMail/go-crypto v0.0.0-20230217124315-7d5c6f04bbb8 // indirect
-	github.com/PuerkitoBio/rehttp v1.2.0 // indirect
 	github.com/agext/levenshtein v1.2.3 // indirect
 	github.com/apparentlymart/go-textseg/v13 v13.0.0 // indirect
 	github.com/armon/go-radix v1.0.0 // indirect
diff --git a/internal/config/config.go b/internal/config/config.go
index 6b2a80643..6edeef712 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -2,9 +2,15 @@ package config
 
 import (
 	"context"
+	"crypto/x509"
 	"fmt"
 	"net/http"
+	"net/url"
+	"regexp"
+	"strconv"
+	"time"
 
+	"github.com/PuerkitoBio/rehttp"
 	"github.com/auth0/go-auth0"
 	"github.com/auth0/go-auth0/management"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
@@ -60,7 +66,8 @@ func ConfigureProvider(terraformVersion *string) schema.ConfigureContextFunc {
 			management.WithDebug(debug),
 			management.WithUserAgent(userAgent(terraformVersion)),
 			management.WithAuth0ClientEnvEntry(providerName, version),
-			management.WithRetries(3, []int{http.StatusTooManyRequests, http.StatusInternalServerError}),
+			management.WithNoRetries(),
+			management.WithClient(customClientWithRetries()),
 		)
 		if err != nil {
 			return nil, diag.FromErr(err)
@@ -106,3 +113,95 @@ func authenticationOption(clientID, clientSecret, apiToken, audience string) man
 
 	return management.WithClientCredentials(ctx, clientID, clientSecret)
 }
+
+func customClientWithRetries() *http.Client {
+	client := &http.Client{
+		Transport: rateLimitTransport(
+			retryableErrorTransport(
+				serverIssueTransport(
+					http.DefaultTransport,
+				),
+			),
+		),
+	}
+
+	return client
+}
+
+func rateLimitTransport(tripper http.RoundTripper) http.RoundTripper {
+	return rehttp.NewTransport(tripper, rateLimitRetry, rateLimitDelay)
+}
+
+func rateLimitRetry(attempt rehttp.Attempt) bool {
+	if attempt.Response == nil {
+		return false
+	}
+
+	return attempt.Response.StatusCode == http.StatusTooManyRequests
+}
+
+func rateLimitDelay(attempt rehttp.Attempt) time.Duration {
+	resetAt := attempt.Response.Header.Get("X-RateLimit-Reset")
+
+	resetAtUnix, err := strconv.ParseInt(resetAt, 10, 64)
+	if err != nil {
+		resetAtUnix = time.Now().Add(5 * time.Second).Unix()
+	}
+
+	return time.Duration(resetAtUnix-time.Now().Unix()) * time.Second
+}
+
+func retryableErrorTransport(tripper http.RoundTripper) http.RoundTripper {
+	return rehttp.NewTransport(
+		tripper,
+		rehttp.RetryIsErr(retryableErrorRetryFunc),
+		rehttp.ExpJitterDelay(250*time.Millisecond, 10*time.Second),
+	)
+}
+
+func retryableErrorRetryFunc(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	if v, ok := err.(*url.Error); ok {
+		// Don't retry if the error was due to too many redirects.
+		if regexp.MustCompile(`stopped after \d+ redirects\z`).MatchString(v.Error()) {
+			return false
+		}
+
+		// Don't retry if the error was due to an invalid protocol scheme.
+		if regexp.MustCompile(`unsupported protocol scheme`).MatchString(v.Error()) {
+			return false
+		}
+
+		// Don't retry if the error was due to TLS cert verification failure.
+		if regexp.MustCompile(`certificate is not trusted`).MatchString(v.Error()) {
+			return false
+		}
+
+		// Don't retry if the certificate issuer is unknown.
+		if _, ok := v.Err.(x509.UnknownAuthorityError); ok {
+			return false
+		}
+	}
+
+	// The error is likely recoverable so retry.
+	return true
+}
+
+func serverIssueTransport(tripper http.RoundTripper) http.RoundTripper {
+	return rehttp.NewTransport(
+		tripper,
+		rehttp.RetryAll(
+			rehttp.RetryMaxRetries(2),
+			rehttp.RetryStatuses(
+				http.StatusServiceUnavailable,
+				http.StatusInternalServerError,
+				http.StatusBadGateway,
+				http.StatusGatewayTimeout,
+			),
+		),
+		rehttp.ExpJitterDelay(500*time.Millisecond, 10*time.Second),
+	)
+}