RFC7253OCB.hs

-- RFC7253OCB.hs: RFC7253 implementation as a
--                workaround until crypton is fixed
-- Copyright © 2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).
{-# LANGUAGE PackageImports #-}

module Codec.Encryption.OpenPGP.Internal.RFC7253OCB
    ( encryptWithOCBRFC7253
    , decryptWithOCBRFC7253
    , decryptWithOCBRFC7253With
    ) where

import Control.Monad (when)
import Crypto.Number.Serialize (i2osp, os2ip)
import Data.Bits
    ( countTrailingZeros
    , shiftL
    , shiftR
    , xor
    , (.&.)
    , (.|.)
    )
import qualified Data.ByteArray as BA
import qualified Data.ByteString as B
import Data.List (foldl')
import Data.Word (Word8)
import qualified "crypton" Crypto.Cipher.Types as CCT

encryptWithOCBRFC7253
    :: CCT.BlockCipher c
    => c
    -> B.ByteString
    -> B.ByteString
    -> B.ByteString
    -> Either String (CCT.AuthTag, B.ByteString)
encryptWithOCBRFC7253 cipher nonce ad plaintext = do
    when (B.length nonce > 15 || B.null nonce) $
        Left "invalid nonce size for OCB"
    offset0 <- ocbOffset0 cipher nonce
    let zeroBlock = B.replicate 16 0
        lStar = CCT.ecbEncrypt cipher zeroBlock
        lDollar = ocbDouble lStar
        lCache = iterate ocbDouble (ocbDouble lDollar)
        hashAd = ocbHash cipher lStar lCache ad
        (fullBlocks, partial) = splitFullAndPartial plaintext
        (cipherBlocks, offsetM, checksum) =
            foldl'
                ( \(accBlocks, offsetPrev, checksumPrev) (idx, pBlock) ->
                    let offsetI = xorBS offsetPrev (lCache !! ntz idx)
                        cipherI = xorBS offsetI (CCT.ecbEncrypt cipher (xorBS offsetI pBlock))
                        checksumI = xorBS checksumPrev pBlock
                     in (accBlocks ++ [cipherI], offsetI, checksumI)
                )
                ([], offset0, zeroBlock)
                (zip [1 ..] fullBlocks)
        (cipherLast, offsetLast, checksumLast) =
            if B.null partial
                then (B.empty, offsetM, checksum)
                else
                    let offsetStar = xorBS offsetM lStar
                        pad = CCT.ecbEncrypt cipher offsetStar
                        cipherPartial = xorBS partial (B.take (B.length partial) pad)
                        checksum' = xorBS checksum (ocbPadPartial partial)
                     in (cipherPartial, offsetStar, checksum')
        tagBytes =
            xorBS
                ( CCT.ecbEncrypt
                    cipher
                    (xorBS (xorBS checksumLast offsetLast) lDollar)
                )
                hashAd
        ciphertext = B.concat cipherBlocks <> cipherLast
     in Right (mkAuthTag (B.take 16 tagBytes), ciphertext)

decryptWithOCBRFC7253
    :: CCT.BlockCipher c
    => c
    -> B.ByteString
    -> B.ByteString
    -> B.ByteString
    -> CCT.AuthTag
    -> Either String B.ByteString
decryptWithOCBRFC7253 =
    decryptWithOCBRFC7253With
        (\_ _ _ _ _ _ -> "OCB authentication failed")

decryptWithOCBRFC7253With
    :: CCT.BlockCipher c
    => ( B.ByteString
         -> B.ByteString
         -> B.ByteString
         -> B.ByteString
         -> B.ByteString
         -> B.ByteString
         -> String
       )
    -> c
    -> B.ByteString
    -> B.ByteString
    -> B.ByteString
    -> CCT.AuthTag
    -> Either String B.ByteString
decryptWithOCBRFC7253With onAuthFailure cipher nonce ad ciphertext authTag = do
    when (B.length nonce > 15 || B.null nonce) $
        Left "invalid nonce size for OCB"
    when (B.length tagBytes /= 16) $
        Left "invalid auth tag size for OCB"
    offset0 <- ocbOffset0 cipher nonce
    let zeroBlock = B.replicate 16 0
        lStar = CCT.ecbEncrypt cipher zeroBlock
        lDollar = ocbDouble lStar
        lCache = iterate ocbDouble (ocbDouble lDollar)
        hashAd = ocbHash cipher lStar lCache ad
        (fullBlocks, partial) = splitFullAndPartial ciphertext
        (plainBlocks, offsetM, checksum) =
            foldl'
                ( \(accBlocks, offsetPrev, checksumPrev) (idx, cBlock) ->
                    let offsetI = xorBS offsetPrev (lCache !! ntz idx)
                        plainI = xorBS offsetI (CCT.ecbDecrypt cipher (xorBS offsetI cBlock))
                        checksumI = xorBS checksumPrev plainI
                     in (accBlocks ++ [plainI], offsetI, checksumI)
                )
                ([], offset0, zeroBlock)
                (zip [1 ..] fullBlocks)
        (plainLast, offsetLast, checksumLast) =
            if B.null partial
                then (B.empty, offsetM, checksum)
                else
                    let offsetStar = xorBS offsetM lStar
                        pad = CCT.ecbEncrypt cipher offsetStar
                        plainPartial = xorBS partial (B.take (B.length partial) pad)
                        checksum' = xorBS checksum (ocbPadPartial plainPartial)
                     in (plainPartial, offsetStar, checksum')
        tagComputed =
            xorBS
                ( CCT.ecbEncrypt
                    cipher
                    (xorBS (xorBS checksumLast offsetLast) lDollar)
                )
                hashAd
        plaintext = B.concat plainBlocks <> plainLast
        computedTag = B.take 16 tagComputed
     in if BA.constEq tagBytes computedTag
            then Right plaintext
            else
                Left
                    (onAuthFailure tagBytes computedTag nonce ad hashAd plaintext)
  where
    tagBytes = BA.convert authTag :: B.ByteString

mkAuthTag :: B.ByteString -> CCT.AuthTag
mkAuthTag = CCT.AuthTag . BA.convert

ocbOffset0
    :: CCT.BlockCipher c
    => c -> B.ByteString -> Either String B.ByteString
ocbOffset0 cipher nonce = do
    let nonceLen = B.length nonce
        prefixLen = 16 - nonceLen
    when (prefixLen <= 0) $
        Left "invalid nonce size for OCB"
    let prefix = B.pack (replicate (prefixLen - 1) 0 <> [1 :: Word8])
        nonceBlock = prefix <> nonce
        bottom = fromIntegral (B.last nonceBlock .&. 0x3f) :: Int
        nonceTop = B.init nonceBlock <> B.singleton (B.last nonceBlock .&. 0xc0)
        kTop = CCT.ecbEncrypt cipher nonceTop
        stretch = kTop <> xorBS (B.take 8 kTop) (B.take 8 (B.drop 1 kTop))
    Right (ocbBitSlice128 stretch bottom)

ocbHash
    :: CCT.BlockCipher c
    => c
    -> B.ByteString
    -> [B.ByteString]
    -> B.ByteString
    -> B.ByteString
ocbHash cipher lStar lCache ad =
    let (fullBlocks, partial) = splitFullAndPartial ad
        (sumBlocks, offsetFinal) =
            foldl'
                ( \(acc, offsetPrev) (idx, block) ->
                    let offsetI = xorBS offsetPrev (lCache !! ntz idx)
                        sumI = xorBS acc (CCT.ecbEncrypt cipher (xorBS offsetI block))
                     in (sumI, offsetI)
                )
                (B.replicate 16 0, B.replicate 16 0)
                (zip [1 ..] fullBlocks)
     in if B.null partial
            then sumBlocks
            else
                let offsetStar = xorBS offsetFinal lStar
                    block = ocbPadPartial partial
                 in xorBS sumBlocks (CCT.ecbEncrypt cipher (xorBS offsetStar block))

ocbPadPartial :: B.ByteString -> B.ByteString
ocbPadPartial bs = bs <> B.singleton 0x80 <> B.replicate (15 - B.length bs) 0

ocbDouble :: B.ByteString -> B.ByteString
ocbDouble bs =
    let shifted = shiftLeftOne bs
        carry = (B.head bs .&. 0x80) /= 0
     in if carry
            then B.init shifted <> B.singleton (B.last shifted `xor` 0x87)
            else shifted

shiftLeftOne :: B.ByteString -> B.ByteString
shiftLeftOne bs = B.pack shifted
  where
    (shifted, _) =
        foldl'
            ( \(acc, carryIn) x ->
                let y = ((x `shiftL` 1) .&. 0xff) .|. carryIn
                    carryOut = if (x .&. 0x80) /= 0 then 1 else 0
                 in (y : acc, carryOut)
            )
            ([], 0)
            (reverse (B.unpack bs))

xorBS :: B.ByteString -> B.ByteString -> B.ByteString
xorBS a b = B.pack (B.zipWith xor a b)

ocbBitSlice128 :: B.ByteString -> Int -> B.ByteString
ocbBitSlice128 stretch startBit =
    let stretchInt = os2ip stretch
        shift = 192 - (startBit + 128)
        mask = (1 `shiftL` (128 :: Int)) - 1
        slice = (stretchInt `shiftR` shift) .&. mask
     in leftPadTo16 (i2osp slice)

leftPadTo16 :: B.ByteString -> B.ByteString
leftPadTo16 bs
    | B.length bs >= 16 = B.drop (B.length bs - 16) bs
    | otherwise = B.replicate (16 - B.length bs) 0 <> bs

splitFullAndPartial
    :: B.ByteString -> ([B.ByteString], B.ByteString)
splitFullAndPartial bs
    | B.null bs = ([], B.empty)
    | otherwise =
        let fullLen = (B.length bs `div` 16) * 16
            (fullPart, rest) = B.splitAt fullLen bs
         in (chunk16 fullPart, rest)

chunk16 :: B.ByteString -> [B.ByteString]
chunk16 bs
    | B.null bs = []
    | otherwise =
        let (h, t) = B.splitAt 16 bs
         in h : chunk16 t

ntz :: Int -> Int
ntz i = fromIntegral (countTrailingZeros (fromIntegral i :: Word))