fix: reject use of pooled connections after release to prevent stale reference escape

ConnectionPool.acquire() yields a raw DqliteConnection. A caller could
stash the reference and use it after the context manager exits, silently
bypassing pool exclusivity. Two tasks could then operate on the same TCP
connection concurrently, corrupting the dqlite wire protocol.

Add a _pool_released flag to DqliteConnection that is checked in
_check_in_use(). The pool sets the flag after cleanup (including any
ROLLBACK for open transactions) completes, and clears it when the
connection is re-acquired. Standalone (non-pooled) connections are not
affected since _pool_released defaults to False.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: antoineleclair

src/dqliteclient/connection.py

@@ -82,6 +82,7 @@ def __init__(
         self._in_use = False
         self._bound_loop: asyncio.AbstractEventLoop | None = None
         self._tx_owner: asyncio.Task[Any] | None = None
+        self._pool_released = False
 
     @property
     def address(self) -> str:
@@ -154,7 +155,12 @@ def _ensure_connected(self) -> tuple[DqliteProtocol, int]:
         return self._protocol, self._db_id
 
     def _check_in_use(self) -> None:
-        """Raise on misuse: wrong event loop or concurrent coroutine access."""
+        """Raise on misuse: wrong event loop, concurrent access, or use after pool release."""
+        if self._pool_released:
+            raise InterfaceError(
+                "This connection has been returned to the pool and can no longer "
+                "be used directly. Acquire a new connection from the pool."
+            )
         if self._bound_loop is not None:
             try:
                 current_loop = asyncio.get_running_loop()

src/dqliteclient/pool.py

@@ -143,17 +143,20 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
                 self._size -= 1
                 conn = await self._create_connection()
 
+        conn._pool_released = False
         try:
             yield conn
         except BaseException:
             if conn.is_connected and not self._closed:
                 # Connection is healthy — user code raised a non-connection error.
                 # Roll back any open transaction, then return to pool.
                 if not await self._reset_connection(conn):
+                    conn._pool_released = True
                     with contextlib.suppress(Exception):
                         await conn.close()
                     self._size -= 1
                     raise
+                conn._pool_released = True
                 try:
                     self._pool.put_nowait(conn)
                 except asyncio.QueueFull:
@@ -162,6 +165,7 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
             else:
                 # Connection is broken (invalidated by execute/fetch error handlers).
                 # Drain other idle connections — they likely point to the same dead server.
+                conn._pool_released = True
                 await self._drain_idle()
                 with contextlib.suppress(BaseException):
                     await conn.close()
@@ -188,16 +192,19 @@ async def _reset_connection(self, conn: DqliteConnection) -> bool:
     async def _release(self, conn: DqliteConnection) -> None:
         """Return a connection to the pool or close it."""
         if self._closed:
+            conn._pool_released = True
             await conn.close()
             self._size -= 1
             return
 
         if not await self._reset_connection(conn):
+            conn._pool_released = True
             with contextlib.suppress(Exception):
                 await conn.close()
             self._size -= 1
             return
 
+        conn._pool_released = True
         try:
             self._pool.put_nowait(conn)
         except asyncio.QueueFull:

tests/test_pool.py

@@ -640,6 +640,148 @@ async def failing_execute(sql, params=None):
 
         await pool.close()
 
+    async def test_escaped_reference_rejected_after_release(self) -> None:
+        """Using a connection after it's returned to the pool must raise InterfaceError."""
+        import asyncio
+
+        from dqliteclient.connection import DqliteConnection
+        from dqliteclient.exceptions import InterfaceError
+
+        pool = ConnectionPool(["localhost:9001"], max_size=1)
+
+        real_conn = DqliteConnection("127.0.0.1:9001")
+        real_conn._protocol = MagicMock()
+        real_conn._db_id = 1
+        real_conn._bound_loop = asyncio.get_running_loop()
+        real_conn._protocol.exec_sql = AsyncMock(return_value=(0, 1))
+        real_conn._protocol.query_sql = AsyncMock(return_value=(["id"], [[1]]))
+        real_conn._protocol.close = MagicMock()
+        real_conn._protocol.wait_closed = AsyncMock()
+
+        with patch.object(pool._cluster, "connect", return_value=real_conn):
+            await pool.initialize()
+
+        # Stash a reference to the connection
+        escaped = None
+        async with pool.acquire() as conn:
+            escaped = conn
+            await conn.execute("SELECT 1")
+
+        # Connection is now back in the pool — escaped reference must be rejected
+        assert escaped is not None
+        with pytest.raises(InterfaceError, match="returned to the pool"):
+            await escaped.execute("SELECT 1")
+
+        await pool.close()
+
+    async def test_escaped_reference_rejected_after_exception(self) -> None:
+        """Escaped reference must be rejected even when user code raises."""
+        import asyncio
+
+        from dqliteclient.connection import DqliteConnection
+        from dqliteclient.exceptions import InterfaceError
+
+        pool = ConnectionPool(["localhost:9001"], max_size=1)
+
+        real_conn = DqliteConnection("127.0.0.1:9001")
+        real_conn._protocol = MagicMock()
+        real_conn._db_id = 1
+        real_conn._bound_loop = asyncio.get_running_loop()
+        real_conn._protocol.exec_sql = AsyncMock(return_value=(0, 1))
+        real_conn._protocol.close = MagicMock()
+        real_conn._protocol.wait_closed = AsyncMock()
+
+        with patch.object(pool._cluster, "connect", return_value=real_conn):
+            await pool.initialize()
+
+        escaped = None
+        with pytest.raises(ValueError, match="app error"):
+            async with pool.acquire() as conn:
+                escaped = conn
+                raise ValueError("app error")
+
+        assert escaped is not None
+        with pytest.raises(InterfaceError, match="returned to the pool"):
+            await escaped.execute("SELECT 1")
+
+        await pool.close()
+
+    async def test_pool_release_rolls_back_transaction_with_real_connection(self) -> None:
+        """_reset_connection must be able to ROLLBACK before _pool_released is set."""
+        import asyncio
+
+        from dqliteclient.connection import DqliteConnection
+
+        pool = ConnectionPool(["localhost:9001"], max_size=1)
+
+        real_conn = DqliteConnection("127.0.0.1:9001")
+        real_conn._protocol = MagicMock()
+        real_conn._db_id = 1
+        real_conn._bound_loop = asyncio.get_running_loop()
+        real_conn._protocol.exec_sql = AsyncMock(return_value=(0, 0))
+        real_conn._protocol.close = MagicMock()
+        real_conn._protocol.wait_closed = AsyncMock()
+
+        with patch.object(pool._cluster, "connect", return_value=real_conn):
+            await pool.initialize()
+
+        # Simulate a connection with an open transaction returned to pool
+        async with pool.acquire() as conn:
+            conn._in_transaction = True
+
+        # The pool should have issued ROLLBACK successfully (not destroyed the conn)
+        assert not real_conn._in_transaction
+        # Connection should be back in the pool (not destroyed)
+        assert pool._pool.qsize() == 1
+        assert pool._size == 1
+
+        await pool.close()
+
+    async def test_escaped_reference_works_when_reacquired(self) -> None:
+        """A connection re-acquired from the pool must work normally."""
+        import asyncio
+
+        from dqliteclient.connection import DqliteConnection
+
+        pool = ConnectionPool(["localhost:9001"], max_size=1)
+
+        real_conn = DqliteConnection("127.0.0.1:9001")
+        real_conn._protocol = MagicMock()
+        real_conn._db_id = 1
+        real_conn._bound_loop = asyncio.get_running_loop()
+        real_conn._protocol.exec_sql = AsyncMock(return_value=(0, 1))
+        real_conn._protocol.close = MagicMock()
+        real_conn._protocol.wait_closed = AsyncMock()
+
+        with patch.object(pool._cluster, "connect", return_value=real_conn):
+            await pool.initialize()
+
+        # First acquire and release
+        async with pool.acquire() as conn:
+            await conn.execute("SELECT 1")
+
+        # Second acquire — same connection from pool must work
+        async with pool.acquire() as conn:
+            await conn.execute("SELECT 2")
+
+        await pool.close()
+
+    async def test_standalone_connection_not_affected_by_pool_guard(self) -> None:
+        """A DqliteConnection used standalone (not from a pool) must not be affected."""
+        import asyncio
+
+        from dqliteclient.connection import DqliteConnection
+
+        conn = DqliteConnection("127.0.0.1:9001")
+        conn._protocol = MagicMock()
+        conn._db_id = 1
+        conn._bound_loop = asyncio.get_running_loop()
+        conn._protocol.exec_sql = AsyncMock(return_value=(0, 1))
+
+        # Standalone connection — no pool involved, should work fine
+        await conn.execute("SELECT 1")
+        await conn.execute("SELECT 2")  # No error — _pool_released is always False
+
 
 class TestConnectionPoolIntegration:
     """Integration tests requiring mocked connections."""