Ed25519BIP32.hs

{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE Rank2Types            #-}
{-# LANGUAGE ViewPatterns          #-}
{-# LANGUAGE GADTs                 #-}
module Crypto.ECC.Ed25519BIP32 where

import qualified Crypto.Hash as C (SHA512, SHA256)
import qualified Crypto.MAC.HMAC as C
import qualified Data.ByteArray as B
import qualified Data.ByteString as BS
import           Data.Bits
import           Data.Kind (Type)
import           Data.Word
import           Data.Proxy
import qualified Crypto.Math.Edwards25519 as ED25519
import           Data.Type.Bool
import           Data.Type.Equality
import           GHC.TypeLits
import           Data.Function (on)
import           Unsafe.Coerce (unsafeCoerce)

import           Crypto.Math.Bits
import           Crypto.Math.Bytes (Bytes)
import qualified Crypto.Math.Bytes as Bytes

-- | A Master secret is a 256 bits random quantity
type MasterSecret = FBits 256

-- | A child key is similar to the key in structure
-- except it has an additional annotation representing
-- the indexes for the hierarchy derivation indexes from
-- a base 'Key' (usually the root key)
type ChildKey (didxs :: DerivationHier) = Key

-- | A key is a 512 bit random value and a chaincode
--
-- Left half need to have:
-- * Lowest 3 bits clear
-- * Highest bit clear
-- * Second highest bit set
-- * Third highest bit clear
--
-- Right half doesn't have any particular structure.
type Key = (FBits 256, FBits 256, ChainCode)

-- | A public part of a key
type Public = (PointCompressed, ChainCode)

-- | A point is 1 bit of x sign and 255 bit of y coordinate (y's 256th bit is always 0)
type PointCompressed = FBits 256

-- | A 256 bits chain code
newtype ChainCode = ChainCode { unChainCode :: Bytes 32 }
    deriving (Eq)

-- | A n bits Digest
newtype Hash n = Hash { unHash :: FBits n }
    deriving (Eq)

-- | A Serialized tag used during HMAC
type Tag = Bytes 1

-- | Serialized Index
newtype SerializedIndex = SerializedIndex (Bytes 4)
    deriving (Eq)

type HMAC_SHA512 = Bytes 64

data DerivationType = Hard | Soft
data DerivationMaterial = ChainCodeMaterial | KeyMaterial

data DerivationIndex (k :: DerivationType) (n :: Nat) = DerivationIndex

data DerivationHier = (:>) Nat DerivationHier | DerivationEnd

type MaxHardIndex = 0xffffffff
type MinHardIndex = 0x80000000
type MaxSoftIndex = MinHardIndex - 1
type MinSoftIndex = 0

data ValidIndex :: Nat -> Type where
    IsValidIndex    :: (ValidDerivationIndex n :~: 'True) -> ValidIndex n
    IsNotValidIndex :: (ValidDerivationIndex n :~: 'False) -> ValidIndex n

data ValidHardIndex :: Nat -> Type where
    IsValidHardIndex    :: (ValidDerivationHardIndex n :~: 'True) -> ValidHardIndex n
    IsNotValidHardIndex :: (ValidDerivationHardIndex n :~: 'False) -> ValidHardIndex n

data ValidSoftIndex :: Nat -> Type where
    IsValidSoftIndex    :: (ValidDerivationSoftIndex n :~: 'True) -> ValidSoftIndex n
    IsNotValidSoftIndex :: (ValidDerivationSoftIndex n :~: 'False) -> ValidSoftIndex n

getValidIndex :: KnownNat n => Proxy n -> Maybe (ValidDerivationIndex n :~: 'True)
getValidIndex n = case isValidIndex n of
                    IsValidIndex Refl -> Just Refl
                    _                 -> Nothing

isValidIndex :: KnownNat n => Proxy n -> ValidIndex n
isValidIndex n
    |  natVal (Proxy @MinSoftIndex) <= natVal n
    && natVal n <= natVal (Proxy @MaxHardIndex) = IsValidIndex (unsafeCoerce Refl)
    | otherwise                                 = IsNotValidIndex (unsafeCoerce Refl)

getValidHardIndex :: KnownNat n => Proxy n -> Maybe (ValidDerivationHardIndex n :~: 'True)
getValidHardIndex n = case isValidHardIndex n of
                    IsValidHardIndex Refl -> Just Refl
                    _                     -> Nothing

isValidHardIndex :: KnownNat n => Proxy n -> ValidHardIndex n
isValidHardIndex n
    |  natVal (Proxy @MinHardIndex) <= natVal n
    && natVal n <= natVal (Proxy @MaxHardIndex) = IsValidHardIndex (unsafeCoerce Refl)
    | otherwise                                 = IsNotValidHardIndex (unsafeCoerce Refl)

getValidSoftIndex :: KnownNat n => Proxy n -> Maybe (ValidDerivationSoftIndex n :~: 'True)
getValidSoftIndex n = case isValidSoftIndex n of
                    IsValidSoftIndex Refl -> Just Refl
                    _                     -> Nothing

isValidSoftIndex :: KnownNat n => Proxy n -> ValidSoftIndex n
isValidSoftIndex n
    |  natVal (Proxy @MinSoftIndex) <= natVal n
    && natVal n <= natVal (Proxy @MaxSoftIndex) = IsValidSoftIndex (unsafeCoerce Refl)
    | otherwise                                 = IsNotValidSoftIndex (unsafeCoerce Refl)

type ValidDerivationIndex     (n :: Nat) = (MinSoftIndex <=? n) && (n <=? MaxHardIndex)
type ValidDerivationHardIndex (n :: Nat) = (MinHardIndex <=? n) && (n <=? MaxHardIndex)
type ValidDerivationSoftIndex (n :: Nat) = (MinSoftIndex <=? n) && (n <=? MaxSoftIndex)

type family ValidDerivationIndexForType (k :: DerivationType) (n :: Nat) :: Bool where
    ValidDerivationIndexForType 'Hard n = ValidDerivationHardIndex n
    ValidDerivationIndexForType 'Soft n = ValidDerivationSoftIndex n

type family DerivationTag (ty :: DerivationType) (material :: DerivationMaterial) :: Nat where
    DerivationTag 'Hard 'KeyMaterial       = 0x0
    DerivationTag 'Hard 'ChainCodeMaterial = 0x1
    DerivationTag 'Soft 'KeyMaterial       = 0x2
    DerivationTag 'Soft 'ChainCodeMaterial = 0x3

-- | Check if the left half is valid
leftHalfValid :: FBits 256 -> Bool
leftHalfValid v =
    and [ testBit v 0 == False
        , testBit v 1 == False
        , testBit v 2 == False
        , testBit v 29 == False
        , testBit v 28 == True
        , testBit v 31 == False
        ]

toPublic :: Key -> Public
toPublic (kl, _, cc) = (kToPoint kl, cc)

kToPoint :: FBits 256 -> PointCompressed
kToPoint k = pointFromRepr $ ED25519.scalarToPoint r
  where r = ED25519.scalarP $ Bytes.fromBits Bytes.LittleEndian k

pointAdd :: PointCompressed -> PointCompressed -> PointCompressed
pointAdd = ((pointFromRepr . ) .  ED25519.pointAdd) `on` pointToRepr

pointToRepr :: PointCompressed -> ED25519.PointCompressed
pointToRepr a = ED25519.pointCompressedP $ Bytes.fromBits Bytes.LittleEndian a

pointFromRepr :: ED25519.PointCompressed -> PointCompressed
pointFromRepr =
      Bytes.toBits Bytes.LittleEndian
    . ED25519.unPointCompressedP

type family BitsToHashScheme (n :: Nat) where
    BitsToHashScheme 256 = C.SHA256
    BitsToHashScheme 512 = C.SHA512

type ValidTag tag = (0 <= tag, tag <= 3)

-- | Compute the HMAC-SHA512 using the ChainCode as the key
fcp :: forall tag idx deriveType deriveMaterial
     . ( KnownNat (DerivationTag deriveType deriveMaterial)
       , KnownNat idx
       , (DerivationTag deriveType deriveMaterial) ~ tag
       , ValidDerivationIndex idx ~ 'True
       , ValidDerivationIndexForType deriveType idx ~ 'True
       )
    => Proxy deriveMaterial
    -> Proxy deriveType
    -> Proxy idx
    -> ChainCode
    -> DerivationIndex deriveType idx
    -> [Word8]
    -> HMAC_SHA512
fcp _ _ pidx c _ input =
    hmacSHA512 key `Bytes.packSome` (Bytes.unpack tagValue ++ input ++ Bytes.unpack idx)
  where
    key = unChainCode c

    (SerializedIndex idx) = indexSerialized pidx

    tagValue :: Tag
    tagValue = Bytes.pack [fromIntegral $ natVal (Proxy :: Proxy tag)]

hmacSHA512 :: Bytes keyLength -> Bytes input -> HMAC_SHA512
hmacSHA512 key ({-Bytes.trace "hmac-input" ->-} msg) =
    Bytes.pack $ B.unpack $ C.hmacGetDigest computed
  where
    computed :: C.HMAC C.SHA512
    computed = C.hmac (BS.pack $ Bytes.unpack key) (BS.pack $ Bytes.unpack msg)

class GetDerivationMaterial (dtype :: DerivationType) mat where
    getDerivationMaterial :: Proxy dtype -> mat -> [Word8]
instance GetDerivationMaterial 'Soft Key where
    getDerivationMaterial p key = getDerivationMaterial p (fst $ toPublic key)
instance GetDerivationMaterial 'Hard Key where
    getDerivationMaterial _ (kl,kr,_) =
        Bytes.unpack $ Bytes.append (Bytes.fromBits Bytes.LittleEndian kl)
                                    (Bytes.fromBits Bytes.LittleEndian kr)
instance GetDerivationMaterial 'Soft PointCompressed where
    getDerivationMaterial _ p = Bytes.unpack $ Bytes.fromBits Bytes.LittleEndian p

derive :: forall dtype idx .
          ( KnownNat (DerivationTag dtype 'KeyMaterial)
          , KnownNat (DerivationTag dtype 'ChainCodeMaterial)
          , KnownNat idx
          , ValidDerivationIndex idx ~ 'True
          , ValidDerivationIndexForType dtype idx ~ 'True
          , GetDerivationMaterial dtype Key)
       => DerivationIndex dtype idx
       -> Key
       -> Key
derive idx key@(kl, kr, c) = (kl', kr', c')
  where
    dtype       = Proxy :: Proxy dtype
    matKeyProxy = Proxy :: Proxy 'KeyMaterial
    matCCProxy  = Proxy :: Proxy 'ChainCodeMaterial
    -- 1) Z
    z = fcp matKeyProxy dtype (Proxy :: Proxy idx) c idx
            (getDerivationMaterial dtype key)

    -- 2) produce kl' and kr'
    (zl8, zr) = step2 z

    kl' = zl8 + kl
    kr' = zr + kr

    -- 3) child chain code
    untrimmedCC = fcp matCCProxy dtype (Proxy :: Proxy idx) c idx
                      (getDerivationMaterial dtype key)
    c' = ChainCode $ Bytes.drop untrimmedCC

derivePublic :: forall idx dtype .
          ( dtype ~ 'Soft -- can only derive public stuff with Soft Derivation
          , KnownNat (DerivationTag dtype 'KeyMaterial)
          , KnownNat (DerivationTag dtype 'ChainCodeMaterial)
          , KnownNat idx
          , ValidDerivationIndex idx ~ 'True
          , ValidDerivationIndexForType dtype idx ~ 'True
          , GetDerivationMaterial dtype PointCompressed)
       => DerivationIndex 'Soft idx
       -> PointCompressed
       -> ChainCode
       -> (PointCompressed, ChainCode)
derivePublic idx p c = (p', c')
  where
    dtype       = Proxy :: Proxy dtype
    matKeyProxy = Proxy :: Proxy 'KeyMaterial
    matCCProxy  = Proxy :: Proxy 'ChainCodeMaterial
    -- 1) Z
    z = fcp matKeyProxy dtype (Proxy :: Proxy idx) c idx (getDerivationMaterial dtype p)

    -- 2) produce kl' and kr'
    (zl8, _) = step2 z

    p' = kToPoint zl8 `pointAdd` p

    -- 3) child chain code
    untrimmedCC = fcp matCCProxy dtype (Proxy :: Proxy idx) c idx
                      (getDerivationMaterial dtype p)
    c' = ChainCode $ Bytes.drop untrimmedCC

-- | Given Z, return 8*ZL(28Bytes) and ZR
step2 :: Bytes 64 -> (FBits 256, FBits 256)
step2 z = (8 * zeroExtendedZl, Bytes.toBits Bytes.LittleEndian zRight)
  where
    (zLeft32, zRight) = Bytes.splitHalf z

    -- step
    -- * take 28 bytes of zLeft32 (zl)
    -- * extend back to 32 bytes
    -- * multiply this number by 8
    zeroExtendedZl = zeroExtender `append` Bytes.toBits Bytes.LittleEndian zl

    zl :: Bytes 28 -- only take 28 bytes
    zl = Bytes.take zLeft32

    zeroExtender :: FBits (4*8) -- re-extend by 4 bytes
    zeroExtender = 0

-- | Serialized index
indexSerialized :: forall idx . (KnownNat idx, ValidDerivationIndex idx ~ 'True)
                => Proxy idx
                -> SerializedIndex
indexSerialized idx = SerializedIndex $ Bytes.fromBits Bytes.LittleEndian (fromInteger n :: FBits 32)
  where n = natVal idx