Merge pull request #61 from SecurityGeekIO/fix-retry-after-case-sensi…

…tive fix: make retry after header value case insensitive
SecurityGeekIO · Apr 17, 2024 · d73b2bf · d73b2bf
2 parents 696bc7c + 2cadda0
commit d73b2bf
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 7 deletions.
diff --git a/zcon/config.go b/zcon/config.go
@@ -296,11 +296,16 @@ func (c *Client) GetContentType() string {
 }
 
 func getRetryAfter(resp *http.Response, l logger.Logger) time.Duration {
-	if s, ok := resp.Header["Retry-After"]; ok {
-		if sleep, err := strconv.ParseInt(s[0], 10, 64); err == nil {
+
+	if s := resp.Header.Get("Retry-After"); s != "" {
+		if sleep, err := strconv.ParseInt(s, 10, 64); err == nil {
 			l.Printf("[INFO] got Retry-After from header:%s\n", s)
 			return time.Second * time.Duration(sleep)
 		} else {
+			dur, err := time.ParseDuration(s)
+			if err == nil {
+				return dur
+			}
 			l.Printf("[INFO] error getting Retry-After from header:%s\n", err)
 		}
 	}

diff --git a/zia/config.go b/zia/config.go
@@ -300,11 +300,15 @@ func (c *Client) GetContentType() string {
 }
 
 func getRetryAfter(resp *http.Response, l logger.Logger) time.Duration {
-	if s, ok := resp.Header["Retry-After"]; ok {
-		if sleep, err := strconv.ParseInt(s[0], 10, 64); err == nil {
+	if s := resp.Header.Get("Retry-After"); s != "" {
+		if sleep, err := strconv.ParseInt(s, 10, 64); err == nil {
 			l.Printf("[INFO] got Retry-After from header:%s\n", s)
 			return time.Second * time.Duration(sleep)
 		} else {
+			dur, err := time.ParseDuration(s)
+			if err == nil {
+				return dur
+			}
 			l.Printf("[INFO] error getting Retry-After from header:%s\n", err)
 		}
 	}

diff --git a/zpa/config.go b/zpa/config.go
@@ -215,10 +215,14 @@ func (c *Config) GetHTTPClient() *http.Client {
 			retryableClient.Backoff = func(min, max time.Duration, attemptNum int, resp *http.Response) time.Duration {
 				if resp != nil {
 					if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode == http.StatusServiceUnavailable {
-						// TODO: ask backend to implement such header, instead of using the logic below
-						if s, ok := resp.Header["Retry-After"]; ok {
-							if sleep, err := strconv.ParseInt(s[0], 10, 64); err == nil {
+						if s := resp.Header.Get("Retry-After"); s != "" {
+							if sleep, err := strconv.ParseInt(s, 10, 64); err == nil {
 								return time.Second * time.Duration(sleep)
+							} else {
+								dur, err := time.ParseDuration(s)
+								if err == nil {
+									return dur
+								}
 							}
 						}
 					}