From 9e0dbb32313b38cb21ce39ac1092e88dd41c0880 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olivier=20Ch=C3=A9ron?= <olivier.cheron@gmail.com>
Date: Sat, 7 Dec 2019 08:35:14 +0100
Subject: [PATCH] Modular square root

---
 Crypto/Number/ModArithmetic.hs | 92 +++++++++++++++++++++++++++++++++-
 tests/Number.hs                | 12 +++++
 2 files changed, 103 insertions(+), 1 deletion(-)

diff --git a/Crypto/Number/ModArithmetic.hs b/Crypto/Number/ModArithmetic.hs
index 9ac1440..edb2679 100644
--- a/Crypto/Number/ModArithmetic.hs
+++ b/Crypto/Number/ModArithmetic.hs
@@ -14,8 +14,10 @@ module Crypto.Number.ModArithmetic
     -- * Inverse computing
     , inverse
     , inverseCoprimes
-    , jacobi
     , inverseFermat
+    -- * Squares
+    , jacobi
+    , squareRoot
     ) where
 
 import Control.Exception (throw, Exception)
@@ -125,3 +127,91 @@ jacobi a n
 -- the modulus is prime but avoids side channels like in 'expSafe'.
 inverseFermat :: Integer -> Integer -> Integer
 inverseFermat g p = expSafe g (p - 2) p
+
+-- | Raised when the assumption about the modulus is invalid.
+data ModulusAssertionError = ModulusAssertionError
+    deriving (Show)
+
+instance Exception ModulusAssertionError
+
+-- | Modular square root of @g@ modulo a prime @p@.
+--
+-- If the modulus is found not to be prime, the function will raise a
+-- 'ModulusAssertionError'.
+--
+-- This implementation is variable time and should be used with public
+-- parameters only.
+squareRoot :: Integer -> Integer -> Maybe Integer
+squareRoot p
+    | p < 2     = throw ModulusAssertionError
+    | otherwise =
+        case p `divMod` 8 of
+           (v, 3) -> method1 (2 * v + 1)
+           (v, 7) -> method1 (2 * v + 2)
+           (u, 5) -> method2 u
+           (_, 1) -> tonelliShanks p
+           (0, 2) -> \a -> Just (if even a then 0 else 1)
+           _      -> throw ModulusAssertionError
+
+  where
+    x `eqMod` y = (x - y) `mod` p == 0
+
+    validate g y | (y * y) `eqMod` g = Just y
+                 | otherwise         = Nothing
+
+    -- p == 4u + 3 and u' == u + 1
+    method1 u' g =
+        let y = expFast g u' p
+         in validate g y
+
+    -- p == 8u + 5
+    method2 u g =
+        let gamma = expFast (2 * g) u p
+            g_gamma = g * gamma
+            i = (2 * g_gamma * gamma) `mod` p
+            y = (g_gamma * (i - 1)) `mod` p
+         in validate g y
+
+tonelliShanks :: Integer -> Integer -> Maybe Integer
+tonelliShanks p a
+    | aa == 0   = Just 0
+    | otherwise =
+        case expFast aa p2 p of
+            b | b == p1   -> Nothing
+              | b == 1    -> Just $ go (expFast aa ((s + 1) `div` 2) p)
+                                       (expFast aa s p)
+                                       (expFast n  s p)
+                                       e
+              | otherwise -> throw ModulusAssertionError
+  where
+    aa = a `mod` p
+    p1 = p - 1
+    p2 = p1 `div` 2
+    n  = findN 2
+
+    x `mul` y = (x * y) `mod` p
+
+    pow2m 0 x = x
+    pow2m i x = pow2m (i - 1) (x `mul` x)
+
+    (e, s) = asPowerOf2AndOdd p1
+
+    -- find a quadratic non-residue
+    findN i
+        | expFast i p2 p == p1 = i
+        | otherwise            = findN (i + 1)
+
+    -- find m such that b^(2^m) == 1 (mod p)
+    findM b i
+        | b == 1    = i
+        | otherwise = findM (b `mul` b) (i + 1)
+
+    go !x b g !r
+        | b == 1    = x
+        | otherwise =
+            let r' = findM b 0
+                z = pow2m (r - r' - 1) g
+                x' = x `mul` z
+                b' = b `mul` g'
+                g' = z `mul` z
+             in go x' b' g' r'
diff --git a/tests/Number.hs b/tests/Number.hs
index 0b8f654..7aa6acf 100644
--- a/tests/Number.hs
+++ b/tests/Number.hs
@@ -10,6 +10,7 @@ import Crypto.Number.Generate
 import qualified Crypto.Number.Serialize    as BE
 import qualified Crypto.Number.Serialize.LE as LE
 import Crypto.Number.Prime
+import Crypto.Number.ModArithmetic
 import Data.Bits
 
 serializationVectors :: [(Int, Integer, ByteString)]
@@ -55,6 +56,17 @@ tests = testGroup "number"
     , testProperty "as-power-of-2-and-odd" $ \n ->
         let (e, a1) = asPowerOf2AndOdd n
          in n == (2^e)*a1
+    , testProperty "squareRoot" $ \testDRG (Int0_2901 baseBits') -> do
+        let baseBits = baseBits' `mod` 500
+            bits = 5 + baseBits -- generating lower than 5 bits causes an error ..
+            p = withTestDRG testDRG $ generatePrime bits
+        g <- choose (1, p - 1)
+        let square x = (x * x) `mod` p
+            r = square <$> squareRoot p g
+        case jacobi g p of
+            Just   1  -> return $ Just g `assertEq` r
+            Just (-1) -> return $ Nothing `assertEq` r
+            _         -> error "invalid jacobi result"
     , testProperty "marshalling-be" $ \qaInt ->
         getQAInteger qaInt == BE.os2ip (BE.i2osp (getQAInteger qaInt) :: Bytes)
     , testProperty "marshalling-le" $ \qaInt ->