Prim.hs

{-# OPTIONS_HADDOCK hide #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}

-- |
-- Module: Lightning.Protocol.BOLT4.Prim
-- Copyright: (c) 2025 Jared Tobin
-- License: MIT
-- Maintainer: Jared Tobin <jared@ppad.tech>
--
-- Keys, secrets and the cryptographic primitives of BOLT #4.

module Lightning.Protocol.BOLT4.Prim (
  -- * Secret keys
    SecretKey
  , secret_key
  , public_key
  , sk_bytes
  , sk_pub

  -- * Shared secrets
  , SharedSecret
  , shared_secret
  , un_shared_secret
  , ecdh

  -- * Derived keys
  , DerivedKey
  , un_derived_key
  , derive_rho
  , derive_mu
  , derive_um
  , derive_ammag
  , derive_pad

  -- * Blinding
  , blinding_factor
  , blinded_node_tweak
  , blind_pub
  , blind_scalar

  -- * Points
  , to_point
  , from_point

  -- * Streams and MACs
  , keystream
  , hmac
  , ct_eq
  , xor_bytes
  ) where

import Control.DeepSeq (NFData(..))
import qualified Crypto.Cipher.ChaCha20 as ChaCha
import qualified Crypto.Curve.Secp256k1 as Secp256k1
import qualified Crypto.Hash.SHA256 as SHA256
import Data.Bits (xor)
import qualified Data.ByteString as BS
import qualified Lightning.Protocol.BOLT1 as BOLT1
import qualified Numeric.Montgomery.Secp256k1.Scalar as S

-- secret keys ----------------------------------------------------------------

-- | A secp256k1 secret key: a session key, a node's private key, or the
--   seed of a blinded path.
--
--   This is secret material: its 'Show' instance is redacted and it
--   has no 'Eq' instance.
data SecretKey = SecretKey
  !BS.ByteString         -- the key
  !Secp256k1.Projective  -- its public key
  !BOLT1.Point           -- its serialized public key

instance Show SecretKey where
  showsPrec d _ = showParen (d > 10) $
    showString "SecretKey <redacted>"

instance NFData SecretKey where
  rnf (SecretKey k p q) = rnf k `seq` p `seq` rnf q

-- | Construct a t'SecretKey' from exactly 32 big-endian bytes encoding an
--   integer in [1, n - 1], where n is the secp256k1 group order.
--
--   >>> secret_key (BS.replicate 32 0x41)
--   Just (SecretKey <redacted>)
--   >>> secret_key (BS.replicate 32 0x00)
--   Nothing
--   >>> secret_key (BS.replicate 31 0x41)
--   Nothing
secret_key :: BS.ByteString -> Maybe SecretKey
secret_key bs
  | BS.length bs /= 32 = Nothing
  | otherwise = do
      pub <- Secp256k1.derive_pub (Secp256k1.unsafe_roll32 bs)
      pt <- to_point pub
      pure $! SecretKey bs pub pt

-- | The public key of a t'SecretKey'.
--
--   >>> let Just sk = secret_key (BS.replicate 32 0x41)
--   >>> BS.length (BOLT1.un_point (public_key sk))
--   33
public_key :: SecretKey -> BOLT1.Point
public_key (SecretKey _ _ p) = p
{-# INLINE public_key #-}

sk_bytes :: SecretKey -> BS.ByteString
sk_bytes (SecretKey k _ _) = k
{-# INLINE sk_bytes #-}

sk_pub :: SecretKey -> Secp256k1.Projective
sk_pub (SecretKey _ p _) = p
{-# INLINE sk_pub #-}

-- shared secrets -------------------------------------------------------------

-- | A 32-byte shared secret, established between the origin of an onion
--   and one of its hops.
--
--   This is secret material: its 'Show' instance is redacted and its
--   'Eq' instance runs in constant time.
newtype SharedSecret = SharedSecret BS.ByteString

instance Eq SharedSecret where
  SharedSecret a == SharedSecret b = ct_eq a b

instance Show SharedSecret where
  showsPrec d _ = showParen (d > 10) $
    showString "SharedSecret <redacted>"

instance NFData SharedSecret where
  rnf (SharedSecret bs) = rnf bs

-- | Construct a t'SharedSecret' from exactly 32 bytes (e.g. to restore one
--   that was stored with 'un_shared_secret').
--
--   >>> let Just ss = shared_secret (BS.replicate 32 0x01)
--   >>> BS.length (un_shared_secret ss)
--   32
--   >>> shared_secret "too short"
--   Nothing
shared_secret :: BS.ByteString -> Maybe SharedSecret
shared_secret bs
  | BS.length bs == 32 = Just (SharedSecret bs)
  | otherwise          = Nothing
{-# INLINE shared_secret #-}

-- | The bytes of a t'SharedSecret'.
un_shared_secret :: SharedSecret -> BS.ByteString
un_shared_secret (SharedSecret bs) = bs
{-# INLINE un_shared_secret #-}

-- | ECDH per BOLT #4: SHA256 of the compressed product of a point and a
--   32-byte scalar. Fails if the scalar is not in [1, n - 1].
ecdh :: BS.ByteString -> Secp256k1.Projective -> Maybe SharedSecret
ecdh k pub = do
  pt <- Secp256k1.mul pub (Secp256k1.unsafe_roll32 k)
  pure $! SharedSecret (SHA256.hash (Secp256k1.serialize_point pt))

-- derived keys ---------------------------------------------------------------

-- A 32-byte key derived from a shared secret.
newtype DerivedKey = DerivedKey BS.ByteString

un_derived_key :: DerivedKey -> BS.ByteString
un_derived_key (DerivedKey k) = k
{-# INLINE un_derived_key #-}

derive :: BS.ByteString -> BS.ByteString -> DerivedKey
derive label ss =
  let SHA256.MAC k = SHA256.hmac label ss
  in  DerivedKey k
{-# INLINE derive #-}

derive_rho :: SharedSecret -> DerivedKey
derive_rho (SharedSecret ss) = derive "rho" ss

derive_mu :: SharedSecret -> DerivedKey
derive_mu (SharedSecret ss) = derive "mu" ss

derive_um :: SharedSecret -> DerivedKey
derive_um (SharedSecret ss) = derive "um" ss

derive_ammag :: SharedSecret -> DerivedKey
derive_ammag (SharedSecret ss) = derive "ammag" ss

-- The pad key is derived from the session key itself.
derive_pad :: SecretKey -> DerivedKey
derive_pad (SecretKey k _ _) = derive "pad" k

-- blinding -------------------------------------------------------------------

-- SHA256(E || ss), the factor by which ephemeral keys are blinded.
blinding_factor :: Secp256k1.Projective -> SharedSecret -> BS.ByteString
blinding_factor e (SharedSecret ss) =
  SHA256.hash (Secp256k1.serialize_point e <> ss)

-- HMAC256("blinded_node_id", ss), the factor by which node ids are
-- blinded in a blinded route.
blinded_node_tweak :: SharedSecret -> BS.ByteString
blinded_node_tweak (SharedSecret ss) =
  let SHA256.MAC t = SHA256.hmac "blinded_node_id" ss
  in  t

-- Multiply a point by a 32-byte factor.
blind_pub
  :: Secp256k1.Projective -> BS.ByteString -> Maybe Secp256k1.Projective
blind_pub p t = Secp256k1.mul p (Secp256k1.unsafe_roll32 t)

-- Multiply a 32-byte scalar by a 32-byte factor, mod n. Fails if the
-- product is zero.
blind_scalar :: BS.ByteString -> BS.ByteString -> Maybe BS.ByteString
blind_scalar k t =
  let !r = S.retr (S.mul (S.to (Secp256k1.unsafe_roll32 k))
                         (S.to (Secp256k1.unsafe_roll32 t)))
  in  if   Secp256k1.ge r
      then Just $! Secp256k1.unroll32 r
      else Nothing

-- points ---------------------------------------------------------------------

to_point :: Secp256k1.Projective -> Maybe BOLT1.Point
to_point = BOLT1.point . Secp256k1.serialize_point
{-# INLINE to_point #-}

from_point :: BOLT1.Point -> Maybe Secp256k1.Projective
from_point = Secp256k1.parse_point . BOLT1.un_point
{-# INLINE from_point #-}

-- streams and MACs -----------------------------------------------------------

-- The ChaCha20 keystream of the given length under a derived key, with
-- an all-zero nonce. The cipher fails only for a key or nonce of the
-- wrong length, or past 256 GiB of output, none of which can occur
-- here; the empty result in that case truncates any output XORed with
-- it, rather than leaving it unencrypted.
keystream :: DerivedKey -> Int -> BS.ByteString
keystream (DerivedKey k) n =
  case ChaCha.cipher k 0 (BS.replicate 12 0) (BS.replicate n 0) of
    Right s -> s
    Left _  -> BS.empty

-- HMAC-SHA256 under a derived key.
hmac :: DerivedKey -> BS.ByteString -> BS.ByteString
hmac (DerivedKey k) m =
  let SHA256.MAC h = SHA256.hmac k m
  in  h
{-# INLINE hmac #-}

-- Constant-time equality (variable-time only in the lengths).
ct_eq :: BS.ByteString -> BS.ByteString -> Bool
ct_eq a b = SHA256.MAC a == SHA256.MAC b
{-# INLINE ct_eq #-}

-- XOR, truncated to the shorter input.
xor_bytes :: BS.ByteString -> BS.ByteString -> BS.ByteString
xor_bytes = BS.packZipWith xor
{-# INLINE xor_bytes #-}