Skip to content

Commit

Permalink
bounded vec
Browse files Browse the repository at this point in the history
  • Loading branch information
andreea-popescu-reef committed Sep 6, 2024
1 parent 2aad6fe commit 6b89f64
Show file tree
Hide file tree
Showing 5 changed files with 53 additions and 17 deletions.
27 changes: 24 additions & 3 deletions pallets/subtensor/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ pub mod pallet {
traits::{
tokens::fungible, OriginTrait, QueryPreimage, StorePreimage, UnfilteredDispatchable,
},
BoundedVec,
};
use frame_system::pallet_prelude::*;
use sp_core::H256;
Expand Down Expand Up @@ -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.
Expand Down
8 changes: 3 additions & 5 deletions pallets/subtensor/src/subnets/serving.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
4 changes: 2 additions & 2 deletions pallets/subtensor/tests/serving.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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"));
});
}

Expand Down
4 changes: 1 addition & 3 deletions pallets/subtensor/tests/swap_hotkey.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
27 changes: 23 additions & 4 deletions pallets/subtensor/tests/uids.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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), ());
});
}

0 comments on commit 6b89f64

Please sign in to comment.