1+ use std:: collections:: { BTreeSet , HashSet } ;
2+
13use digest:: { Digest , FixedOutput } ;
2- use std:: collections:: BTreeSet ;
34
45use crate :: {
5- Parameters , RegisterError , SignerIndex , Stake , StmResult ,
6+ Parameters , RegisterError , SignerIndex , Stake , StmResult , VerificationKeyForConcatenation ,
67 VerificationKeyProofOfPossessionForConcatenation ,
78 membership_commitment:: { MerkleTree , MerkleTreeLeaf } ,
89 protocol:: key_registration:: ClosedRegistrationEntry ,
@@ -14,16 +15,22 @@ use crate::VerificationKeyForSnark;
1415use super :: RegistrationEntry ;
1516
1617/// Key Registration
17- #[ derive( Clone , Default , PartialEq , Eq , PartialOrd , Ord , Debug ) ]
18+ #[ derive( Clone , Default , PartialEq , Eq , Debug ) ]
1819pub struct KeyRegistration {
1920 registration_entries : BTreeSet < RegistrationEntry > ,
21+ registered_keys_for_concatenation : HashSet < VerificationKeyForConcatenation > ,
22+ #[ cfg( feature = "future_snark" ) ]
23+ registered_keys_for_snark : HashSet < VerificationKeyForSnark > ,
2024}
2125
2226impl KeyRegistration {
2327 /// Initialize an empty registration
2428pub fn initialize ( ) -> Self {
2529 Self {
2630 registration_entries : Default :: default ( ) ,
31+ registered_keys_for_concatenation : Default :: default ( ) ,
32+ #[ cfg( feature = "future_snark" ) ]
33+ registered_keys_for_snark : Default :: default ( ) ,
2734 }
2835 }
2936
@@ -32,11 +39,28 @@ impl KeyRegistration {
3239/// # Error
3340/// The function fails when the entry is already registered.
3441pub fn register_by_entry ( & mut self , entry : & RegistrationEntry ) -> StmResult < ( ) > {
35- if !self . registration_entries . contains ( entry) {
36- self . registration_entries . insert ( * entry) ;
37- return Ok ( ( ) ) ;
42+ let vk_concatenation = entry. get_verification_key_for_concatenation ( ) ;
43+ let is_already_registered =
44+ self . registered_keys_for_concatenation . contains ( & vk_concatenation) ;
45+
46+ #[ cfg( feature = "future_snark" ) ]
47+ let is_already_registered = is_already_registered
48+ || entry
49+ . get_verification_key_for_snark ( )
50+ . is_some_and ( |vk_snark| self . registered_keys_for_snark . contains ( & vk_snark) ) ;
51+
52+ if is_already_registered {
53+ return Err ( RegisterError :: EntryAlreadyRegistered ( Box :: new ( * entry) ) . into ( ) ) ;
3854 }
39- Err ( RegisterError :: EntryAlreadyRegistered ( Box :: new ( * entry) ) . into ( ) )
55+
56+ self . registered_keys_for_concatenation . insert ( vk_concatenation) ;
57+ #[ cfg( feature = "future_snark" ) ]
58+ if let Some ( vk_snark) = entry. get_verification_key_for_snark ( ) {
59+ self . registered_keys_for_snark . insert ( vk_snark) ;
60+ }
61+ self . registration_entries . insert ( * entry) ;
62+
63+ Ok ( ( ) )
4064 }
4165
4266 /// Registers a new signer with the given verification key proof of possession and stake.
@@ -308,6 +332,66 @@ mod tests {
308332 }
309333 }
310334
335+ #[ test]
336+ fn register_by_entry_rejects_same_verification_key_with_different_stake ( ) {
337+ let mut rng = ChaCha20Rng :: from_seed ( [ 0u8 ; 32 ] ) ;
338+ let mut kr = KeyRegistration :: initialize ( ) ;
339+ let vk_pop = VerificationKeyProofOfPossessionForConcatenation :: from (
340+ & BlsSigningKey :: generate ( & mut rng) ,
341+ ) ;
342+
343+ let first_entry = RegistrationEntry :: new (
344+ vk_pop,
345+ 100 ,
346+ #[ cfg( feature = "future_snark" ) ]
347+ None ,
348+ )
349+ . unwrap ( ) ;
350+ kr. register_by_entry ( & first_entry)
351+ . expect ( "registering a new verification key should succeed" ) ;
352+
353+ let second_entry = RegistrationEntry :: new (
354+ vk_pop,
355+ 200 ,
356+ #[ cfg( feature = "future_snark" ) ]
357+ None ,
358+ )
359+ . unwrap ( ) ;
360+ let result = kr. register_by_entry ( & second_entry) ;
361+
362+ assert ! ( matches!(
363+ result. unwrap_err( ) . downcast_ref:: <RegisterError >( ) ,
364+ Some ( RegisterError :: EntryAlreadyRegistered ( _) )
365+ ) ) ;
366+ }
367+
368+ #[ cfg( feature = "future_snark" ) ]
369+ #[ test]
370+ fn register_by_entry_rejects_same_snark_key_with_different_concatenation_key ( ) {
371+ let mut rng = ChaCha20Rng :: from_seed ( [ 0u8 ; 32 ] ) ;
372+ let mut kr = KeyRegistration :: initialize ( ) ;
373+ let schnorr_vk =
374+ SchnorrVerificationKey :: new_from_signing_key ( SchnorrSigningKey :: generate ( & mut rng) ) ;
375+
376+ let first_vk_pop = VerificationKeyProofOfPossessionForConcatenation :: from (
377+ & BlsSigningKey :: generate ( & mut rng) ,
378+ ) ;
379+ let first_entry = RegistrationEntry :: new ( first_vk_pop, 100 , Some ( schnorr_vk) ) . unwrap ( ) ;
380+ kr. register_by_entry ( & first_entry)
381+ . expect ( "registering a new verification key pair should succeed" ) ;
382+
383+ let second_vk_pop = VerificationKeyProofOfPossessionForConcatenation :: from (
384+ & BlsSigningKey :: generate ( & mut rng) ,
385+ ) ;
386+ let second_entry = RegistrationEntry :: new ( second_vk_pop, 200 , Some ( schnorr_vk) ) . unwrap ( ) ;
387+ let result = kr. register_by_entry ( & second_entry) ;
388+
389+ assert ! ( matches!(
390+ result. unwrap_err( ) . downcast_ref:: <RegisterError >( ) ,
391+ Some ( RegisterError :: EntryAlreadyRegistered ( _) )
392+ ) ) ;
393+ }
394+
311395 proptest ! {
312396 #[ test]
313397 fn test_keyreg( stake in vec( 1 ..1u64 << 60 , 2 ..=10 ) ,
@@ -333,8 +417,10 @@ mod tests {
333417 VerificationKeyProofOfPossessionForConcatenation :: from( & sk)
334418 } ;
335419
336- // Record successful registrations
420+ // Record successful registrations, keyed by verification key since that's
421+ // the uniqueness criterion enforced by register_by_entry
337422 let mut keys = BTreeSet :: new( ) ;
423+ let mut registered_entries = BTreeSet :: new( ) ;
338424
339425 for ( i, & stake) in stake. iter( ) . enumerate( ) {
340426 let mut pk = gen_keys[ i % gen_keys. len( ) ] ;
@@ -350,15 +436,17 @@ mod tests {
350436
351437 match entry_result {
352438 Ok ( entry) => {
439+ let vk = entry. get_verification_key_for_concatenation( ) ;
353440 let reg = kr. register_by_entry( & entry) ;
354441 match reg {
355442 Ok ( _) => {
356- assert!( keys. insert( entry) ) ;
443+ assert!( keys. insert( vk) ) ;
444+ assert!( registered_entries. insert( entry) ) ;
357445 } ,
358446 Err ( error) => match error. downcast_ref:: <RegisterError >( ) {
359447 Some ( RegisterError :: EntryAlreadyRegistered ( e1) ) => {
360448 assert!( e1. as_ref( ) == & entry) ;
361- assert!( keys. contains( & entry ) ) ;
449+ assert!( keys. contains( & vk ) ) ;
362450 } ,
363451 _ => { panic!( "Unexpected error: {error}" ) }
364452 }
@@ -380,7 +468,7 @@ mod tests {
380468 let retrieved_keys = closed. closed_registration_entries. iter( )
381469 . map( |entry| ( * entry) . clone( ) . into( ) )
382470 . collect:: <BTreeSet <RegistrationEntry >>( ) ;
383- assert!( retrieved_keys == keys ) ;
471+ assert!( retrieved_keys == registered_entries ) ;
384472 }
385473 }
386474 }
0 commit comments