Skip to content

Commit

Permalink
tlock commitments unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@JohnReedV
JohnReedV committed Mar 3, 2025
1 parent be5bf5e commit 423fe2c
Show file tree
Hide file tree
Showing 5 changed files with 380 additions and 106 deletions.
3 changes: 3 additions & 0 deletions Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 4 additions & 0 deletions pallets/commitments/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,10 @@ pallet-drand = { path = "../drand", default-features = false }
tle = { workspace = true, default-features = false }
ark-serialize = { workspace = true, default-features = false }
w3f-bls = { workspace = true, default-features = false }
rand_chacha = { workspace = true }
hex = { workspace = true }
sha2 = { workspace = true }

log = { workspace = true }

[dev-dependencies]
Expand Down
179 changes: 78 additions & 101 deletions pallets/commitments/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ pub use weights::WeightInfo;

use ark_serialize::CanonicalDeserialize;
use frame_support::{BoundedVec, traits::Currency};
use frame_system::pallet_prelude::BlockNumberFor;
use sp_runtime::{Saturating, traits::Zero};
use sp_std::{boxed::Box, vec::Vec};
use tle::{
Expand Down Expand Up @@ -398,123 +397,101 @@ impl<T: Config> Pallet<T> {
.iter()
.find(|data| matches!(data, Data::TimelockEncrypted { .. }))
{
// Calculate reveal block
let reveal_block = Self::calculate_reveal_block(*reveal_round, registration.block)?;

// Check if the current block has reached or exceeded the reveal block
if current_block >= reveal_block {
// Deserialize the encrypted commitment into a TLECiphertext
let reader = &mut &encrypted[..];
let commit = TLECiphertext::<TinyBLS381>::deserialize_compressed(reader)
// Check if the corresponding Drand round data exists
let pulse = match pallet_drand::Pulses::<T>::get(*reveal_round) {
Some(p) => p,
None => continue,
};

// Prepare the signature bytes
let signature_bytes = pulse
.signature
.strip_prefix(b"0x")
.unwrap_or(&pulse.signature);
let sig_reader = &mut &signature_bytes[..];
let sig =
<TinyBLS381 as EngineBLS>::SignatureGroup::deserialize_compressed(sig_reader)
.map_err(|e| {
log::warn!("Failed to deserialize TLECiphertext for {:?}: {:?}", who, e)
log::warn!(
"Failed to deserialize drand signature for {:?}: {:?}",
who,
e
)
})
.ok();

let commit = match commit {
Some(c) => c,
None => continue,
};
let sig = match sig {
Some(s) => s,
None => continue,
};

// Get the drand pulse for the reveal round
let pulse = match pallet_drand::Pulses::<T>::get(*reveal_round) {
Some(p) => p,
None => {
log::warn!(
"Failed to reveal commit for subnet {} by {:?}: missing drand round {}",
netuid,
who,
reveal_round
);
continue;
}
};

// Prepare the signature bytes
let signature_bytes = pulse
.signature
.strip_prefix(b"0x")
.unwrap_or(&pulse.signature);
let sig_reader = &mut &signature_bytes[..];
let sig = <TinyBLS381 as EngineBLS>::SignatureGroup::deserialize_compressed(
sig_reader,
)
// Attempt to deserialize the encrypted commitment
let reader = &mut &encrypted[..];
let commit = TLECiphertext::<TinyBLS381>::deserialize_compressed(reader)
.map_err(|e| {
log::warn!(
"Failed to deserialize drand signature for {:?}: {:?}",
who,
e
)
log::warn!("Failed to deserialize TLECiphertext for {:?}: {:?}", who, e)
})
.ok();

let sig = match sig {
Some(s) => s,
None => continue,
};

// Decrypt the timelock commitment
let decrypted_bytes: Vec<u8> =
tld::<TinyBLS381, AESGCMStreamCipherProvider>(commit, sig)
.map_err(|e| {
log::warn!("Failed to decrypt timelock for {:?}: {:?}", who, e)
})
.ok()
.unwrap_or_default();

if decrypted_bytes.is_empty() {
continue;
}
let commit = match commit {
Some(c) => c,
None => continue,
};

// Decode the decrypted bytes into CommitmentInfo (assuming it’s SCALE-encoded CommitmentInfo)
let mut reader = &decrypted_bytes[..];
let revealed_info: CommitmentInfo<T::MaxFields> = Decode::decode(&mut reader)
// Decrypt the timelock commitment
let decrypted_bytes: Vec<u8> =
tld::<TinyBLS381, AESGCMStreamCipherProvider>(commit, sig)
.map_err(|e| {
log::warn!("Failed to decode decrypted data for {:?}: {:?}", who, e)
log::warn!("Failed to decrypt timelock for {:?}: {:?}", who, e)
})
.ok()
.unwrap_or_else(|| CommitmentInfo {
fields: BoundedVec::default(),
});

// Create RevealedData for storage
let revealed_data = RevealedData {
info: revealed_info,
revealed_block: current_block,
deposit: registration.deposit,
};

// Store the revealed data in RevealedCommitments
<RevealedCommitments<T>>::insert(netuid, &who, revealed_data);

// Remove the TimelockEncrypted field from the original commitment
let filtered_fields: Vec<Data> = registration.info.fields.into_iter()
.filter(|data| !matches!(data, Data::TimelockEncrypted { reveal_round: r, .. } if r == reveal_round))
.collect();
registration.info.fields = BoundedVec::try_from(filtered_fields)
.map_err(|_| "Failed to filter timelock fields")?;

Self::deposit_event(Event::CommitmentRevealed { netuid, who });
.unwrap_or_default();

if decrypted_bytes.is_empty() {
continue;
}

// Decode the decrypted bytes into CommitmentInfo
let mut reader = &decrypted_bytes[..];
let revealed_info: CommitmentInfo<T::MaxFields> = match Decode::decode(&mut reader)
{
Ok(info) => info,
Err(e) => {
log::warn!("Failed to decode decrypted data for {:?}: {:?}", who, e);
continue;
}
};

// Store the revealed data
let revealed_data = RevealedData {
info: revealed_info,
revealed_block: current_block,
deposit: registration.deposit,
};
<RevealedCommitments<T>>::insert(netuid, &who, revealed_data);

// Remove the TimelockEncrypted field from the original commitment
let filtered_fields: Vec<Data> = registration
.info
.fields
.into_iter()
.filter(|data| {
!matches!(
data,
Data::TimelockEncrypted {
reveal_round: r, ..
} if r == reveal_round
)
})
.collect();

registration.info.fields = BoundedVec::try_from(filtered_fields)
.map_err(|_| "Failed to filter timelock fields")?;

Self::deposit_event(Event::CommitmentRevealed { netuid, who });
}
}

Ok(())
}

fn calculate_reveal_block(
reveal_round: u64,
commit_block: BlockNumberFor<T>,
) -> Result<BlockNumberFor<T>, &'static str> {
let last_drand_round = pallet_drand::LastStoredRound::<T>::get();
let blocks_per_round = 12_u64.checked_div(3).unwrap_or(0); // 4 blocks per round (12s blocktime / 3s round)
let rounds_since_last = reveal_round.saturating_sub(last_drand_round);
let blocks_to_reveal = rounds_since_last.saturating_mul(blocks_per_round);
let reveal_block = commit_block.saturating_add(
blocks_to_reveal
.try_into()
.map_err(|_| "Block number conversion failed")?,
);
Ok(reveal_block)
}
}
76 changes: 76 additions & 0 deletions pallets/commitments/src/mock.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -181,3 +181,79 @@ pub fn new_test_ext() -> sp_io::TestExternalities {
ext.execute_with(|| System::set_block_number(1));
ext
}

