From 0b27584de8f94311f8022e3e025961d1ddf96d32 Mon Sep 17 00:00:00 2001
From: Ian Good <ian@icgood.net>
Date: Mon, 3 Aug 2020 20:20:23 -0400
Subject: [PATCH] Use a thread pool for auth hashing

---
 pymap/backend/redis/__init__.py |  2 +-
 pymap/config.py                 | 13 +++++++++++++
 pymap/user.py                   | 15 +++++++++++----
 setup.py                        |  2 +-
 4 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/pymap/backend/redis/__init__.py b/pymap/backend/redis/__init__.py
index c1598ecf..2c542878 100644
--- a/pymap/backend/redis/__init__.py
+++ b/pymap/backend/redis/__init__.py
@@ -315,7 +315,7 @@ async def _check_user(self, redis: Redis,
             data = await self._get_user(redis, user)
             if data is None:
                 raise InvalidAuth()
-            data.check_password(credentials)
+            await data.check_password(credentials)
             if user != credentials.identity:
                 raise InvalidAuth(authorization=True)
         return credentials.identity
diff --git a/pymap/config.py b/pymap/config.py
index cadb98d2..8ef26c03 100644
--- a/pymap/config.py
+++ b/pymap/config.py
@@ -7,6 +7,7 @@
 from abc import abstractmethod, ABCMeta
 from argparse import Namespace
 from collections import OrderedDict
+from concurrent.futures import ThreadPoolExecutor
 from ssl import SSLContext
 from typing import Any, TypeVar, Type, Union, Optional, Iterable, Iterator, \
     Sequence, Mapping, Dict
@@ -97,6 +98,8 @@ class IMAPConfig(metaclass=ABCMeta):
             ``key_file`` and an SSL context will be created.
         starttls_enabled: True if opportunistic TLS should be supported.
         hash_context: The hash to use for passwords.
+        cpu_subsystem: The subsystem to use for CPU-heavy operations,
+            defaulting to a small thread pool.
         preauth_credentials: If given, clients will pre-authenticate on
             connection using these credentials.
         proxy_protocol: The PROXY protocol implementation to use.
@@ -123,6 +126,7 @@ def __init__(self, args: Namespace, *,
                  preauth_credentials: AuthenticationCredentials = None,
                  proxy_protocol: ProxyProtocol = None,
                  hash_context: HashInterface = None,
+                 cpu_subsystem: Subsystem = None,
                  max_append_len: Optional[int] = 1000000000,
                  bad_command_limit: Optional[int] = 5,
                  disable_search_keys: Iterable[bytes] = None,
@@ -138,6 +142,8 @@ def __init__(self, args: Namespace, *,
         self.disable_search_keys: Final = disable_search_keys or []
         self.hash_context: Final = hash_context or \
             get_hash(passlib_config=args.passlib_cfg)
+        self.cpu_subsystem: Final = cpu_subsystem or \
+            self._get_cpu_subsystem()
         self._ssl_context = ssl_context or self._load_certs(extra)
         self._tls_enabled = tls_enabled
         self._preauth_credentials = preauth_credentials
@@ -188,6 +194,13 @@ def backend_capability(self) -> BackendCapability:
         """
         ...
 
+    @classmethod
+    def _get_cpu_subsystem(cls) -> Subsystem:
+        cpu_count = os.cpu_count() or 1
+        cpus_minus_one = max(1, cpu_count - 1)
+        executor = ThreadPoolExecutor(max_workers=cpus_minus_one)
+        return Subsystem.for_threading(executor)
+
     @classmethod
     def _load_certs(cls, extra: Mapping[str, Any]) -> SSLContext:
         cert_file: Optional[str] = extra.get('cert_file')
diff --git a/pymap/user.py b/pymap/user.py
index 8b82ee74..fe19287d 100644
--- a/pymap/user.py
+++ b/pymap/user.py
@@ -4,7 +4,7 @@
 from typing import Any, Optional, Mapping, Dict
 from typing_extensions import Final
 
-from pysasl import AuthenticationCredentials
+from pysasl import AuthenticationCredentials, HashInterface
 
 from .config import IMAPConfig
 from .exceptions import InvalidAuth
@@ -61,7 +61,7 @@ def to_dict(self) -> Mapping[str, Any]:
             data['password'] = self.password
         return data
 
-    def check_password(self, creds: AuthenticationCredentials) -> None:
+    async def check_password(self, creds: AuthenticationCredentials) -> None:
         """Check the given credentials against the known password comparison
         data. If the known data used a hash, then the equivalent hash of the
         provided secret is compared.
@@ -73,8 +73,15 @@ def check_password(self, creds: AuthenticationCredentials) -> None:
             :class:`~pymap.exceptions.InvalidAuth`
 
         """
-        config = self.config
         if self.password is None:
             raise InvalidAuth()
-        elif not creds.check_secret(self.password, hash=config.hash_context):
+        hash_context = self.config.hash_context
+        cpu_subsystem = self.config.cpu_subsystem
+        fut = self._check_secret(creds, self.password, hash_context)
+        if not await cpu_subsystem.execute(fut):
             raise InvalidAuth()
+
+    async def _check_secret(self, creds: AuthenticationCredentials,
+                            password: str,
+                            hash_context: HashInterface) -> bool:
+        return creds.check_secret(password, hash=hash_context)
diff --git a/setup.py b/setup.py
index 3d8f63ab..0c5eb978 100644
--- a/setup.py
+++ b/setup.py
@@ -28,7 +28,7 @@
     license = f.read()
 
 setup(name='pymap',
-      version='0.19.0',
+      version='0.20.0',
       author='Ian Good',
       author_email='icgood@gmail.com',
       description='Lightweight, asynchronous IMAP serving in Python.',