diff --git a/pallets/subtensor/rpc/src/lib.rs b/pallets/subtensor/rpc/src/lib.rs index 2445a5eda..d99388193 100644 --- a/pallets/subtensor/rpc/src/lib.rs +++ b/pallets/subtensor/rpc/src/lib.rs @@ -41,7 +41,6 @@ pub trait SubtensorCustomApi { fn get_neurons(&self, netuid: u16, at: Option) -> RpcResult>; #[method(name = "neuronInfo_getNeuron")] fn get_neuron(&self, netuid: u16, uid: u16, at: Option) -> RpcResult>; - #[method(name = "subnetInfo_getSubnetInfo")] fn get_subnet_info(&self, netuid: u16, at: Option) -> RpcResult>; #[method(name = "subnetInfo_getSubnetsInfo")] diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index 2985736c8..cc3d7d025 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -67,12 +67,14 @@ pub mod pallet { traits::{ tokens::fungible, OriginTrait, QueryPreimage, StorePreimage, UnfilteredDispatchable, }, + BoundedVec, }; use frame_system::pallet_prelude::*; use sp_core::H256; use sp_runtime::traits::{Dispatchable, TrailingZeroInput}; use sp_std::vec; use sp_std::vec::Vec; + use subtensor_macros::freeze_struct; #[cfg(not(feature = "std"))] use alloc::boxed::Box; @@ -129,6 +131,36 @@ pub mod pallet { pub placeholder2: u8, } + /// Struct for NeuronCertificate. + pub type NeuronCertificateOf = NeuronCertificate; + /// Data structure for NeuronCertificate information. + #[freeze_struct("1c232be200d9ec6c")] + #[derive(Decode, Encode, Default, TypeInfo, PartialEq, Eq, Clone, Debug)] + pub struct NeuronCertificate { + /// The neuron TLS public key + pub public_key: BoundedVec>, + /// The algorithm used to generate the public key + pub algorithm: u8, + } + + impl TryFrom> for NeuronCertificate { + type Error = (); + + fn try_from(value: Vec) -> Result { + if value.len() > 65 { + return Err(()); + } + // take the first byte as the algorithm + let algorithm = value.first().ok_or(())?; + // and the rest as the public_key + let certificate = value.get(1..).ok_or(())?.to_vec(); + Ok(Self { + public_key: BoundedVec::try_from(certificate).map_err(|_| ())?, + algorithm: *algorithm, + }) + } + } + /// Struct for Prometheus. pub type PrometheusInfoOf = PrometheusInfo; @@ -1162,6 +1194,17 @@ pub mod pallet { /// --- MAP ( netuid, hotkey ) --> axon_info pub type Axons = StorageDoubleMap<_, Identity, u16, Blake2_128Concat, T::AccountId, AxonInfoOf, OptionQuery>; + /// --- MAP ( netuid, hotkey ) --> certificate + #[pallet::storage] + pub type NeuronCertificates = StorageDoubleMap< + _, + Identity, + u16, + Blake2_128Concat, + T::AccountId, + NeuronCertificateOf, + OptionQuery, + >; #[pallet::storage] /// --- MAP ( netuid, hotkey ) --> prometheus_info pub type Prometheus = StorageDoubleMap< @@ -1538,6 +1581,10 @@ where let transaction_fee = 0; Ok((CallType::Serve, transaction_fee, who.clone())) } + Some(Call::serve_axon_tls { .. }) => { + let transaction_fee = 0; + Ok((CallType::Serve, transaction_fee, who.clone())) + } Some(Call::register_network { .. }) => { let transaction_fee = 0; Ok((CallType::RegisterNetwork, transaction_fee, who.clone())) diff --git a/pallets/subtensor/src/macros/dispatches.rs b/pallets/subtensor/src/macros/dispatches.rs index a97e4494d..c4b985a49 100644 --- a/pallets/subtensor/src/macros/dispatches.rs +++ b/pallets/subtensor/src/macros/dispatches.rs @@ -435,7 +435,7 @@ mod dispatches { Self::do_remove_stake(origin, hotkey, amount_unstaked) } - /// Serves or updates axon /promethteus information for the neuron associated with the caller. If the caller is + /// Serves or updates axon /prometheus information for the neuron associated with the caller. If the caller is /// already registered the metadata is updated. If the caller is not registered this call throws NotRegistered. /// /// # Args: @@ -511,6 +511,92 @@ mod dispatches { protocol, placeholder1, placeholder2, + None, + ) + } + + /// Same as `serve_axon` but takes a certificate as an extra optional argument. + /// Serves or updates axon /prometheus information for the neuron associated with the caller. If the caller is + /// already registered the metadata is updated. If the caller is not registered this call throws NotRegistered. + /// + /// # Args: + /// * 'origin': (Origin): + /// - The signature of the caller. + /// + /// * 'netuid' (u16): + /// - The u16 network identifier. + /// + /// * 'version' (u64): + /// - The bittensor version identifier. + /// + /// * 'ip' (u64): + /// - The endpoint ip information as a u128 encoded integer. + /// + /// * 'port' (u16): + /// - The endpoint port information as a u16 encoded integer. + /// + /// * 'ip_type' (u8): + /// - The endpoint ip version as a u8, 4 or 6. + /// + /// * 'protocol' (u8): + /// - UDP:1 or TCP:0 + /// + /// * 'placeholder1' (u8): + /// - Placeholder for further extra params. + /// + /// * 'placeholder2' (u8): + /// - Placeholder for further extra params. + /// + /// * 'certificate' (Vec): + /// - TLS certificate for inter neuron communitation. + /// + /// # Event: + /// * AxonServed; + /// - On successfully serving the axon info. + /// + /// # Raises: + /// * 'SubNetworkDoesNotExist': + /// - Attempting to set weights on a non-existent network. + /// + /// * 'NotRegistered': + /// - Attempting to set weights from a non registered account. + /// + /// * 'InvalidIpType': + /// - The ip type is not 4 or 6. + /// + /// * 'InvalidIpAddress': + /// - The numerically encoded ip address does not resolve to a proper ip. + /// + /// * 'ServingRateLimitExceeded': + /// - Attempting to set prometheus information withing the rate limit min. + /// + #[pallet::call_index(40)] + #[pallet::weight((Weight::from_parts(46_000_000, 0) + .saturating_add(T::DbWeight::get().reads(4)) + .saturating_add(T::DbWeight::get().writes(1)), DispatchClass::Normal, Pays::No))] + pub fn serve_axon_tls( + origin: OriginFor, + netuid: u16, + version: u32, + ip: u128, + port: u16, + ip_type: u8, + protocol: u8, + placeholder1: u8, + placeholder2: u8, + certificate: Vec, + ) -> DispatchResult { + Self::do_serve_axon( + origin, + netuid, + version, + ip, + port, + ip_type, + protocol, + placeholder1, + placeholder2, + Some(certificate), ) } diff --git a/pallets/subtensor/src/subnets/serving.rs b/pallets/subtensor/src/subnets/serving.rs index 1a9240c36..7e2b9a0f0 100644 --- a/pallets/subtensor/src/subnets/serving.rs +++ b/pallets/subtensor/src/subnets/serving.rs @@ -31,6 +31,9 @@ impl Pallet { /// * 'placeholder2' (u8): /// - Placeholder for further extra params. /// + /// * 'certificate' (Option>): + /// - Certificate for mutual Tls connection between neurons + /// /// # Event: /// * AxonServed; /// - On successfully serving the axon info. @@ -61,6 +64,7 @@ impl Pallet { protocol: u8, placeholder1: u8, placeholder2: u8, + certificate: Option>, ) -> dispatch::DispatchResult { // We check the callers (hotkey) signature. let hotkey_id = ensure_signed(origin)?; @@ -86,6 +90,13 @@ impl Pallet { Error::::ServingRateLimitExceeded ); + // Check certificate + if let Some(certificate) = certificate { + if let Ok(certificate) = NeuronCertificateOf::try_from(certificate) { + NeuronCertificates::::insert(netuid, hotkey_id.clone(), certificate) + } + } + // We insert the axon meta. prev_axon.block = Self::get_current_block_as_u64(); prev_axon.version = version; diff --git a/pallets/subtensor/src/subnets/uids.rs b/pallets/subtensor/src/subnets/uids.rs index fff358f1c..2a5ceedb4 100644 --- a/pallets/subtensor/src/subnets/uids.rs +++ b/pallets/subtensor/src/subnets/uids.rs @@ -45,6 +45,9 @@ impl Pallet { Uids::::insert(netuid, new_hotkey.clone(), uid_to_replace); // Make uid - hotkey association. BlockAtRegistration::::insert(netuid, uid_to_replace, block_number); // Fill block at registration. IsNetworkMember::::insert(new_hotkey.clone(), netuid, true); // Fill network is member. + + // 4. Clear neuron certificates + NeuronCertificates::::remove(netuid, old_hotkey.clone()); } /// Appends the uid to the network. diff --git a/pallets/subtensor/src/swap/swap_hotkey.rs b/pallets/subtensor/src/swap/swap_hotkey.rs index 793e34bff..ca3d0b5a7 100644 --- a/pallets/subtensor/src/swap/swap_hotkey.rs +++ b/pallets/subtensor/src/swap/swap_hotkey.rs @@ -276,6 +276,18 @@ impl Pallet { weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2)); } } + + // 9.7. Swap neuron TLS certificates. + // NeuronCertificates( netuid, hotkey ) -> Vec -- the neuron certificate for the hotkey. + if is_network_member { + if let Ok(old_neuron_certificates) = + NeuronCertificates::::try_get(netuid, old_hotkey) + { + NeuronCertificates::::remove(netuid, old_hotkey); + NeuronCertificates::::insert(netuid, new_hotkey, old_neuron_certificates); + weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2)); + } + } } // 10. Swap Stake. diff --git a/pallets/subtensor/tests/serving.rs b/pallets/subtensor/tests/serving.rs index 49a963951..6bc30c76f 100644 --- a/pallets/subtensor/tests/serving.rs +++ b/pallets/subtensor/tests/serving.rs @@ -99,6 +99,64 @@ fn test_serving_ok() { }); } +#[test] +fn test_serving_tls_ok() { + new_test_ext(1).execute_with(|| { + let hotkey_account_id = U256::from(1); + let netuid: u16 = 1; + let tempo: u16 = 13; + let version: u32 = 2; + let ip: u128 = 1676056785; + let port: u16 = 128; + let ip_type: u8 = 4; + let modality: u16 = 0; + let protocol: u8 = 0; + let placeholder1: u8 = 0; + let placeholder2: u8 = 0; + let certificate: Vec = "CERT".as_bytes().to_vec(); + add_network(netuid, tempo, modality); + register_ok_neuron(netuid, hotkey_account_id, U256::from(66), 0); + assert_ok!(SubtensorModule::serve_axon_tls( + <::RuntimeOrigin>::signed(hotkey_account_id), + netuid, + version, + ip, + port, + ip_type, + protocol, + placeholder1, + placeholder2, + certificate.clone() + )); + + let stored_certificate = NeuronCertificates::::get(netuid, hotkey_account_id) + .expect("Certificate should exist"); + assert_eq!( + stored_certificate.public_key.clone().into_inner(), + certificate.get(1..).expect("Certificate should exist") + ); + let new_certificate = "UPDATED_CERT".as_bytes().to_vec(); + assert_ok!(SubtensorModule::serve_axon_tls( + <::RuntimeOrigin>::signed(hotkey_account_id), + netuid, + version, + ip, + port, + ip_type, + protocol, + placeholder1, + placeholder2, + new_certificate.clone() + )); + let stored_certificate = NeuronCertificates::::get(netuid, hotkey_account_id) + .expect("Certificate should exist"); + assert_eq!( + stored_certificate.public_key.clone().into_inner(), + new_certificate.get(1..).expect("Certificate should exist") + ); + }); +} + #[test] fn test_serving_set_metadata_update() { new_test_ext(1).execute_with(|| { diff --git a/pallets/subtensor/tests/swap_hotkey.rs b/pallets/subtensor/tests/swap_hotkey.rs index bff738b86..89938e3eb 100644 --- a/pallets/subtensor/tests/swap_hotkey.rs +++ b/pallets/subtensor/tests/swap_hotkey.rs @@ -311,6 +311,38 @@ fn test_swap_axons() { }); } +// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_certificates --exact --nocapture +#[test] +fn test_swap_certificates() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey = U256::from(3); + let netuid = 0u16; + let certificate = NeuronCertificate::try_from(vec![1, 2, 3]).unwrap(); + let mut weight = Weight::zero(); + + add_network(netuid, 0, 1); + IsNetworkMember::::insert(old_hotkey, netuid, true); + NeuronCertificates::::insert(netuid, old_hotkey, certificate.clone()); + + assert_ok!(SubtensorModule::perform_hotkey_swap( + &old_hotkey, + &new_hotkey, + &coldkey, + &mut weight + )); + + assert!(!NeuronCertificates::::contains_key( + netuid, old_hotkey + )); + assert_eq!( + NeuronCertificates::::get(netuid, new_hotkey), + Some(certificate) + ); + }); +} + // SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_weight_commits --exact --nocapture #[test] fn test_swap_weight_commits() { diff --git a/pallets/subtensor/tests/uids.rs b/pallets/subtensor/tests/uids.rs index 82adc6b8a..6b4c00328 100644 --- a/pallets/subtensor/tests/uids.rs +++ b/pallets/subtensor/tests/uids.rs @@ -1,8 +1,9 @@ #![allow(clippy::unwrap_used)] use crate::mock::*; -use frame_support::assert_ok; +use frame_support::{assert_err, assert_ok}; use frame_system::Config; +use pallet_subtensor::*; use sp_core::U256; mod mock; @@ -32,6 +33,7 @@ fn test_replace_neuron() { let new_hotkey_account_id = U256::from(2); let _new_colkey_account_id = U256::from(12345); + let certificate = NeuronCertificate::try_from(vec![1, 2, 3]).unwrap(); //add network add_network(netuid, tempo, 0); @@ -51,6 +53,9 @@ fn test_replace_neuron() { let neuron_uid = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &hotkey_account_id); assert_ok!(neuron_uid); + // Set a neuron certificate for it + NeuronCertificates::::insert(netuid, hotkey_account_id, certificate); + // Replace the neuron. SubtensorModule::replace_neuron( netuid, @@ -77,6 +82,10 @@ fn test_replace_neuron() { &new_hotkey_account_id )); assert_eq!(curr_hotkey.unwrap(), new_hotkey_account_id); + + // Check neuron certificate was reset + let certificate = NeuronCertificates::::get(netuid, hotkey_account_id); + assert_eq!(certificate, None); }); } @@ -371,3 +380,24 @@ fn test_replace_neuron_multiple_subnets_unstake_all() { ); }); } + +#[test] +fn test_neuron_certificate() { + new_test_ext(1).execute_with(|| { + // 512 bits key + let mut data = [0; 65].to_vec(); + assert_ok!(NeuronCertificate::try_from(data)); + + // 256 bits key + data = [1; 33].to_vec(); + assert_ok!(NeuronCertificate::try_from(data)); + + // too much data + data = [8; 88].to_vec(); + assert_err!(NeuronCertificate::try_from(data), ()); + + // no data + data = vec![]; + assert_err!(NeuronCertificate::try_from(data), ()); + }); +}