use alloy_primitives::Address;
use alloy_signer::k256::ecdsa::VerifyingKey;
use alloy_signer::utils::public_key_to_address;
use rayon::prelude::*;
use crate::{Stamp, StampDigest, StampError};
use nectar_primitives::SwarmAddress;
#[derive(Debug, Clone)]
pub struct VerifyResult {
pub index: usize,
pub result: Result<Address, StampError>,
}
pub fn verify_stamps_parallel(stamps: &[(&Stamp, &SwarmAddress)]) -> Vec<VerifyResult> {
stamps
.par_iter()
.enumerate()
.map(|(index, (stamp, address))| {
let result = recover_stamp_signer(stamp, address);
VerifyResult { index, result }
})
.collect()
}
pub fn verify_stamps_parallel_with_owner(
stamps: &[(&Stamp, &SwarmAddress)],
expected_owner: Address,
) -> Vec<VerifyResult> {
stamps
.par_iter()
.enumerate()
.map(|(index, (stamp, address))| {
let result = verify_stamp_owner(stamp, address, expected_owner);
VerifyResult { index, result }
})
.collect()
}
pub fn verify_stamps_parallel_with_pubkey(
stamps: &[(&Stamp, &SwarmAddress)],
owner_pubkey: &VerifyingKey,
) -> Vec<VerifyResult> {
let owner_address = public_key_to_address(owner_pubkey);
stamps
.par_iter()
.enumerate()
.map(|(index, (stamp, address))| {
let result = match stamp.verify_with_pubkey(address, owner_pubkey) {
Ok(()) => Ok(owner_address),
Err(e) => Err(e),
};
VerifyResult { index, result }
})
.collect()
}
fn recover_stamp_signer(stamp: &Stamp, address: &SwarmAddress) -> Result<Address, StampError> {
let digest = StampDigest::new(
*address,
stamp.batch(),
stamp.stamp_index(),
stamp.timestamp(),
);
let prehash = digest.to_prehash();
stamp
.signature()
.recover_address_from_msg(prehash.as_slice())
.map_err(|_| StampError::InvalidSignature)
}
fn verify_stamp_owner(
stamp: &Stamp,
address: &SwarmAddress,
expected_owner: Address,
) -> Result<Address, StampError> {
let recovered = recover_stamp_signer(stamp, address)?;
if recovered != expected_owner {
return Err(StampError::OwnerMismatch {
expected: expected_owner,
actual: recovered,
});
}
Ok(recovered)
}
#[cfg(test)]
mod tests {
use super::*;
use alloy_primitives::B256;
use alloy_signer::SignerSync;
use alloy_signer_local::PrivateKeySigner;
use crate::{Stamp, StampIndex, current_timestamp};
fn create_test_stamp(
signer: &PrivateKeySigner,
chunk_address: &SwarmAddress,
batch_id: B256,
) -> Stamp {
let index = StampIndex::new(0, 0);
let timestamp = current_timestamp();
let digest = StampDigest::new(*chunk_address, batch_id, index, timestamp);
let prehash = digest.to_prehash();
let sig = signer.sign_message_sync(prehash.as_slice()).unwrap();
Stamp::with_index(batch_id, index, timestamp, sig)
}
#[test]
fn test_parallel_verification() {
let signer = PrivateKeySigner::random();
let expected_owner = signer.address();
let batch_id = B256::ZERO;
let addresses: Vec<_> = (0..50)
.map(|_| SwarmAddress::from(B256::random()))
.collect();
let stamps: Vec<_> = addresses
.iter()
.map(|addr| create_test_stamp(&signer, addr, batch_id))
.collect();
let verify_input: Vec<_> = stamps.iter().zip(addresses.iter()).collect();
let verify_results = verify_stamps_parallel_with_owner(&verify_input, expected_owner);
assert_eq!(verify_results.len(), 50);
for result in &verify_results {
assert!(result.result.is_ok());
assert_eq!(result.result.as_ref().unwrap(), &expected_owner);
}
}
#[test]
fn test_verify_wrong_signer() {
let signer = PrivateKeySigner::random();
let wrong_owner = Address::repeat_byte(0xFF);
let batch_id = B256::ZERO;
let address = SwarmAddress::from(B256::random());
let stamp = create_test_stamp(&signer, &address, batch_id);
let verify_input = [(&stamp, &address)];
let verify_results = verify_stamps_parallel_with_owner(&verify_input, wrong_owner);
assert!(matches!(
verify_results[0].result,
Err(StampError::OwnerMismatch { .. })
));
}
#[test]
fn test_verify_stamps_parallel_basic() {
let signer = PrivateKeySigner::random();
let expected_owner = signer.address();
let batch_id = B256::ZERO;
let address = SwarmAddress::from(B256::random());
let stamp = create_test_stamp(&signer, &address, batch_id);
let verify_input = [(&stamp, &address)];
let results = verify_stamps_parallel(&verify_input);
assert_eq!(results.len(), 1);
assert_eq!(results[0].result.as_ref().unwrap(), &expected_owner);
}
#[test]
fn test_verify_stamps_parallel_with_pubkey() {
let signer = PrivateKeySigner::random();
let expected_owner = signer.address();
let batch_id = B256::ZERO;
let addresses: Vec<_> = (0..50)
.map(|_| SwarmAddress::from(B256::random()))
.collect();
let stamps: Vec<_> = addresses
.iter()
.map(|addr| create_test_stamp(&signer, addr, batch_id))
.collect();
let pubkey = stamps[0].recover_pubkey(&addresses[0]).unwrap();
let verify_input: Vec<_> = stamps.iter().zip(addresses.iter()).collect();
let verify_results = verify_stamps_parallel_with_pubkey(&verify_input, &pubkey);
assert_eq!(verify_results.len(), 50);
for result in &verify_results {
assert!(result.result.is_ok());
assert_eq!(result.result.as_ref().unwrap(), &expected_owner);
}
}
}