{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-cse -fno-full-laziness #-}

-- |
-- Module      : Crypto.PubKey.RSA.Blinded.Internal
-- License     : BSD-style
-- Stability   : experimental
-- Portability : Good
--
-- A private key with a blinding pair that is reused by squaring, as OpenSSL
-- does, rather than generated afresh for every operation.  The key is a value
-- threaded through the operations, as a hash 'Crypto.Hash.Context' is.
module Crypto.PubKey.RSA.Blinded.Internal (
    BlindedKey (..),
    blindedKey,
    dpBlinded,
    Draw,
    drawFor,
    withBlindingFor,
    dpWithPair,
) where

import Data.IORef
import System.IO.Unsafe (unsafePerformIO)

import Crypto.Internal.ByteArray (ByteArray, ScrubbedBytes)
import Crypto.Number.Generate (generateMax)
import Crypto.Number.ModArithmetic (expFast, inverseCoprimes)
import Crypto.Number.Serialize (i2ospOf_, os2ip)
import Crypto.PubKey.RSA.Prim (dp)
import Crypto.PubKey.RSA.Types
import Crypto.Random.Types

-- | A private key together with the blinding its next operation uses.
--
-- Each operation multiplies the input by @r^e@ and the output by @r^-1@, and
-- returns the key with both squared for the next one, so a pair costs two
-- multiplications after the first.  A new @r@ is drawn every 32 operations.
-- Use the key an operation returns: handing the same one to two operations
-- gives both the same pair.
data BlindedKey = BlindedKey
    { BlindedKey -> PrivateKey
blindedPrivateKey :: PrivateKey
    -- ^ the key it was made from
    , BlindedKey -> Pair
blindedPair :: Pair
    }

-- | @r^e@ and @r^-1@ modulo n, and how many more uses they have after this.
data Pair = Pair !Int !Integer !Integer | Spent

-- | Uses of one @r@ before another is drawn, OpenSSL's number.
refresh :: Int
refresh :: Int
refresh = Int
32

-- | A t'BlindedKey' for a private key.  The first pair is drawn when it is
-- first used.
blindedKey :: PrivateKey -> BlindedKey
blindedKey :: PrivateKey -> BlindedKey
blindedKey PrivateKey
pk = PrivateKey -> Pair -> BlindedKey
BlindedKey PrivateKey
pk Pair
Spent

-- | The random numbers a new pair is made from: drawing them is cheap, and
-- making the pair, an inverse and an exponentiation, is done only if needed.
data Draw = Draw !Integer !Integer

-- | Draw for a key, as 'Crypto.PubKey.RSA.generateBlinder' does.
drawFor :: forall m. MonadRandom m => PrivateKey -> m Draw
drawFor :: forall (m :: * -> *). MonadRandom m => PrivateKey -> m Draw
drawFor PrivateKey
pk = do
    Integer
r <- Integer -> m Integer
forall (m :: * -> *). MonadRandom m => Integer -> m Integer
generateMax (PrivateKey -> Integer
private_n PrivateKey
pk)
    Integer
u <- ScrubbedBytes -> Integer
forall ba. ByteArrayAccess ba => ba -> Integer
os2ip (ScrubbedBytes -> Integer) -> m ScrubbedBytes -> m Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> m ScrubbedBytes
forall byteArray. ByteArray byteArray => Int -> m byteArray
forall (m :: * -> *) byteArray.
(MonadRandom m, ByteArray byteArray) =>
Int -> m byteArray
getRandomBytes Int
16 :: m ScrubbedBytes)
    Draw -> m Draw
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Integer -> Integer -> Draw
Draw Integer
r Integer
u)

-- | @r^e@ and @r^-1@.  The inverse is taken of @r*u@ and multiplied by @u@,
-- so the extended Euclidean algorithm never sees @r@, as in
-- 'Crypto.PubKey.RSA.generateBlinder'.
makePair :: PrivateKey -> Draw -> (Integer, Integer)
makePair :: PrivateKey -> Draw -> (Integer, Integer)
makePair PrivateKey
pk (Draw Integer
r Integer
u) =
    (Integer -> Integer -> Integer -> Integer
expFast Integer
r (PrivateKey -> Integer
private_e PrivateKey
pk) Integer
n, (Integer -> Integer -> Integer
inverseCoprimes ((Integer
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
u) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n) Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
u) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n)
  where
    n :: Integer
