fradrive/src/Handler/Utils/Zip.hs

{-# OPTIONS_GHC -fno-warn-missing-fields #-} -- This concerns zipEntrySize in produceZip
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Handler.Utils.Zip
  ( typeZip, extensionZip
  , ZipError(..)
  , ZipInfo(..)
  , produceZip
  , consumeZip
  , modifyFileTitle
  , receiveFiles, acceptFile
  ) where

import Import.NoFoundation

import Handler.Utils.Files (acceptFile)
import Handler.Utils.DateTime (localTimeToUTCSimple, utcToLocalTime)

import Codec.Archive.Zip.Conduit.Types
import Codec.Archive.Zip.Conduit.UnZip
import Codec.Archive.Zip.Conduit.Zip

import System.FilePath

import Data.List (dropWhileEnd)

import qualified Data.Map as Map

import qualified Data.Conduit.List as C


import Data.Encoding ( decodeStrictByteStringExplicit
                     , encodeStrictByteStringExplicit
                     )
import Data.Encoding.CP437
import qualified Data.Char as Char

import Control.Monad.Trans.Cont
import Control.Monad.Trans.State.Strict (evalStateT)
import qualified Control.Monad.State.Class as State


typeZip :: ContentType
typeZip = "application/zip"

extensionZip :: Extension
extensionZip = fromMaybe "zip" $ listToMaybe [ ext | (ext, mime) <- Map.toList mimeMap, mime == typeZip ]



instance Default ZipInfo where
  def = ZipInfo
    { zipComment = mempty
    }


data ConsumeZipException
  = ConsumeZipUnZipException SomeException
  | ConsumeZipUnexpectedContent
  deriving (Show, Generic, Typeable)
  deriving anyclass (Exception)


consumeZip :: forall m m'.
              ( MonadThrow m
              , PrimMonad m
              , MonadUnliftIO m
              , MonadResource m
              , MonadIO m'
              , MonadThrow m'
              )
           => ConduitT () ByteString m () -> ConduitT () (File m') m ZipInfo
consumeZip inpBS = do
  inpChunk <- liftIO newEmptyTMVarIO
  zipAsync <- lift . allocateLinkedAsync $
    runConduit $ (inpBS .| unZipStream) `fuseUpstream` C.mapM_ (atomically . putTMVar inpChunk)

  flip evalStateT Nothing . evalContT . callCC $ \finishConsume -> forever $ do
    inpChunk' <- atomically $
          Right <$> takeTMVar inpChunk
      <|> Left  <$> waitCatchSTM zipAsync

    fileSink <- State.get
    case (fileSink, inpChunk') of
      (mFSink    , Left  (Left unzipExc)    ) -> do
        for_ mFSink $ \fSink' -> atomically $ do
          writeTMChan fSink' $ Left unzipExc
          closeTMChan fSink'
        throwM unzipExc
  
      (mFSink    , Left  (Right zInfo)      ) -> do
        for_ mFSink $ atomically . closeTMChan
        finishConsume zInfo
        
      (Just fSink, Right (Right bs)         ) ->
        atomically . writeTMChan fSink $ Right bs

      (Nothing   , Right (Right _)          ) ->
        throwM ConsumeZipUnexpectedContent

      (mFSink    , Right (Left ZipEntry{..})) -> do
        for_ mFSink $ atomically . closeTMChan
        State.put Nothing

        zipEntryName' <- decodeZipEntryName zipEntryName
        let
          fileTitle = "." <//> zipEntryName'
                    & normalise
                    & makeValid
                    & dropWhile isPathSeparator
                    & dropWhileEnd isPathSeparator
                    & normalise
                    & makeValid
          fileModified = localTimeToUTCSimple zipEntryTime
          isDirectory = hasTrailingPathSeparator zipEntryName'
        fileContent <- if
          | isDirectory -> return Nothing
          | otherwise   -> do
              fileChan <- liftIO newTMChanIO
              State.put $ Just fileChan
              return . Just . evalContT . callCC $ \finishFileContent -> forever $ do
                nextVal <- atomically $ asum
                  [ readTMChan fileChan
                  , do
                      inpChunk'' <- Right <$> takeTMVar inpChunk
                                <|> Left  <$> waitCatchSTM zipAsync
                      case inpChunk'' of
                        Left  (Left  unzipExc) -> return . Just $ Left unzipExc
                        Left  (Right _       ) -> return Nothing
                        Right (Left  zInfo   ) -> Nothing <$ putTMVar inpChunk (Left zInfo)
                        Right (Right bs      ) -> return . Just $ Right bs
                  ]
                case nextVal of
                  Nothing -> finishFileContent ()
                  Just (Right bs) -> lift $ yield bs
                  Just (Left exc) -> throwM $ ConsumeZipUnZipException exc
        lift . lift $ yield File{..}
  
produceZip :: forall m.
            ( MonadThrow m
            , PrimMonad m
            )
           => ZipInfo
           -> ConduitT (File m) ByteString m ()
produceZip info = C.map toZipData .| void (zipStream zipOptions)
  where
    zipOptions = ZipOptions
      { zipOpt64 = False
      , zipOptCompressLevel = defaultCompression
      , zipOptInfo = info
      }

    -- toZipData :: forall v. File m -> ConduitT v (ZipEntry, ZipData m) m ()
    -- toZipData f
    --   | Just fc <- fileContent f = do
    --       outpChunk <- newEmptyTMVarIO
    --       outpAsync <- lift . allocateLinkedAsync $
    --         runConduit $ fc .| C.mapM_ (atomically . putTMVar outpChunk)
    --       yield ( toZipEntry f
    --             , ZipDataSource . evalContT . callCC $ \finishContent -> forever $ do
    --                 nextVal <- atomically $
    --                       Right <$> takeTMVar outpChunk
    --                   <|> Left  <$> waitCatchSTM outpAsync
    --                 case nextVal of
    --                   Right  chunk      -> lift $ yield chunk
    --                   Left  (Right () ) -> finishContent ()
    --                   Left  (Left  exc) -> throwM exc
    --             )
    --   | otherwise = yield (toZipEntry f, mempty)

    toZipData :: File m -> (ZipEntry, ZipData m)
    toZipData f@File{..}
       = (toZipEntry f, maybe mempty ZipDataSource fileContent)

    toZipEntry :: File m -> ZipEntry
    toZipEntry File{..} = ZipEntry{..}
      where
        isDir = is _Nothing fileContent

        zipEntryName = "." <//> fileTitle
                     & normalise
                     & makeValid
                     & dropWhile isPathSeparator
                     & dropWhileEnd isPathSeparator
                     & bool id addTrailingPathSeparator isDir
                     & normalise
                     & makeValid
                     & encodeZipEntryName
        zipEntryTime = utcToLocalTime fileModified
        zipEntrySize = Nothing
        zipEntryExternalAttributes = Nothing

modifyFileTitle :: Monad m => (FilePath -> FilePath) -> ConduitT (File m') (File m') m ()
modifyFileTitle f = mapC $ \x@File{..} -> x{ fileTitle = f fileTitle }

-- Takes FileInfo and if it is a ZIP-Archive, extract files, otherwiese yield fileinfo
receiveFiles :: (MonadLogger m, MonadResource m, MonadThrow m, PrimMonad m, MonadUnliftIO m, MonadResource m', MonadThrow m') => FileInfo -> ConduitT () (File m') m ()
receiveFiles fInfo
    | ((==) `on` simpleContentType) mimeType typeZip = do
        $logInfoS "sourceFiles" "Unpacking ZIP"
        void . consumeZip $ fileSource fInfo
    | otherwise = do
        $logDebugS "sourceFiles" [st|Not unpacking file of type #{decodeUtf8 mimeType}|]
        yieldM $ acceptFile fInfo
  where
    mimeType = mimeLookup $ fileName fInfo


decodeZipEntryName :: MonadThrow m => Either Text ByteString -> m FilePath
-- ^ Extract the filename from a 'ZipEntry' doing decoding along the way.
--
-- Throws 'Data.Encoding.Exception.DecodingException's.
decodeZipEntryName = \case
  Left  t     -> return $ unpack t
  Right cp437 -> throwLeft $ decodeStrictByteStringExplicit CP437 cp437

encodeZipEntryName :: FilePath -> Either Text ByteString
-- ^ Encode a filename for use in a 'ZipEntry', encodes as
--   'Data.Encoding.UTF8.UTF8' iff the given path contains non-ascii characters.
--
-- Does not do any normalisation (in particular this function does not ensure
-- that the 'FilePath' does not start with a slash).
encodeZipEntryName path = fromMaybe (Left $ pack path) $ do
  guard $ all Char.isAscii path

  either (const mzero) (return . Right) $ encodeStrictByteStringExplicit CP437 path