fixed

commit 1ed06104158a825aeee9ae94a30c499f24cea862
parent 7578f945e2d6d54e91f036887432b11114349c83
Author: Jared Tobin <jared@jtobin.io>
Date:   Wed, 24 Dec 2025 07:52:37 -0330

lib: mark sqrt as vartime

We branch at the end based on whether or not the input is a quadratic
residue. This is typically benign, but worth being strict about in our
naming.

Diffstat:
M
bench/Main.hs
 | 
4
++--
M
bench/Weight.hs
 | 
4
++--
M
lib/Numeric/Montgomery/Secp256k1/Curve.hs
 | 
29
+++++++++++++++--------------

3 files changed, 19 insertions(+), 18 deletions(-)
diff --git a/bench/Main.hs b/bench/Main.hs
@@ -99,8 +99,8 @@ sqrt =
   let !c2 = 2 :: C.Montgomery
       !c_big = (2 ^ 255 - 19) :: C.Montgomery
   in  bgroup "sqrt" [
-          bench "curve:  sqrt M(2)" $ nf C.sqrt c2
-        , bench "curve:  sqrt M(2 ^ 255 - 19)" $ nf C.sqrt c_big
+          bench "curve:  sqrt M(2)" $ nf C.sqrt_vartime c2
+        , bench "curve:  sqrt M(2 ^ 255 - 19)" $ nf C.sqrt_vartime c_big
         ]
 
 exp :: Benchmark
diff --git a/bench/Weight.hs b/bench/Weight.hs
@@ -125,8 +125,8 @@ sqrt =
   let !c2 = 2 :: C.Montgomery
       !c_big = (2 ^ 255 - 19) :: C.Montgomery
   in  wgroup "sqrt" $ do
-        func "curve:  sqrt M(2)" C.sqrt c2
-        func "curve:  sqrt M(2 ^ 255 - 19)" C.sqrt c_big
+        func "curve:  sqrt M(2)" C.sqrt_vartime c2
+        func "curve:  sqrt M(2 ^ 255 - 19)" C.sqrt_vartime c_big
 
 redc :: Weigh ()
 redc =
diff --git a/lib/Numeric/Montgomery/Secp256k1/Curve.hs b/lib/Numeric/Montgomery/Secp256k1/Curve.hs
@@ -51,7 +51,7 @@ module Numeric.Montgomery.Secp256k1.Curve (
   , neg#
   , inv
   , inv#
-  , sqrt
+  , sqrt_vartime
   , sqrt#
   , exp
   , odd#
@@ -989,24 +989,27 @@ inv (Montgomery w) = Montgomery (inv# w)
 
 -- | Square root (Tonelli-Shanks) in the Montgomery domain.
 --
---   For a, return x such that a = x x mod p. Returns nothing if no such
---   square root exists.
+--   Returns 'Nothing' if the square root doesn't exist.
 --
---   >>> sqrt 4
+--   Note that the square root calculation itself is performed in
+--   constant time; we branch only when casting to 'Maybe' at the end.
+--
+--   >>> sqrt_vartime 4
 --   Just 2
---   >>> sqrt 15
+--   >>> sqrt_vartime 15
 --   Just 69211104694897500952317515077652022726490027694212560352756646854116994689233
---   >>> (*) <$> sqrt 15 <*> sqrt 15
+--   >>> (*) <$> sqrt_vartime 15 <*> sqrt_vartime 15
 --   Just 15
-sqrt :: Montgomery -> Maybe Montgomery
-sqrt (Montgomery n) = case sqrt# n of
-  (# a | #) -> Just $! Montgomery a
-  _         -> Nothing
+sqrt_vartime :: Montgomery -> Maybe Montgomery
+sqrt_vartime (Montgomery n) = case sqrt# n of
+  (# a, c #)
+    | C.decide c -> Just $! Montgomery a
+    | otherwise  -> Nothing
 
 -- generated by etc/generate_sqrt.sh
 sqrt#
   :: (# Limb, Limb, Limb, Limb #)
-  -> (# (# Limb, Limb, Limb, Limb #) | () #)
+  -> (# (# Limb, Limb, Limb, Limb #), C.Choice #)
 sqrt# a =
   let !t0 = (# Limb 0x1000003D1##, Limb 0##, Limb 0##, Limb 0## #)
       !t1 = sqr# t0
@@ -1513,9 +1516,7 @@ sqrt# a =
       !t502 = sqr# t501
       !t503 = sqr# t502
       !r = t503
-  in  if   C.decide (WW.eq# (sqr# r) a)
-      then (# r | #)
-      else (# | () #)
+  in  (# r, WW.eq# (sqr# r) a #)
 {-# INLINE sqrt# #-}
 
 -- | Exponentiation in the Montgomery domain.