Make openFile exception safe wrt #8

System/File/OsPath.hs

@@ -1,9 +1,32 @@
-module System.File.OsPath where
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+
+module System.File.OsPath (
+  openBinaryFile
+, withFile
+, withBinaryFile
+, withFile'
+, withBinaryFile'
+, readFile
+, readFile'
+, writeFile
+, writeFile'
+, appendFile
+, appendFile'
+, openFile
+, openExistingFile
+) where
 
 import qualified System.File.Platform as P
 
-import Control.Exception (bracket)
-import System.IO (IOMode(..), Handle, hSetBinaryMode, hClose)
+import Control.Concurrent.MVar
+import Control.Exception.Base
+import Control.Monad (void)
+import GHC.IO.Handle (hClose_help)
+import GHC.IO.Handle.Internals (debugIO, hClose_impl)
+import GHC.IO.Handle.Types (Handle__, Handle(..))
+import Prelude (IO, FilePath, Bool(..), (++), ($), pure, return, show)
+import System.IO (IOMode(..), hSetBinaryMode)
 import System.OsPath
 import System.OsString.Internal.Types
 
@@ -31,16 +54,20 @@ openBinaryFile fp iomode = do
 --
 -- The 'Handle' is automatically closed afther the action.
 withFile :: OsPath -> IOMode -> (Handle -> IO r) -> IO r
-withFile fp iomode action = bracket
-  (openFile fp iomode)
-  hClose
-  action
+withFile fp iomode action = mask $ \restore -> do
+  hndl <- openFile fp iomode
+  addHandleFinalizer hndl handleFinalizer
+  r <- restore (action hndl) `onException` hClose_impl hndl
+  hClose_impl hndl
+  pure r
 
 withBinaryFile :: OsPath -> IOMode -> (Handle -> IO r) -> IO r
-withBinaryFile fp iomode action = bracket
-  (openBinaryFile fp iomode)
-  hClose
-  action
+withBinaryFile fp iomode action = mask $ \restore -> do
+  hndl <- openBinaryFile fp iomode
+  addHandleFinalizer hndl handleFinalizer
+  r <- restore (action hndl) `onException` hClose_impl hndl
+  hClose_impl hndl
+  pure r
 
 -- | Run an action on a file.
 --
@@ -98,3 +125,30 @@ openFile (OsString fp) = P.openFile fp
 -- | Open an existing file and return the 'Handle'.
 openExistingFile :: OsPath -> IOMode -> IO Handle
 openExistingFile (OsString fp) = 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)
+