From f77f51c9596e8cdf7468c56329ad30d9fe335816 Mon Sep 17 00:00:00 2001
From: Zekun Li <52300039+zekun000@users.noreply.github.com>
Date: Fri, 1 Mar 2024 15:10:12 -0800
Subject: [PATCH] [dkg] remove aptos-dkg unsafe code (#12328)

Co-authored-by: Zekun Li <zekunli@Zekuns-MBP.localdomain>
Co-authored-by: Alin Tomescu <tomescu.alin@gmail.com>
---
 crates/aptos-dkg/src/utils/mod.rs             |  5 ++--
 .../src/utils/parallel_multi_pairing.rs       | 28 +++++--------------
 2 files changed, 9 insertions(+), 24 deletions(-)

diff --git a/crates/aptos-dkg/src/utils/mod.rs b/crates/aptos-dkg/src/utils/mod.rs
index b8b07e0d37717..01e3716ad1442 100644
--- a/crates/aptos-dkg/src/utils/mod.rs
+++ b/crates/aptos-dkg/src/utils/mod.rs
@@ -1,8 +1,7 @@
 // Copyright © Aptos Foundation
 
 use crate::utils::{
-    parallel_multi_pairing::parallel_multi_miller_loop_and_final_exp,
-    random::random_scalar_from_uniform_bytes,
+    parallel_multi_pairing::parallel_multi_pairing_slice, random::random_scalar_from_uniform_bytes,
 };
 use blstrs::{
     pairing, Bls12, G1Affine, G1Projective, G2Affine, G2Prepared, G2Projective, Gt, Scalar,
@@ -112,7 +111,7 @@ where
     I1: Iterator<Item = &'a G1Projective>,
     I2: Iterator<Item = &'a G2Projective>,
 {
-    parallel_multi_miller_loop_and_final_exp(
+    parallel_multi_pairing_slice(
         lhs.zip(rhs)
             .map(|(g1, g2)| (g1.to_affine(), g2.to_affine()))
             .collect::<Vec<(G1Affine, G2Affine)>>()
diff --git a/crates/aptos-dkg/src/utils/parallel_multi_pairing.rs b/crates/aptos-dkg/src/utils/parallel_multi_pairing.rs
index 44d8458b743a5..d93890267148e 100644
--- a/crates/aptos-dkg/src/utils/parallel_multi_pairing.rs
+++ b/crates/aptos-dkg/src/utils/parallel_multi_pairing.rs
@@ -1,13 +1,12 @@
 // Copyright © Aptos Foundation
 
-use blst::{blst_final_exp, blst_fp12, blst_fp12_mul, blst_fp12_one, blst_miller_loop};
+use blst::blst_fp12;
 use blstrs::{Fp12, G1Affine, G2Affine, Gt};
 use group::prime::PrimeCurveAffine;
 use rayon::{prelude::*, ThreadPool};
 
-/// Computes $$\sum_{i=1}^n \textbf{ML}(a_i, b_i)$$ given a series of terms
-/// $$(a_1, b_1), (a_2, b_2), ..., (a_n, b_n).$$
-pub fn parallel_multi_miller_loop_and_final_exp(
+/// Computes a multi-pairing $$\prod_{i=1}^n e(a_i, b_i)$$ using multiple threads from `pool`.
+pub fn parallel_multi_pairing_slice(
     terms: &[(&G1Affine, &G2Affine)],
     pool: &ThreadPool,
     min_length: usize,
@@ -19,27 +18,14 @@ pub fn parallel_multi_miller_loop_and_final_exp(
             .map(|(p, q)| {
                 if (p.is_identity() | q.is_identity()).into() {
                     // Define pairing with zero as one, matching what `pairing` does.
-                    unsafe { *blst_fp12_one() }
+                    blst_fp12::default()
                 } else {
-                    unsafe {
-                        let mut tmp = blst_fp12::default();
-                        blst_miller_loop(&mut tmp, q.as_ref(), p.as_ref());
-                        tmp
-                    }
+                    blst_fp12::miller_loop(q.as_ref(), p.as_ref())
                 }
             })
-            .reduce(
-                || unsafe { *blst_fp12_one() },
-                |mut acc, val| {
-                    unsafe {
-                        blst_fp12_mul(&mut acc, &acc, &val);
-                    }
-                    acc
-                },
-            )
+            .reduce(|| blst_fp12::default(), |acc, val| acc * val)
     });
 
-    let mut out = blst_fp12::default();
-    unsafe { blst_final_exp(&mut out, &res) };
+    let out = blst_fp12::final_exp(&res);
     Fp12::from(out).into()
 }