use super::*;
use crate::{EngineBLS, MAX_TIMELOCK_COMMITMENT_SIZE_BYTES, TinyBLS381};
use ark_serialize::CanonicalSerialize;
use frame_support::BoundedVec;
use rand_chacha::{ChaCha20Rng, rand_core::SeedableRng};
use sha2::Digest;
use tle::{ibe::fullident::Identity, stream_ciphers::AESGCMStreamCipherProvider, tlock::tle};

// Drand Quicknet public key and signature for round=1000:
pub const DRAND_QUICKNET_PUBKEY_HEX: &str = "83cf0f2896adee7eb8b5f01fcad3912212c437e0073e911fb90022d3e760183c8c4b450b6\
a0a6c3ac6a5776a2d1064510d1fec758c921cc22b0e17e63aaf4bcb5ed66304de9cf809b\
d274ca73bab4af5a6e9c76a4bc09e76eae8991ef5ece45a";
pub const DRAND_QUICKNET_SIG_HEX: &str = "b44679b9a59af2ec876b1a6b1ad52ea9b1615fc3982b19576350f93447cb1125e342b73a8dd2bacbe47e4b6b63ed5e39";

/// Inserts a Drand pulse for `round` with the given `signature_bytes`.
pub fn insert_drand_pulse(round: u64, signature_bytes: &[u8]) {
let sig_bounded: BoundedVec<u8, ConstU32<144>> = signature_bytes
.to_vec()
.try_into()
.expect("Signature within 144 bytes");

let randomness_bounded: BoundedVec<u8, ConstU32<32>> = vec![0u8; 32]
.try_into()
.expect("Randomness must be exactly 32 bytes");

pallet_drand::Pulses::<Test>::insert(
round,
pallet_drand::types::Pulse {
round,
randomness: randomness_bounded,
signature: sig_bounded,
},
);
}

/// Produces a **real** ciphertext by TLE-encrypting `plaintext` for Drand Quicknet `round`.
///
/// The returned `BoundedVec<u8, ConstU32<MAX_TIMELOCK_COMMITMENT_SIZE_BYTES>>`
/// will decrypt if you pass in the valid signature for the same round.
pub fn produce_ciphertext(
plaintext: &[u8],
round: u64,
) -> BoundedVec<u8, ConstU32<MAX_TIMELOCK_COMMITMENT_SIZE_BYTES>> {
// 1) Deserialize the known Drand Quicknet public key:
let pub_key_bytes = hex::decode(DRAND_QUICKNET_PUBKEY_HEX).expect("decode pubkey");
let pub_key =
<TinyBLS381 as EngineBLS>::PublicKeyGroup::deserialize_compressed(&pub_key_bytes[..])
.expect("bad pubkey bytes");

// 2) Prepare the identity for that round
// by hashing round.to_be_bytes() with SHA256:
let msg = {
let mut hasher = sha2::Sha256::new();
hasher.update(round.to_be_bytes());
hasher.finalize().to_vec()
};
let identity = Identity::new(b"", vec![msg]);

// 3) Actually encrypt
// (just an example ephemeral secret key & RNG seed)
let esk = [2u8; 32];
let rng = ChaCha20Rng::seed_from_u64(0);

let ct = tle::<TinyBLS381, AESGCMStreamCipherProvider, ChaCha20Rng>(
pub_key, esk, plaintext, identity, rng,
)
.expect("Encryption failed in produce_real_ciphertext");

// 4) Serialize the ciphertext to BoundedVec
let mut ct_bytes = Vec::new();
ct.serialize_compressed(&mut ct_bytes)
.expect("serialize TLECiphertext");

ct_bytes.try_into().expect("Ciphertext is within max size")
}
Loading

0 comments on commit 423fe2c

Please sign in to comment.