From 24e2f4a033055b532c42d965674f1e55f0255c57 Mon Sep 17 00:00:00 2001 From: Andreea Popescu Date: Thu, 8 Aug 2024 15:55:34 +0100 Subject: [PATCH] add neuron certificate handling --- pallets/subtensor/rpc/src/lib.rs | 27 +++++- pallets/subtensor/runtime-api/src/lib.rs | 1 + pallets/subtensor/src/certificate.rs | 4 + pallets/subtensor/src/lib.rs | 24 ++++++ pallets/subtensor/src/macros/dispatches.rs | 84 ++++++++++++++++++- pallets/subtensor/src/rpc_info/neuron_info.rs | 17 ++++ pallets/subtensor/src/subnets/serving.rs | 20 +++++ pallets/subtensor/tests/serving.rs | 50 +++++++++++ runtime/src/lib.rs | 46 +++------- 9 files changed, 235 insertions(+), 38 deletions(-) create mode 100644 pallets/subtensor/src/certificate.rs diff --git a/pallets/subtensor/rpc/src/lib.rs b/pallets/subtensor/rpc/src/lib.rs index 2f71e9c21..fe0ce3f51 100644 --- a/pallets/subtensor/rpc/src/lib.rs +++ b/pallets/subtensor/rpc/src/lib.rs @@ -41,7 +41,13 @@ 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 = "neuronInfo_getNeuronCertificate")] + fn get_neuron_certificate( + &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")] @@ -183,6 +189,25 @@ where .map_err(|e| Error::RuntimeError(format!("Unable to get neuron info: {:?}", e)).into()) } + fn get_neuron_certificate( + &self, + netuid: u16, + uid: u16, + at: Option<::Hash>, + ) -> RpcResult> { + let api = self.client.runtime_api(); + let at = at.unwrap_or_else(|| self.client.info().best_hash); + + api.get_neuron_certificate(at, netuid, uid).map_err(|e| { + CallError::Custom(ErrorObject::owned( + Error::RuntimeError.into(), + "Unable to get neuron certificate.", + Some(e.to_string()), + )) + .into() + }) + } + fn get_subnet_info( &self, netuid: u16, diff --git a/pallets/subtensor/runtime-api/src/lib.rs b/pallets/subtensor/runtime-api/src/lib.rs index 9095ad54a..d31e95b4a 100644 --- a/pallets/subtensor/runtime-api/src/lib.rs +++ b/pallets/subtensor/runtime-api/src/lib.rs @@ -14,6 +14,7 @@ sp_api::decl_runtime_apis! { pub trait NeuronInfoRuntimeApi { fn get_neurons(netuid: u16) -> Vec; fn get_neuron(netuid: u16, uid: u16) -> Vec; + fn get_neuron_certificate(netuid: u16, uid: u16) -> Vec; fn get_neurons_lite(netuid: u16) -> Vec; fn get_neuron_lite(netuid: u16, uid: u16) -> Vec; } diff --git a/pallets/subtensor/src/certificate.rs b/pallets/subtensor/src/certificate.rs new file mode 100644 index 000000000..f0f39845d --- /dev/null +++ b/pallets/subtensor/src/certificate.rs @@ -0,0 +1,4 @@ +use sp_std::vec::Vec; + +// TODO Certificate Encryption Pubkey +pub type Certificate = Vec; diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index b305661f7..63777fc8c 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -60,6 +60,7 @@ extern crate alloc; #[frame_support::pallet] pub mod pallet { + use crate::certificate::Certificate; use crate::migrations; use frame_support::{ dispatch::GetDispatchInfo, @@ -116,6 +117,14 @@ pub mod pallet { pub placeholder2: u8, } + /// Struct for NeuronCertificate. + pub type NeuronCertificateOf = NeuronCertificate; + + #[derive(Decode, Encode, Default, TypeInfo, PartialEq, Eq, Clone, Debug)] + pub struct NeuronCertificate { + pub certificate: Certificate, + } + /// Struct for Prometheus. pub type PrometheusInfoOf = PrometheusInfo; /// Data structure for Prometheus information. @@ -1065,6 +1074,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(super) type NeuronCertificates = StorageDoubleMap< + _, + Identity, + u16, + Blake2_128Concat, + T::AccountId, + NeuronCertificateOf, + OptionQuery, + >; #[pallet::storage] /// --- MAP ( netuid, hotkey ) --> prometheus_info pub type Prometheus = StorageDoubleMap< @@ -1395,6 +1415,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 00865f8db..a0b45c55b 100644 --- a/pallets/subtensor/src/macros/dispatches.rs +++ b/pallets/subtensor/src/macros/dispatches.rs @@ -431,7 +431,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: @@ -510,6 +510,88 @@ mod dispatches { ) } + /// 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. + /// + /// # 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: Certificate, + ) -> DispatchResult { + Self::do_serve_axon( + origin, + netuid, + version, + ip, + port, + ip_type, + protocol, + placeholder1, + placeholder2, + Some(certificate), + ) + } + /// ---- Set prometheus information for the neuron. /// # Args: /// * 'origin': (Origin): diff --git a/pallets/subtensor/src/rpc_info/neuron_info.rs b/pallets/subtensor/src/rpc_info/neuron_info.rs index cadd4b6e3..8cc0b02a9 100644 --- a/pallets/subtensor/src/rpc_info/neuron_info.rs +++ b/pallets/subtensor/src/rpc_info/neuron_info.rs @@ -147,6 +147,23 @@ impl Pallet { Some(neuron) } + fn get_neuron_certificate(netuid: u16, uid: u16) -> Option { + if !Self::if_subnet_exist(netuid) { + return None; + } + + let hotkey = match Self::get_hotkey_for_net_and_uid(netuid, uid) { + Ok(h) => h, + Err(_) => return None, + }; + + if Self::has_neuron_certificate(netuid, &hotkey) { + NeuronCertificates::::get(netuid, hotkey) + } else { + None + } + } + pub fn get_neuron(netuid: u16, uid: u16) -> Option> { if !Self::if_subnet_exist(netuid) { return None; diff --git a/pallets/subtensor/src/subnets/serving.rs b/pallets/subtensor/src/subnets/serving.rs index eb7fa4369..035de7da7 100644 --- a/pallets/subtensor/src/subnets/serving.rs +++ b/pallets/subtensor/src/subnets/serving.rs @@ -1,4 +1,7 @@ use super::*; +use crate::certificate::Certificate; +use frame_support::inherent::Vec; +use frame_support::sp_std::vec; impl Pallet { /// ---- The implementation for the extrinsic serve_axon which sets the ip endpoint information for a uid on a network. @@ -31,6 +34,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 +67,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 +93,15 @@ impl Pallet { Error::::ServingRateLimitExceeded ); + // Check certificate + if let Some(certificate) = certificate { + NeuronCertificates::::insert( + netuid, + hotkey_id.clone(), + NeuronCertificate { certificate }, + ) + } + // We insert the axon meta. prev_axon.block = Self::get_current_block_as_u64(); prev_axon.version = version; @@ -239,6 +255,10 @@ impl Pallet { Axons::::contains_key(netuid, hotkey) } + pub fn has_neuron_certificate(netuid: u16, hotkey: &T::AccountId) -> bool { + return NeuronCertificates::::contains_key(netuid, hotkey); + } + pub fn has_prometheus_info(netuid: u16, hotkey: &T::AccountId) -> bool { Prometheus::::contains_key(netuid, hotkey) } diff --git a/pallets/subtensor/tests/serving.rs b/pallets/subtensor/tests/serving.rs index b0eada8e6..003947ddf 100644 --- a/pallets/subtensor/tests/serving.rs +++ b/pallets/subtensor/tests/serving.rs @@ -99,6 +99,56 @@ fn test_serving_ok() { }); } +#[test] +fn test_serving_tls_ok() { + new_test_ext().execute_with(|| { + let hotkey_account_id = U256::from(1); + let uid: u16 = 0; + 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 = SubtensorModule::get_neuron_certificate(netuid, uid); + assert_eq!(stored_certificate.unwrap().certificate, certificate); + 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 = SubtensorModule::get_neuron_certificate(netuid, uid); + assert_eq!(stored_certificate.unwrap().certificate, new_certificate) + }); +} + #[test] fn test_serving_set_metadata_update() { new_test_ext(1).execute_with(|| { diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index d047288dd..612908835 100644 --- a/runtime/src/lib.rs +++ b/runtime/src/lib.rs @@ -141,7 +141,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion { // `spec_version`, and `authoring_version` are the same between Wasm and native. // This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use // the compatible custom types. - spec_version: 192, + spec_version: 193, impl_version: 1, apis: RUNTIME_API_VERSIONS, transaction_version: 1, @@ -1598,13 +1598,7 @@ impl_runtime_apis! { } fn get_delegate(delegate_account_vec: Vec) -> Vec { - let _result = SubtensorModule::get_delegate(delegate_account_vec); - if _result.is_some() { - let result = _result.expect("Could not get DelegateInfo"); - result.encode() - } else { - vec![] - } + SubtensorModule::get_delegate(delegate_account_vec).map(|r| r.encode()).unwrap_or(vec![]) } fn get_delegated(delegatee_account_vec: Vec) -> Vec { @@ -1620,13 +1614,7 @@ impl_runtime_apis! { } fn get_neuron_lite(netuid: u16, uid: u16) -> Vec { - let _result = SubtensorModule::get_neuron_lite(netuid, uid); - if _result.is_some() { - let result = _result.expect("Could not get NeuronInfoLite"); - result.encode() - } else { - vec![] - } + SubtensorModule::get_neuron_lite(netuid, uid).map(|r| r.encode()).unwrap_or(vec![]) } fn get_neurons(netuid: u16) -> Vec { @@ -1635,25 +1623,17 @@ impl_runtime_apis! { } fn get_neuron(netuid: u16, uid: u16) -> Vec { - let _result = SubtensorModule::get_neuron(netuid, uid); - if _result.is_some() { - let result = _result.expect("Could not get NeuronInfo"); - result.encode() - } else { - vec![] - } + SubtensorModule::get_neuron(netuid, uid).map(|r| r.encode()).unwrap_or(vec![]) + } + + fn get_neuron_certificate(netuid: u16, uid: u16) -> Vec { + SubtensorModule::get_neuron_certificate(netuid, uid).map(|r| r.encode()).unwrap_or(vec![]) } } impl subtensor_custom_rpc_runtime_api::SubnetInfoRuntimeApi for Runtime { fn get_subnet_info(netuid: u16) -> Vec { - let _result = SubtensorModule::get_subnet_info(netuid); - if _result.is_some() { - let result = _result.expect("Could not get SubnetInfo"); - result.encode() - } else { - vec![] - } + SubtensorModule::get_subnet_info(netuid).map(|r| r.encode()).unwrap_or(vec![]) } fn get_subnets_info() -> Vec { @@ -1662,13 +1642,7 @@ impl_runtime_apis! { } fn get_subnet_hyperparams(netuid: u16) -> Vec { - let _result = SubtensorModule::get_subnet_hyperparams(netuid); - if _result.is_some() { - let result = _result.expect("Could not get SubnetHyperparams"); - result.encode() - } else { - vec![] - } + SubtensorModule::get_subnet_hyperparams(netuid).map(|r| r.encode()).unwrap_or(vec![]) } }