From 23685aa5cc9a9cd062717d909a6b66e0892f8321 Mon Sep 17 00:00:00 2001
From: Edsko de Vries <edsko@well-typed.com>
Date: Fri, 19 Jul 2024 12:13:47 +0200
Subject: [PATCH] Fix treatment of async exceptions

In https://github.com/kazu-yamamoto/http2/pull/92 we added an exception handler
that was meant to catch _all_ exceptions (sync and async). This got changed in
https://github.com/kazu-yamamoto/http2/pull/114 (specifically,
https://github.com/kazu-yamamoto/http2/pull/114/commits/52a9619ba95b67d469205cb0dea546ada8489baa):
when we moved from `Control.Exception` to `UnliftIO.Exception`, we got a
different behaviour for `catch` and friends (see
https://github.com/well-typed/grapesy/issues/193#issuecomment-2238704595) for a
full list. This commit fixes some unintended consequences of this change.
---
 Network/HTTP2/Client/Run.hs | 16 ++++++----------
 Network/HTTP2/H2/Manager.hs | 15 +++++++++++----
 Network/HTTP2/Server/Run.hs | 10 ++--------
 3 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/Network/HTTP2/Client/Run.hs b/Network/HTTP2/Client/Run.hs
index 0db39014..4b466d18 100644
--- a/Network/HTTP2/Client/Run.hs
+++ b/Network/HTTP2/Client/Run.hs
@@ -140,16 +140,8 @@ setup ClientConfig{..} conf@Config{..} = do
 
 runH2 :: Config -> Context -> IO a -> IO a
 runH2 conf ctx runClient = do
-    stopAfter mgr (race runBackgroundThreads runClient) $ \res -> do
-        closeAllStreams (oddStreamTable ctx) (evenStreamTable ctx) $
-            either Just (const Nothing) res
-        case res of
-            Left err ->
-                throwIO err
-            Right (Left ()) ->
-                undefined -- never reach
-            Right (Right x) ->
-                return x
+    stopAfter mgr (clientResult <$> race runBackgroundThreads runClient) $ \res ->
+        closeAllStreams (oddStreamTable ctx) (evenStreamTable ctx) res
   where
     mgr = threadManager ctx
     runReceiver = frameReceiver ctx conf
@@ -158,6 +150,10 @@ runH2 conf ctx runClient = do
         labelMe "H2 runBackgroundThreads"
         concurrently_ runReceiver runSender
 
+    clientResult :: Either () a -> a
+    clientResult (Left ()) = undefined -- unreachable
+    clientResult (Right a) = a
+
 sendRequest
     :: Config
     -> Context
diff --git a/Network/HTTP2/H2/Manager.hs b/Network/HTTP2/H2/Manager.hs
index a60117d8..ce4b7ea4 100644
--- a/Network/HTTP2/H2/Manager.hs
+++ b/Network/HTTP2/H2/Manager.hs
@@ -58,12 +58,19 @@ start timmgr = do
                  in go q tset
 
 -- | Stopping the manager.
-stopAfter :: Manager -> IO a -> (Either SomeException a -> IO b) -> IO b
+--
+-- The action is run in the scope of an exception handler that catches all
+-- exceptions (including asynchronous ones); this allows the cleanup handler
+-- to cleanup in all circumstances. If an exception is caught, it is rethrown
+-- after the cleanup is complete.
+stopAfter :: Manager -> IO a -> (Maybe SomeException -> IO ()) -> IO a
 stopAfter (Manager q _ _) action cleanup = do
     mask $ \unmask -> do
-        ma <- try $ unmask action
+        ma <- trySyncOrAsync $ unmask action
         atomically $ writeTQueue q $ Stop (either Just (const Nothing) ma)
-        cleanup ma
+        case ma of
+          Left err -> cleanup (Just err) >> throwIO err
+          Right a  -> cleanup Nothing    >> return a
 
 ----------------------------------------------------------------
 
@@ -86,7 +93,7 @@ forkManagedUnmask mgr label io =
         incCounter mgr
         -- We catch the exception and do not rethrow it: we don't want the
         -- exception printed to stderr.
-        io unmask `catch` \(_e :: SomeException) -> return ()
+        io unmask `catchSyncOrAsync` \(_e :: SomeException) -> return ()
         deleteMyId mgr
         decCounter mgr
   where
diff --git a/Network/HTTP2/Server/Run.hs b/Network/HTTP2/Server/Run.hs
index 7d4d64e5..f833a446 100644
--- a/Network/HTTP2/Server/Run.hs
+++ b/Network/HTTP2/Server/Run.hs
@@ -129,14 +129,8 @@ runH2 conf ctx = do
         runReceiver = frameReceiver ctx conf
         runSender = frameSender ctx conf
         runBackgroundThreads = concurrently_ runReceiver runSender
-    stopAfter mgr runBackgroundThreads $ \res -> do
-        closeAllStreams (oddStreamTable ctx) (evenStreamTable ctx) $
-            either Just (const Nothing) res
-        case res of
-            Left err ->
-                throwIO err
-            Right x ->
-                return x
+    stopAfter mgr runBackgroundThreads $ \res ->
+        closeAllStreams (oddStreamTable ctx) (evenStreamTable ctx) res
 
 -- connClose must not be called here since Run:fork calls it
 goaway :: Config -> ErrorCode -> ByteString -> IO ()