Decouple woker_type and worker_mode (#2476)

* add worker_mode * adjust test * update extrinsic param
litentry · Feb 11, 2024 · aaccf7d · aaccf7d
1 parent 5ce842f
commit aaccf7d
Show file tree

Hide file tree

Showing 9 changed files with 58 additions and 22 deletions.
diff --git a/pallets/teebag/src/lib.rs b/pallets/teebag/src/lib.rs
@@ -135,17 +135,18 @@ pub mod pallet {
 		EnclaveIdentifierNotExist,
 		/// The enclave identifier already exists.
 		EnclaveIdentifierAlreadyExist,
-		/// when we try to re-register an existing enclave with a differnet worker type
-		WorkerTypeNotAllowed,
+		/// The worker type is unexpected, becuase e.g. when we try to re-register an
+		/// existing enclave with a differnet worker type
+		UnexpectedWorkerType,
 		/// The shard doesn't match the enclave.
 		WrongMrenclaveForShard,
 		/// The worker url is too long.
 		EnclaveUrlTooLong,
 		/// The raw attestation data is too long.
 		AttestationTooLong,
-		/// The worker type is unexpected, because e.g. a non-sidechain worker calls sidechain
-		/// specific extrinsic
-		UnexpectedWorkerType,
+		/// The worker mode is unexpected, because e.g. a non-sidechain worker calls
+		/// sidechain specific extrinsic
+		UnexpectedWorkerMode,
 		/// Can not found the desired scheduled enclave.
 		ScheduledEnclaveNotExist,
 		/// Enclave not in the scheduled list, therefore unexpected.
@@ -412,6 +413,7 @@ pub mod pallet {
 		pub fn register_enclave(
 			origin: OriginFor<T>,
 			worker_type: WorkerType,
+			worker_mode: WorkerMode,
 			attestation: Vec<u8>,
 			worker_url: Vec<u8>,
 			shielding_pubkey: Option<Vec<u8>>,
@@ -423,6 +425,7 @@ pub mod pallet {
 			ensure!(worker_url.len() <= MAX_URL_LEN, Error::<T>::EnclaveUrlTooLong);
 
 			let mut enclave = Enclave::new(worker_type)
+				.with_worker_mode(worker_mode)
 				.with_url(worker_url)
 				.with_shielding_pubkey(shielding_pubkey)
 				.with_vc_pubkey(vc_pubkey)
@@ -564,7 +567,10 @@ pub mod pallet {
 				Error::<T>::WrongMrenclaveForShard
 			);
 
-			ensure!(sender_enclave.worker_type.is_sidechain(), Error::<T>::UnexpectedWorkerType,);
+			ensure!(
+				sender_enclave.worker_mode == WorkerMode::Sidechain,
+				Error::<T>::UnexpectedWorkerMode
+			);
 
 			sender_enclave.last_seen_timestamp = Self::now().saturated_into();
 
@@ -626,7 +632,7 @@ impl<T: Config> Pallet<T> {
 		match EnclaveRegistry::<T>::get(sender) {
 			Some(old_enclave) => ensure!(
 				old_enclave.worker_type == enclave.worker_type,
-				Error::<T>::WorkerTypeNotAllowed
+				Error::<T>::UnexpectedWorkerType
 			),
 			None => Self::add_enclave_identifier(enclave.worker_type, sender)?,
 		};

diff --git a/pallets/teebag/src/tests.rs b/pallets/teebag/src/tests.rs
@@ -90,6 +90,7 @@ fn register_enclave_dev_works_with_no_scheduled_enclave() {
 		assert_ok!(Teebag::register_enclave(
 			RuntimeOrigin::signed(alice()),
 			Default::default(),
+			Default::default(),
 			TEST4_MRENCLAVE.to_vec(),
 			URL.to_vec(),
 			None,
@@ -114,6 +115,7 @@ fn register_enclave_dev_works_with_sgx_build_mode_debug() {
 		assert_ok!(Teebag::register_enclave(
 			RuntimeOrigin::signed(signer4.clone()),
 			Default::default(),
+			Default::default(),
 			TEST4_CERT.to_vec(),
 			URL.to_vec(),
 			None,
@@ -152,6 +154,7 @@ fn parentchain_block_processed_works() {
 		assert_ok!(Teebag::register_enclave(
 			RuntimeOrigin::signed(signer7.clone()),
 			WorkerType::BitAcross,
+			Default::default(),
 			TEST7_CERT.to_vec(),
 			URL.to_vec(),
 			None,
@@ -195,6 +198,7 @@ fn register_dcap_enclave_works() {
 		assert_ok!(Teebag::register_enclave(
 			RuntimeOrigin::signed(signer.clone()),
 			WorkerType::Identity,
+			Default::default(),
 			TEST1_DCAP_QUOTE.to_vec(),
 			URL.to_vec(),
 			None,
@@ -227,6 +231,7 @@ fn register_enclave_prod_works_with_sgx_build_mode_debug() {
 		assert_ok!(Teebag::register_enclave(
 			RuntimeOrigin::signed(signer4.clone()),
 			Default::default(),
+			Default::default(),
 			TEST4_CERT.to_vec(),
 			URL.to_vec(),
 			None,
@@ -261,6 +266,7 @@ fn register_enclave_prod_works_with_sgx_build_mode_production() {
 		assert_ok!(Teebag::register_enclave(
 			RuntimeOrigin::signed(signer8.clone()),
 			Default::default(),
+			Default::default(),
 			TEST8_CERT.to_vec(),
 			URL.to_vec(),
 			None,
@@ -287,6 +293,7 @@ fn register_enclave_prod_fails_with_wrong_attestation_type() {
 			Teebag::register_enclave(
 				RuntimeOrigin::signed(alice()),
 				Default::default(),
+				Default::default(),
 				TEST4_MRENCLAVE.to_vec(),
 				URL.to_vec(),
 				None,
@@ -307,6 +314,7 @@ fn register_enclave_prod_fails_with_no_scheduled_enclave() {
 			Teebag::register_enclave(
 				RuntimeOrigin::signed(signer),
 				Default::default(),
+				Default::default(),
 				TEST4_CERT.to_vec(),
 				URL.to_vec(),
 				None,
@@ -333,6 +341,7 @@ fn register_enclave_prod_fails_with_max_limit_reached() {
 		assert_ok!(Teebag::register_enclave(
 			RuntimeOrigin::signed(signer5.clone()),
 			WorkerType::BitAcross,
+			Default::default(),
 			TEST5_CERT.to_vec(),
 			URL.to_vec(),
 			None,
@@ -344,6 +353,7 @@ fn register_enclave_prod_fails_with_max_limit_reached() {
 			Teebag::register_enclave(
 				RuntimeOrigin::signed(signer6.clone()),
 				WorkerType::BitAcross,
+				Default::default(),
 				TEST6_CERT.to_vec(),
 				URL.to_vec(),
 				None,
@@ -359,13 +369,14 @@ fn register_enclave_prod_fails_with_max_limit_reached() {
 			Teebag::register_enclave(
 				RuntimeOrigin::signed(signer5.clone()),
 				WorkerType::Identity,
+				Default::default(),
 				TEST5_CERT.to_vec(),
 				URL.to_vec(),
 				None,
 				None,
 				AttestationType::Ias,
 			),
-			Error::<Test>::WorkerTypeNotAllowed
+			Error::<Test>::UnexpectedWorkerType
 		);
 
 		// remove and re-register it should work
@@ -374,6 +385,7 @@ fn register_enclave_prod_fails_with_max_limit_reached() {
 		assert_ok!(Teebag::register_enclave(
 			RuntimeOrigin::signed(signer5),
 			WorkerType::Identity,
+			Default::default(),
 			TEST5_CERT.to_vec(),
 			URL.to_vec(),
 			None,
@@ -386,6 +398,7 @@ fn register_enclave_prod_fails_with_max_limit_reached() {
 			Teebag::register_enclave(
 				RuntimeOrigin::signed(signer6),
 				WorkerType::Identity,
+				Default::default(),
 				TEST6_CERT.to_vec(),
 				URL.to_vec(),
 				None,

diff --git a/pallets/teebag/src/types.rs b/pallets/teebag/src/types.rs
@@ -74,10 +74,12 @@ pub enum WorkerType {
 	BitAcross,
 }
 
-impl WorkerType {
-	pub fn is_sidechain(&self) -> bool {
-		self == &Self::Identity
-	}
+#[derive(Encode, Decode, Clone, Copy, Default, PartialEq, Eq, RuntimeDebug, TypeInfo)]
+pub enum WorkerMode {
+	#[default]
+	OffChainWorker,
+	Sidechain,
+	Teeracle,
 }
 
 #[derive(Encode, Decode, Copy, Clone, Default, PartialEq, Eq, RuntimeDebug, TypeInfo)]
@@ -99,6 +101,7 @@ pub struct SidechainBlockConfirmation {
 #[derive(Encode, Decode, Clone, Default, PartialEq, Eq, RuntimeDebug, TypeInfo)]
 pub struct Enclave {
 	pub worker_type: WorkerType,
+	pub worker_mode: WorkerMode,
 	pub mrenclave: MrEnclave,
 	pub last_seen_timestamp: u64, // unix epoch in milliseconds when it's last seen
 	pub url: Vec<u8>,             // utf8 encoded url
@@ -113,6 +116,11 @@ impl Enclave {
 		Enclave { worker_type, ..Default::default() }
 	}
 
+	pub fn with_worker_mode(mut self, worker_mode: WorkerMode) -> Self {
+		self.worker_mode = worker_mode;
+		self
+	}
+
 	pub fn with_mrenclave(mut self, mrenclave: MrEnclave) -> Self {
 		self.mrenclave = mrenclave;
 		self

diff --git a/tee-worker/Cargo.lock b/tee-worker/Cargo.lock
diff --git a/tee-worker/core-primitives/settings/Cargo.toml b/tee-worker/core-primitives/settings/Cargo.toml
@@ -5,10 +5,12 @@ authors = ['Trust Computing GmbH <info@litentry.com>', 'Integritee AG <hello@int
 edition = "2021"
 
 [dependencies]
-
+litentry-primitives = { path = "../../litentry/primitives", default-features = false }
 
 [features]
-production = []
+production = [
+    "litentry-primitives/production",
+]
 sidechain = []
 offchain-worker = []
 teeracle = []
diff --git a/tee-worker/core-primitives/settings/src/worker_mode.rs b/tee-worker/core-primitives/settings/src/worker_mode.rs
@@ -15,12 +15,7 @@
 
 */
 
-#[derive(Eq, PartialEq, Debug, Clone)]
-pub enum WorkerMode {
-	OffChainWorker,
-	Sidechain,
-	Teeracle,
-}
+pub use litentry_primitives::WorkerMode;
 
 pub trait ProvideWorkerMode {
 	fn worker_mode() -> WorkerMode;

diff --git a/tee-worker/enclave-runtime/Cargo.lock b/tee-worker/enclave-runtime/Cargo.lock
diff --git a/tee-worker/enclave-runtime/src/attestation.rs b/tee-worker/enclave-runtime/src/attestation.rs
@@ -47,7 +47,10 @@ use itp_node_api::metadata::{
 	Error as MetadataError,
 };
 use itp_node_api_metadata::NodeMetadata;
-use itp_settings::worker::MR_ENCLAVE_SIZE;
+use itp_settings::{
+	worker::MR_ENCLAVE_SIZE,
+	worker_mode::{ProvideWorkerMode, WorkerModeProvider},
+};
 use itp_sgx_crypto::{
 	ed25519_derivation::DeriveEd25519, key_repository::AccessKey, Error as SgxCryptoError,
 };
@@ -332,6 +335,7 @@ pub fn generate_dcap_ra_extrinsic_from_quote_internal(
 	let call = OpaqueCall::from_tuple(&(
 		call_ids,
 		WorkerType::Identity,
+		WorkerModeProvider::worker_mode(),
 		quote,
 		url,
 		shielding_pubkey,
@@ -361,6 +365,7 @@ pub fn generate_dcap_skip_ra_extrinsic_from_mr_enclave(
 	let call = OpaqueCall::from_tuple(&(
 		call_ids,
 		WorkerType::Identity,
+		WorkerModeProvider::worker_mode(),
 		quote,
 		url,
 		shielding_pubkey,
@@ -401,6 +406,7 @@ pub fn generate_ias_ra_extrinsic_from_der_cert_internal(
 	let call = OpaqueCall::from_tuple(&(
 		call_ids,
 		WorkerType::Identity,
+		WorkerModeProvider::worker_mode(),
 		cert_der,
 		url,
 		shielding_pubkey,

diff --git a/tee-worker/litentry/primitives/src/lib.rs b/tee-worker/litentry/primitives/src/lib.rs
@@ -45,7 +45,7 @@ use litentry_hex_utils::hex_encode;
 use log::error;
 pub use pallet_teebag::{
 	decl_rsa_request, extract_tcb_info_from_raw_dcap_quote, AttestationType, Enclave,
-	EnclaveFingerprint, MrEnclave, ShardIdentifier, SidechainBlockNumber, WorkerType,
+	EnclaveFingerprint, MrEnclave, ShardIdentifier, SidechainBlockNumber, WorkerMode, WorkerType,
 };
 pub use parentchain_primitives::{
 	all_bitcoin_web3networks, all_evm_web3networks, all_substrate_web3networks, all_web3networks,