diff --git a/Codec/Archive/Zip/Conduit/Zip.hs b/Codec/Archive/Zip/Conduit/Zip.hs --- a/Codec/Archive/Zip/Conduit/Zip.hs +++ b/Codec/Archive/Zip/Conduit/Zip.hs @@ -1,6 +1,7 @@ -- |Stream the creation of a zip file, e.g., as it's being uploaded. {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ViewPatterns #-} +{-# LANGUAGE BangPatterns #-} module Codec.Archive.Zip.Conduit.Zip ( zipStream , ZipOptions(..) @@ -13,6 +14,7 @@ import qualified Codec.Compression.Zlib.Raw as Z import Control.Arrow ((&&&), (+++), left) +import Control.DeepSeq (force) import Control.Monad (when) import Control.Monad.Catch (MonadThrow) import Control.Monad.Primitive (PrimMonad) @@ -35,16 +37,16 @@ import qualified Data.Text as T import qualified Data.Text.Encoding as TE import Data.Time (LocalTime(..), TimeOfDay(..), toGregorian) -import Data.Word (Word16, Word64) +import Data.Word (Word16, Word32, Word64) import Codec.Archive.Zip.Conduit.Types import Codec.Archive.Zip.Conduit.Internal -- |Options controlling zip file parameters and features data ZipOptions = ZipOptions - { zipOpt64 :: Bool -- ^Allow 'ZipDataSource's over 4GB (reduces compatibility in some cases); this is automatically enabled for any files of known size (e.g., 'zipEntrySize') - , zipOptCompressLevel :: Int -- ^Compress zipped files (0 = store only, 1 = minimal, 9 = best; non-zero improves compatibility, since some unzip programs don't supported stored, streamed files, including the one in this package) - , zipOptInfo :: ZipInfo -- ^Other parameters to store in the zip file + { zipOpt64 :: !Bool -- ^Allow 'ZipDataSource's over 4GB (reduces compatibility in some cases); this is automatically enabled for any files of known size (e.g., 'zipEntrySize') + , zipOptCompressLevel :: !Int -- ^Compress zipped files (0 = store only, 1 = minimal, 9 = best; non-zero improves compatibility, since some unzip programs don't supported stored, streamed files, including the one in this package) + , zipOptInfo :: !ZipInfo -- ^Other parameters to store in the zip file } defaultZipOptions :: ZipOptions @@ -88,6 +90,67 @@ maxBound16 :: Integral n => n maxBound16 = fromIntegral (maxBound :: Word16) +data CommonFileHeaderInfo = CommonFileHeaderInfo + { cfhiIsStreamingEntry :: !Bool + , cfhiHasUtf8Filename :: !Bool + , cfhiIsCompressed :: !Bool + , cfhiTime :: !Word16 + , cfhiDate :: !Word16 + } deriving (Eq, Ord, Show) + +putCommonFileHeaderPart :: CommonFileHeaderInfo -> P.PutM () +putCommonFileHeaderPart CommonFileHeaderInfo{..} = do + P.putWord16le $ cfhiIsStreamingEntry ?* bit 3 .|. cfhiHasUtf8Filename ?* bit 11 + P.putWord16le $ cfhiIsCompressed ?* 8 + P.putWord16le $ cfhiTime + P.putWord16le $ cfhiDate + +-- | This is retained in memory until the end of the archive is written. +-- +-- To avoid space leaks, this should contain only strict data. +data CentralDirectoryInfo = CentralDirectoryInfo + { cdiOff :: !Word64 + , cdiZ64 :: !Bool + , cdiCommonFileHeaderInfo :: !CommonFileHeaderInfo + , cdiCrc :: !Word32 + , cdiUsz :: !Word64 + , cdiName :: !BSC.ByteString + , cdiCsz :: !Word64 + , cdiZipEntryExternalAttributes :: !(Maybe Word32) -- lazy Maybe must be e.g. via `force` at creation + } deriving (Eq, Ord, Show) + +putCentralDirectory :: CentralDirectoryInfo -> P.PutM () +putCentralDirectory CentralDirectoryInfo{..} = do + -- central directory + let o64 = cdiOff >= maxBound32 + l64 = cdiZ64 ?* 16 + o64 ?* 8 + a64 = cdiZ64 || o64 + P.putWord32le 0x02014b50 + P.putWord8 zipVersion + P.putWord8 osVersion + P.putWord8 $ if a64 then 45 else 20 + P.putWord8 osVersion + putCommonFileHeaderPart cdiCommonFileHeaderInfo + P.putWord32le cdiCrc + P.putWord32le $ if cdiZ64 then maxBound32 else fromIntegral cdiCsz + P.putWord32le $ if cdiZ64 then maxBound32 else fromIntegral cdiUsz + P.putWord16le $ fromIntegral (BS.length cdiName) + P.putWord16le $ a64 ?* (4 + l64) + P.putWord16le 0 -- comment length + P.putWord16le 0 -- disk number + P.putWord16le 0 -- internal file attributes + P.putWord32le $ fromMaybe 0 cdiZipEntryExternalAttributes + P.putWord32le $ if o64 then maxBound32 else fromIntegral cdiOff + P.putByteString cdiName + when a64 $ do + P.putWord16le 0x0001 + P.putWord16le l64 + when cdiZ64 $ do + P.putWord64le cdiUsz + P.putWord64le cdiCsz + when o64 $ + P.putWord64le cdiOff + -- |Stream produce a zip file, reading a sequence of entries with data. -- Although file data is never kept in memory (beyond a single 'ZipDataByteString'), the format of zip files requires producing a final directory of entries at the end of the file, consuming an additional ~100 bytes of state per entry during streaming. -- The final result is the total size of the zip file. @@ -112,92 +175,73 @@ next (succ cnt) $ dir >> d) entry (ZipEntry{..}, zipData -> dat) = do let usiz = dataSize dat - sdat = left ((C..| sizeCRC) . C.toProducer) dat - comp = zipOptCompressLevel /= 0 + sdat = left (\x -> C.toProducer x C..| sizeCRC) dat + cfhiIsCompressed = zipOptCompressLevel /= 0 && all (0 /=) usiz && all (0 /=) zipEntrySize + cfhiIsStreamingEntry = isLeft dat + compressPlainBs = + Z.compressWith + Z.defaultCompressParams + { Z.compressLevel = + if zipOptCompressLevel == -1 + then Z.defaultCompression + else Z.compressionLevel zipOptCompressLevel + } (cdat, csiz) - | comp = + | cfhiIsCompressed = ( ((`C.fuseBoth` (outputSize $ CZ.compress zipOptCompressLevel deflateWindowBits)) - +++ Z.compress) sdat -- level for Z.compress? + +++ compressPlainBs) sdat , dataSize cdat) | otherwise = (left (fmap (id &&& fst)) sdat, usiz) - z64 = maybe (zipOpt64 || any (maxBound32 <) zipEntrySize) + cdiZ64 = maybe (zipOpt64 || any (maxBound32 <) zipEntrySize) (maxBound32 <) (max <$> usiz <*> csiz) - name = either TE.encodeUtf8 id zipEntryName - namelen = BS.length name - (time, date) = toDOSTime zipEntryTime + cfhiHasUtf8Filename = isLeft zipEntryName + cdiName = either TE.encodeUtf8 id zipEntryName + namelen = BS.length cdiName + (cfhiTime, cfhiDate) = toDOSTime zipEntryTime mcrc = either (const Nothing) (Just . crc32) dat + !cdiCommonFileHeaderInfo = CommonFileHeaderInfo{..} when (namelen > maxBound16) $ zipError $ either T.unpack BSC.unpack zipEntryName ++ ": entry name too long" - let common = do - P.putWord16le $ isLeft dat ?* bit 3 .|. isLeft zipEntryName ?* bit 11 - P.putWord16le $ comp ?* 8 - P.putWord16le $ time - P.putWord16le $ date - off <- get + cdiOff <- get output $ do P.putWord32le 0x04034b50 - P.putWord8 $ if z64 then 45 else 20 + P.putWord8 $ if cdiZ64 then 45 else 20 P.putWord8 osVersion - common + putCommonFileHeaderPart cdiCommonFileHeaderInfo P.putWord32le $ fromMaybe 0 mcrc - P.putWord32le $ if z64 then maxBound32 else maybe 0 fromIntegral csiz - P.putWord32le $ if z64 then maxBound32 else maybe 0 fromIntegral usiz + P.putWord32le $ if cdiZ64 then maxBound32 else maybe 0 fromIntegral csiz + P.putWord32le $ if cdiZ64 then maxBound32 else maybe 0 fromIntegral usiz P.putWord16le $ fromIntegral namelen - P.putWord16le $ z64 ?* 20 - P.putByteString name - when z64 $ do + P.putWord16le $ cdiZ64 ?* 20 + P.putByteString cdiName + when cdiZ64 $ do P.putWord16le 0x0001 P.putWord16le 16 P.putWord64le $ fromMaybe 0 usiz P.putWord64le $ fromMaybe 0 csiz - let outsz c = stateC $ \o -> (id &&& (o +) . snd) <$> c - ((usz, crc), csz) <- either + let outsz c = stateC $ \(!o) -> (id &&& (o +) . snd) <$> c + ((cdiUsz, cdiCrc), cdiCsz) <- either (\cd -> do r@((usz, crc), csz) <- outsz cd -- write compressed data - when (not z64 && (usz > maxBound32 || csz > maxBound32)) $ zipError $ either T.unpack BSC.unpack zipEntryName ++ ": file too large and zipOpt64 disabled" + when (not cdiZ64 && (usz > maxBound32 || csz > maxBound32)) $ zipError $ either T.unpack BSC.unpack zipEntryName ++ ": file too large and zipOpt64 disabled" output $ do P.putWord32le 0x08074b50 P.putWord32le crc let putsz - | z64 = P.putWord64le + | cdiZ64 = P.putWord64le | otherwise = P.putWord32le . fromIntegral putsz csz putsz usz return r) (\b -> outsz $ ((fromJust usiz, fromJust mcrc), fromJust csiz) <$ CB.sourceLbs b) cdat - when (any (usz /=) zipEntrySize) $ zipError $ either T.unpack BSC.unpack zipEntryName ++ ": incorrect zipEntrySize" - return $ do - -- central directory - let o64 = off >= maxBound32 - l64 = z64 ?* 16 + o64 ?* 8 - a64 = z64 || o64 - P.putWord32le 0x02014b50 - P.putWord8 zipVersion - P.putWord8 osVersion - P.putWord8 $ if a64 then 45 else 20 - P.putWord8 osVersion - common - P.putWord32le crc - P.putWord32le $ if z64 then maxBound32 else fromIntegral csz - P.putWord32le $ if z64 then maxBound32 else fromIntegral usz - P.putWord16le $ fromIntegral namelen - P.putWord16le $ a64 ?* (4 + l64) - P.putWord16le 0 -- comment length - P.putWord16le 0 -- disk number - P.putWord16le 0 -- internal file attributes - P.putWord32le $ fromMaybe 0 zipEntryExternalAttributes - P.putWord32le $ if o64 then maxBound32 else fromIntegral off - P.putByteString name - when a64 $ do - P.putWord16le 0x0001 - P.putWord16le l64 - when z64 $ do - P.putWord64le usz - P.putWord64le csz - when o64 $ - P.putWord64le off + when (any (cdiUsz /=) zipEntrySize) $ zipError $ either T.unpack BSC.unpack zipEntryName ++ ": incorrect zipEntrySize" + let !centralDirectoryInfo = CentralDirectoryInfo + { cdiZipEntryExternalAttributes = force zipEntryExternalAttributes + , .. } + return $ putCentralDirectory centralDirectoryInfo + endDirectory cdoff cdlen cnt = do let z64 = zipOpt64 || cdoff > maxBound32 || cnt > maxBound16 when z64 $ output $ do diff --git a/tests/Main.hs b/tests/Main.hs new file mode 100644 --- /dev/null +++ b/tests/Main.hs @@ -0,0 +1,59 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE OverloadedStrings #-} + +module Main (main) where + +import Control.Monad (when, void) +import Control.Monad.IO.Class (liftIO) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Lazy as BSL +import qualified Data.Conduit as C +import Data.Conduit.Combinators (sinkNull) +import Data.Foldable (for_) +import qualified Data.Text as T +import Data.Time.LocalTime (utc, utcToLocalTime) +import Data.Time.Clock.POSIX (posixSecondsToUTCTime) +import GHC.Stats (getRTSStats, RTSStats(..), GCDetails(..)) +import System.Mem (performMajorGC) +import Test.Hspec (hspec, describe, it) + +import Codec.Archive.Zip.Conduit.Zip + + +main :: IO () +main = hspec $ do + describe "zipping" $ do + it "ZipDataByteString streams in constant memory" $ do + C.runConduitRes $ + (do + -- Stream 1000 * 4 MiB = 4 GiB + for_ [(1::Int)..1024] $ \i -> do + -- `bs` needs to depend on loop variable `i`, otherwise GHC may hoist + -- it out of the loop ("floating"), making the memory constant + -- even for incorrect implementations, thus making the test useless. + let !bs = BS.replicate (4 * 1024 * 1024) (fromIntegral i) -- 4 MiB + + C.yield + ( ZipEntry + { zipEntryName = Left ("file-" <> T.pack (show i) <> ".bin") + , zipEntryTime = utcToLocalTime utc (posixSecondsToUTCTime 0) + , zipEntrySize = Nothing + , zipEntryExternalAttributes = Nothing + } + , ZipDataByteString (BSL.fromStrict bs) -- `copy` to avoid sharing + ) + + liftIO $ do + -- GC every 40 MB to make it easy to observe constant memory. + when (i `mod` 10 == 0) performMajorGC + + RTSStats{ gc = GCDetails{ gcdetails_live_bytes } } <- getRTSStats + when (gcdetails_live_bytes > 3 * 1024 * 1024 * 1024) $ do -- 3 GiB + error $ "Memory usage too high (" ++ show gcdetails_live_bytes ++ " B), probably streaming is not constant-memory" + + ) + C..| void (zipStream defaultZipOptions{ zipOptCompressLevel = 0 }) + C..| sinkNull + :: IO () + diff --git a/zip-stream.cabal b/zip-stream.cabal --- a/zip-stream.cabal +++ b/zip-stream.cabal @@ -1,5 +1,5 @@ name: zip-stream -version: 0.2.1.0 +version: 0.2.2.0 synopsis: ZIP archive streaming using conduits description: Process (extract and create) zip files as streams (e.g., over the network), accessing contained files without having to write the zip file to disk (unlike zip-conduit). license: BSD3 @@ -31,6 +31,7 @@ bytestring, conduit >= 1.3, conduit-extra, + deepseq, digest, exceptions, mtl, @@ -75,3 +76,19 @@ time, transformers, zip-stream + +test-suite tests + type: exitcode-stdio-1.0 + hs-source-dirs: tests + default-language: Haskell2010 + main-is: Main.hs + -- `-T` is needed for https://hackage.haskell.org/package/base-4.17.0.0/docs/GHC-Stats.html + ghc-options: -Wall -threaded -with-rtsopts=-T + build-depends: + base + , zip-stream + , bytestring + , conduit + , hspec + , text + , time