TempDir.hs

{-# LANGUAGE ForeignFunctionInterface,ViewPatterns,RecordWildCards #-}
-- Various routines for dealing with temporary directories and files.
module Support.TempDir(
    getTempDir,
    createTempFile,
    fileInTempDir,
    cleanTempDir,
    setTempDir,
    addAtExit,
    withStackStatus,
    wrapMain
   ) where

import Control.Exception as E
import Control.Monad
import Data.IORef
import Data.Maybe
import System.Directory
import System.Exit
import System.FilePath as FP
import System.IO
import System.IO.Unsafe
import System.IO.Temp
import Text.Printf
import qualified Data.Set as Set
import GenUtil (iocatch)
import Support.CompatMingw32

data TempDir = TempDir {
    tempDirClean   :: Bool,  -- ^ whether to delete the directory afterwords.
    tempDirDump    :: Bool,
    tempDirPath    :: Maybe String,
    tempDirAtExit  :: [IO ()],
    tempDirCleanup :: Set.Set FilePath
    }

putLog :: String -> IO ()
--putLog = putStrLn
putLog _ = return ()

cleanTempDir :: Bool -> IO ()
cleanTempDir b = modifyIORef tdRef $ \x -> x { tempDirClean = b }

setTempDir :: FilePath -> IO ()
setTempDir (FP.normalise -> fp) = do
    TempDir {..} <- readIORef tdRef
    when (isJust $ tempDirPath) $ do
        fail $ printf
          "Cannot set temp directory to '%s'; it is already set to '%s'."
          fp (fromJust tempDirPath)
    putLog $ printf "Setting work directory to '%s'" fp
    createDirectoryIfMissing False fp
    writeIORef tdRef TempDir { tempDirPath = Just fp,  .. }
    cleanTempDir False

getTempDir :: IO FilePath
getTempDir = do
    td <- readIORef tdRef
    case tempDirPath td of
        Just fp -> return fp
        Nothing -> do
            tmpdir <- getTemporaryDirectory
            fp <- createTempDirectory tmpdir "jhc_"
            putLog $ printf "Created work directory '%s'" fp
            writeIORef tdRef td { tempDirPath = Just fp }
            return fp

addAtExit :: IO () -> IO ()
addAtExit action = do
    td <- readIORef tdRef
    writeIORef tdRef td { tempDirAtExit = action:tempDirAtExit td }

createTempFile :: FilePath -> IO (FilePath, Handle)
createTempFile (FP.normalise -> fp) = do
    unless (filePathSafe fp) $
        fail $ "createTempFile: unsafe path " ++ fp
    dir <- getTempDir
    (fp,h) <- openBinaryTempFile dir (if null fp then "temp.tmp" else fp)
    putLog $ printf "Created temporary file '%s'" fp
    addCleanup fp
    return (fp,h)

-- make sure nothing is sneaky about the file path
filePathSafe fp = FP.isRelative fp &&
        ".." `notElem` FP.splitPath fp && not (hasDrive fp)

fileInTempDir :: FilePath -> (FilePath -> IO ()) -> IO FilePath
fileInTempDir (FP.normalise -> fp) action = do
    unless (filePathSafe fp) $
        fail $ "fileinTempDir: unsafe path " ++ fp
    let (FP.normalise -> dpart,_) = FP.splitFileName fp
    tdir <- getTempDir
    let f ("./":ps) cp = f ps cp
        f (".":ps) cp = f ps cp
        f (p:ps) cp = do
            putLog $ printf "Creating directory '%s' '%s' '%s' '%s' '%s'" tdir cp p dpart fp
            createDirectoryIfMissing False (tdir </> cp </> p)
            let cp' = FP.normalise (cp </> p)
            addCleanup cp'
            f ps cp'
        f [] _ = return ()
    f (FP.splitPath dpart) ""
    --unless (null $ FP.normalise dpart) $
    --    fold (FP.splitPath dpart) $ addCleanup
    --    createDirectoryIfMissing True (tdir </> dpart)
    let nfp = FP.normalise (tdir </> fp)
    b <- addCleanup fp
    when b $ action nfp
    return $ noEscapePath nfp

cleanUp :: IO ()
cleanUp = do
    td <- readIORef tdRef
    sequence_ (tempDirAtExit td)
    if not (tempDirClean td) ||
        isNothing (tempDirPath td) then return () else do
    dir <- getTempDir
    forM_ (reverse . Set.toList $ tempDirCleanup td) $ \fp -> do
        putLog $ printf "Removing '%s'" (dir </> fp)
        ignoreError (removeDirectory $ dir </> fp)
        ignoreError (removeFile $ dir </> fp)
    putLog $ printf "Removing '%s'" dir
    ignoreError (removeDirectory dir)

addCleanup :: FilePath -> IO Bool
addCleanup fp = do
    td <- readIORef tdRef
    if fp `Set.member` tempDirCleanup td then return False else do
    writeIORef tdRef td { tempDirCleanup = fp `Set.insert` tempDirCleanup td }
    return True

wrapMain :: IO () -> IO ()
wrapMain main = E.catch (main >> cleanUp) f where
    panic = raiseSigIntCompat
    f (fromException -> Just code) = cleanUp >> exitWith code
    f (fromException -> Just UserInterrupt) = cleanUp >> panic
    f e = do
        ss <- readIORef stackRef
        td <- readIORef tdRef
        case tempDirPath td of
            Just td -> hPutStrLn stderr $
                printf "Exiting abnormally. Work directory is '%s'" td
            _ -> return ()
        unless (null ss) $
            forM_ ("Stack:":ss) (hPutStrLn stderr)
        throwIO e

-------------------
-- support routines
-------------------

ignoreError :: IO () -> IO ()
ignoreError action = iocatch action (\_ -> return ())

{-# NOINLINE tdRef #-}
tdRef :: IORef TempDir
tdRef = unsafePerformIO $ newIORef TempDir {
    tempDirClean   = True,
    tempDirDump    = False,
    tempDirPath    = Nothing,
    tempDirAtExit  = [],
    tempDirCleanup = Set.empty
    }

{-# NOINLINE stackRef #-}
stackRef :: IORef [String]
stackRef = unsafePerformIO $ newIORef []

withStackStatus :: String -> IO a -> IO a
withStackStatus s action = do
    cs <- readIORef stackRef
    writeIORef stackRef (s:cs)
    r <- action
    writeIORef stackRef cs
    return r