Compression.hs

-- Compression.hs: OpenPGP (RFC9580) compression and decompression
-- Copyright © 2012-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).

module Codec.Encryption.OpenPGP.Compression
    ( decompressPkt
    , compressPkts
    ) where

import qualified Codec.Compression.BZip as BZip
import qualified Codec.Compression.Zlib as Zlib
import qualified Codec.Compression.Zlib.Raw as ZlibRaw
import Data.Binary (get, put)
import Data.Binary.Get (runGetOrFail)
import Data.Binary.Put (runPut)
import qualified Data.ByteString.Lazy as BL

import Codec.Encryption.OpenPGP.Serialize ()
import Codec.Encryption.OpenPGP.Types
import Codec.Encryption.OpenPGP.Types.Internal.Errors
    ( CompressionError (..)
    , renderCompressionError
    )

{- | Decompress an OpenPGP Compressed Data packet, classifying structural
failures as 'CompressionErrorReason.  Non-Compressed-Data packets are returned
unchanged in @Right [p]@.  Corrupt compressed streams may still throw
'IOException' from the underlying decompression library.
-}
decompressPkt :: Pkt -> Either CompressionError [Pkt]
decompressPkt compressed@(CompressedDataPkt (OtherCA _) _) = Right [compressed]
decompressPkt (CompressedDataPkt algo bs)
    | BL.null bs = Left (EmptyCompressedPayload algo)
    | otherwise =
        case runGetOrFail get (dfunc algo bs) of
            Left (_, _, err) -> Left (InnerPacketParseFailed algo err)
            Right (_, _, packs) ->
                let pkts = unBlock packs
                 in case pkts of
                        [] -> Left (ZeroLengthDecompressedPayload algo)
                        _ | all isMarkerPkt pkts -> Left (MarkerOnlyPayload algo)
                        _ -> Right pkts
  where
    dfunc Uncompressed = id
    dfunc ZIP = ZlibRaw.decompress
    dfunc ZLIB = Zlib.decompress
    dfunc BZip2 = BZip.decompress
    dfunc (OtherCA _) = id
decompressPkt p = Right [p]

isMarkerPkt :: Pkt -> Bool
isMarkerPkt (MarkerPkt _) = True
isMarkerPkt _ = False

compressPkts :: CompressionAlgorithm -> [Pkt] -> Pkt
compressPkts ca packs =
    let bs = runPut $ put (Block packs)
        cbs = cfunc ca bs
        outAlgo =
            if isSupportedCompressionAlgorithm ca then ca else Uncompressed
     in CompressedDataPkt outAlgo cbs
  where
    cfunc Uncompressed = id
    cfunc ZIP = ZlibRaw.compress
    cfunc ZLIB = Zlib.compress
    cfunc BZip2 = BZip.compress
    cfunc _ = id

isSupportedCompressionAlgorithm :: CompressionAlgorithm -> Bool
isSupportedCompressionAlgorithm Uncompressed = True
isSupportedCompressionAlgorithm ZIP = True
isSupportedCompressionAlgorithm ZLIB = True
isSupportedCompressionAlgorithm BZip2 = True
isSupportedCompressionAlgorithm (OtherCA _) = False