From 0cec622ddf7d344669797f460806cec737b812ac Mon Sep 17 00:00:00 2001
From: Colin Atkinson <colin.william.atkinson@gmail.com>
Date: Sat, 24 Dec 2016 17:43:54 -0500
Subject: [PATCH] Fix generate(Safe)Prime to guarantee prime size

Add check for size in generatePrime

Add size test in generateSafePrime

Require only that top bit is set, instead of top 2

This is the general standard, see e.g. OpenSSL

Add an error for too few bits being supplied to prime generator, and add documentation

Add some documentation and require highest two bits set

Simplify return syntax in generatePrime and generateSafePrime

Switch exponent to bit-shift for small performance boost
---
 .gitignore             |  2 ++
 Crypto/Error/Types.hs  |  2 ++
 Crypto/Number/Prime.hs | 35 +++++++++++++++++++++++++++++------
 cryptonite.cabal       |  2 +-
 4 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/.gitignore b/.gitignore
index 33ee2ff..a213015 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,3 +10,5 @@ QA
 benchs/Bench
 benchs/Hash
 *.sublime-workspace
+.cabal-sandbox/
+cabal.sandbox.config
diff --git a/Crypto/Error/Types.hs b/Crypto/Error/Types.hs
index 4aaf4e0..cb512d2 100644
--- a/Crypto/Error/Types.hs
+++ b/Crypto/Error/Types.hs
@@ -43,6 +43,8 @@ data CryptoError =
     -- Message authentification error
     | CryptoError_MacKeyInvalid
     | CryptoError_AuthenticationTagSizeInvalid
+    -- Prime generation error
+    | CryptoError_PrimeSizeInvalid
     deriving (Show,Eq,Enum,Data,Typeable)
 
 instance E.Exception CryptoError
diff --git a/Crypto/Number/Prime.hs b/Crypto/Number/Prime.hs
index 89a92da..c993648 100644
--- a/Crypto/Number/Prime.hs
+++ b/Crypto/Number/Prime.hs
@@ -27,6 +27,7 @@ import Crypto.Number.Basic (sqrti, gcde)
 import Crypto.Number.ModArithmetic (expSafe)
 import Crypto.Random.Types
 import Crypto.Random.Probabilistic
+import Crypto.Error
 
 import Data.Bits
 
@@ -40,22 +41,44 @@ isProbablyPrime !n
     | primalityTestFermat 50 (n`div`2) n                   = primalityTestMillerRabin 30 n
     | otherwise                                            = False
 
--- | generate a prime number of the required bitsize
+-- | generate a prime number of the required bitsize (i.e. in the range
+--   [2^(b-1)+2^(b-2), 2^b)).
+--
+--   May throw a CryptoError_PrimeSizeInvalid if the requested size is less
+--   than 5 bits, as the smallest prime meeting these conditions is 29.
+--   This function requires that the two highest bits are set, so that when
+--   multiplied with another prime to create a key, it is guaranteed to be of
+--   the proper size.
 generatePrime :: MonadRandom m => Int -> m Integer
 generatePrime bits = do
-    sp <- generateParams bits (Just SetTwoHighest) True
-    return $ findPrimeFrom sp
+    if bits < 5 then
+        throwCryptoError $ CryptoFailed $ CryptoError_PrimeSizeInvalid
+    else do
+        sp <- generateParams bits (Just SetTwoHighest) True
+        let prime = findPrimeFrom sp
+        if prime < 1 `shiftL` bits then
+            return $ prime
+        else generatePrime bits
 
 -- | generate a prime number of the form 2p+1 where p is also prime.
 -- it is also knowed as a Sophie Germaine prime or safe prime.
 --
 -- The number of safe prime is significantly smaller to the number of prime,
 -- as such it shouldn't be used if this number is supposed to be kept safe.
+--
+-- May throw a CryptoError_PrimeSizeInvalid if the requested size is less than
+-- 6 bits, as the smallest safe prime with the two highest bits set is 59.
 generateSafePrime :: MonadRandom m => Int -> m Integer
 generateSafePrime bits = do
-    sp <- generateParams bits (Just SetTwoHighest) True
-    let p = findPrimeFromWith (\i -> isProbablyPrime (2*i+1)) (sp `div` 2)
-    return (2*p+1)
+    if bits < 6 then
+        throwCryptoError $ CryptoFailed $ CryptoError_PrimeSizeInvalid
+    else do
+        sp <- generateParams bits (Just SetTwoHighest) True
+        let p = findPrimeFromWith (\i -> isProbablyPrime (2*i+1)) (sp `div` 2)
+        let val = 2 * p + 1
+        if val < 1 `shiftL` bits then
+            return $ val
+        else generateSafePrime bits
 
 -- | find a prime from a starting point where the property hold.
 findPrimeFromWith :: (Integer -> Bool) -> Integer -> Integer
diff --git a/cryptonite.cabal b/cryptonite.cabal
index 1ea74f7..e82f1e5 100644
--- a/cryptonite.cabal
+++ b/cryptonite.cabal
@@ -205,7 +205,7 @@ Library
                    , bytestring
                    , memory >= 0.8
                    , ghc-prim
-  ghc-options:       -Wall -fwarn-tabs -optc-O3 -fno-warn-unused-imports
+  ghc-options:       -Wall -fwarn-tabs -optc-O3 -fno-warn-unused-imports -fobject-code
   default-language:  Haskell2010
   cc-options:        -std=gnu99
   if flag(old_toolchain_inliner)