From 53286f535c90357c6c29fb9e284c560708b34327 Mon Sep 17 00:00:00 2001
From: arnaucube <root@arnaucube.com>
Date: Tue, 9 Mar 2021 18:39:05 +0100
Subject: [PATCH] Add blake lib support for feature aarch64

---
 .github/workflows/test.yml |  4 +++-
 Cargo.toml                 |  8 ++++++--
 src/lib.rs                 | 31 ++++++++++++++++++++++++++-----
 3 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 9635bb9..58b7ded 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -10,4 +10,6 @@ jobs:
     - name: Build
       run: cargo build --verbose
     - name: Run tests
-      run: cargo test --verbose
+      run: |
+        cargo test --verbose
+        cargo test --verbose --no-default-features --features=aarch64
diff --git a/Cargo.toml b/Cargo.toml
index de411e9..652a5ef 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "babyjubjub-rs"
-version = "0.0.7"
+version = "0.0.8"
 authors = ["arnaucube <root@arnaucube.com>"]
 edition = "2018"
 license = "GPL-3.0"
@@ -15,7 +15,8 @@ rand6 = {package="rand", version="0.6.5"}
 num = "0.2.0"
 num-bigint = {version = "0.2.2", features = ["rand"]}
 num-traits = "0.2.8"
-blake-hash = "0.4.0"
+blake-hash = {version="0.4.0", optional=true}
+blake = {version="2.0.1", optional=true}
 generic-array = "0.13.2"
 tiny-keccak = "1.5"
 rustc-hex = "1.0.0"
@@ -30,3 +31,6 @@ criterion = "0.3"
 name = "bench_babyjubjub"
 harness = false
 
+[features]
+default = ["blake-hash"]
+aarch64 = ["blake"]
diff --git a/src/lib.rs b/src/lib.rs
index 392069e..c687c1c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,3 +1,5 @@
+// For LICENSE check https://github.com/arnaucube/babyjubjub-rs
+
 extern crate rand;
 #[macro_use]
 extern crate ff;
@@ -17,10 +19,16 @@ extern crate rand6;
 use rand6::Rng;
 
 // use blake2::{Blake2b, Digest};
+
+#[cfg(feature = "default")]
 extern crate blake_hash; // compatible version with Blake used at circomlib
+#[cfg(feature = "default")]
 #[macro_use]
 use blake_hash::Digest;
 
+#[cfg(feature = "aarch64")]
+extern crate blake; // compatible version with Blake used at circomlib
+
 use std::cmp::min;
 
 use num_bigint::{BigInt, RandBigInt, Sign, ToBigInt};
@@ -232,6 +240,19 @@ pub fn decompress_point(bb: [u8; 32]) -> Result<Point, String> {
     Ok(Point { x: x_fr, y: y_fr })
 }
 
+#[cfg(feature = "default")]
+fn blh(b: &Vec<u8>) -> Vec<u8> {
+    let hash = blake_hash::Blake512::digest(&b);
+    hash.to_vec()
+}
+
+#[cfg(feature = "aarch64")]
+fn blh(b: &Vec<u8>) -> Vec<u8> {
+    let mut hash = [0; 64];
+    blake::hash(512, b, &mut hash).unwrap();
+    hash.to_vec()
+}
+
 #[derive(Debug, Clone)]
 pub struct Signature {
     r_b8: Point,
@@ -287,7 +308,7 @@ impl PrivateKey {
         // let mut h = hasher.finalize();
 
         // compatible with circomlib implementation
-        let hash = blake_hash::Blake512::digest(&self.key.to_vec());
+        let mut hash: Vec<u8> = blh(&self.key.to_vec());
         let mut h: Vec<u8> = hash[..32].to_vec();
 
         h[0] = h[0] & 0xF8;
@@ -312,7 +333,7 @@ impl PrivateKey {
         // let mut hasher = Blake2b::new();
         // hasher.update(sk_bytes);
         // let mut h = hasher.finalize(); // h: hash(sk), s: h[32:64]
-        let mut h = blake_hash::Blake512::digest(&self.key);
+        let mut h: Vec<u8> = blh(&self.key.to_vec());
 
         let (_, msg_bytes) = msg.to_bytes_le();
         let mut msg32: [u8; 32] = [0; 32];
@@ -322,7 +343,7 @@ impl PrivateKey {
         // https://tools.ietf.org/html/rfc8032#section-5.1.6
         let s = GenericArray::<u8, generic_array::typenum::U32>::from_mut_slice(&mut h[32..64]);
         let r_bytes = utils::concatenate_arrays(s, &msg32);
-        let r_hashed = blake_hash::Blake512::digest(&r_bytes);
+        let r_hashed: Vec<u8> = blh(&r_bytes);
         let mut r = BigInt::from_bytes_le(Sign::Plus, &r_hashed[..]);
         r = utils::modulus(&r, &SUBORDER);
         let r8: Point = B8.mul_scalar(&r);
@@ -644,7 +665,7 @@ mod tests {
             let random_bytes = rand6::thread_rng().gen::<[u8; 32]>();
             let sk_raw: BigInt = BigInt::from_bytes_le(Sign::Plus, &random_bytes[..]);
             let (_, sk_raw_bytes) = sk_raw.to_bytes_be();
-            let mut h = blake_hash::Blake512::digest(&sk_raw_bytes);
+            let mut h: Vec<u8> = blh(&sk_raw_bytes);
 
             h[0] = h[0] & 0xF8;
             h[31] = h[31] & 0x7F;
@@ -699,7 +720,7 @@ mod tests {
                 .unwrap();
 
         // test blake compatible with circomlib implementation
-        let h = blake_hash::Blake512::digest(&sk_raw_bytes);
+        let mut h: Vec<u8> = blh(&sk_raw_bytes);
         assert_eq!(h.to_hex(), "c992db23d6290c70ffcc02f7abeb00b9d00fa8b43e55d7949c28ba6be7545d3253882a61bd004a236ef1cdba01b27ba0aedfb08eefdbfb7c19657c880b43ddf1");
 
         // test private key