From a3400670480fcc6d31e3c33706dd1368009376bd Mon Sep 17 00:00:00 2001 From: disqualifier Date: Thu, 2 Jul 2026 16:40:50 -0400 Subject: [PATCH] fix: reconnect-on-stale re-selects folder; lock connect() against races ensure_connection() reconnected via connect(), which only reaches IMAP state AUTH, not SELECTED - every search/fetch/store after a mid-pass link drop silently failed (swallowed into []/None/False) because aioimaplib rejects those commands outside SELECTED. Track the currently-selected folder and re-select it after reconnecting; note pre-reconnect sequence-number ids are invalid post-reconnect unless use_uid=True. Also add an asyncio.Lock around connect()/close()/ensure_connection(): concurrent callers on one instance used to race inside connect(), where task B's `await self.close()` tore down task A's mid-handshake connection, leaking sockets and orphaning server sessions. Superseded connections are now always logged out instead of silently overwritten. Bump to 0.1.6, update README install pins. Signed-off-by: disqualifier --- README.md | 10 ++--- pyproject.toml | 2 +- src/aiomail/client.py | 90 +++++++++++++++++++++++++++++++++++++------ 3 files changed, 84 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 22722b8..9ab456e 100644 --- a/README.md +++ b/README.md @@ -11,22 +11,22 @@ This reads codes from email; it does not generate them (that is `pyotp`'s job). `requirements.txt`: ``` -aiomail @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.5 +aiomail @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.6 # OAuth token providers (Microsoft / Google) need the extra: -aiomail[oauth] @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.5 +aiomail[oauth] @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.6 ``` Direct: ```bash -pip install "aiomail @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.5" -pip install "aiomail[oauth] @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.5" +pip install "aiomail @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.6" +pip install "aiomail[oauth] @ git+ssh://git@git.rethinkstudios.io/rethink-public/aiomail.git@v0.1.6" ``` Requires `aioimaplib` and `beautifulsoup4` (pulled transitively). The `oauth` extra adds `aiohttp` for the refresh-token providers. -Drop the `@v0.1.5` suffix from the line above to install the latest unpinned. +Drop the `@v0.1.6` suffix from the line above to install the latest unpinned. ## Password auth diff --git a/pyproject.toml b/pyproject.toml index 4c863f3..02614da 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "aiomail" -version = "0.1.5" +version = "0.1.6" description = "async IMAP one-time-code retrieval with password/OAuth2 auth and dynamic matching" requires-python = ">=3.10" dependencies = [ diff --git a/src/aiomail/client.py b/src/aiomail/client.py index 05222c2..43f23cb 100644 --- a/src/aiomail/client.py +++ b/src/aiomail/client.py @@ -4,6 +4,16 @@ a thin, provider-agnostic client: it owns the connection lifecycle (connect with retries, reconnect-on-stale, close) and exposes the handful of operations the OTP flow needs (folders, search, fetch, mark-seen). auth is injected, so the same client serves password and OAuth accounts. + +reconnect-on-stale re-selects whatever folder was selected before the drop, so +search/fetch/store keep working against a reconnected session. sequence-number ids +from before a reconnect are not valid afterward (a fresh SELECT can renumber the +mailbox) — pass `use_uid=True` if ids need to survive a reconnect. + +concurrency: one `IMAPClient` instance is not safe to drive from multiple +concurrent tasks/coroutines without external serialization; connect/reconnect +internally uses a lock to avoid corrupting `_mail`, but overlapping calls to the +same instance are not the intended usage pattern. """ import asyncio import email @@ -62,6 +72,8 @@ class IMAPClient: self.timeout = timeout self.max_retries = max_retries self._mail = None + self._selected_folder: Optional[str] = None + self._lock = asyncio.Lock() async def __aenter__(self) -> "IMAPClient": await self.ensure_connection() @@ -71,8 +83,21 @@ class IMAPClient: await self.close() async def connect(self) -> bool: - """open a connection and authenticate, retrying with linear backoff""" - await self.close() + """open a connection and authenticate, retrying with linear backoff + + serialized by an internal lock: concurrent callers on the same instance queue + up rather than tearing down each other's in-progress handshake. a connection + superseded while a caller waited is logged out, never silently overwritten. + """ + async with self._lock: + return await self._connect_locked() + + async def _connect_locked(self) -> bool: + """connect()'s body; caller must hold self._lock""" + superseded, self._mail = self._mail, None + self._selected_folder = None + if superseded is not None: + await self._discard_mail(superseded) for attempt in range(self.max_retries): try: if self.use_ssl: @@ -114,22 +139,60 @@ class IMAPClient: async def close(self) -> None: """log out and drop the connection, swallowing teardown errors""" + async with self._lock: + await self._close_locked() + + async def _close_locked(self) -> None: + """close()'s body; caller must hold self._lock""" if self._mail is not None: + mail, self._mail = self._mail, None try: - await self._mail.logout() + await mail.logout() except Exception as exc: log.debug("logout error ignored: %s", exc) - self._mail = None + self._selected_folder = None async def ensure_connection(self) -> bool: - """return a live connection, reconnecting if the link is stale""" - if self._mail is None: - return await self.connect() - try: - await self._mail.noop() + """return a live, SELECTED-if-applicable connection, reconnecting if the link is stale + + a reconnect only re-authenticates (state AUTH); if a folder was selected before + the drop, it is re-selected here so search/fetch/store keep working afterward. + note: sequence-number ids from before the reconnect are NOT valid against the + new session (a fresh SELECT can renumber/re-EXISTS the mailbox) unless + use_uid=True, in which case UIDs remain stable across the reconnect. + + serialized by an internal lock, so concurrent callers on the same instance + never race each other's connect/reconnect; a caller that waits behind another + rechecks liveness first instead of tearing down a connection that just came up. + """ + async with self._lock: + if self._mail is not None: + try: + await self._mail.noop() + return True + except Exception: + pass + return await self._connect_and_reselect_locked() + + async def _connect_and_reselect_locked(self) -> bool: + """connect() then re-select the previously-selected folder; caller must hold self._lock""" + folder = self._selected_folder + if not await self._connect_locked(): + return False + if folder is None: return True - except Exception: - return await self.connect() + try: + result, _ = await self._mail.select(f'"{folder}"') + except Exception as exc: + log.debug("re-select %s after reconnect failed: %s", folder, exc) + self._selected_folder = None + return False + if result != "OK": + log.debug("re-select %s after reconnect failed: %s", folder, result) + self._selected_folder = None + return False + self._selected_folder = folder + return True async def get_folders(self) -> List[str]: """list mailbox folder names""" @@ -154,10 +217,13 @@ class IMAPClient: return False try: result, _ = await self._mail.select(f'"{folder}"') - return result == "OK" except Exception as exc: log.debug("select %s failed: %s", folder, exc) return False + if result == "OK": + self._selected_folder = folder + return True + return False async def search(self, query: str) -> List[int]: """search the selected folder, returning ids newest-first"""