diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index 4f553ed7a..cc3d7d025 100644 --- a/pallets/subtensor/src/lib.rs +++ b/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, + /// 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. diff --git a/pallets/subtensor/src/subnets/serving.rs b/pallets/subtensor/src/subnets/serving.rs index 22550fb93..7e2b9a0f0 100644 --- a/pallets/subtensor/src/subnets/serving.rs +++ b/pallets/subtensor/src/subnets/serving.rs @@ -92,11 +92,9 @@ impl Pallet { // Check certificate if let Some(certificate) = certificate { - NeuronCertificates::::insert( - netuid, - hotkey_id.clone(), - NeuronCertificate { certificate }, - ) + if let Ok(certificate) = NeuronCertificateOf::try_from(certificate) { + NeuronCertificates::::insert(netuid, hotkey_id.clone(), certificate) + } } // We insert the axon meta. diff --git a/pallets/subtensor/tests/serving.rs b/pallets/subtensor/tests/serving.rs index 17b8d2144..4516c3b05 100644 --- a/pallets/subtensor/tests/serving.rs +++ b/pallets/subtensor/tests/serving.rs @@ -131,7 +131,7 @@ fn test_serving_tls_ok() { let stored_certificate = NeuronCertificates::::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( <::RuntimeOrigin>::signed(hotkey_account_id), @@ -147,7 +147,7 @@ fn test_serving_tls_ok() { )); let stored_certificate = NeuronCertificates::::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")); }); } diff --git a/pallets/subtensor/tests/swap_hotkey.rs b/pallets/subtensor/tests/swap_hotkey.rs index 845338be9..89938e3eb 100644 --- a/pallets/subtensor/tests/swap_hotkey.rs +++ b/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); diff --git a/pallets/subtensor/tests/uids.rs b/pallets/subtensor/tests/uids.rs index 827d4ec1a..6b4c00328 100644 --- a/pallets/subtensor/tests/uids.rs +++ b/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), ()); + }); +}