From 2b5937ec0ae88cd0b4cc0c8534f21c435ee94662 Mon Sep 17 00:00:00 2001
From: "Miss Islington (bot)"
 <31488909+miss-islington@users.noreply.github.com>
Date: Fri, 2 Feb 2018 15:14:38 -0800
Subject: [PATCH] bpo-32734: Fix asyncio.Lock multiple acquire safety issue
 (GH-5466) (#5501)

(cherry picked from commit 2f79c014931cbb23b08a7d16c534a3cc9607ae14)

Co-authored-by: Bar Harel <bzvi7919@gmail.com>
---
 Lib/asyncio/locks.py                          | 32 ++++++++----
 Lib/test/test_asyncio/test_locks.py           | 50 +++++++++++++++++++
 Misc/ACKS                                     |  1 +
 .../2018-02-01-01-34-47.bpo-32734.gCV9AD.rst  |  2 +
 4 files changed, 75 insertions(+), 10 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2018-02-01-01-34-47.bpo-32734.gCV9AD.rst

diff --git a/Lib/asyncio/locks.py b/Lib/asyncio/locks.py
index 61938373509..508a2142d84 100644
--- a/Lib/asyncio/locks.py
+++ b/Lib/asyncio/locks.py
@@ -183,16 +183,22 @@ class Lock(_ContextManagerMixin):
 
         fut = self._loop.create_future()
         self._waiters.append(fut)
+
+        # Finally block should be called before the CancelledError
+        # handling as we don't want CancelledError to call
+        # _wake_up_first() and attempt to wake up itself.
         try:
-            await fut
-            self._locked = True
-            return True
+            try:
+                await fut
+            finally:
+                self._waiters.remove(fut)
         except futures.CancelledError:
             if not self._locked:
                 self._wake_up_first()
             raise
-        finally:
-            self._waiters.remove(fut)
+
+        self._locked = True
+        return True
 
     def release(self):
         """Release a lock.
@@ -212,11 +218,17 @@ class Lock(_ContextManagerMixin):
             raise RuntimeError('Lock is not acquired.')
 
     def _wake_up_first(self):
-        """Wake up the first waiter who isn't cancelled."""
-        for fut in self._waiters:
-            if not fut.done():
-                fut.set_result(True)
-                break
+        """Wake up the first waiter if it isn't done."""
+        try:
+            fut = next(iter(self._waiters))
+        except StopIteration:
+            return
+
+        # .done() necessarily means that a waiter will wake up later on and
+        # either take the lock, or, if it was cancelled and lock wasn't
+        # taken already, will hit this again and wake up a new waiter.
+        if not fut.done():
+            fut.set_result(True)
 
 
 class Event:
diff --git a/Lib/test/test_asyncio/test_locks.py b/Lib/test/test_asyncio/test_locks.py
index 3c5069778b5..3e3dd799273 100644
--- a/Lib/test/test_asyncio/test_locks.py
+++ b/Lib/test/test_asyncio/test_locks.py
@@ -200,6 +200,56 @@ class LockTests(test_utils.TestCase):
         self.assertTrue(tb.cancelled())
         self.assertTrue(tc.done())
 
+    def test_cancel_release_race(self):
+        # Issue 32734
+        # Acquire 4 locks, cancel second, release first
+        # and 2 locks are taken at once.
+        lock = asyncio.Lock(loop=self.loop)
+        lock_count = 0
+        call_count = 0
+
+        async def lockit():
+            nonlocal lock_count
+            nonlocal call_count
+            call_count += 1
+            await lock.acquire()
+            lock_count += 1
+  
+        async def lockandtrigger():
+            await lock.acquire()
+            self.loop.call_soon(trigger)
+          
+        def trigger():
+            t1.cancel()
+            lock.release()
+
+        t0 = self.loop.create_task(lockandtrigger())
+        t1 = self.loop.create_task(lockit())
+        t2 = self.loop.create_task(lockit())
+        t3 = self.loop.create_task(lockit())
+
+        # First loop acquires all
+        test_utils.run_briefly(self.loop)
+        self.assertTrue(t0.done())
+
+        # Second loop calls trigger
+        test_utils.run_briefly(self.loop)
+        # Third loop calls cancellation
+        test_utils.run_briefly(self.loop)
+
+        # Make sure only one lock was taken
+        self.assertEqual(lock_count, 1)
+        # While 3 calls were made to lockit()
+        self.assertEqual(call_count, 3)
+        self.assertTrue(t1.cancelled() and t2.done())
+
+        # Cleanup the task that is stuck on acquire.
+        t3.cancel()
+        test_utils.run_briefly(self.loop)
+        self.assertTrue(t3.cancelled())
+
+
+
     def test_finished_waiter_cancelled(self):
         lock = asyncio.Lock(loop=self.loop)
 
diff --git a/Misc/ACKS b/Misc/ACKS
index 3fdfa3041f8..c5eadc5ba01 100644
--- a/Misc/ACKS
+++ b/Misc/ACKS
@@ -602,6 +602,7 @@ Milton L. Hankins
 Stephen Hansen
 Barry Hantman
 Lynda Hardman
+Bar Harel
 Derek Harland
 Jason Harper
 David Harrigan
diff --git a/Misc/NEWS.d/next/Library/2018-02-01-01-34-47.bpo-32734.gCV9AD.rst b/Misc/NEWS.d/next/Library/2018-02-01-01-34-47.bpo-32734.gCV9AD.rst
new file mode 100644
index 00000000000..14d4bbdade7
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2018-02-01-01-34-47.bpo-32734.gCV9AD.rst
@@ -0,0 +1,2 @@
+Fixed ``asyncio.Lock()`` safety issue which allowed acquiring and locking
+the same lock multiple times, without it being free. Patch by Bar Harel.