Update the ink! contract to match docs

contracts/extension/src/lib.rs

@@ -3,22 +3,32 @@
 use ink_lang as ink;
 use ink_env::{Environment, DefaultEnvironment, AccountId};
 
+pub type Curve = u16;
+pub type Coin = u32;
+pub type GlobalValidatorSetId = u32;
+pub type ValidatorSetIndex = u8;
+pub type Key = Vec<u8>;
+
 #[ink::chain_extension]
 pub trait SeraiExtension {
   type ErrorCode = ();
 
-  /// Returns the amount of active validators on the current chain.
+  /// Returns the ID for the current global validator set.
   #[ink(extension = 0, handle_status = false, returns_result = false)]
-  fn active_validators_len() -> u16;
+  fn global_validator_set_id() -> GlobalValidatorSetId;
 
-  /// Returns the ID for the current validator set for the current chain.
-  // TODO: Decide if this should be an increasing unsigned integer instead of a hash.
+  /// Returns the amount of active validator sets within the global validator set.
   #[ink(extension = 1, handle_status = false, returns_result = false)]
-  fn validator_set_id() -> [u8; 32];
+  fn validator_sets() -> u8;
 
-  /// Returns if the specified account is an active validator for the current chain.
+  /// Returns the amount of key shares used within the specified validator set.
   #[ink(extension = 2, handle_status = false, returns_result = false)]
-  fn is_active_validator(account: &AccountId) -> bool;
+  fn validator_set_shares(set: ValidatorSetIndex) -> u16;
+
+  /// Returns the validator set the specified account is in, along with their amount of shares in
+  /// that validator set, if they are in a current validator
+  #[ink(extension = 3, handle_status = false, returns_result = false)]
+  fn active_validator(account: &AccountId) -> Option<(ValidatorSetIndex, u16)>;
 }
 
 pub struct SeraiEnvironment;
