Issue #2334 - HTTPDigestAuth - All state now in thread local storage

Following feedback from tardyp and @vincentxb.

1 files changed, 38 insertions, 27 deletions

diff --git a/requests/auth.py b/requests/auth.py
index 5c5332c8..7555114e 100644
--- a/requests/auth.py
+++ b/requests/auth.py
@@ -64,19 +64,24 @@ class HTTPDigestAuth(AuthBase):
     def __init__(self, username, password):
         self.username = username
         self.password = password
-        self.last_nonce = ''
-        self.nonce_count = 0
-        self.chal = {}
-        self.pos = None
-        self.num_401_calls = threading.local()
+        # Keep state in per-thread local storage
+        self.tl = threading.local()
+        self.init_per_thread_state()
+
+    def init_per_thread_state(self):
+        self.tl.last_nonce = ''
+        self.tl.nonce_count = 0
+        self.tl.chal = {}
+        self.tl.pos = None
+        self.tl.num_401_calls = None
 
     def build_digest_header(self, method, url):
 
-        realm = self.chal['realm']
-        nonce = self.chal['nonce']
-        qop = self.chal.get('qop')
-        algorithm = self.chal.get('algorithm')
-        opaque = self.chal.get('opaque')
+        realm = self.tl.chal['realm']
+        nonce = self.tl.chal['nonce']
+        qop = self.tl.chal.get('qop')
+        algorithm = self.tl.chal.get('algorithm')
+        opaque = self.tl.chal.get('opaque')
 
         if algorithm is None:
             _algorithm = 'MD5'
@@ -114,12 +119,12 @@ class HTTPDigestAuth(AuthBase):
         HA1 = hash_utf8(A1)
         HA2 = hash_utf8(A2)
 
-        if nonce == self.last_nonce:
-            self.nonce_count += 1
+        if nonce == self.tl.last_nonce:
+            self.tl.nonce_count += 1
         else:
-            self.nonce_count = 1
-        ncvalue = '%08x' % self.nonce_count
-        s = str(self.nonce_count).encode('utf-8')
+            self.tl.nonce_count = 1
+        ncvalue = '%08x' % self.tl.nonce_count
+        s = str(self.tl.nonce_count).encode('utf-8')
         s += nonce.encode('utf-8')
         s += time.ctime().encode('utf-8')
         s += os.urandom(8)
@@ -139,7 +144,7 @@ class HTTPDigestAuth(AuthBase):
             # XXX handle auth-int.
             return None
 
-        self.last_nonce = nonce
+        self.tl.last_nonce = nonce
 
         # XXX should the partial digests be encoded too?
         base = 'username="%s", realm="%s", nonce="%s", uri="%s", ' \
@@ -158,23 +163,23 @@ class HTTPDigestAuth(AuthBase):
     def handle_redirect(self, r, **kwargs):
         """Reset num_401_calls counter on redirects."""
         if r.is_redirect:
-            self.num_401_calls.value = 1
+            self.tl.num_401_calls = 1
 
     def handle_401(self, r, **kwargs):
         """Takes the given response and tries digest-auth, if needed."""
 
-        if self.pos is not None:
+        if self.tl.pos is not None:
             # Rewind the file position indicator of the body to where
             # it was to resend the request.
-            r.request.body.seek(self.pos)
-        num_401_calls = self.num_401_calls.value
+            r.request.body.seek(self.tl.pos)
+        num_401_calls = self.tl.num_401_calls
         s_auth = r.headers.get('www-authenticate', '')
 
         if 'digest' in s_auth.lower() and num_401_calls < 2:
 
-            self.num_401_calls.value += 1
+            self.tl.num_401_calls += 1
             pat = re.compile(r'digest ', flags=re.IGNORECASE)
-            self.chal = parse_dict_header(pat.sub('', s_auth, count=1))
+            self.tl.chal = parse_dict_header(pat.sub('', s_auth, count=1))
 
             # Consume content and release the original connection
             # to allow our new request to reuse the same one.
@@ -192,23 +197,29 @@ class HTTPDigestAuth(AuthBase):
 
             return _r
 
-        self.num_401_calls.value = 1
+        self.tl.num_401_calls = 1
         return r
 
     def __call__(self, r):
+        # When called from a thread other than the one that __init__'ed us
+        # per-thread state may be missing: initialize it if that's the case.
+        try:
+            self.tl.last_nonce
+        except AttributeError:
+            self.init_per_thread_state()
         # If we have a saved nonce, skip the 401
-        if self.last_nonce:
+        if self.tl.last_nonce:
             r.headers['Authorization'] = self.build_digest_header(r.method, r.url)
         try:
-            self.pos = r.body.tell()
+            self.tl.pos = r.body.tell()
         except AttributeError:
             # In the case of HTTPDigestAuth being reused and the body of
             # the previous request was a file-like object, pos has the
             # file position of the previous body. Ensure it's set to
             # None.
-            self.pos = None
+            self.tl.pos = None
         r.register_hook('response', self.handle_401)
         r.register_hook('response', self.handle_redirect)
-        self.num_401_calls.value = 1
+        self.tl.num_401_calls = 1
 
         return r