add neuron certificate handling

pallets/subtensor/rpc/src/lib.rs

@@ -41,7 +41,13 @@ pub trait SubtensorCustomApi<BlockHash> {
     fn get_neurons(&self, netuid: u16, at: Option<BlockHash>) -> RpcResult<Vec<u8>>;
     #[method(name = "neuronInfo_getNeuron")]
     fn get_neuron(&self, netuid: u16, uid: u16, at: Option<BlockHash>) -> RpcResult<Vec<u8>>;
-
+    #[method(name = "neuronInfo_getNeuronCertificate")]
+    fn get_neuron_certificate(
+        &self,
+        netuid: u16,
+        uid: u16,
+        at: Option<BlockHash>,
+    ) -> RpcResult<Vec<u8>>;
     #[method(name = "subnetInfo_getSubnetInfo")]
     fn get_subnet_info(&self, netuid: u16, at: Option<BlockHash>) -> RpcResult<Vec<u8>>;
     #[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<<Block as BlockT>::Hash>,
+    ) -> RpcResult<Vec<u8>> {
+        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,

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<u8>;
         fn get_neuron(netuid: u16, uid: u16) -> Vec<u8>;
+        fn get_neuron_certificate(netuid: u16, uid: u16) -> Vec<u8>;
         fn get_neurons_lite(netuid: u16) -> Vec<u8>;
         fn get_neuron_lite(netuid: u16, uid: u16) -> Vec<u8>;
     }

pallets/subtensor/src/certificate.rs

@@ -0,0 +1,4 @@
+use sp_std::vec::Vec;
+
+// TODO Certificate Encryption Pubkey
+pub type Certificate = Vec<u8>;

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<T: Config> =
         StorageDoubleMap<_, Identity, u16, Blake2_128Concat, T::AccountId, AxonInfoOf, OptionQuery>;
+    /// --- MAP ( netuid, hotkey ) --> certificate
+    #[pallet::storage]
+    pub(super) type NeuronCertificates<T: Config> = StorageDoubleMap<
+        _,
+        Identity,
+        u16,
+        Blake2_128Concat,
+        T::AccountId,
+        NeuronCertificateOf,
+        OptionQuery,
+    >;
     #[pallet::storage]
     /// --- MAP ( netuid, hotkey ) --> prometheus_info
     pub type Prometheus<T: Config> = 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()))

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': (<T as frame_system::Config>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<T>,
+            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': (<T as frame_system::Config>Origin):

pallets/subtensor/src/rpc_info/neuron_info.rs

@@ -147,6 +147,23 @@ impl<T: Config> Pallet<T> {
         Some(neuron)
     }
 
+    fn get_neuron_certificate(netuid: u16, uid: u16) -> Option<NeuronCertificate> {
+        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::<T>::get(netuid, hotkey)
+        } else {
+            None
+        }
+    }
+
     pub fn get_neuron(netuid: u16, uid: u16) -> Option<NeuronInfo<T>> {
         if !Self::if_subnet_exist(netuid) {
             return None;

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<T: Config> Pallet<T> {
     /// ---- The implementation for the extrinsic serve_axon which sets the ip endpoint information for a uid on a network.
@@ -31,6 +34,9 @@ impl<T: Config> Pallet<T> {
     /// * 'placeholder2' (u8):
     ///     - Placeholder for further extra params.
     ///
+    ///	* 'certificate' (Option<Vec<u8>>):
+    ///		- Certificate for mutual Tls connection between neurons
+    ///
     /// # Event:
     /// * AxonServed;
     ///     - On successfully serving the axon info.
@@ -61,6 +67,7 @@ impl<T: Config> Pallet<T> {
         protocol: u8,
         placeholder1: u8,
         placeholder2: u8,
+        certificate: Option<Certificate>,
     ) -> dispatch::DispatchResult {
         // We check the callers (hotkey) signature.
         let hotkey_id = ensure_signed(origin)?;
@@ -86,6 +93,15 @@ impl<T: Config> Pallet<T> {
             Error::<T>::ServingRateLimitExceeded
         );
 
+        // Check certificate
+        if let Some(certificate) = certificate {
+            NeuronCertificates::<T>::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<T: Config> Pallet<T> {
         Axons::<T>::contains_key(netuid, hotkey)
     }
 
+    pub fn has_neuron_certificate(netuid: u16, hotkey: &T::AccountId) -> bool {
+        return NeuronCertificates::<T>::contains_key(netuid, hotkey);
+    }
+
     pub fn has_prometheus_info(netuid: u16, hotkey: &T::AccountId) -> bool {
         Prometheus::<T>::contains_key(netuid, hotkey)
     }

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<u8> = "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(
+            <<Test as Config>::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(
+            <<Test as Config>::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(|| {