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
  ( CompressionError(..)
  , renderCompressionError
  , 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 Codec.Encryption.OpenPGP.Serialize ()
import Codec.Encryption.OpenPGP.Types
import Data.Binary (get, put)
import Data.Binary.Get (runGetOrFail)
import Data.Binary.Put (runPut)
import qualified Data.ByteString.Lazy as BL

-- | Errors that can arise during decompression of an OpenPGP Compressed Data
-- packet.  Note that corrupt stream exceptions from the underlying
-- zlib\/bzip2 library are not captured here; they propagate as 'IOException'
-- through the call stack.
data CompressionError
  = -- | The compressed payload bytes are empty; nothing to decompress.
    EmptyCompressedPayload CompressionAlgorithm
  | -- | Decompression succeeded but the binary parse of the inner packet
    -- sequence failed.
    InnerPacketParseFailed CompressionAlgorithm String
  | -- | Decompression and parse succeeded but produced no packets at all
    -- (zero-length uncompressed payload).
    ZeroLengthDecompressedPayload CompressionAlgorithm
  | -- | The decompressed content contains only Marker packets (RFC4880 §5.8
    -- marker packets carry no semantic content).
    MarkerOnlyPayload CompressionAlgorithm
  deriving (Eq, Show)

renderCompressionError :: CompressionError -> String
renderCompressionError (EmptyCompressedPayload algo) =
  "Compressed Data packet (" ++ show algo ++ "): empty compressed payload"
renderCompressionError (InnerPacketParseFailed algo err) =
  "Compressed Data packet (" ++ show algo ++ "): inner packet parse failed: " ++ err
renderCompressionError (ZeroLengthDecompressedPayload algo) =
  "Compressed Data packet (" ++ show algo ++ "): zero-length decompressed payload"
renderCompressionError (MarkerOnlyPayload algo) =
  "Compressed Data packet (" ++ show algo ++ "): decompressed content contains only Marker packets"

-- | Decompress an OpenPGP Compressed Data packet, classifying structural
-- failures as 'CompressionError'.  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