From bf694c67ea3322f77b440a12b3e21e2c3940fe92 Mon Sep 17 00:00:00 2001
From: Ludo Pulles <ludo.pulles@gmail.com>
Date: Wed, 16 Oct 2024 13:24:08 +0200
Subject: [PATCH] Fix MitM Babai e-admissability computation!

It was simply wrong. The e^(-x^2) forgot to take the square in the exponent, and the sigma was not
really the standard deviation of the Gaussian sample (which they use in the proof).
---
 estimator/lwe_primal.py |  2 +-
 estimator/prob.py       | 24 +++++++-----------------
 2 files changed, 8 insertions(+), 18 deletions(-)

diff --git a/estimator/lwe_primal.py b/estimator/lwe_primal.py
index 6f38b24..80fed53 100644
--- a/estimator/lwe_primal.py
+++ b/estimator/lwe_primal.py
@@ -420,7 +420,7 @@ def ssf(x):
 
         if mitm and zeta > 0:
             if babai:
-                probability *= mitm_babai_probability(r, params.Xe.stddev, params.q)
+                probability *= mitm_babai_probability(r, params.Xe.stddev)
             else:
                 # TODO: the probability in this case needs to be analysed
                 probability *= 1
diff --git a/estimator/prob.py b/estimator/prob.py
index 345cff8..041dd2b 100644
--- a/estimator/prob.py
+++ b/estimator/prob.py
@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 from sage.all import binomial, ZZ, log, ceil, RealField, oo, exp, pi
 from sage.all import RealDistribution, RR, sqrt, prod, erf
-from .nd import sigmaf
 from .conf import max_n_cache
 
 
@@ -78,35 +77,26 @@ def gaussian_cdf(mu, sigma, t):
     return RR((1/2)*(1 + erf((t - mu)/(sqrt(2)*sigma))))
 
 
-def mitm_babai_probability(r, stddev, q, fast=False):
+def mitm_babai_probability(r, stddev, fast=False):
     """
     Compute the "e-admissibility" probability associated to the mitm step, according to
     [EPRINT:SonChe19]_
 
     :params r: the squared GSO lengths
     :params stddev: the std.dev of the error distribution
-    :params q: the LWE modulus
     :param fast: toggle for setting p = 1 (faster, but underestimates security)
     :return: probability for the mitm process
-
-    # NOTE: the model sometimes outputs negative probabilities, we set p = 0 in this case
     """
-
     if fast:
         # overestimate the probability -> underestimate security
         return 1
 
-    # get non-squared norms
-    alphaq = sigmaf(stddev)
-    probs = (
-        RR(
-            erf(s * sqrt(RR(pi)) / alphaq)
-            + (alphaq / s) * ((exp(-s * sqrt(RR(pi)) / alphaq) - 1) / RR(pi))
-        )
-        for s in map(sqrt, r)
-    )
-    p = RR(prod(probs))
-    return p if 0 <= p <= 1 else 0.0
+    # Note: `r` contains *square norms*, so convert to non-square norms.
+    # Follow the proof of Lemma 4.2 [EPRINT_SonChe19]_, because that one uses standard deviation.
+    xs = [sqrt(.5 * ri) / stddev for ri in r]
+    p = prod(RR(erf(x) - (1 - exp(-x**2)) / (x * sqrt(pi))) for x in xs)
+    assert 0.0 <= p <= 1.0
+    return p
 
 
 def babai(r, norm):