Merge branch 'auth-digest-multi-thread' of https://github.com/exvito/requests into proposed/2.8.0

requests/auth.py

@@ -11,6 +11,7 @@ import os
 import re
 import time
 import hashlib
+import threading
 
 from base64 import b64encode
 
@@ -63,19 +64,26 @@ 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 = 1
+        # Keep state in per-thread local storage
+        self._thread_local = threading.local()
+
+    def init_per_thread_state(self):
+        # Ensure state is initialized just once per-thread
+        if not hasattr(self._thread_local, 'init'):
+            self._thread_local.init = True
+            self._thread_local.last_nonce = ''
+            self._thread_local.nonce_count = 0
+            self._thread_local.chal = {}
+            self._thread_local.pos = None
+            self._thread_local.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._thread_local.chal['realm']
+        nonce = self._thread_local.chal['nonce']
+        qop = self._thread_local.chal.get('qop')
+        algorithm = self._thread_local.chal.get('algorithm')
+        opaque = self._thread_local.chal.get('opaque')
 
         if algorithm is None:
             _algorithm = 'MD5'
@@ -114,12 +122,12 @@ class HTTPDigestAuth(AuthBase):
         HA1 = hash_utf8(A1)
         HA2 = hash_utf8(A2)
 
-        if nonce == self.last_nonce:
-            self.nonce_count += 1
+        if nonce == self._thread_local.last_nonce:
+            self._thread_local.nonce_count += 1
         else:
-            self.nonce_count = 1
-        ncvalue = '%08x' % self.nonce_count
-        s = str(self.nonce_count).encode('utf-8')
+            self._thread_local.nonce_count = 1
+        ncvalue = '%08x' % self._thread_local.nonce_count
+        s = str(self._thread_local.nonce_count).encode('utf-8')
         s += nonce.encode('utf-8')
         s += time.ctime().encode('utf-8')
         s += os.urandom(8)
@@ -139,7 +147,7 @@ class HTTPDigestAuth(AuthBase):
             # XXX handle auth-int.
             return None
 
-        self.last_nonce = nonce
+        self._thread_local.last_nonce = nonce
 
         # XXX should the partial digests be encoded too?
         base = 'username="%s", realm="%s", nonce="%s", uri="%s", ' \
@@ -158,23 +166,22 @@ class HTTPDigestAuth(AuthBase):
     def handle_redirect(self, r, **kwargs):
         """Reset num_401_calls counter on redirects."""
         if r.is_redirect:
-            self.num_401_calls = 1
+            self._thread_local.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._thread_local.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 = getattr(self, 'num_401_calls', 1)
+            r.request.body.seek(self._thread_local.pos)
         s_auth = r.headers.get('www-authenticate', '')
 
-        if 'digest' in s_auth.lower() and num_401_calls < 2:
+        if 'digest' in s_auth.lower() and self._thread_local.num_401_calls < 2:
 
-            self.num_401_calls += 1
+            self._thread_local.num_401_calls += 1
             pat = re.compile(r'digest ', flags=re.IGNORECASE)
-            self.chal = parse_dict_header(pat.sub('', s_auth, count=1))
+            self._thread_local.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,21 +199,25 @@ class HTTPDigestAuth(AuthBase):
 
             return _r
 
-        self.num_401_calls = 1
+        self._thread_local.num_401_calls = 1
         return r
 
     def __call__(self, r):
+        # Initialize per-thread state, if needed
+        self.init_per_thread_state()
         # If we have a saved nonce, skip the 401
-        if self.last_nonce:
+        if self._thread_local.last_nonce:
             r.headers['Authorization'] = self.build_digest_header(r.method, r.url)
         try:
-            self.pos = r.body.tell()
+            self._thread_local.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._thread_local.pos = None
         r.register_hook('response', self.handle_401)
         r.register_hook('response', self.handle_redirect)
+        self._thread_local.num_401_calls = 1
+
         return r

test_requests.py

@@ -33,6 +33,11 @@ try:
 except ImportError:
     import io as StringIO
 
+try:
+    from multiprocessing.pool import ThreadPool
+except ImportError:
+    ThreadPool = None
+
 if is_py3:
     def u(s):
         return s
@@ -412,6 +417,21 @@ class RequestsTestCase(unittest.TestCase):
         r = requests.get(url, auth=auth)
         assert '"auth"' in r.request.headers['Authorization']
 
+    def test_DIGESTAUTH_THREADED(self):
+
+        auth = HTTPDigestAuth('user', 'pass')
+        url = httpbin('digest-auth', 'auth', 'user', 'pass')
+        session = requests.Session()
+        session.auth=auth
+
+        def do_request(i):
+            r = session.get(url)
+            assert '"auth"' in r.request.headers['Authorization']
+            return 1
+        if ThreadPool is not None:
+            pool = ThreadPool(processes=50)
+            pool.map(do_request, range(100))
+
     def test_POSTBIN_GET_POST_FILES(self):
 
         url = httpbin('post')