Unpack.hs

module Codec.Archive.Internal.Unpack ( hsEntriesAbs
                                     , unpackEntriesFp
                                     , unpackArchive
                                     , readArchiveFile
                                     , readArchiveBS
                                     , archiveFile
                                     , bsToArchive
                                     , unpackToDir
                                     , readBS
                                     , readBSL
                                     , readEntry
                                     , readContents
                                     , getHsEntry
                                     , hsEntries
                                     , hsEntriesST
                                     , hsEntriesSTLazy
                                     , hsEntriesSTAbs
                                     ) where

import           Codec.Archive.Foreign
import           Codec.Archive.Internal.Monad
import           Codec.Archive.Types
import           Control.Monad                (when, (<=<))
import           Control.Monad.Except         (throwError)
import           Control.Monad.IO.Class       (liftIO)
import qualified Control.Monad.ST.Lazy        as LazyST
import qualified Control.Monad.ST.Lazy.Unsafe as LazyST
import           Data.Bifunctor               (first)
import qualified Data.ByteString              as BS
import qualified Data.ByteString.Lazy         as BSL
import           Data.Functor                 (void, ($>))
import           Foreign.C.String
import           Foreign.Concurrent           (newForeignPtr)
import           Foreign.ForeignPtr           (castForeignPtr, newForeignPtr_)
import           Foreign.Marshal.Alloc        (allocaBytes, free, mallocBytes)
import           Foreign.Marshal.Utils        (copyBytes)
import           Foreign.Ptr                  (castPtr, nullPtr)
import           System.FilePath              (normalise, splitDirectories, (</>))
import           System.IO.Unsafe             (unsafeDupablePerformIO)

