Merge pull request #750 from backend-developers-ltd/encrypt

add neuron certificate handling

Author: unconst

pallets/subtensor/rpc/src/lib.rs

@@ -41,7 +41,6 @@ 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 = "subnetInfo_getSubnetInfo")]
     fn get_subnet_info(&self, netuid: u16, at: Option<BlockHash>) -> RpcResult<Vec<u8>>;
     #[method(name = "subnetInfo_getSubnetsInfo")]

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<u8, ConstU32<64>>,
+        ///  The algorithm used to generate the public key
+        pub algorithm: u8,
+    }
+
+    impl TryFrom<Vec<u8>> for NeuronCertificate {
+        type Error = ();
+
+        fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
+            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<T: Config> =
         StorageDoubleMap<_, Identity, u16, Blake2_128Concat, T::AccountId, AxonInfoOf, OptionQuery>;
+    /// --- MAP ( netuid, hotkey ) --> certificate
+    #[pallet::storage]
+    pub 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<
@@ -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()))

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': (<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.
+        ///
+        /// * 'certificate' (Vec<u8>):
+        ///     - 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<T>,
+            netuid: u16,
+            version: u32,
+            ip: u128,
+            port: u16,
+            ip_type: u8,
+            protocol: u8,
+            placeholder1: u8,
+            placeholder2: u8,
+            certificate: Vec<u8>,
+        ) -> DispatchResult {
+            Self::do_serve_axon(
+                origin,
+                netuid,
+                version,
+                ip,
+                port,
+                ip_type,
+                protocol,
+                placeholder1,
+                placeholder2,
+                Some(certificate),
             )
         }
 

pallets/subtensor/src/subnets/serving.rs

@@ -31,6 +31,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 +64,7 @@ impl<T: Config> Pallet<T> {
         protocol: u8,
         placeholder1: u8,
         placeholder2: u8,
+        certificate: Option<Vec<u8>>,
     ) -> dispatch::DispatchResult {
         // We check the callers (hotkey) signature.
         let hotkey_id = ensure_signed(origin)?;
@@ -86,6 +90,13 @@ impl<T: Config> Pallet<T> {
             Error::<T>::ServingRateLimitExceeded
         );
 
+        // Check certificate
+        if let Some(certificate) = certificate {
+            if let Ok(certificate) = NeuronCertificateOf::try_from(certificate) {
+                NeuronCertificates::<T>::insert(netuid, hotkey_id.clone(), certificate)
+            }
+        }
+
         // We insert the axon meta.
         prev_axon.block = Self::get_current_block_as_u64();
         prev_axon.version = version;

pallets/subtensor/src/subnets/uids.rs

@@ -45,6 +45,9 @@ impl<T: Config> Pallet<T> {
         Uids::<T>::insert(netuid, new_hotkey.clone(), uid_to_replace); // Make uid - hotkey association.
         BlockAtRegistration::<T>::insert(netuid, uid_to_replace, block_number); // Fill block at registration.
         IsNetworkMember::<T>::insert(new_hotkey.clone(), netuid, true); // Fill network is member.
+
+        // 4. Clear neuron certificates
+        NeuronCertificates::<T>::remove(netuid, old_hotkey.clone());
     }
 
     /// Appends the uid to the network.

pallets/subtensor/src/swap/swap_hotkey.rs

@@ -276,6 +276,18 @@ impl<T: Config> Pallet<T> {
                     weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
                 }
             }
+
+            // 9.7. Swap neuron TLS certificates.
+            // NeuronCertificates( netuid, hotkey ) -> Vec<u8> -- the neuron certificate for the hotkey.
+            if is_network_member {
+                if let Ok(old_neuron_certificates) =
+                    NeuronCertificates::<T>::try_get(netuid, old_hotkey)
+                {
+                    NeuronCertificates::<T>::remove(netuid, old_hotkey);
+                    NeuronCertificates::<T>::insert(netuid, new_hotkey, old_neuron_certificates);
+                    weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
+                }
+            }
         }
 
         // 10. Swap Stake.

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<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 = NeuronCertificates::<Test>::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(
+            <<Test as Config>::RuntimeOrigin>::signed(hotkey_account_id),
+            netuid,
+            version,
+            ip,
+            port,
+            ip_type,
+            protocol,
+            placeholder1,
+            placeholder2,
+            new_certificate.clone()
+        ));
+        let stored_certificate = NeuronCertificates::<Test>::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(|| {

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::<Test>::insert(old_hotkey, netuid, true);
+        NeuronCertificates::<Test>::insert(netuid, old_hotkey, certificate.clone());
+
+        assert_ok!(SubtensorModule::perform_hotkey_swap(
+            &old_hotkey,
+            &new_hotkey,
+            &coldkey,
+            &mut weight
+        ));
+
+        assert!(!NeuronCertificates::<Test>::contains_key(
+            netuid, old_hotkey
+        ));
+        assert_eq!(
+            NeuronCertificates::<Test>::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() {