small refactor and check if fp is canonical

crates/precompile/src/bls12_381/g1.rs

@@ -1,9 +1,6 @@
-use super::utils::{fp_to_bytes, remove_padding, PADDED_FP_LENGTH};
-use blst::{
-    blst_bendian_from_fp, blst_fp, blst_fp_from_bendian, blst_p1_affine, blst_p1_affine_in_g1,
-    blst_p1_affine_on_curve,
-};
-use revm_primitives::{Bytes, PrecompileError};
+use super::utils::{fp_from_bendian, fp_to_bytes, remove_padding, PADDED_FP_LENGTH};
+use crate::primitives::{Bytes, PrecompileError};
+use blst::{blst_p1_affine, blst_p1_affine_in_g1, blst_p1_affine_on_curve};
 
 /// Length of each of the elements in a g1 operation input.
 pub(super) const G1_INPUT_ITEM_LENGTH: usize = 128;
@@ -33,35 +30,13 @@ pub(super) fn decode_and_check_g1(
     p0_y: &[u8; 48],
 ) -> Result<blst_p1_affine, PrecompileError> {
     let out = blst_p1_affine {
-        x: check_canonical_fp(p0_x)?,
-        y: check_canonical_fp(p0_y)?,
+        x: fp_from_bendian(p0_x)?,
+        y: fp_from_bendian(p0_y)?,
     };
 
     Ok(out)
 }
 
-/// Checks whether or not the input represents a canonical field element, returning the field
-/// element if successful.
-pub(super) fn check_canonical_fp(input: &[u8; 48]) -> Result<blst_fp, PrecompileError> {
-    let mut fp = blst_fp::default();
-    let mut out = [0; 48];
-    // SAFETY: input has fixed length, and fp is a blst value.
-    unsafe {
-        // TODO: cursed and horrible, we need to ensure the input is canonical. It would be nice if
-        // `blst_fp_from_bendian` checked for canonical
-        //
-        // This performs the check for canonical field elements
-        blst_fp_from_bendian(&mut fp, input.as_ptr());
-        blst_bendian_from_fp(out.as_mut_ptr(), &fp);
-    }
-
-    if *input != out {
-        return Err(PrecompileError::Other("non-canonical fp value".to_string()));
-    }
-
-    Ok(fp)
-}
-
 /// Extracts a G1 point in Affine format from a 128 byte slice representation.
 ///
 /// NOTE: This function will perform a G1 subgroup check if `subgroup_check` is set to `true`.

crates/precompile/src/bls12_381/g2.rs

@@ -1,9 +1,6 @@
-use super::{
-    g1::check_canonical_fp,
-    utils::{fp_to_bytes, remove_padding, FP_LENGTH, PADDED_FP_LENGTH},
-};
+use super::utils::{fp_from_bendian, fp_to_bytes, remove_padding, FP_LENGTH, PADDED_FP_LENGTH};
+use crate::primitives::{Bytes, PrecompileError};
 use blst::{blst_fp2, blst_p2_affine, blst_p2_affine_in_g2, blst_p2_affine_on_curve};
-use revm_primitives::{Bytes, PrecompileError};
 
 /// Length of each of the elements in a g2 operation input.
 pub(super) const G2_INPUT_ITEM_LENGTH: usize = 256;
