From 17336857c5fec89429c4d81853e2d2e463195d88 Mon Sep 17 00:00:00 2001
From: Will Song <will.song@trailofbits.com>
Date: Mon, 1 Jun 2020 20:56:42 -0500
Subject: [PATCH] implement square roots in f2m

---
 Crypto/Number/F2m.hs | 52 +++++++++++++++++++++++++++++++++++++++++---
 tests/Number/F2m.hs  |  8 +++++++
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/Crypto/Number/F2m.hs b/Crypto/Number/F2m.hs
index 93b1f48..ad0941f 100644
--- a/Crypto/Number/F2m.hs
+++ b/Crypto/Number/F2m.hs
@@ -16,7 +16,10 @@ module Crypto.Number.F2m
     , mulF2m
     , squareF2m'
     , squareF2m
+    , powF2m
+    , powF2m'
     , modF2m
+    , sqrtF2m
     , invF2m
     , divF2m
     ) where
@@ -66,8 +69,8 @@ mulF2m :: BinaryPolynomial -- ^ Modulus
 mulF2m fx n1 n2
     |    fx < 0
       || n1 < 0
-      || n2 < 0 = error "mulF2m: negative number represent no binary binary polynomial"
-    | fx == 0   = error "modF2m: cannot multiply modulo zero polynomial"
+      || n2 < 0 = error "mulF2m: negative number represent no binary polynomial"
+    | fx == 0   = error "mulF2m: cannot multiply modulo zero polynomial"
     | otherwise = modF2m fx $ go (if n2 `mod` 2 == 1 then n1 else 0) (log2 n2)
       where
         go n s | s == 0  = n
@@ -96,10 +99,53 @@ squareF2m fx = modF2m fx . squareF2m'
 squareF2m' :: Integer
            -> Integer
 squareF2m' n
-    | n < 0     = error "mulF2m: negative number represent no binary binary polynomial"
+    | n < 0     = error "mulF2m: negative number represent no binary polynomial"
     | otherwise = foldl' (\acc s -> if testBit n s then setBit acc (2 * s) else acc) 0 [0 .. log2 n]
 {-# INLINE squareF2m' #-}
 
+-- | Exponentiation in F₂m by computing @a^b mod fx@.
+--
+-- This implements an exponentiation by squaring based solution. It inherits the
+-- same restrictions as 'squareF2m'. Negative exponents are disallowed. See
+-- 'powF2m'' for one that handles this case
+powF2m :: BinaryPolynomial -- ^Modulus
+       -> Integer          -- ^a
+       -> Integer          -- ^b
+       -> Integer
+powF2m fx a b
+  | b == 0 = 1
+  | b > 0  = squareF2m fx x * if even b then 1 else a
+  | b < 0  = error "powF2m: negative exponents disallowed"
+  | otherwise = error "powF2m: impossible"
+  where x = powF2m fx a (b `div` 2)
+
+-- | Exponentiation in F₂m by computing @a^b mod fx@.
+--
+-- This implements an exponentiation by squaring based solution. It inherits the
+-- same restrictions as 'squareF2m'. 'Nothing' is returned when a negative
+-- exponent is given and @a@ is not invertible.
+powF2m' :: BinaryPolynomial -- ^Modulus
+        -> Integer          -- ^a
+        -> Integer          -- ^b
+        -> Maybe Integer
+powF2m' fx a b
+  | b == 0 = Just 1
+  | b > 0  = Just $ powF2m fx a b
+  | b < 0  = case invF2m fx a of
+               Just inv -> Just $ powF2m fx inv (-b)
+               Nothing  -> Nothing
+  | otherwise = error "impossible"
+
+-- | Square rooot in F₂m.
+--
+-- We exploit the fact that @a^(2^m) = a@, or in particular, @a^(2^m - 1) = 1@
+-- from a classical result by Lagrange. Thus the square root is simply @a^(2^(m
+-- - 1))@.
+sqrtF2m :: BinaryPolynomial -- ^Modulus
+        -> Integer          -- ^a
+        -> Integer
+sqrtF2m fx a = powF2m fx a (2 ^ (log2 fx - 1))
+
 -- | Extended GCD algorithm for polynomials. For @a@ and @b@ returns @(g, u, v)@ such that @a * u + b * v == g@.
 --
 -- Reference: https://en.wikipedia.org/wiki/Polynomial_greatest_common_divisor#B.C3.A9zout.27s_identity_and_extended_GCD_algorithm
diff --git a/tests/Number/F2m.hs b/tests/Number/F2m.hs
index afa6e50..4434a9b 100644
--- a/tests/Number/F2m.hs
+++ b/tests/Number/F2m.hs
@@ -52,6 +52,14 @@ mulTests = testGroup "mulF2m"
 squareTests = testGroup "squareF2m"
     [ testProperty "sqr(a) == a * a"
         $ \(Positive m) (NonNegative a) -> mulF2m m a a == squareF2m m a
+    -- disabled because we require @m@ to be a suitable modulus and there is no
+    -- way to guarantee this
+    -- , testProperty "sqrt(a) * sqrt(a) = a"
+    --     $ \(Positive m) (NonNegative aa) -> let a = sqrtF2m m aa in mulF2m m a a == modF2m m aa
+    , testProperty "sqrt(a) * sqrt(a) = a in GF(2^16)"
+        $ let m = 65581 :: Integer -- x^16 + x^5 + x^3 + x^2 + 1
+              nums = [0 .. 65535 :: Integer]
+          in  nums == [let y = sqrtF2m m x in squareF2m m y | x <- nums]
     ]
 
 invTests = testGroup "invF2m"