bounded vec

Author: andreea-popescu-reef

pallets/subtensor/src/lib.rs

@@ -67,6 +67,7 @@ pub mod pallet {
         traits::{
             tokens::fungible, OriginTrait, QueryPreimage, StorePreimage, UnfilteredDispatchable,
         },
+        BoundedVec,
     };
     use frame_system::pallet_prelude::*;
     use sp_core::H256;
@@ -133,11 +134,31 @@ pub mod pallet {
     /// Struct for NeuronCertificate.
     pub type NeuronCertificateOf = NeuronCertificate;
     /// Data structure for NeuronCertificate information.
-    #[freeze_struct("e6193a76002d491")]
+    #[freeze_struct("1c232be200d9ec6c")]
     #[derive(Decode, Encode, Default, TypeInfo, PartialEq, Eq, Clone, Debug)]
     pub struct NeuronCertificate {
-        ///  The neuron certificate.
-        pub certificate: Vec<u8>,
+        ///  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.

pallets/subtensor/src/subnets/serving.rs

@@ -92,11 +92,9 @@ impl<T: Config> Pallet<T> {
 
         // Check certificate
         if let Some(certificate) = certificate {
-            NeuronCertificates::<T>::insert(
-                netuid,
-                hotkey_id.clone(),
-                NeuronCertificate { certificate },
-            )
+            if let Ok(certificate) = NeuronCertificateOf::try_from(certificate) {
+                NeuronCertificates::<T>::insert(netuid, hotkey_id.clone(), certificate)
+            }
         }
 
         // We insert the axon meta.

pallets/subtensor/tests/serving.rs

@@ -131,7 +131,7 @@ fn test_serving_tls_ok() {
 
         let stored_certificate = NeuronCertificates::<Test>::get(netuid, hotkey_account_id)
             .expect("Certificate should exist");
-        assert_eq!(stored_certificate.certificate, certificate);
+        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),
@@ -147,7 +147,7 @@ fn test_serving_tls_ok() {
         ));
         let stored_certificate = NeuronCertificates::<Test>::get(netuid, hotkey_account_id)
             .expect("Certificate should exist");
-        assert_eq!(stored_certificate.certificate, new_certificate)
+        assert_eq!(stored_certificate.public_key.clone().into_inner(), new_certificate.get(1..).expect("Certificate should exist"));
     });
 }
 

pallets/subtensor/tests/swap_hotkey.rs

@@ -319,9 +319,7 @@ fn test_swap_certificates() {
         let new_hotkey = U256::from(2);
         let coldkey = U256::from(3);
         let netuid = 0u16;
-        let certificate = NeuronCertificate {
-            certificate: vec![1, 2, 3],
-        };
+        let certificate = NeuronCertificate::try_from(vec![1, 2, 3]).unwrap();
         let mut weight = Weight::zero();
 
         add_network(netuid, 0, 1);

pallets/subtensor/tests/uids.rs

@@ -1,7 +1,7 @@
 #![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;
@@ -33,9 +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 {
-            certificate: vec![1, 2, 3],
-        };
+        let certificate = NeuronCertificate::try_from(vec![1, 2, 3]).unwrap();
 
         //add network
         add_network(netuid, tempo, 0);
@@ -382,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), ());
+    });
+}