From 0fa83e32d8c57470fc47e08469ad1dafc220bd30 Mon Sep 17 00:00:00 2001
From: Vincent Hanquez <vincent@snarc.org>
Date: Sat, 9 Apr 2016 13:45:05 +0100
Subject: [PATCH] [ECDH][DH] change SharedKey representation to be the usual
 bytes-like representation

Prevent mistake when the serialization is not done properly, for example missing
the padding when necessary.
---
 CHANGELOG.md            | 3 +++
 Crypto/PubKey/DH.hs     | 8 +++++---
 Crypto/PubKey/ECC/DH.hs | 6 ++++--
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2a0711e..1431c9c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,9 @@
 * Add a Seed capability to the main DRG, to be able to debug/reproduce randomized program
   where you would want to disable the randomness.
 * Add support for Cipher-based Message Authentication Code (CMAC) (Kei Hibino)
+* *CHANGE* Change the `SharedKey` for `Crypto.PubKey.DH` and `Crypto.PubKey.ECC.DH`,
+  from an Integer newtype to a ScrubbedBytes newtype. Prevent mistake where the
+  bytes representation is generated without the right padding (when needed).
 * *CHANGE* Keep The field size in bits, in the `Params` in `Crypto.PubKey.DH`,
   moving from 2 elements to 3 elements in the structure.
 
diff --git a/Crypto/PubKey/DH.hs b/Crypto/PubKey/DH.hs
index 30f5879..8cfa5c8 100644
--- a/Crypto/PubKey/DH.hs
+++ b/Crypto/PubKey/DH.hs
@@ -23,7 +23,9 @@ import Crypto.Internal.Imports
 import Crypto.Number.ModArithmetic (expSafe)
 import Crypto.Number.Prime (generateSafePrime)
 import Crypto.Number.Generate (generateMax)
+import Crypto.Number.Serialize (i2ospOf_)
 import Crypto.Random.Types
+import Data.ByteArray (ByteArrayAccess, ScrubbedBytes)
 import Data.Data
 
 -- | Represent Diffie Hellman parameters namely P (prime), and G (generator).
@@ -42,8 +44,8 @@ newtype PrivateNumber = PrivateNumber Integer
     deriving (Show,Read,Eq,Enum,Real,Num,Ord)
 
 -- | Represent Diffie Hellman shared secret.
-newtype SharedKey = SharedKey Integer
-    deriving (Show,Read,Eq,Enum,Real,Num,Ord)
+newtype SharedKey = SharedKey ScrubbedBytes
+    deriving (Show,Eq,ByteArrayAccess)
 
 -- | generate params from a specific generator (2 or 5 are common values)
 -- we generate a safe prime (a prime number of the form 2p+1 where p is also prime)
@@ -74,4 +76,4 @@ generatePublic = calculatePublic
 
 -- | generate a shared key using our private number and the other party public number
 getShared :: Params -> PrivateNumber -> PublicNumber -> SharedKey
-getShared (Params p _ bits) (PrivateNumber x) (PublicNumber y) = SharedKey $ expSafe y x p
+getShared (Params p _ bits) (PrivateNumber x) (PublicNumber y) = SharedKey $ i2ospOf_ (bits + 7 `div` 8) $ expSafe y x p
diff --git a/Crypto/PubKey/ECC/DH.hs b/Crypto/PubKey/ECC/DH.hs
index fbe3cbe..4a3e727 100644
--- a/Crypto/PubKey/ECC/DH.hs
+++ b/Crypto/PubKey/ECC/DH.hs
@@ -19,10 +19,11 @@ module Crypto.PubKey.ECC.DH
     ) where
 
 import Crypto.Number.Generate (generateMax)
+import Crypto.Number.Serialize (i2ospOf_)
 import Crypto.PubKey.ECC.Prim (pointMul)
 import Crypto.Random.Types
 import Crypto.PubKey.DH (SharedKey(..))
-import Crypto.PubKey.ECC.Types (PublicPoint, PrivateNumber, Curve, Point(..))
+import Crypto.PubKey.ECC.Types (PublicPoint, PrivateNumber, Curve, Point(..), curveSizeBits)
 import Crypto.PubKey.ECC.Types (ecc_n, ecc_g, common_curve)
 
 -- | Generating a private number d.
@@ -41,6 +42,7 @@ calculatePublic curve d = q
 -- | Generating a shared key using our private number and
 --   the other party public point.
 getShared :: Curve -> PrivateNumber -> PublicPoint -> SharedKey
-getShared curve db qa = SharedKey x
+getShared curve db qa = SharedKey $ i2ospOf_ (nbBits + 7 `div` 8) x
   where
     Point x _ = pointMul curve db qa
+    nbBits    = curveSizeBits curve