From 3fd7b028461c18f2c9fcfae6caacc3a4d6519f48 Mon Sep 17 00:00:00 2001
From: Simon Hengel <sol@typeful.net>
Date: Thu, 11 May 2023 16:22:34 +0700
Subject: [PATCH] Ensure that `waitForProcess` is never called multiple times
 (fixes #69)

---
 ChangeLog.md                     |  5 +++++
 package.yaml                     |  2 +-
 src/System/Process/Typed.hs      | 35 ++++++++++++--------------------
 test/System/Process/TypedSpec.hs | 10 +++++++++
 4 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/ChangeLog.md b/ChangeLog.md
index f44b4a4..ca59f25 100644
--- a/ChangeLog.md
+++ b/ChangeLog.md
@@ -1,5 +1,10 @@
 # ChangeLog for typed-process
 
+## 0.2.12.0
+
+* Ensure that `waitForProcess` is never called multiple times
+  [#69](https://github.com/fpco/typed-process/pull/69)
+
 ## 0.2.11.0
 
 * Expose more from `System.Process.Typed.Internal`
diff --git a/package.yaml b/package.yaml
index 5b6a7e5..0339e8a 100644
--- a/package.yaml
+++ b/package.yaml
@@ -1,5 +1,5 @@
 name:        typed-process
-version:     0.2.11.0
+version:     0.2.12.0
 synopsis:    Run external processes, with strong typing of streams
 description: Please see the tutorial at <https://github.com/fpco/typed-process#readme>
 category:    System
diff --git a/src/System/Process/Typed.hs b/src/System/Process/Typed.hs
index 731c0f9..53c8f9d 100644
--- a/src/System/Process/Typed.hs
+++ b/src/System/Process/Typed.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE DeriveDataTypeable #-}
 {-# LANGUAGE RecordWildCards #-}
@@ -135,8 +136,9 @@ import qualified System.Process as P
 import System.IO (hClose)
 import System.IO.Error (isPermissionError)
 import Control.Concurrent (threadDelay)
-import Control.Concurrent.Async (asyncWithUnmask, cancel, waitCatch)
-import Control.Concurrent.STM (newEmptyTMVarIO, atomically, putTMVar, TMVar, readTMVar, tryReadTMVar, STM, tryPutTMVar, throwSTM, catchSTM)
+import Control.Concurrent.Async (asyncWithUnmask)
+import qualified Control.Concurrent.Async as Async
+import Control.Concurrent.STM (newEmptyTMVarIO, atomically, putTMVar, TMVar, readTMVar, tryReadTMVar, STM, throwSTM, catchSTM)
 import System.Exit (ExitCode (ExitSuccess, ExitFailure))
 import System.Process.Typed.Internal
 import qualified Data.ByteString.Lazy as L
@@ -239,23 +241,12 @@ startProcess pConfig'@ProcessConfig {..} = liftIO $ do
               atomically $ putTMVar pExitCode ec
               return ec
 
+          let waitForProcess = Async.wait waitingThread :: IO ExitCode
+
           let pCleanup = pCleanup1 `finally` do
-                  -- First: stop calling waitForProcess, so that we can
-                  -- avoid race conditions where the process is removed from
-                  -- the system process table while we're trying to
-                  -- terminate it.
-                  cancel waitingThread
-
-                  -- Now check if the process had already exited
-                  eec <- waitCatch waitingThread
-
-                  case eec of
-                      -- Process already exited, nothing to do
-                      Right _ec -> return ()
-
-                      -- Process didn't exit yet, let's terminate it and
-                      -- then call waitForProcess ourselves
-                      Left _ -> do
+                  _ :: ExitCode <- Async.poll waitingThread >>= \ case
+                      Just r -> either throwIO return r
+                      Nothing -> do
                           eres <- try $ P.terminateProcess pHandle
                           ec <-
                             case eres of
@@ -272,11 +263,11 @@ startProcess pConfig'@ProcessConfig {..} = liftIO $ do
                                 -- Recommendation: always use the multi-threaded
                                 -- runtime!
                                 | isPermissionError e && not multiThreadedRuntime && isWindows ->
-                                  P.waitForProcess pHandle
+                                  waitForProcess
                                 | otherwise -> throwIO e
-                              Right () -> P.waitForProcess pHandle
-                          success <- atomically $ tryPutTMVar pExitCode ec
-                          evaluate $ assert success ()
+                              Right () -> waitForProcess
+                          return ec
+                  return ()
 
           return Process {..}
   where
diff --git a/test/System/Process/TypedSpec.hs b/test/System/Process/TypedSpec.hs
index 45003f9..0e8a015 100644
--- a/test/System/Process/TypedSpec.hs
+++ b/test/System/Process/TypedSpec.hs
@@ -5,6 +5,7 @@ module System.Process.TypedSpec (spec) where
 import System.Process.Typed
 import System.Process.Typed.Internal
 import System.IO
+import Control.Exception
 import Control.Concurrent.Async (Concurrently (..))
 import Control.Concurrent.STM (atomically)
 import Test.Hspec
@@ -170,3 +171,12 @@ spec = do
     it "empty param are showed" $
       let expected = "Raw command: podman exec --detach-keys \"\" ctx bash\n"
        in show (proc "podman" ["exec", "--detach-keys", "", "ctx", "bash"]) `shouldBe` expected
+
+    describe "stopProcess" $ do
+        it "never calls waitForProcess multiple times (fix for #69)" $ do
+            -- https://github.com/fpco/typed-process/issues/69
+            let config = setStdout createPipe (proc "echo" ["foo"])
+            withProcessWait config $ \p -> do
+              _ <- S.hGetContents (getStdout p)
+              throwIO DivideByZero
+            `shouldThrow` (== DivideByZero)