@@ -45,41 +55,57 @@ pub fn test_validators() -> Vec<AccountId> {
 }
 
 pub fn test_register() {
-  struct ExtensionLen;
-  impl ink_env::test::ChainExtension for ExtensionLen {
-    fn func_id(&self) -> u32 {
-      0
-    }
-
-    fn call(&mut self, _: &[u8], output: &mut Vec<u8>) -> u32 {
-      scale::Encode::encode_to(&u16::try_from(test_validators().len()).unwrap(), output);
-      0
-    }
-  }
-  ink_env::test::register_chain_extension(ExtensionLen);
-
   struct ExtensionId;
   impl ink_env::test::ChainExtension for ExtensionId {
     fn func_id(&self) -> u32 {
-      1
+      0
     }
 
     fn call(&mut self, _: &[u8], output: &mut Vec<u8>) -> u32 {
-      scale::Encode::encode_to(&[0xffu8; 32], output);
+      // Non-0 global validator set ID
+      scale::Encode::encode_to(&1u32, output);
       0
     }
   }
   ink_env::test::register_chain_extension(ExtensionId);
 
-  struct ExtensionActive;
-  impl ink_env::test::ChainExtension for ExtensionActive {
+  struct ExtensionSets;
+  impl ink_env::test::ChainExtension for ExtensionSets {
+    fn func_id(&self) -> u32 {
+      1
+    }
+
+    fn call(&mut self, _: &[u8], output: &mut Vec<u8>) -> u32 {
+      // 1 validator set
+      scale::Encode::encode_to(&1u8, output);
+      0
+    }
+  }
+  ink_env::test::register_chain_extension(ExtensionSets);
+
+  struct ExtensionShares;
+  impl ink_env::test::ChainExtension for ExtensionShares {
     fn func_id(&self) -> u32 {
       2
     }
 
+    fn call(&mut self, _: &[u8], output: &mut Vec<u8>) -> u32 {
+      // 1 key share per validator
+      scale::Encode::encode_to(&u16::try_from(test_validators().len()).unwrap(), output);
+      0
+    }
+  }
+  ink_env::test::register_chain_extension(ExtensionShares);
+
+  struct ExtensionActive;
+  impl ink_env::test::ChainExtension for ExtensionActive {
+    fn func_id(&self) -> u32 {
+      3
+    }
+
     fn call(&mut self, input: &[u8], output: &mut Vec<u8>) -> u32 {
       use scale::Decode;
-      let potential = AccountId::decode(&mut &input[1 ..]).unwrap(); // TODO: Why is this 1 ..?
+      let potential = AccountId::decode(&mut &input[1 ..]).unwrap(); // TODO: Why is this [1 ..]?
 
       let mut presence = false;
       for validator in test_validators() {
@@ -87,7 +113,8 @@ pub fn test_register() {
           presence = true;
         }
       }
-      scale::Encode::encode_to(&presence, output);
+      // Validator set 0, 1 key share
+      scale::Encode::encode_to(&Some((0u8, 1u16)).filter(|_| presence), output);
       0
     }
   }

contracts/multisig/lib.rs

@@ -2,6 +2,10 @@
 
 use ink_lang as ink;
 
+use serai_extension::{Curve, GlobalValidatorSetId, ValidatorSetIndex, Key};
+
+type KeysHash = [u8; 32];
+
 #[ink::contract(env = serai_extension::SeraiEnvironment)]
 mod multisig {
   use scale::Encode;
@@ -9,54 +13,74 @@ mod multisig {
   use ink_storage::{traits::SpreadAllocate, Mapping};
   use ink_env::{hash::Blake2x256, hash_encoded};
 
+  use super::*;
+
   /// A contract which tracks the current multisig keys.
+  /// Mapping of each validator set to their multisigs.
   #[ink(storage)]
   #[derive(SpreadAllocate)]
   pub struct Multisig {
-    /// Validator set currently holding the multisig.
-    validator_set: [u8; 32],
-    /// Mapping from a curve's index to the multisig's current public key for it.
+    /// Global validator set ID under which this multisig was updated.
+    /// Used to track if the multisig has been updated to the latest instantiation of a validator
+    /// set or not.
+    /// May be behind, and still healthy, if a validator set didn't change despite the global
+    /// validator set doing so.
+    updated_at: Mapping<ValidatorSetIndex, GlobalValidatorSetId>,
+    /// Mapping from a curve's index to the multisig's current public key for it, if it has one.
     // This is a mapping due to ink's eager loading. Considering we're right now only considering
-    // secp256k1 and Ed25519, it may be notably more efficient to use a Vec here.
-    keys: Mapping<u8, Vec<u8>>,
-    /// Voter + Keys -> Voted already or not
-    voted: Mapping<(AccountId, [u8; 32]), ()>,
-    /// Validator Set + Keys -> Vote Count
-    votes: Mapping<([u8; 32], [u8; 32]), u16>,
+    // Secp256k1 and Ed25519, it may be notably more efficient to use a Vec here.
+    // In practice, we're likely discussing up to 7 curves in total, so it may always be better to
+    // simply use a Vec here, especially since it'd be Vec<Option<Key>>.
+    keys: Mapping<(ValidatorSetIndex, Curve), Key>,
+    /// Validator + Keys -> Voted already or not.
+    /// Prevents voting multiple times on the same set of keys.
+    voted: Mapping<(AccountId, KeysHash), ()>,
+    /// Global Validator Set ID + Validator + Keys -> Vote Count.
+    /// Including the GVSID locks it to a specific time period, preventing a validator from joining
+    /// a set, voting on old keys, and then moving their bond to a new account to vote again.
+    votes: Mapping<(GlobalValidatorSetId, ValidatorSetIndex, KeysHash), u16>,
   }
 
-  /// Event emitted when a new set of multisig keys is voted on. Only for the first vote on a set
-  // of keys will they be present in this event.
+  /// Event emitted when a new set of multisig keys is voted on.
   #[ink(event)]
   pub struct Vote {
     /// Validator who issued the vote.
     #[ink(topic)]
     validator: AccountId,
+    /// Global validator set ID under which keys are being generated.
+    #[ink(topic)]
+    global_validator_set: GlobalValidatorSetId,
     /// Validator set for which keys are being generated.
     #[ink(topic)]
-    validator_set: [u8; 32],
+    validator_set: ValidatorSetIndex,
     /// Hash of the keys voted on.
     #[ink(topic)]
-    hash: [u8; 32],
-    /// Keys voted on.
-    keys: Option<Vec<Vec<u8>>>,
+    hash: KeysHash,
+    /// Keys voted on. Only present in the first event for a given set of keys.
+    keys: Option<Vec<Option<Key>>>,
   }
 
-  /// Event emitted when the new keys are fully generated for all curves, having been fully voted
-  /// on.
+  /// Event emitted when the new keys are fully generated for a validator set, having been fully
+  /// voted on.
   #[ink(event)]
   pub struct KeyGen {
     #[ink(topic)]
-    validator_set: [u8; 32],
+    global_validator_set: GlobalValidatorSetId,
     #[ink(topic)]
-    hash: [u8; 32],
+    validator_set: ValidatorSetIndex,
+    #[ink(topic)]
+    hash: KeysHash,
   }
 
   /// The Multisig error types.
   #[derive(Debug, PartialEq, Eq, scale::Encode, scale::Decode)]
   #[cfg_attr(feature = "std", derive(scale_info::TypeInfo))]
   pub enum Error {
-    /// Returned if a curve index doesn't have a key registered for it.
+    /// Returned if a validator set hasn't had keys registered for it yet.
+    NonExistentValidatorSet,
+    /// Returned if a validator set and curve index doesn't have a key registered for it.
+    NonExistentKey,
+    /// Returned if a curve index doesn't exist.
     NonExistentCurve,
     /// Returned if a non-validator is voting.
     NotValidator,
@@ -76,76 +100,85 @@ mod multisig {
       ink_lang::utils::initialize_contract(|_| {})
     }
 
-    /// Validator set currently holding the multisig.
+    /// Global validator set ID under which a validator set updated their multisig.
     #[ink(message)]
-    pub fn validator_set(&self) -> [u8; 32] {
-      self.validator_set
+    pub fn updated_at(&self, validator_set: ValidatorSetIndex) -> Result<GlobalValidatorSetId> {
+      self.updated_at.get(validator_set).ok_or(Error::NonExistentValidatorSet)
     }
 
-    /// Returns the key currently in-use for a given curve ID. This is then bound to a given chain
-    /// by applying a personalized additive offset, as done by the processor. Each chain then has
-    /// its own way of receiving funds to these keys, leaving this not for usage by wallets, nor
-    /// the processor which is expected to track events for this information. This is really solely
-    /// for debugging purposes.
+    /// Returns the key currently in-use for a given validator set and curve.
+    /// This is then bound to a given chain by applying a network-specific additive offset, as done
+    /// by the processor. Each chain then has its own way of receiving funds to these keys, leaving
+    /// this not for usage by wallets, nor the processor which is expected to track events for this
+    /// information. This is really solely for debugging purposes.
     #[ink(message)]
-    pub fn key(&self, curve: u8) -> Result<Vec<u8>> {
-      self.keys.get(curve).ok_or(Error::NonExistentCurve)
+    pub fn key(&self, validator_set: ValidatorSetIndex, curve: Curve) -> Result<Key> {
+      self.keys.get((validator_set, curve)).ok_or(Error::NonExistentKey)
     }
 
     // TODO: voted
     // TODO: votes
 
-    fn hash<T: Encode>(value: &T) -> [u8; 32] {
-      let mut output = [0; 32];
+    fn hash<T: Encode>(value: &T) -> KeysHash {
+      let mut output = KeysHash::default();
       hash_encoded::<Blake2x256, _>(value, &mut output);
       output
     }
 
     /// Vote for a given set of keys.
     #[ink(message)]
-    pub fn vote(&mut self, keys: Vec<Vec<u8>>) -> Result<()> {
+    pub fn vote(&mut self, keys: Vec<Option<Key>>) -> Result<()> {
       if keys.len() > 256 {
         Err(Error::NonExistentCurve)?;
       }
 
+      // Make sure they're a valid validator.
       let validator = self.env().caller();
-      if !self.env().extension().is_active_validator(&validator) {
+      let active_validator = self.env().extension().active_validator(&validator);
+      if active_validator.is_none() {
         Err(Error::NotValidator)?;
       }
+      let (validator_set, shares) = active_validator.unwrap();
 
-      let validator_set = self.env().extension().validator_set_id();
-      if self.validator_set == validator_set {
+      // Prevent a validator set from generating keys multiple times. Only the first-voted-in keys
+      // should be acknowledged.
+      let global_validator_set = self.env().extension().global_validator_set_id();
+      if self.updated_at.get(validator_set) == Some(global_validator_set) {
         Err(Error::AlreadyGeneratedKeys)?;
       }
 
+      // Prevent a validator from voting on keys multiple times.
       let keys_hash = Self::hash(&keys);
       if self.voted.get((validator, keys_hash)).is_some() {
         Err(Error::AlreadyVoted)?;
       }
       self.voted.insert((validator, keys_hash), &());
 
-      let votes = if let Some(votes) = self.votes.get((validator_set, keys_hash)) {
-        self.env().emit_event(Vote { validator, validator_set, hash: keys_hash, keys: None });
-        votes + 1
+      let votes = if let Some(votes) = self.votes.get((global_validator_set, validator_set, keys_hash)) {
+        self.env().emit_event(Vote { validator, global_validator_set, validator_set, hash: keys_hash, keys: None });
+        votes + shares
       } else {
         self.env().emit_event(Vote {
           validator,
+          global_validator_set,
           validator_set,
           hash: keys_hash,
           keys: Some(keys.clone()),
         });
-        1
+        shares
       };
       // We could skip writing this if we've reached consensus, yet best to keep our ducks in a row
-      self.votes.insert((validator_set, keys_hash), &votes);
+      self.votes.insert((global_validator_set, validator_set, keys_hash), &votes);
 
       // If we've reached consensus, action this.
-      if votes == self.env().extension().active_validators_len() {
-        self.validator_set = validator_set;
+      if votes == self.env().extension().validator_set_shares(validator_set) {
+        self.updated_at.insert(validator_set, &global_validator_set);
         for (k, key) in keys.iter().enumerate() {
-          self.keys.insert(u8::try_from(k).unwrap(), key);
+          if let Some(key) = key {
+            self.keys.insert((validator_set, Curve::try_from(k).unwrap()), key);
+          }
         }
-        self.env().emit_event(KeyGen { validator_set, hash: keys_hash });
+        self.env().emit_event(KeyGen { global_validator_set, validator_set, hash: keys_hash });
       }
 
       Ok(())
@@ -170,11 +203,11 @@ mod multisig {
     type Event = <Multisig as ::ink_lang::reflect::ContractEventBase>::Type;
 
     lazy_static! {
-      static ref EXPECTED_VALIDATOR_SET: [u8; 32] = [0xff; 32];
-
-      static ref KEYS: Vec<Vec<u8>> = vec![vec![0, 1], vec![2, 3]];
-      static ref EXPECTED_HASH: [u8; 32] = {
-        let mut hash = [0; 32];
+      static ref EXPECTED_GLOBAL_VALIDATOR_SET: GlobalValidatorSetId = 1;
+      static ref EXPECTED_VALIDATOR_SET: ValidatorSetIndex = 0;
+      static ref KEYS: Vec<Option<Key>> = vec![Some(vec![0, 1]), Some(vec![2, 3])];
+      static ref EXPECTED_HASH: KeysHash = {
+        let mut hash = KeysHash::default();
         ink_env::hash_encoded::<Blake2x256, _>(&*KEYS, &mut hash);
         hash
       };
@@ -182,7 +215,7 @@ mod multisig {
 
     fn hash_prefixed<T: scale::Encode>(prefixed: PrefixedValue<T>) -> [u8; 32] {
       let encoded = prefixed.encode();
-      let mut hash = [0; 32];
+      let mut hash = KeysHash::default();
       if encoded.len() < 32 {
         hash[.. encoded.len()].copy_from_slice(&encoded);
       } else {
@@ -199,9 +232,10 @@ mod multisig {
       let decoded_event = <Event as scale::Decode>::decode(&mut &event.data[..])
         .expect("encountered invalid contract event data buffer");
 
-      if let Event::Vote(Vote { validator, validator_set, hash, keys: actual_keys }) = decoded_event
+      if let Event::Vote(Vote { validator, global_validator_set, validator_set, hash, keys: actual_keys }) = decoded_event
       {
         assert_eq!(validator, expected_validator);
+        assert_eq!(global_validator_set, *EXPECTED_GLOBAL_VALIDATOR_SET);
         assert_eq!(validator_set, *EXPECTED_VALIDATOR_SET);
         assert_eq!(hash, *EXPECTED_HASH);
         assert_eq!(actual_keys.as_ref(), expected_keys.map(|_| &*KEYS));
@@ -215,6 +249,10 @@ mod multisig {
           prefix: b"Multisig::Vote::validator",
           value: &expected_validator,
         }),
+        hash_prefixed(PrefixedValue {
+          prefix: b"Multisig::Vote::global_validator_set",
+          value: &*EXPECTED_GLOBAL_VALIDATOR_SET,
+        }),
         hash_prefixed(PrefixedValue {
           prefix: b"Multisig::Vote::validator_set",
           value: &*EXPECTED_VALIDATOR_SET,
@@ -233,8 +271,8 @@ mod multisig {
       let decoded_event = <Event as scale::Decode>::decode(&mut &event.data[..])
         .expect("encountered invalid contract event data buffer");
 
-      if let Event::KeyGen(KeyGen { validator_set, hash }) = decoded_event
-      {
+      if let Event::KeyGen(KeyGen { global_validator_set, validator_set, hash }) = decoded_event {
+        assert_eq!(global_validator_set, *EXPECTED_GLOBAL_VALIDATOR_SET);
         assert_eq!(validator_set, *EXPECTED_VALIDATOR_SET);
         assert_eq!(hash, *EXPECTED_HASH);
       } else {
@@ -243,14 +281,15 @@ mod multisig {
 
       let expected_topics = vec![
         hash_prefixed(PrefixedValue { prefix: b"", value: b"Multisig::KeyGen" }),
+        hash_prefixed(PrefixedValue {
+          prefix: b"Multisig::KeyGen::global_validator_set",
+          value: &*EXPECTED_GLOBAL_VALIDATOR_SET,
+        }),
         hash_prefixed(PrefixedValue {
           prefix: b"Multisig::KeyGen::validator_set",
           value: &*EXPECTED_VALIDATOR_SET,
         }),
-        hash_prefixed(PrefixedValue {
-          prefix: b"Multisig::KeyGen::hash",
-          value: &*EXPECTED_HASH,
-        })
+        hash_prefixed(PrefixedValue { prefix: b"Multisig::KeyGen::hash", value: &*EXPECTED_HASH }),
       ];
 
       for (n, (actual_topic, expected_topic)) in
@@ -264,13 +303,13 @@ mod multisig {
     #[ink::test]
     fn new() {
       let multisig = Multisig::new();
-      assert_eq!(multisig.validator_set(), [0; 32]);
+      assert_eq!(multisig.updated_at(0), Err(Error::NonExistentValidatorSet));
     }
 
-    /// Non-existent curves error accordingly.
+    /// Non-existent keys error accordingly.
     #[ink::test]
-    fn non_existent_curve() {
-      assert_eq!(Multisig::new().key(0), Err(Error::NonExistentCurve));
+    fn non_existent_key() {
+      assert_eq!(Multisig::new().key(0, 0), Err(Error::NonExistentKey));
     }
 
     #[ink::test]
@@ -296,7 +335,7 @@ mod multisig {
       }
 
       // Since this should have key gen'd, verify that
-      assert_eq!(multisig.validator_set(), *EXPECTED_VALIDATOR_SET);
+      assert_eq!(multisig.updated_at(0).unwrap(), *EXPECTED_GLOBAL_VALIDATOR_SET);
       assert_key_gen(&emitted_events[test_validators().len()]);
     }
   }