Create a cache for BLS

README.md

@@ -25,6 +25,8 @@ cargo test --all --release
 You may need a python virtual environment activated for the tests to link.
 This seems to be caused by the pyo3 dependency in the `wheel`.
 
+To run the python tests, check out the README in the `wheel` folder.
+
 # Benchmarks
 
 To run benchmarks for a specific crate:

crates/chia-bls/Cargo.toml

@@ -23,7 +23,9 @@ blst = { version = "0.3.11", git = "https://github.com/supranational/blst.git",
 hex = "0.4.3"
 thiserror = "1.0.44"
 pyo3 = { version = "0.19.0", features = ["multiple-pymethods"], optional = true }
-arbitrary = { version = "1.3.0", optional = true }
+arbitrary = { version = "1.3.0" , optional = true}
+lru = "0.12.2"
+
 
 [dev-dependencies]
 rand = "0.8.5"

crates/chia-bls/src/cached_bls.rs

@@ -0,0 +1,234 @@
+// This cache is a bit weird because it's trying to account for validating
+// mempool signatures versus block signatures. When validating block signatures,
+// there's not much point in caching the pairings because we're probably not going
+// to see them again unless there's a reorg. However, a spend in the mempool
+// is likely to reappear in a block later, so we can save having to do the pairing
+// again. So caching is primarily useful when synced and monitoring the mempool in real-time.
+
+use crate::aggregate_verify_gt as agg_ver_gt;
+use crate::gtelement::GTElement;
+use crate::hash_to_g2;
+use crate::PublicKey;
+use crate::Signature;
+use lru::LruCache;
+use sha2::{Digest, Sha256};
+use std::num::NonZeroUsize;
+
+#[cfg(feature = "py-bindings")]
+use pyo3::types::{PyBool, PyInt, PyList};
+#[cfg(feature = "py-bindings")]
+use pyo3::{pyclass, pymethods, PyResult};
+
+#[cfg_attr(feature = "py-bindings", pyclass(name = "BLSCache"))]
+pub struct BLSCache {
+    cache: LruCache<[u8; 32], GTElement>, // LRUCache of hash(pubkey + message) -> GTElement
+}
+
+impl Default for BLSCache {
+    fn default() -> Self {
+        Self::new(50000)
+    }
+}
+
+impl BLSCache {
+    pub fn new(cache_size: usize) -> BLSCache {
+        let cache: LruCache<[u8; 32], GTElement> =
+            LruCache::new(NonZeroUsize::new(cache_size).unwrap());
+        Self { cache }
+    }
+
+    pub fn generator(cache_size: Option<usize>) -> Self {
+        let cache: LruCache<[u8; 32], GTElement> =
+            LruCache::new(NonZeroUsize::new(cache_size.unwrap_or(50000)).unwrap());
+        Self { cache }
+    }
+
+    pub fn aggregate_verify<M: AsRef<[u8]>>(
+        &mut self,
+        pks: &[PublicKey],
+        msgs: &[M],
+        sig: &Signature,
+    ) -> bool {
+        let iter = pks.iter().zip(msgs.iter()).map(|(pk, msg)| -> GTElement {
+            let mut hasher = Sha256::new();
+            hasher.update(pk.to_bytes());
+            hasher.update(msg); // pk + msg
+            let h: [u8; 32] = hasher.finalize().into();
+
+            let pairing: Option<GTElement> = self.cache.get(&h).cloned();
+
+            if let Some(pairing) = pairing {
+                // equivalent to `if pairing is not None`
+                return pairing;
+            }
+            // if pairing is None then make pairing and add to cache
+            let mut aug_msg = pk.to_bytes().to_vec();
+            aug_msg.extend_from_slice(msg.as_ref()); // pk + msg
+            let aug_hash: Signature = hash_to_g2(&aug_msg);
+
+            let pairing: GTElement = aug_hash.pair(pk);
+            let mut hasher = Sha256::new();
+            hasher.update(&aug_msg);
+            let h: [u8; 32] = hasher.finalize().into();
+            self.cache.put(h, pairing.clone());
+            pairing
+        });
+        agg_ver_gt(sig, iter)
+    }
+}
+
+// Python Functions
+
+// Commented out for now as we may remove these
+// as the python consensus code that uses it is being ported to rust.
+
+// #[cfg(feature = "py-bindings")]
+// #[pymethods]
+// impl BLSCache {
+//     #[new]
+//     pub fn init() -> Self {
+//         Self::default()
+//     }
+
+//     #[staticmethod]
+//     #[pyo3(name = "generator")]
+//     pub fn py_generator(size: Option<&PyInt>) -> Self {
+//         size.map(|s| Self::new(s.extract().unwrap()))
+//             .unwrap_or_default()
+//     }
+
+//     #[pyo3(name = "aggregate_verify")]
+//     pub fn py_aggregate_verify(
+//         &mut self,
+//         pks: &PyList,
+//         msgs: &PyList,
+//         sig: &Signature,
+//         force_cache: &PyBool,
+//     ) -> PyResult<bool> {
+//         let pks_r: Vec<[u8; 48]> = pks
+//             .iter()
+//             .map(|item| item.extract::<[u8; 48]>())
+//             .collect::<PyResult<_>>()?;
+//         let msgs_r: Vec<&[u8]> = msgs
+//             .iter()
+//             .map(|item| item.extract::<&[u8]>())
+//             .collect::<PyResult<_>>()?;
+//         let force_cache_bool = force_cache.extract::<bool>()?;
+//         Ok(self.aggregate_verify(&pks_r, &msgs_r, sig, force_cache_bool))
+//     }
+
+//     #[pyo3(name = "len")]
+//     pub fn py_len(&self) -> PyResult<usize> {
+//         Ok(self.cache.len())
+//     }
+// }
+
+#[cfg(test)]
+pub mod tests {
+    use super::*;
+    use crate::aggregate;
+    use crate::sign;
+    use crate::SecretKey;
+
+    #[test]
+    pub fn test_instantiation() {
+        let mut bls_cache: BLSCache = BLSCache::default();
+        let byte_array: [u8; 32] = [0; 32];
+        let sk: SecretKey = SecretKey::from_seed(&byte_array);
+        let pk: PublicKey = sk.public_key();
+        let msg: [u8; 32] = [106; 32];
+        let mut aug_msg: Vec<u8> = pk.clone().to_bytes().to_vec();
+        aug_msg.extend_from_slice(&msg); // pk + msg
+        let aug_hash = hash_to_g2(&aug_msg);
+        let pairing = aug_hash.pair(&pk);
+        let mut hasher = Sha256::new();
+        hasher.update(&aug_msg);
+        let h: [u8; 32] = hasher.finalize().into();
+        bls_cache.cache.put(h, pairing.clone());
+        assert_eq!(*bls_cache.cache.get(&h).unwrap(), pairing);
+    }
+
+    #[test]
+    pub fn test_aggregate_verify() {
+        let mut bls_cache: BLSCache = BLSCache::default();
+        assert_eq!(bls_cache.cache.len(), 0);
+        let byte_array: [u8; 32] = [0; 32];
+        let sk: SecretKey = SecretKey::from_seed(&byte_array);
+        let pk: PublicKey = sk.public_key();
+        let msg: &[u8] = &[106; 32];
+        let sig: Signature = sign(&sk, msg);
+        let pk_list: Vec<PublicKey> = [pk].to_vec();
+        let msg_list: Vec<&[u8]> = [msg].to_vec();
+        assert!(bls_cache.aggregate_verify(&pk_list, &msg_list, &sig));
+        assert_eq!(bls_cache.cache.len(), 1);
+        // try again with (pk, msg) cached
+        assert!(bls_cache.aggregate_verify(&pk_list, &msg_list, &sig));
+        assert_eq!(bls_cache.cache.len(), 1);
+    }
+
+    #[test]
+    pub fn test_cache() {
+        let mut bls_cache: BLSCache = BLSCache::default();
+        assert_eq!(bls_cache.cache.len(), 0);
+        let byte_array: [u8; 32] = [0; 32];
+        let sk: SecretKey = SecretKey::from_seed(&byte_array);
+        let pk: PublicKey = sk.public_key();
+        let msg: &[u8] = &[106; 32];
+        let sig: Signature = sign(&sk, msg);
+        let mut pk_list: Vec<PublicKey> = [pk].to_vec();
+        let mut msg_list: Vec<&[u8]> = [msg].to_vec();
+        // add first to cache
+        // try one cached, one not cached
+        assert!(bls_cache.aggregate_verify(&pk_list, &msg_list, &sig));
+        assert_eq!(bls_cache.cache.len(), 1);
+        let byte_array: [u8; 32] = [1; 32];
+        let sk: SecretKey = SecretKey::from_seed(&byte_array);
+        let pk: PublicKey = sk.public_key();
+        let msg: &[u8] = &[107; 32];
+        let sig = aggregate([sig, sign(&sk, msg)]);
+        pk_list.push(pk);
+        msg_list.push(msg);
+        assert!(bls_cache.aggregate_verify(&pk_list, &msg_list, &sig));
+        assert_eq!(bls_cache.cache.len(), 2);
+        // try reusing a pubkey
+        let pk: PublicKey = sk.public_key();
+        let msg: &[u8] = &[108; 32];
+        let sig = aggregate([sig, sign(&sk, msg)]);
+        pk_list.push(pk);
+        msg_list.push(msg);
+        // check verification
+        assert!(bls_cache.aggregate_verify(&pk_list, &msg_list, &sig));
+        assert_eq!(bls_cache.cache.len(), 3);
+    }
+
+    #[test]
+    pub fn test_cache_limit() {
+        // set cache size to 3
+        let mut bls_cache: BLSCache = BLSCache::new(3);
+        assert_eq!(bls_cache.cache.len(), 0);
+        // create 5 pk/msg combos
+        for i in 1..=5 {
+            let byte_array: [u8; 32] = [i as u8; 32];
+            let sk: SecretKey = SecretKey::from_seed(&byte_array);
+            let pk: PublicKey = sk.public_key();
+            let msg: &[u8] = &[106; 32];
+            let sig: Signature = sign(&sk, msg);
+            let pk_list: Vec<PublicKey> = [pk].to_vec();
+            let msg_list: Vec<&[u8]> = vec![msg];
+            assert!(bls_cache.aggregate_verify(&pk_list, &msg_list, &sig));
+        }
+        assert_eq!(bls_cache.cache.len(), 3);
+        // recreate first key
+        let byte_array: [u8; 32] = [1; 32];
+        let sk: SecretKey = SecretKey::from_seed(&byte_array);
+        let pk: PublicKey = sk.public_key();
+        let msg: Vec<u8> = vec![106; 32];
+        let mut aug_msg = pk.to_bytes().to_vec();
+        aug_msg.extend_from_slice(&msg); // pk + msg
+        let mut hasher = Sha256::new();
+        hasher.update(aug_msg);
+        let h: [u8; 32] = hasher.finalize().into();
+        // assert first key has been removed
+        assert!(bls_cache.cache.get(&h).is_none());
+    }
+}

crates/chia-bls/src/lib.rs

@@ -1,3 +1,4 @@
+pub mod cached_bls;
 pub mod derivable_key;
 pub mod derive_keys;
 pub mod error;
@@ -7,6 +8,7 @@ pub mod public_key;
 pub mod secret_key;
 pub mod signature;
 
+pub use cached_bls::BLSCache;
 pub use derivable_key::DerivableKey;
 pub use error::{Error, Result};
 pub use gtelement::GTElement;

tests/test_blscache.py

@@ -0,0 +1,58 @@
+from chia_rs import G1Element, PrivateKey, AugSchemeMPL, G2Element, BLSCache
+from chia.util.ints import uint64
+from chia.types.blockchain_format.sized_bytes import bytes48
+import pytest
+from typing import List
+
+
+def test_instantiation() -> None:
+    bls_cache = BLSCache()
+    assert bls_cache.len() == 0
+    assert BLSCache is not None
+    seed: bytes = bytes.fromhex(
+        "003206f418c701193458c013120c5906dc12663ad1520c3e596eb6092c14fe16"
+    )
+
+    sk: PrivateKey = AugSchemeMPL.key_gen(seed)
+    pk: G1Element = sk.get_g1()
+    msg = b"hello"
+    sig: G2Element = AugSchemeMPL.sign(sk, msg)
+    pks: List[bytes48] = [bytes48(pk.to_bytes())]
+    msgs: List[bytes] = [msg]
+    result = bls_cache.aggregate_verify(pks, msgs, sig, True)
+    assert result
+    assert bls_cache.len() == 1
+    result = bls_cache.aggregate_verify(pks, msgs, sig, True)
+    assert result
+    assert bls_cache.len() == 1
+    pks.append(bytes48(pk.to_bytes()))
+
+    msg = b"world"
+    msgs.append(msg)
+    sig = AugSchemeMPL.aggregate([sig, AugSchemeMPL.sign(sk, msg)])
+    result = bls_cache.aggregate_verify(pks, msgs, sig, True)
+    assert result
+    assert bls_cache.len() == 2
+
+
+def test_cache_limit() -> None:
+    bls_cache = BLSCache.generator(3)
+    assert bls_cache.len() == 0
+    assert BLSCache is not None
+    seed: bytes = bytes.fromhex(
+        "003206f418c701193458c013120c5906dc12663ad1520c3e596eb6092c14fe16"
+    )
+
+    sk: PrivateKey = AugSchemeMPL.key_gen(seed)
+    pk: G1Element = sk.get_g1()
+    pks: List[bytes48] = []
+    msgs: List[bytes] = []
+    pk_bytes = bytes48(pk.to_bytes())
+    sigs: List[G2Element] = []
+    for i in [0xCAFE, 0xF00D, 0xABCD, 0x1234]:
+        msgs.append(i.to_bytes(8, byteorder="little"))
+        pks.append(pk_bytes)
+        sigs.append(AugSchemeMPL.sign(sk, i.to_bytes(8, byteorder="little")))
+    result = bls_cache.aggregate_verify(pks, msgs, AugSchemeMPL.aggregate(sigs), True)
+    assert result
+    assert bls_cache.len() == 3

wheel/README.md

@@ -1 +1,2 @@
 The `chia_rs` wheel contains python bindings for code from the `chia` crate.
+

wheel/generate_type_stubs.py

@@ -338,6 +338,12 @@ def serialized_length(program: ReadableBuffer) -> int: ...
 def tree_hash(program: ReadableBuffer) -> bytes32: ...
 def get_puzzle_and_solution_for_coin(program: ReadableBuffer, args: ReadableBuffer, max_cost: int, find_parent: bytes32, find_amount: int, find_ph: bytes32, flags: int) -> Tuple[bytes, bytes]: ...
 
+class BLSCache:
+    def len(self) -> uint128: ...
+    def aggregate_verify(self, pks: List[bytes48], msgs: List[bytes], sig: G2Element, force_cache: bool) -> bool: ...
+    @staticmethod
+    def generator(size: Optional[int]) -> BLSCache: ...
+
 class AugSchemeMPL:
     @staticmethod
     def sign(pk: PrivateKey, msg: bytes, prepend_pk: Optional[G1Element] = None) -> G2Element: ...

wheel/python/chia_rs/chia_rs.pyi

@@ -81,6 +81,12 @@ def serialized_length(program: ReadableBuffer) -> int: ...
 def tree_hash(program: ReadableBuffer) -> bytes32: ...
 def get_puzzle_and_solution_for_coin(program: ReadableBuffer, args: ReadableBuffer, max_cost: int, find_parent: bytes32, find_amount: int, find_ph: bytes32, flags: int) -> Tuple[bytes, bytes]: ...
 
+class BLSCache:
+    def len(self) -> uint128: ...
+    def aggregate_verify(self, pks: List[bytes48], msgs: List[bytes], sig: G2Element, force_cache: bool) -> bool: ...
+    @staticmethod
+    def generator(size: Optional[int]) -> BLSCache: ...
+
 class AugSchemeMPL:
     @staticmethod
     def sign(pk: PrivateKey, msg: bytes, prepend_pk: Optional[G1Element] = None) -> G2Element: ...

wheel/src/api.rs

@@ -63,7 +63,8 @@ use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs};
 use clvmr::ChiaDialect;
 
 use chia_bls::{
-    hash_to_g2 as native_hash_to_g2, DerivableKey, GTElement, PublicKey, SecretKey, Signature,
+    hash_to_g2 as native_hash_to_g2, BLSCache, DerivableKey, GTElement, PublicKey, SecretKey,
+    Signature,
 };
 
 #[pyfunction]
@@ -512,6 +513,7 @@ pub fn chia_rs(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<GTElement>()?;
     m.add_class::<SecretKey>()?;
     m.add_class::<AugSchemeMPL>()?;
+    m.add_class::<BLSCache>()?;
 
     Ok(())
 }