From cea18d17d36dd84c09a0e6b0d8b17bc067ab4437 Mon Sep 17 00:00:00 2001
From: Samuel Schlesinger <sgschlesinger@gmail.com>
Date: Fri, 2 Oct 2026 23:33:00 -0400
Subject: [PATCH] fix(DiffieHellman): reject mismatched moduli

---
 .../DiffieHellman/Basic.lean                  |  7 ++++---
 CslibTests/DiffieHellman.lean                 | 20 +++++++++----------
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/Cslib/Algorithms/StatefulProcesses/DiffieHellman/Basic.lean b/Cslib/Algorithms/StatefulProcesses/DiffieHellman/Basic.lean
index f0d04272c..2b7187b98 100644
--- a/Cslib/Algorithms/StatefulProcesses/DiffieHellman/Basic.lean
+++ b/Cslib/Algorithms/StatefulProcesses/DiffieHellman/Basic.lean
@@ -139,12 +139,13 @@ def bob : Process Pid Var params.Val FunId params.SelLabel params.ProcName :=
 
 /-! ## Semantics -/
 
-/-- Implementation of local function call evaluation. -/
+/-- Implementation of local function call evaluation.
+Calls with a mismatched modulus are rejected. -/
 def funEval : FunCallEval FunId params.Val
   | .computePublicMessage, [.nat p, .zMod g, .nat privateExp], v =>
-    (h : p = params.p) → v = (.zMod <| h ▸ computePublicMessage p (h ▸ g) privateExp)
+    p = params.p ∧ v = .zMod (computePublicMessage params.p g privateExp)
   | .computeSharedSecret, [.nat p, .zMod msg, .nat privateExp], v =>
-    (h : p = params.p) → v = (.zMod <| h ▸ computeSharedSecret p (h ▸ msg) privateExp)
+    p = params.p ∧ v = .zMod (computeSharedSecret params.p msg privateExp)
   | _, _, _ => False
 
 /-- DH network. -/
diff --git a/CslibTests/DiffieHellman.lean b/CslibTests/DiffieHellman.lean
index 9e8039258..076a6a5e3 100644
--- a/CslibTests/DiffieHellman.lean
+++ b/CslibTests/DiffieHellman.lean
@@ -16,30 +16,30 @@ variable {Pid Var : Type*} (params : Params Pid Var)
 example (σ : LocalStore Var params.Val) :
     (funEval params).EvalExpr σ (aliceComputeMesg params) (.zMod (params.g ^ params.a)) := by
   apply FunCallEval.EvalExpr.call (.cons .val (.cons .val (.cons .val .nil)))
-  intro hmod
-  rfl
+  exact ⟨rfl, rfl⟩
 
 example (σ : LocalStore Var params.Val) :
     (funEval params).EvalExpr σ (bobComputeMesg params) (.zMod (params.g ^ params.b)) := by
   apply FunCallEval.EvalExpr.call (.cons .val (.cons .val (.cons .val .nil)))
-  intro hmod
-  rfl
+  exact ⟨rfl, rfl⟩
 
 -- Each role must also use its own exponent on the received message.
 example (σ : LocalStore Var params.Val) (message : ZMod params.p)
     (h : σ params.y = .zMod message) :
     (funEval params).EvalExpr σ (aliceComputeSharedSecret params) (.zMod (message ^ params.a)) := by
   apply FunCallEval.EvalExpr.call (.cons .val (.cons .var (.cons .val .nil)))
-  simp only [h, funEval]
-  intro hmod
-  rfl
+  simp [h, funEval, computeSharedSecret]
 
 example (σ : LocalStore Var params.Val) (message : ZMod params.p)
     (h : σ params.x = .zMod message) :
     (funEval params).EvalExpr σ (bobComputeSharedSecret params) (.zMod (message ^ params.b)) := by
   apply FunCallEval.EvalExpr.call (.cons .val (.cons .var (.cons .val .nil)))
-  simp only [h, funEval]
-  intro hmod
-  rfl
+  simp [h, funEval, computeSharedSecret]
+
+-- A mismatched modulus must reject every result for either function.
+example (f : FunId) (p : ℕ) (hp : p ≠ params.p) (message : ZMod params.p)
+    (privateExp : ℕ) (v : params.Val) :
+    ¬ funEval params f [.nat p, .zMod message, .nat privateExp] v := by
+  cases f <;> exact fun h => hp h.1
 
 end CslibTests.DiffieHellman