@@ -49,8 +46,8 @@ pub(super) fn check_canonical_fp2(
     input_1: &[u8; 48],
     input_2: &[u8; 48],
 ) -> Result<blst_fp2, PrecompileError> {
-    let fp_1 = check_canonical_fp(input_1)?;
-    let fp_2 = check_canonical_fp(input_2)?;
+    let fp_1 = fp_from_bendian(input_1)?;
+    let fp_2 = fp_from_bendian(input_2)?;
 
     let fp2 = blst_fp2 { fp: [fp_1, fp_2] };
 

crates/precompile/src/bls12_381/map_fp_to_g1.rs

@@ -1,6 +1,6 @@
 use super::{
-    g1::{check_canonical_fp, encode_g1_point},
-    utils::{remove_padding, PADDED_FP_LENGTH},
+    g1::encode_g1_point,
+    utils::{fp_from_bendian, remove_padding, PADDED_FP_LENGTH},
 };
 use crate::{u64_to_address, PrecompileWithAddress};
 use blst::{blst_map_to_g1, blst_p1, blst_p1_affine, blst_p1_to_affine};
@@ -32,7 +32,7 @@ pub(super) fn map_fp_to_g1(input: &Bytes, gas_limit: u64) -> PrecompileResult {
     }
 
     let input_p0 = remove_padding(input)?;
-    let fp = check_canonical_fp(input_p0)?;
+    let fp = fp_from_bendian(input_p0)?;
 
     let mut p = blst_p1::default();
     // SAFETY: p and fp are blst values.

crates/precompile/src/bls12_381/utils.rs

@@ -1,4 +1,8 @@
-use blst::{blst_bendian_from_fp, blst_fp, blst_scalar, blst_scalar_from_bendian};
+use core::cmp::Ordering;
+
+use blst::{
+    blst_bendian_from_fp, blst_fp, blst_fp_from_bendian, blst_scalar, blst_scalar_from_bendian,
+};
 use revm_primitives::PrecompileError;
 
 /// Number of bits used in the BLS12-381 curve finite field elements.
@@ -13,6 +17,12 @@ pub(super) const PADDED_FP2_LENGTH: usize = 128;
 pub(super) const PADDING_LENGTH: usize = 16;
 /// Scalar length.
 pub(super) const SCALAR_LENGTH: usize = 32;
+// Big-endian non-Montgomery form.
+pub(super) const MODULUS_REPR: [u8; 48] = [
+    0x1a, 0x01, 0x11, 0xea, 0x39, 0x7f, 0xe6, 0x9a, 0x4b, 0x1b, 0xa7, 0xb6, 0x43, 0x4b, 0xac, 0xd7,
+    0x64, 0x77, 0x4b, 0x84, 0xf3, 0x85, 0x12, 0xbf, 0x67, 0x30, 0xd2, 0xa0, 0xf6, 0xb0, 0xf6, 0x24,
+    0x1e, 0xab, 0xff, 0xfe, 0xb1, 0x53, 0xff, 0xff, 0xb9, 0xfe, 0xff, 0xff, 0xff, 0xff, 0xaa, 0xab,
+];
 
 /// Encodes a single finite field element into byte slice with padding.
 pub(super) fn fp_to_bytes(out: &mut [u8], input: *const blst_fp) {
@@ -57,3 +67,32 @@ pub(super) fn extract_scalar_input(input: &[u8]) -> Result<blst_scalar, Precompi
 
     Ok(out)
 }
+
+/// Checks if the input is a valid big-endian representation of a field element.
+fn is_valid_be(input: &[u8; 48]) -> bool {
+    for (i, modul) in input.iter().zip(MODULUS_REPR.iter()) {
+        match i.cmp(modul) {
+            Ordering::Greater => return false,
+            Ordering::Less => return true,
+            Ordering::Equal => continue,
+        }
+    }
+    // false if matching the modulus
+    false
+}
+
+/// Checks whether or not the input represents a canonical field element, returning the field
+/// element if successful.
+pub(super) fn fp_from_bendian(input: &[u8; 48]) -> Result<blst_fp, PrecompileError> {
+    if !is_valid_be(input) {
+        return Err(PrecompileError::Other("non-canonical fp value".to_string()));
+    }
+    let mut fp = blst_fp::default();
+    // SAFETY: input has fixed length, and fp is a blst value.
+    unsafe {
+        // This performs the check for canonical field elements
+        blst_fp_from_bendian(&mut fp, input.as_ptr());
+    }
+
+    Ok(fp)
+}