From 54fe9c8ed9d493197eae5da63fccb67577d8ab1c Mon Sep 17 00:00:00 2001
From: Alex Vaghin <1977446+desdeel2d0m@users.noreply.github.com>
Date: Sat, 25 Feb 2017 19:24:21 -0800
Subject: [PATCH] acme: add function to check rate limits

This change exposes a function to extract rate limit duration
from a client error using Retry-After response header.
Author: David Calavera <david.calavera@gmail.com>.

Fixes golang/go#19304.
Change-Id: Iec9cfab398b84c6f216b95d3265ffad1ce2f29a7
Reviewed-on: https://go-review.googlesource.com/37463
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
---
 acme/types.go      | 26 +++++++++++++++++++
 acme/types_test.go | 63 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+)
 create mode 100644 acme/types_test.go

diff --git a/acme/types.go b/acme/types.go
index ea0d235..ab4de0b 100644
--- a/acme/types.go
+++ b/acme/types.go
@@ -1,3 +1,7 @@
+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
 package acme
 
 import (
@@ -5,6 +9,7 @@ import (
 	"fmt"
 	"net/http"
 	"strings"
+	"time"
 )
 
 // ACME server response statuses used to describe Authorization and Challenge states.
@@ -79,6 +84,27 @@ func (a *AuthorizationError) Error() string {
 	return fmt.Sprintf("acme: authorization error for %s: %s", a.Identifier, strings.Join(e, "; "))
 }
 
+// RateLimit reports whether err represents a rate limit error and
+// any Retry-After duration returned by the server.
+//
+// See the following for more details on rate limiting:
+// https://tools.ietf.org/html/draft-ietf-acme-acme-05#section-5.6
+func RateLimit(err error) (time.Duration, bool) {
+	e, ok := err.(*Error)
+	if !ok {
+		return 0, false
+	}
+	// Some CA implementations may return incorrect values.
+	// Use case-insensitive comparison.
+	if !strings.HasSuffix(strings.ToLower(e.ProblemType), ":ratelimited") {
+		return 0, false
+	}
+	if e.Header == nil {
+		return 0, true
+	}
+	return retryAfter(e.Header.Get("Retry-After"), 0), true
+}
+
 // Account is a user account. It is associated with a private key.
 type Account struct {
 	// URI is the account unique ID, which is also a URL used to retrieve
diff --git a/acme/types_test.go b/acme/types_test.go
new file mode 100644
index 0000000..a7553e6
--- /dev/null
+++ b/acme/types_test.go
@@ -0,0 +1,63 @@
+// Copyright 2017 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package acme
+
+import (
+	"errors"
+	"net/http"
+	"testing"
+	"time"
+)
+
+func TestRateLimit(t *testing.T) {
+	now := time.Date(2017, 04, 27, 10, 0, 0, 0, time.UTC)
+	f := timeNow
+	defer func() { timeNow = f }()
+	timeNow = func() time.Time { return now }
+
+	h120, hTime := http.Header{}, http.Header{}
+	h120.Set("Retry-After", "120")
+	hTime.Set("Retry-After", "Tue Apr 27 11:00:00 2017")
+
+	err1 := &Error{
+		ProblemType: "urn:ietf:params:acme:error:nolimit",
+		Header:      h120,
+	}
+	err2 := &Error{
+		ProblemType: "urn:ietf:params:acme:error:rateLimited",
+		Header:      h120,
+	}
+	err3 := &Error{
+		ProblemType: "urn:ietf:params:acme:error:rateLimited",
+		Header:      nil,
+	}
+	err4 := &Error{
+		ProblemType: "urn:ietf:params:acme:error:rateLimited",
+		Header:      hTime,
+	}
+
+	tt := []struct {
+		err error
+		res time.Duration
+		ok  bool
+	}{
+		{nil, 0, false},
+		{errors.New("dummy"), 0, false},
+		{err1, 0, false},
+		{err2, 2 * time.Minute, true},
+		{err3, 0, true},
+		{err4, time.Hour, true},
+	}
+	for i, test := range tt {
+		res, ok := RateLimit(test.err)
+		if ok != test.ok {
+			t.Errorf("%d: RateLimit(%+v): ok = %v; want %v", i, test.err, ok, test.ok)
+			continue
+		}
+		if res != test.res {
+			t.Errorf("%d: RateLimit(%+v) = %v; want %v", i, test.err, res, test.res)
+		}
+	}
+}