n = PrivateKey -> Integer
private_n PrivateKey
pk

-- | The pair for this operation, and the key holding the next one.  The
-- draw is used only when a new pair is due.
takePair :: Draw -> BlindedKey -> ((Integer, Integer), BlindedKey)
takePair :: Draw -> BlindedKey -> ((Integer, Integer), BlindedKey)
takePair Draw
d (BlindedKey PrivateKey
pk Pair
pair) = case Pair
pair of
    Pair Int
k Integer
a Integer
ai -> ((Integer
a, Integer
ai), PrivateKey -> Pair -> BlindedKey
BlindedKey PrivateKey
pk (Int -> Integer -> Integer -> Pair
next Int
k Integer
a Integer
ai))
    Pair
Spent -> case PrivateKey -> Draw -> (Integer, Integer)
makePair PrivateKey
pk Draw
d of
        (Integer
a, Integer
ai) -> ((Integer
a, Integer
ai), PrivateKey -> Pair -> BlindedKey
BlindedKey PrivateKey
pk (Int -> Integer -> Integer -> Pair
next Int
refresh Integer
a Integer
ai))
  where
    n :: Integer
n = PrivateKey -> Integer
private_n PrivateKey
pk
    next :: Int -> Integer -> Integer -> Pair
next Int
k Integer
a Integer
ai
        | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 = Int -> Integer -> Integer -> Pair
Pair (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ((Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
a) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n) ((Integer
ai Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
ai) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n)
        | Bool
otherwise = Pair
Spent

-- | The decryption primitive, blinded with a given pair.
dpWithPair :: ByteArray ba => (Integer, Integer) -> PrivateKey -> ba -> ba
dpWithPair :: forall ba.
ByteArray ba =>
(Integer, Integer) -> PrivateKey -> ba -> ba
dpWithPair (Integer
a, Integer
ai) PrivateKey
pk ba
c = Int -> Integer -> ba
forall ba. ByteArray ba => Int -> Integer -> ba
i2ospOf_ Int
k ((Integer
ai Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* ba -> Integer
forall ba. ByteArrayAccess ba => ba -> Integer
os2ip ba
m) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n)
  where
    n :: Integer
n = PrivateKey -> Integer
private_n PrivateKey
pk
    k :: Int
k = PrivateKey -> Int
private_size PrivateKey
pk
    m :: ba
m = Maybe Blinder -> PrivateKey -> ba -> ba
forall ba. ByteArray ba => Maybe Blinder -> PrivateKey -> ba -> ba
dp Maybe Blinder
forall a. Maybe a
Nothing PrivateKey
pk (Int -> Integer -> ba
forall ba. ByteArray ba => Int -> Integer -> ba
i2ospOf_ Int
k ((Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* ba -> Integer
forall ba. ByteArrayAccess ba => ba -> Integer
os2ip ba
c) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
n) ba -> ba -> ba
forall a. a -> a -> a
`asTypeOf` ba
c)

-- | The decryption primitive, blinded with the key's pair, and the key for
-- the next operation.  Randomness is drawn only when a new pair is due.
dpBlinded
    :: (MonadRandom m, ByteArray ba) => BlindedKey -> ba -> m (ba, BlindedKey)
dpBlinded :: forall (m :: * -> *) ba.
(MonadRandom m, ByteArray ba) =>
BlindedKey -> ba -> m (ba, BlindedKey)
dpBlinded BlindedKey
bk ba
c = do
    Draw
d <- case BlindedKey -> Pair
blindedPair BlindedKey
bk of
        Pair {} -> Draw -> m Draw
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Integer -> Integer -> Draw
Draw Integer
0 Integer
0)
        Pair
Spent -> PrivateKey -> m Draw
forall (m :: * -> *). MonadRandom m => PrivateKey -> m Draw
drawFor (BlindedKey -> PrivateKey
blindedPrivateKey BlindedKey
bk)
    let ((Integer, Integer)
pair, BlindedKey
bk') = Draw -> BlindedKey -> ((Integer, Integer), BlindedKey)
takePair Draw
d BlindedKey
bk
    (ba, BlindedKey) -> m (ba, BlindedKey)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Integer, Integer) -> PrivateKey -> ba -> ba
