From 4d04d14b2038d0364c26e49164eff6cf7624a593 Mon Sep 17 00:00:00 2001
From: Jakub Warmuz <jakub@warmuz.org>
Date: Thu, 7 Jan 2016 20:25:07 +0000
Subject: [PATCH] Fix "global" max_attempt bug (#1719)

---
 acme/acme/client.py | 30 ++++++++++++++++++++----------
 acme/acme/errors.py | 17 ++++++++---------
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/acme/acme/client.py b/acme/acme/client.py
index c3e28ef47..9970e2046 100644
--- a/acme/acme/client.py
+++ b/acme/acme/client.py
@@ -1,4 +1,5 @@
 """ACME client API."""
+import collections
 import datetime
 import heapq
 import logging
@@ -336,8 +337,9 @@ class Client(object):  # pylint: disable=too-many-instance-attributes
         :param authzrs: `list` of `.AuthorizationResource`
         :param int mintime: Minimum time before next attempt, used if
             ``Retry-After`` is not present in the response.
-        :param int max_attempts: Maximum number of attempts before
-            `PollError` with non-empty ``waiting`` is raised.
+        :param int max_attempts: Maximum number of attempts (per
+            authorization) before `PollError` with non-empty ``waiting``
+            is raised.
 
         :returns: ``(cert, updated_authzrs)`` `tuple` where ``cert`` is
             the issued certificate (`.messages.CertificateResource`),
@@ -351,6 +353,11 @@ class Client(object):  # pylint: disable=too-many-instance-attributes
             was marked by the CA as invalid
 
         """
+        # pylint: disable=too-many-locals
+        assert max_attempts > 0
+        attempts = collections.defaultdict(int)
+        exhausted = set()
+
         # priority queue with datetime (based on Retry-After) as key,
         # and original Authorization Resource as value
         waiting = [(datetime.datetime.now(), authzr) for authzr in authzrs]
@@ -358,8 +365,7 @@ class Client(object):  # pylint: disable=too-many-instance-attributes
         # recently updated one
         updated = dict((authzr, authzr) for authzr in authzrs)
 
-        while waiting and max_attempts:
-            max_attempts -= 1
+        while waiting:
             # find the smallest Retry-After, and sleep if necessary
             when, authzr = heapq.heappop(waiting)
             now = datetime.datetime.now()
@@ -373,16 +379,20 @@ class Client(object):  # pylint: disable=too-many-instance-attributes
             updated_authzr, response = self.poll(updated[authzr])
             updated[authzr] = updated_authzr
 
+            attempts[authzr] += 1
             # pylint: disable=no-member
             if updated_authzr.body.status not in (
                     messages.STATUS_VALID, messages.STATUS_INVALID):
-                # push back to the priority queue, with updated retry_after
-                heapq.heappush(waiting, (self.retry_after(
-                    response, default=mintime), authzr))
+                if attempts[authzr] < max_attempts:
+                    # push back to the priority queue, with updated retry_after
+                    heapq.heappush(waiting, (self.retry_after(
+                        response, default=mintime), authzr))
+                else:
+                    exhausted.add(authzr)
 
-        if not max_attempts or any(authzr.body.status == messages.STATUS_INVALID
-                                   for authzr in six.itervalues(updated)):
-            raise errors.PollError(waiting, updated)
+        if exhausted or any(authzr.body.status == messages.STATUS_INVALID
+                            for authzr in six.itervalues(updated)):
+            raise errors.PollError(exhausted, updated)
 
         updated_authzrs = tuple(updated[authzr] for authzr in authzrs)
         return self.request_issuance(csr, updated_authzrs), updated_authzrs
diff --git a/acme/acme/errors.py b/acme/acme/errors.py
index 0385667c7..77d47c522 100644
--- a/acme/acme/errors.py
+++ b/acme/acme/errors.py
@@ -56,26 +56,25 @@ class MissingNonce(NonceError):
 class PollError(ClientError):
     """Generic error when polling for authorization fails.
 
-    This might be caused by either timeout (`waiting` will be non-empty)
+    This might be caused by either timeout (`exhausted` will be non-empty)
     or by some authorization being invalid.
 
-    :ivar waiting: Priority queue with `datetime.datatime` (based on
-        ``Retry-After``) as key, and original `.AuthorizationResource`
-        as value.
+    :ivar exhausted: Set of `.AuthorizationResource` that didn't finish
+        within max allowed attempts.
     :ivar updated: Mapping from original `.AuthorizationResource`
         to the most recently updated one
 
     """
-    def __init__(self, waiting, updated):
-        self.waiting = waiting
+    def __init__(self, exhausted, updated):
+        self.exhausted = exhausted
         self.updated = updated
         super(PollError, self).__init__()
 
     @property
     def timeout(self):
         """Was the error caused by timeout?"""
-        return bool(self.waiting)
+        return bool(self.exhausted)
 
     def __repr__(self):
-        return '{0}(waiting={1!r}, updated={2!r})'.format(
-            self.__class__.__name__, self.waiting, self.updated)
+        return '{0}(exhausted={1!r}, updated={2!r})'.format(
+            self.__class__.__name__, self.exhausted, self.updated)