Create a cache for BLS

crates/chia-bls/src/cached_bls.rs

@@ -19,7 +19,7 @@ use std::num::NonZeroUsize;
 #[cfg(feature = "py-bindings")]
 use pyo3::exceptions::PyValueError;
 #[cfg(feature = "py-bindings")]
-use pyo3::types::{PyInt, PyList};
+use pyo3::types::PyList;
 #[cfg(feature = "py-bindings")]
 use pyo3::{pyclass, pymethods, PyResult};
 
@@ -37,8 +37,9 @@ impl Default for BLSCache {
 
 impl BLSCache {
     pub fn new(cache_size: Option<NonZeroUsize>) -> BLSCache {
-        let cache: LruCache<[u8; 32], GTElement> =
-            LruCache::new(cache_size.unwrap_or(NonZeroUsize::new(50000).unwrap()));
+        let cache: LruCache<[u8; 32], GTElement> = LruCache::new(
+            cache_size.unwrap_or(NonZeroUsize::new(50000).expect("50000 should be non-zero")),
+        );
         Self { cache }
     }
 
@@ -88,18 +89,15 @@ impl BLSCache {
 #[pymethods]
 impl BLSCache {
     #[new]
-    pub fn init(size: Option<&PyInt>) -> PyResult<Self> {
+    pub fn init(size: Option<u32>) -> PyResult<Self> {
         match size {
             Some(p_size) => {
-                let r_size = p_size.extract::<isize>().unwrap();
-                if r_size < 1 {
+                if p_size < 1 {
                     Err(PyValueError::new_err(
                         "Cannot have a cache size less than one.",
                     ))
                 } else {
-                    Ok(Self::new(NonZeroUsize::new(
-                        p_size.extract::<usize>().unwrap(),
-                    )))
+                    Ok(Self::new(NonZeroUsize::new(p_size as usize)))
                 }
             }
             None => Ok(Self::default()),
@@ -216,6 +214,7 @@ pub mod tests {
             let sig: Signature = sign(&sk, msg);
             let pk_list: Vec<PublicKey> = [pk].to_vec();
             let msg_list: Vec<&[u8]> = vec![msg];
+            // add to cache by validating them one at a time
             assert!(bls_cache.aggregate_verify(pk_list.iter(), msg_list.iter(), &sig));
         }
         assert_eq!(bls_cache.cache.len(), 3);
@@ -232,4 +231,12 @@ pub mod tests {
         // assert first key has been removed
         assert!(bls_cache.cache.get(&h).is_none());
     }
+
+    pub fn test_empty_sig() {
+        let mut bls_cache: BLSCache = BLSCache::default();
+        let sig: Signature = aggregate(&[]);
+        let pk_list: [PublicKey; 0] = [];
+        let msg_list: Vec<&[u8]> = vec![];
+        assert!(bls_cache.aggregate_verify(pk_list, msg_list, &sig));
+    }
 }

tests/test_blscache.py

@@ -135,17 +135,30 @@ def test_cached_bls_repeat_pk():
     assert cached_bls_old.aggregate_verify(pks_bytes, msgs, agg_sig, True)
 
 
+def test_empty_sig():
+    sig = AugSchemeMPL.aggregate([])
+    cached_bls = BLSCache()
+    assert cached_bls.aggregate_verify([], [], sig)
+    assert cached_bls_old.aggregate_verify([], [], sig)
+
+
 def test_bad_cache_size():
     with pytest.raises(ValueError) as exc_info:
         bls_cache = BLSCache(0)
 
     assert str(exc_info.value) == "Cannot have a cache size less than one."
-    with pytest.raises(ValueError) as exc_info:
+
+    with pytest.raises(OverflowError) as exc_info:
         bls_cache = BLSCache(-1)
 
-    assert str(exc_info.value) == "Cannot have a cache size less than one."
+    assert str(exc_info.value) == "can't convert negative int to unsigned"
 
-    with pytest.raises(ValueError) as exc_info:
+    with pytest.raises(OverflowError) as exc_info:
         bls_cache = BLSCache(-100000)
 
-    assert str(exc_info.value) == "Cannot have a cache size less than one."
+    assert str(exc_info.value) == "can't convert negative int to unsigned"
+
+    with pytest.raises(OverflowError) as exc_info:
+        bls_cache = BLSCache(-9223372036854775809)
+
+    assert str(exc_info.value) == "can't convert negative int to unsigned"