forall ba.
ByteArray ba =>
(Integer, Integer) -> PrivateKey -> ba -> ba
dpWithPair (Integer, Integer)
pair (BlindedKey -> PrivateKey
blindedPrivateKey BlindedKey
bk) ba
c, BlindedKey
bk')

-- | Run an operation with the pair kept for a key, for the functions that
-- take a bare t'PrivateKey'.  One t'BlindedKey' is kept per public key, for
-- the last few keys used; taking its pair and leaving the next is atomic, so
-- no two operations get the same pair.  The caller draws every time, since
-- whether a new pair is due is known only in here.
withBlindingFor :: PrivateKey -> Draw -> ((Integer, Integer) -> a) -> a
withBlindingFor :: forall a. PrivateKey -> Draw -> ((Integer, Integer) -> a) -> a
withBlindingFor PrivateKey
pk Draw
d (Integer, Integer) -> a
f = IO a -> a
forall a. IO a -> a
unsafePerformIO (IO a -> a) -> IO a -> a
forall a b. (a -> b) -> a -> b
$ do
    (Integer, Integer)
pair <- IORef [BlindedKey]
-> ([BlindedKey] -> ([BlindedKey], (Integer, Integer)))
-> IO (Integer, Integer)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef [BlindedKey]
table [BlindedKey] -> ([BlindedKey], (Integer, Integer))
step
    a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Integer, Integer) -> a
f (Integer, Integer)
pair)
  where
    pub :: PublicKey
pub = PrivateKey -> PublicKey
private_pub PrivateKey
pk
    step :: [BlindedKey] -> ([BlindedKey], (Integer, Integer))
step [BlindedKey]
ks = case (BlindedKey -> Bool)
-> [BlindedKey] -> ([BlindedKey], [BlindedKey])
forall a. (a -> Bool) -> [a] -> ([a], [a])
break ((PublicKey -> PublicKey -> Bool
forall a. Eq a => a -> a -> Bool
== PublicKey
pub) (PublicKey -> Bool)
-> (BlindedKey -> PublicKey) -> BlindedKey -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrivateKey -> PublicKey
private_pub (PrivateKey -> PublicKey)
-> (BlindedKey -> PrivateKey) -> BlindedKey -> PublicKey
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BlindedKey -> PrivateKey
blindedPrivateKey) [BlindedKey]
ks of
        ([BlindedKey]
before, BlindedKey
bk : [BlindedKey]
after) -> BlindedKey -> [BlindedKey] -> ([BlindedKey], (Integer, Integer))
found BlindedKey
bk ([BlindedKey]
before [BlindedKey] -> [BlindedKey] -> [BlindedKey]
forall a. [a] -> [a] -> [a]
++ [BlindedKey]
after)
        ([BlindedKey], [BlindedKey])
_ -> BlindedKey -> [BlindedKey] -> ([BlindedKey], (Integer, Integer))
found (PrivateKey -> BlindedKey
blindedKey PrivateKey
pk) [BlindedKey]
ks
    found :: BlindedKey -> [BlindedKey] -> ([BlindedKey], (Integer, Integer))
found BlindedKey
bk [BlindedKey]
rest = case Draw -> BlindedKey -> ((Integer, Integer), BlindedKey)
takePair Draw
d BlindedKey
bk of
        (pair :: (Integer, Integer)
pair@(!Integer
_, !Integer
_), BlindedKey
bk') -> (Int -> [BlindedKey] -> [BlindedKey]
forall a. Int -> [a] -> [a]
take Int
tableSize (BlindedKey
bk' BlindedKey -> [BlindedKey] -> [BlindedKey]
forall a. a -> [a] -> [a]
: [BlindedKey]
rest), (Integer, Integer)
pair)
{-# NOINLINE withBlindingFor #-}

-- | How many keys 'withBlindingFor' remembers.
tableSize :: Int
tableSize :: Int
tableSize = Int
4

table :: IORef [BlindedKey]
table :: IORef [BlindedKey]
table = IO (IORef [BlindedKey]) -> IORef [BlindedKey]
forall a. IO a -> a
unsafePerformIO ([BlindedKey] -> IO (IORef [BlindedKey])
forall a. a -> IO (IORef a)
newIORef [])
{-# NOINLINE table #-}