-- | Read an archive contained in a 'BS.ByteString'. The format of the archive is
-- automatically detected.
--
-- @since 1.0.0.0
readArchiveBS :: BS.ByteString -> Either ArchiveResult [Entry FilePath BS.ByteString]
readArchiveBS = unsafeDupablePerformIO . runArchiveM . (go hsEntries <=< bsToArchive)
    where go f (y, act) = f y <* liftIO act
{-# NOINLINE readArchiveBS #-}

bsToArchive :: BS.ByteString -> ArchiveM (ArchivePtr, IO ())
bsToArchive bs = do
    preA <- liftIO archiveReadNew
    a <- liftIO $ castForeignPtr <$> newForeignPtr (castPtr preA) (void $ archiveFree preA)
    ignore $ archiveReadSupportFormatAll a
    bufPtr <- useAsCStringLenArchiveM bs $
        \(buf, sz) -> do
            buf' <- liftIO $ mallocBytes sz
            _ <- liftIO $ copyBytes buf' buf sz
            handle $ archiveReadOpenMemory a buf (fromIntegral sz)
            pure buf'
    pure (a, free bufPtr)

-- | Read an archive from a file. The format of the archive is automatically
-- detected.
--
-- @since 1.0.0.0
readArchiveFile :: FilePath -> ArchiveM [Entry FilePath BS.ByteString]
readArchiveFile fp = act =<< liftIO (do
    pre <- archiveReadNew
    castForeignPtr <$> newForeignPtr (castPtr pre) (void $ archiveFree pre))

    where act a =
            archiveFile fp a $> LazyST.runST (hsEntriesST a)

{-# INLINE archiveAbs #-}
archiveAbs :: (ArchivePtr -> IO ArchiveResult) -- ^ Function to set format support
           -> FilePath
           -> ArchivePtr
           -> ArchiveM ()
archiveAbs support fp a = withCStringArchiveM fp $ \cpath ->
    ignore (support a) *>
    handle (archiveReadOpenFilename a cpath 10240)

-- TODO: general function for format
archiveFile :: FilePath -> ArchivePtr -> ArchiveM ()
archiveFile = archiveAbs archiveReadSupportFormatAll

-- | This is more efficient than
--
-- @
-- unpackToDir "llvm" =<< BS.readFile "llvm.tar"
-- @
unpackArchive :: FilePath -- ^ Filepath pointing to archive
              -> FilePath -- ^ Dirctory to unpack in
              -> ArchiveM ()
unpackArchive tarFp dirFp = do
    preA <- liftIO archiveReadNew
    a <- liftIO $ castForeignPtr <$> newForeignPtr (castPtr preA) (void $ archiveFree preA)
    act a

    where act a =
            archiveFile tarFp a *>
            unpackEntriesFp a dirFp

readEntry :: Integral a
          => (ArchivePtr -> a -> IO e)
          -> ArchivePtr
          -> ArchiveEntryPtr
          -> IO (Entry FilePath e)
readEntry read' a entry =
    Entry
        <$> (peekCString =<< archiveEntryPathname entry)
        <*> readContents read' a entry
        <*> archiveEntryPerm entry
        <*> readOwnership entry
        <*> readTimes entry

-- | Yield the next entry in an archive
getHsEntry :: Integral a
           => (ArchivePtr -> a -> IO e)
           -> ArchivePtr
           -> IO (Maybe (Entry FilePath e))
getHsEntry read' a = do
    entry <- getEntry a
    case entry of
        Nothing -> pure Nothing
        Just x  -> Just <$> readEntry read' a x

-- | Return a list of 'Entry's.
hsEntries :: ArchivePtr -> ArchiveM [Entry FilePath BS.ByteString]
hsEntries = hsEntriesAbs readBS

hsEntriesAbs :: Integral a
             => (ArchivePtr -> a -> IO e)
             -> ArchivePtr
             -> ArchiveM [Entry FilePath e]
hsEntriesAbs read' p = pure (LazyST.runST $ hsEntriesSTAbs read' p)

-- | Return a list of 'Entry's.
hsEntriesST :: ArchivePtr -> LazyST.ST s [Entry FilePath BS.ByteString]
hsEntriesST = hsEntriesSTAbs readBS

hsEntriesSTLazy :: ArchivePtr -> LazyST.ST s [Entry FilePath BSL.ByteString]
hsEntriesSTLazy = hsEntriesSTAbs readBSL

hsEntriesSTAbs :: Integral a
               => (ArchivePtr -> a -> IO e)
               -> ArchivePtr
               -> LazyST.ST s [Entry FilePath e]
hsEntriesSTAbs read' a = do
    next <- LazyST.unsafeIOToST (getHsEntry read' a)
    case next of
        Nothing -> pure []
        Just x  -> (x:) <$> hsEntriesSTAbs read' a

abortIfSuspect :: FilePath -> ArchiveM ()
abortIfSuspect fp | ("..":_) <- splitDirectories (normalise fp) = throwError ArchiveFatal
                  | otherwise = pure ()

-- | Unpack an archive in a given directory.
--
-- Throws 'ArchiveFatal' if any entry's path begins with @".."@
unpackEntriesFp :: ArchivePtr -> FilePath -> ArchiveM ()
unpackEntriesFp a fp = do
    res <- liftIO $ getEntry a
    case res of
        Nothing -> pure ()
        Just x  -> do
            preFile <- liftIO $ archiveEntryPathname x
            file <- liftIO $ peekCString preFile
            abortIfSuspect file
            let file' = fp </> file
            liftIO $ withCString file' $ \fileC ->
                archiveEntrySetPathname x fileC
            ft <- liftIO $ archiveEntryFiletype x
            case ft of
                Just{} ->
                    ignore $ archiveReadExtract a x archiveExtractTime
                Nothing -> do
                    hardlink <- liftIO $ peekCString =<< archiveEntryHardlink x
                    let hardlink' = fp </> hardlink
                    liftIO $ withCString hardlink' $ \hl ->
                        archiveEntrySetHardlink x hl
                    ignore $ archiveReadExtract a x archiveExtractTime
            ignore $ archiveReadDataSkip a
            unpackEntriesFp a fp

{-# INLINE readBS #-}
readBS :: ArchivePtr -> Int -> IO BS.ByteString
readBS a sz =
    allocaBytes sz $ \buff ->
        archiveReadData a buff (fromIntegral sz) *>
        BS.packCStringLen (buff, sz)

-- TODO: sanity check by comparing to archiveEntrySize?
readBSL :: ArchivePtr -> Int -> IO BSL.ByteString
readBSL a _ = BSL.fromChunks <$> loop
    where step =
            allocaBytes bufSz $ \bufPtr -> do
                bRead <- archiveReadData a bufPtr (fromIntegral bufSz)
                if bRead == 0
                    then pure Nothing
                    else Just <$> BS.packCStringLen (bufPtr, fromIntegral bRead)

          loop = do
            res <- step
            case res of
                Just b  -> (b:) <$> loop
                Nothing -> pure []

          bufSz = 32 * 1024 -- read in 32k blocks

readContents :: Integral a
             => (ArchivePtr -> a -> IO e)
             -> ArchivePtr
             -> ArchiveEntryPtr
             -> IO (EntryContent FilePath e)
readContents read' a entry = go =<< archiveEntryFiletype entry
    where go Nothing            = Hardlink <$> (peekCString =<< archiveEntryHardlink entry)
          go (Just FtRegular)   = NormalFile <$> (read' a =<< sz)
          go (Just FtLink)      = Symlink <$> (peekCString =<< archiveEntrySymlink entry) <*> archiveEntrySymlinkType entry
          go (Just FtDirectory) = pure Directory
          go (Just _)           = error "Unsupported filetype"
          sz = fromIntegral <$> archiveEntrySize entry

archiveGetterHelper :: (ArchiveEntryPtr -> IO a) -> (ArchiveEntryPtr -> IO Bool) -> ArchiveEntryPtr -> IO (Maybe a)
archiveGetterHelper get check entry = do
    check' <- check entry
    if check'
        then Just <$> get entry
        else pure Nothing

archiveGetterNull :: (ArchiveEntryPtr -> IO CString) -> ArchiveEntryPtr -> IO (Maybe String)
archiveGetterNull get entry = do
    res <- get entry
    if res == nullPtr
        then pure Nothing
        else fmap Just (peekCString res)

readOwnership :: ArchiveEntryPtr -> IO Ownership
readOwnership entry =
    Ownership
        <$> archiveGetterNull archiveEntryUname entry
        <*> archiveGetterNull archiveEntryGname entry
        <*> (fromIntegral <$> archiveEntryUid entry)
        <*> (fromIntegral <$> archiveEntryGid entry)

readTimes :: ArchiveEntryPtr -> IO (Maybe ModTime)
readTimes = archiveGetterHelper go archiveEntryMtimeIsSet
    where go entry =
            (,) <$> archiveEntryMtime entry <*> archiveEntryMtimeNsec entry

-- | Get the next 'ArchiveEntry' in an 'Archive'
getEntry :: ArchivePtr -> IO (Maybe ArchiveEntryPtr)
getEntry a = do
    let done ArchiveOk    = False
        done ArchiveRetry = False
        done _            = True
    (stop, res) <- first done <$> archiveReadNextHeader a
    if stop
        then pure Nothing
        else Just . castForeignPtr <$> newForeignPtr_ (castPtr res)

-- | Throws 'ArchiveFatal' if any entry's path begins with @".."@
unpackToDir :: FilePath -- ^ Directory to unpack in
            -> BS.ByteString -- ^ 'BS.ByteString' containing archive
            -> ArchiveM ()
unpackToDir fp bs = do
    (a, act) <- bsToArchive bs
    unpackEntriesFp a fp
    liftIO act