{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ViewPatterns #-}

-- |
-- Module: Crypto.MAC.Poly1305
-- Copyright: (c) 2025 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- A pure Poly1305 MAC implementation, as specified by
-- [RFC 8439](https://datatracker.ietf.org/doc/html/rfc8439).

module Crypto.MAC.Poly1305 (
    -- * Poly1305 message authentication code
    mac

    -- testing
  , _poly1305_loop
  , _roll
  ) where

import Data.Bits ((.&.), (.|.), (.<<.), (.>>.))
import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BI

fi :: (Integral a, Num b) => a -> b
fi :: forall a b. (Integral a, Num b) => a -> b
fi = a -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral
{-# INLINE fi #-}

-- arbitrary-size little-endian bytestring decoding
_roll :: BS.ByteString -> Integer
_roll :: ByteString -> Integer
_roll = (Word8 -> Integer -> Integer) -> Integer -> ByteString -> Integer
forall a. (Word8 -> a -> a) -> a -> ByteString -> a
BS.foldr Word8 -> Integer -> Integer
forall {a} {a}. (Integral a, Bits a, Num a) => a -> a -> a
alg Integer
0 where
  alg :: a -> a -> a
alg (a -> a
forall a b. (Integral a, Num b) => a -> b
fi -> !a
b) !a
a = (a
a a -> Int -> a
forall a. Bits a => a -> Int -> a
.<<. Int
8) a -> a -> a
forall a. Bits a => a -> a -> a
.|. a
b
{-# INLINE _roll #-}

-- little-endian bytestring encoding
unroll :: Integer -> BS.ByteString
unroll :: Integer -> ByteString
unroll Integer
i = case Integer
i of
    Integer
0 -> Word8 -> ByteString
BS.singleton Word8
0
    Integer
_ -> (Integer -> Maybe (Word8, Integer)) -> Integer -> ByteString
forall a. (a -> Maybe (Word8, a)) -> a -> ByteString
BS.unfoldr Integer -> Maybe (Word8, Integer)
coalg Integer
i
  where
    coalg :: Integer -> Maybe (Word8, Integer)
coalg = \case
      Integer
0 -> Maybe (Word8, Integer)
forall a. Maybe a
Nothing
      Integer
m -> (Word8, Integer) -> Maybe (Word8, Integer)
forall a. a -> Maybe a
Just ((Word8, Integer) -> Maybe (Word8, Integer))
-> (Word8, Integer) -> Maybe (Word8, Integer)
forall a b. (a -> b) -> a -> b
$! (Integer -> Word8
forall a b. (Integral a, Num b) => a -> b
fi Integer
m, Integer
m Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
.>>. Int
8)
{-# INLINE unroll #-}

-- little-endian bytestring encoding for 128-bit ints, right-padding
-- with zeros
unroll16 :: Integer -> BS.ByteString
unroll16 :: Integer -> ByteString
unroll16 (Integer -> ByteString
unroll -> u :: ByteString
u@(BI.PS ForeignPtr Word8
_ Int
_ Int
l))
  | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
16 = ByteString
u ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Int -> Word8 -> ByteString
BS.replicate (Int
16 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l) Word8
0
  | Bool
otherwise = ByteString
u
{-# INLINE unroll16 #-}

clamp :: Integer -> Integer
clamp :: Integer -> Integer
clamp Integer
r = Integer
r Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. Integer
0x0ffffffc0ffffffc0ffffffc0fffffff
{-# INLINE clamp #-}

-- | Produce a Poly1305 MAC for the provided message, given the provided
--   key.
--
--   Per RFC8439: the key, which is essentially a /one-time/ key, should
--   be unique, and MUST be unpredictable for each invocation.
--
--   The key must be exactly 256 bits in length. Providing an invalid
--   key will cause the function to throw an ErrorCall exception.
--
--   >>> mac "i'll never use this key again!!!" "a message needing authentication"
--   "O'\231Z\224\149\148\246\203[}\210\203\b\200\207"
mac
  :: BS.ByteString -- ^ 256-bit one-time key
  -> BS.ByteString -- ^ arbitrary-length message
  -> BS.ByteString -- ^ 128-bit message authentication code
mac :: ByteString -> ByteString -> ByteString
mac key :: ByteString
key@(BI.PS ForeignPtr Word8
_ Int
_ Int
kl) ByteString
msg
  | Int
kl Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
32  = [Char] -> ByteString
forall a. HasCallStack => [Char] -> a
error [Char]
"ppad-poly1305 (mac): invalid key"
  | Bool
otherwise =
      let (Integer -> Integer
clamp (Integer -> Integer)
-> (ByteString -> Integer) -> ByteString -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Integer
_roll -> Integer
r, ByteString -> Integer
_roll -> Integer
s) = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
16 ByteString
key
      in  Integer -> Integer -> ByteString -> ByteString
_poly1305_loop Integer
r Integer
s ByteString
msg

_poly1305_loop :: Integer -> Integer -> BS.ByteString -> BS.ByteString
_poly1305_loop :: Integer -> Integer -> ByteString -> ByteString
_poly1305_loop !Integer
r !Integer
s !ByteString
msg =
    let loop :: Integer -> ByteString -> ByteString
loop !Integer
acc !ByteString
bs = case Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
16 ByteString
bs of
          (chunk :: ByteString
chunk@(BI.PS ForeignPtr Word8
_ Int
_ Int
l), ByteString
etc)
            | Int
l Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 -> Int -> ByteString -> ByteString
BS.take Int
16 (Integer -> ByteString
unroll16 (Integer
acc Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
s))
            | Bool
otherwise ->
                let !n :: Integer
n = ByteString -> Integer
_roll ByteString
chunk Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.|. (Integer
0x01 Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
.<<. (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
l))
                    !nacc :: Integer
nacc = Integer
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Integer
acc Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
n) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`rem` Integer
p
                in  Integer -> ByteString -> ByteString
loop Integer
nacc ByteString
etc
    in  Integer -> ByteString -> ByteString
loop Integer
0 ByteString
msg
  where
    p :: Integer
p = Integer
1361129467683753853853498429727072845819 -- (1 << 130) - 5
{-# INLINE _poly1305_loop #-}