From bc3a7f662e97f0211c5cc367abf68b12dac2d3f2 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 more than once
 (fixes #69)

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

diff --git a/ChangeLog.md b/ChangeLog.md
index f44b4a4..24bb133 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 more than once
+  [#70](https://github.com/fpco/typed-process/pull/70)
+
 ## 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 c94701a..e4212f6 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 #-}
@@ -130,13 +131,15 @@ module System.Process.Typed
     ) where
 
 import Control.Exception hiding (bracket, finally)
+import Control.Monad ((>=>), guard)
 import Control.Monad.IO.Class
 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,27 +242,18 @@ 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
+                  _ :: ExitCode <- Async.poll waitingThread >>= \ case
                       -- Process already exited, nothing to do
-                      Right _ec -> return ()
+                      Just r -> either throwIO return r
 
                       -- Process didn't exit yet, let's terminate it and
                       -- then call waitForProcess ourselves
-                      Left _ -> do
+                      Nothing -> do
                           terminateProcess pHandle
-                          ec <- P.waitForProcess pHandle
-                          success <- atomically $ tryPutTMVar pExitCode ec
-                          evaluate $ assert success ()
+                          waitForProcess
+                  return ()
 
           return Process {..}
   where
diff --git a/test/System/Process/TypedSpec.hs b/test/System/Process/TypedSpec.hs
index 45003f9..fec29be 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 more than once (fix for #69)" $ do
+            -- https://github.com/fpco/typed-process/issues/70
+            let config = setStdout createPipe (proc "echo" ["foo"])
+            withProcessWait config $ \p -> do
+              _ <- S.hGetContents (getStdout p)
+              throwIO DivideByZero
+            `shouldThrow` (== DivideByZero)