From 7b16d96d138e59292a7095e8a53af5c08c816142 Mon Sep 17 00:00:00 2001
From: Julian Ospald <hasufell@posteo.de>
Date: Sun, 14 Jan 2024 18:28:55 +0800
Subject: [PATCH] Make openFile exception safe wrt #8

---
 System/File/OsPath.hs | 60 ++++++++++++++++++++++++++++++++++---------
 1 file changed, 48 insertions(+), 12 deletions(-)

diff --git a/System/File/OsPath.hs b/System/File/OsPath.hs
index 56a4e16..fa66c95 100644
--- a/System/File/OsPath.hs
+++ b/System/File/OsPath.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE BangPatterns #-}
 
 module System.File.OsPath (
   openBinaryFile
@@ -16,16 +17,21 @@ module System.File.OsPath (
 , openExistingFile
 ) where
 
+
 import qualified System.File.Platform as P
 
-import Prelude ((.), ($), String, IO, pure, either, const, flip, Maybe(..), fmap, (<$>), id, ioError, (=<<), Bool(..))
+import Prelude ((.), ($), String, IO, ioError, pure, either, const, flip, Maybe(..), fmap, (<$>), id, (=<<), Bool(..), FilePath, (++), return, show, (>>=))
 import GHC.IO (catchException)
 import GHC.IO.Exception (IOException(..))
+import GHC.IO.Handle (hClose_help)
+import GHC.IO.Handle.Internals (debugIO)
+import GHC.IO.Handle.Types (Handle__, Handle(..))
+import Control.Concurrent.MVar
+import Control.Monad (void)
 import Control.DeepSeq (force)
-import Control.Exception (SomeException, try, evaluate, bracket)
-import System.IO (IOMode(..), Handle)
+import Control.Exception (SomeException, try, evaluate, mask, onException)
+import System.IO (IOMode(..), hSetBinaryMode, hClose)
 import System.IO.Unsafe (unsafePerformIO)
-import System.IO (hSetBinaryMode, hClose)
 import System.OsPath as OSP
 import System.OsString.Internal.Types
 
@@ -56,16 +62,10 @@ openBinaryFile' (OsString fp) iomode =do
 --
 -- The 'Handle' is automatically closed afther the action.
 withFile :: OsPath -> IOMode -> (Handle -> IO r) -> IO r
-withFile osfp@(OsString fp) iomode action = either ioError pure =<< (augmentError "withFile" osfp $ bracket
-  (P.openFile fp iomode)
-  hClose
-  (try . action))
+withFile osfp iomode act = (augmentError "withFile" osfp $ withOpenFile' openFile' osfp iomode (try . act)) >>= either ioError pure
 
 withBinaryFile :: OsPath -> IOMode -> (Handle -> IO r) -> IO r
-withBinaryFile osfp iomode action = either ioError pure =<< (augmentError "withBinaryFile" osfp $ bracket
-  (openBinaryFile' osfp iomode)
-  hClose
-  (try . action))
+withBinaryFile osfp iomode act = (augmentError "withBinaryFile" osfp $ withOpenFile' openBinaryFile' osfp iomode (try . act)) >>= either ioError pure
 
 -- | Run an action on a file.
 --
@@ -120,10 +120,46 @@ appendFile' fp contents = withFile fp AppendMode (`BS.hPut` contents)
 openFile :: OsPath -> IOMode -> IO Handle
 openFile osfp@(OsString fp) = augmentError "openFile" osfp . P.openFile fp
 
+openFile' :: OsPath -> IOMode -> IO Handle
+openFile' (OsString fp) = P.openFile fp
+
 -- | Open an existing file and return the 'Handle'.
 openExistingFile :: OsPath -> IOMode -> IO Handle
 openExistingFile osfp@(OsString fp) = augmentError "openExistingFile" osfp . P.openExistingFile fp
 
+
+-- ---------------------------------------------------------------------------
+-- Internals
+
+handleFinalizer :: FilePath -> MVar Handle__ -> IO ()
+handleFinalizer _fp m = do
+  handle_ <- takeMVar m
+  (handle_', _) <- hClose_help handle_
+  putMVar m handle_'
+  return ()
+
+type HandleFinalizer = FilePath -> MVar Handle__ -> IO ()
+
+-- | Add a finalizer to a 'Handle'. Specifically, the finalizer
+-- will be added to the 'MVar' of a file handle or the write-side
+-- 'MVar' of a duplex handle. See Handle Finalizers for details.
+addHandleFinalizer :: Handle -> HandleFinalizer -> IO ()
+addHandleFinalizer hndl finalizer = do
+  debugIO $ "Registering finalizer: " ++ show filepath
+  void $ mkWeakMVar mv (finalizer filepath mv)
+  where
+    !(filepath, !mv) = case hndl of
+      FileHandle fp m -> (fp, m)
+      DuplexHandle fp _ write_m -> (fp, write_m)
+
+withOpenFile' :: (OsPath -> IOMode -> IO Handle) -> OsPath -> IOMode -> (Handle -> IO r) -> IO r
+withOpenFile' acquire fp iomode action = mask $ \restore -> do
+  hndl <- acquire fp iomode
+  addHandleFinalizer hndl handleFinalizer
+  r <- restore (action hndl) `onException` hClose hndl
+  hClose hndl
+  pure r
+
 addFilePathToIOError :: String -> OsPath -> IOException -> IOException
 addFilePathToIOError fun fp ioe = unsafePerformIO $ do
   fp'  <- either (const (fmap OSP.toChar . OSP.unpack $ fp)) id <$> try @SomeException (OSP.decodeFS fp)