Fix state provisioning bug when reading large byte buffers (#1083)

* Fix state provisioning when reading large states
* initialize state regardless of whether it was already initialzed.
Required for state provisioning in case only shielding key is exchanged
(offchain-worker)

Cargo.lock

@@ -3307,6 +3307,7 @@ dependencies = [
  "parity-scale-codec",
  "rust-base58 0.0.4 (registry+https://github.com/rust-lang/crates.io-index)",
  "rust-base58 0.0.4 (git+https://github.com/mesalock-linux/rust-base58-sgx?rev=sgx_1.1.3)",
+ "sgx_crypto_helper",
  "sgx_tcrypto",
  "sgx_tstd",
  "sgx_types",

core-primitives/stf-state-handler/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2021"
 
 [dependencies]
 # sgx dependencies
+sgx_crypto_helper = { default-features = false, optional = true, features = ["mesalock_sgx"], version = "1.1.5", git = "https://github.com/apache/incubator-teaclave-sgx-sdk", tag = "v1.1.5" }
 sgx_tcrypto = { branch = "master", git = "https://github.com/apache/teaclave-sgx-sdk.git", optional = true }
 sgx_tstd = { branch = "master", git = "https://github.com/apache/teaclave-sgx-sdk.git", optional = true }
 sgx_types = { branch = "master", git = "https://github.com/apache/teaclave-sgx-sdk.git" }
@@ -71,4 +72,5 @@ sgx = [
 test = [
     "itp-sgx-crypto/mocks",
     "itp-stf-interface/mocks",
+    "sgx_crypto_helper",
 ]

core-primitives/stf-state-handler/src/file_io.rs

@@ -93,34 +93,40 @@ pub mod sgx {
 	use codec::Decode;
 	use core::fmt::Debug;
 	use ita_stf::AccountId;
-	use itp_sgx_crypto::{key_repository::AccessKey, StateCrypto};
+	use itp_sgx_crypto::{
+		ed25519_derivation::DeriveEd25519, key_repository::AccessKey, StateCrypto,
+	};
 	use itp_sgx_externalities::SgxExternalitiesTrait;
 	use itp_sgx_io::{read as io_read, write as io_write};
 	use itp_stf_interface::InitState;
 	use itp_types::H256;
 	use log::*;
 	use sgx_tcrypto::rsgx_sha256_slice;
+	use sp_core::Pair;
 	use std::{fs, marker::PhantomData, path::Path, sync::Arc};
 
 	/// SGX state file I/O.
-	pub struct SgxStateFileIo<StateKeyRepository, Stf, State> {
+	pub struct SgxStateFileIo<StateKeyRepository, ShieldingKeyRepository, Stf, State> {
 		state_key_repository: Arc<StateKeyRepository>,
-		enclave_account: AccountId,
+		shielding_key_repository: Arc<ShieldingKeyRepository>,
 		_phantom: PhantomData<(State, Stf)>,
 	}
 
-	impl<StateKeyRepository, Stf, State> SgxStateFileIo<StateKeyRepository, Stf, State>
+	impl<StateKeyRepository, ShieldingKeyRepository, Stf, State>
+		SgxStateFileIo<StateKeyRepository, ShieldingKeyRepository, Stf, State>
 	where
 		StateKeyRepository: AccessKey,
 		<StateKeyRepository as AccessKey>::KeyType: StateCrypto,
+		ShieldingKeyRepository: AccessKey,
+		<ShieldingKeyRepository as AccessKey>::KeyType: DeriveEd25519,
 		Stf: InitState<State, AccountId>,
 		State: SgxExternalitiesTrait,
 	{
 		pub fn new(
 			state_key_repository: Arc<StateKeyRepository>,
-			enclave_account: AccountId,
+			shielding_key_repository: Arc<ShieldingKeyRepository>,
 		) -> Self {
-			SgxStateFileIo { state_key_repository, enclave_account, _phantom: PhantomData }
+			SgxStateFileIo { state_key_repository, shielding_key_repository, _phantom: PhantomData }
 		}
 
 		fn read(&self, path: &Path) -> Result<Vec<u8>> {
@@ -157,10 +163,13 @@ pub mod sgx {
 		}
 	}
 
-	impl<StateKeyRepository, Stf, State> StateFileIo for SgxStateFileIo<StateKeyRepository, Stf, State>
+	impl<StateKeyRepository, ShieldingKeyRepository, Stf, State> StateFileIo
+		for SgxStateFileIo<StateKeyRepository, ShieldingKeyRepository, Stf, State>
 	where
 		StateKeyRepository: AccessKey,
 		<StateKeyRepository as AccessKey>::KeyType: StateCrypto,
+		ShieldingKeyRepository: AccessKey,
+		<ShieldingKeyRepository as AccessKey>::KeyType: DeriveEd25519,
 		Stf: InitState<State, AccountId>,
 		State: SgxExternalitiesTrait + Debug,
 		<State as SgxExternalitiesTrait>::SgxExternalitiesType: Encode + Decode,
@@ -219,7 +228,8 @@ pub mod sgx {
 			state_id: StateId,
 		) -> Result<Self::HashType> {
 			init_shard(&shard_identifier)?;
-			let state = Stf::init_state(self.enclave_account.clone());
+			let enclave_account = self.shielding_key_repository.retrieve_key()?.derive_ed25519()?;
+			let state = Stf::init_state(enclave_account.public().into());
 			self.write(shard_identifier, state_id, state)
 		}
 

core-primitives/stf-state-handler/src/state_snapshot_repository.rs

@@ -51,6 +51,8 @@ pub trait VersionedStateAccess {
 	) -> Result<Self::StateType>;
 
 	/// Initialize a new shard.
+	///
+	/// If the shard already exists, it will re-initialize it.
 	fn initialize_new_shard(&mut self, shard_identifier: ShardIdentifier)
 		-> Result<Self::HashType>;
 
@@ -248,11 +250,6 @@ where
 		&mut self,
 		shard_identifier: ShardIdentifier,
 	) -> Result<Self::HashType> {
-		if let Some(state_snapshots) = self.snapshot_history.get(&shard_identifier) {
-			warn!("Shard ({:?}) already exists, will not initialize again", shard_identifier);
-			return state_snapshots.front().map(|s| s.state_hash).ok_or(Error::EmptyRepository)
-		}
-
 		let snapshot_metadata =
 			initialize_shard_with_snapshot(&shard_identifier, self.file_io.as_ref())?;
 

core-primitives/stf-state-handler/src/test/sgx_tests.rs

@@ -31,21 +31,27 @@ use crate::{
 };
 use codec::{Decode, Encode};
 use ita_sgx_runtime::Runtime;
-use ita_stf::{AccountId, State as StfState, StateType as StfStateType, Stf};
+use ita_stf::{State as StfState, StateType as StfStateType, Stf};
 use itp_sgx_crypto::{mocks::KeyRepositoryMock, Aes, AesSeal, StateCrypto};
 use itp_sgx_externalities::{SgxExternalities, SgxExternalitiesTrait};
 use itp_sgx_io::{write, StaticSealedIO};
 use itp_stf_interface::mocks::{CallExecutorMock, GetterExecutorMock};
 use itp_stf_state_observer::state_observer::StateObserver;
 use itp_types::{ShardIdentifier, H256};
+use sgx_crypto_helper::{rsa3072::Rsa3072KeyPair, RsaKeyPair};
 use sp_core::hashing::blake2_256;
 use std::{sync::Arc, thread, vec::Vec};
 
 const STATE_SNAPSHOTS_CACHE_SIZE: usize = 3;
 
 type TestStf = Stf<CallExecutorMock, GetterExecutorMock, SgxExternalities, Runtime>;
 type StateKeyRepositoryMock = KeyRepositoryMock<Aes>;
-type TestStateFileIo = SgxStateFileIo<StateKeyRepositoryMock, TestStf, SgxExternalities>;
+type TestStateFileIo = SgxStateFileIo<
+	StateKeyRepositoryMock,
+	KeyRepositoryMock<Rsa3072KeyPair>,
+	TestStf,
+	SgxExternalities,
+>;
 type TestStateRepository = StateSnapshotRepository<TestStateFileIo>;
 type TestStateRepositoryLoader = StateSnapshotRepositoryLoader<TestStateFileIo>;
 type TestStateObserver = StateObserver<StfState>;
@@ -240,8 +246,10 @@ pub fn test_file_io_get_state_hash_works() {
 	let _shard_dir_handle = ShardDirectoryHandle::new(shard).unwrap();
 	let state_key_access =
 		Arc::new(StateKeyRepositoryMock::new(AesSeal::unseal_from_static_file().unwrap()));
+	let shielding_key_repository =
+		Arc::new(KeyRepositoryMock::<Rsa3072KeyPair>::new(Rsa3072KeyPair::new().unwrap()));
 
-	let file_io = TestStateFileIo::new(state_key_access, AccountId::new([1u8; 32]));
+	let file_io = TestStateFileIo::new(state_key_access, shielding_key_repository);
 
 	let state_id = 1234u128;
 	let state_hash = file_io.create_initialized(&shard, state_id).unwrap();
@@ -284,8 +292,10 @@ pub fn test_list_state_ids_ignores_files_not_matching_the_pattern() {
 	let _shard_dir_handle = ShardDirectoryHandle::new(shard).unwrap();
 	let state_key_access =
 		Arc::new(StateKeyRepositoryMock::new(AesSeal::unseal_from_static_file().unwrap()));
+	let shielding_key_repository =
+		Arc::new(KeyRepositoryMock::<Rsa3072KeyPair>::new(Rsa3072KeyPair::new().unwrap()));
 
-	let file_io = TestStateFileIo::new(state_key_access, AccountId::new([1u8; 32]));
+	let file_io = TestStateFileIo::new(state_key_access, shielding_key_repository);
 
 	let mut invalid_state_file_path = shard_path(&shard);
 	invalid_state_file_path.push("invalid-state.bin");
@@ -320,7 +330,10 @@ fn initialize_state_handler_with_directory_handle(
 fn initialize_state_handler() -> Arc<TestStateHandler> {
 	let state_key_access =
 		Arc::new(StateKeyRepositoryMock::new(AesSeal::unseal_from_static_file().unwrap()));
-	let file_io = Arc::new(TestStateFileIo::new(state_key_access, AccountId::new([1u8; 32])));
+	let shielding_key_repository =
+		Arc::new(KeyRepositoryMock::<Rsa3072KeyPair>::new(Rsa3072KeyPair::new().unwrap()));
+
+	let file_io = Arc::new(TestStateFileIo::new(state_key_access, shielding_key_repository));
 	let state_repository_loader = TestStateRepositoryLoader::new(file_io);
 	let state_observer = Arc::new(TestStateObserver::default());
 	let state_snapshot_repository = state_repository_loader

core-primitives/test/src/mock/handle_state_mock.rs

@@ -54,18 +54,7 @@ impl HandleState for HandleStateMock {
 	type HashType = H256;
 
 	fn initialize_shard(&self, shard: ShardIdentifier) -> Result<Self::HashType> {
-		let maybe_state = self.state_map.read().unwrap().get(&shard).cloned();
-
-		return match maybe_state {
-			// Initialize with default state, if it doesn't exist yet.
-			None => {
-				let state = StfState::default();
-				let state_hash = state.using_encoded(blake2_256).into();
-				self.state_map.write().unwrap().insert(shard, state);
-				Ok(state_hash)
-			},
-			Some(s) => Ok(s.using_encoded(blake2_256).into()),
-		}
+		self.reset(StfState::default(), &shard)
 	}
 
 	fn load(&self, shard: &ShardIdentifier) -> Result<StfState> {

enclave-runtime/Cargo.lock

@@ -1998,6 +1998,7 @@ dependencies = [
  "log",
  "parity-scale-codec",
  "rust-base58",
+ "sgx_crypto_helper",
  "sgx_tcrypto",
  "sgx_tstd",
  "sgx_types",

enclave-runtime/src/global_components.rs

@@ -84,7 +84,8 @@ pub type EnclaveTrustedCallSigned = TrustedCallSigned;
 pub type EnclaveStf = Stf<EnclaveTrustedCallSigned, EnclaveGetter, StfState, Runtime>;
 pub type EnclaveStateKeyRepository = KeyRepository<Aes, AesSeal>;
 pub type EnclaveShieldingKeyRepository = KeyRepository<Rsa3072KeyPair, Rsa3072Seal>;
-pub type EnclaveStateFileIo = SgxStateFileIo<EnclaveStateKeyRepository, EnclaveStf, StfState>;
+pub type EnclaveStateFileIo =
+	SgxStateFileIo<EnclaveStateKeyRepository, EnclaveShieldingKeyRepository, EnclaveStf, StfState>;
 pub type EnclaveStateSnapshotRepository = StateSnapshotRepository<EnclaveStateFileIo>;
 pub type EnclaveStateObserver = StateObserver<StfState>;
 pub type EnclaveStateHandler = StateHandler<EnclaveStateSnapshotRepository, EnclaveStateObserver>;
@@ -179,12 +180,8 @@ pub type EnclaveSidechainBlockImportQueueWorker = BlockImportQueueWorker<
 	EnclaveSidechainBlockImportQueue,
 	EnclaveSidechainBlockSyncer,
 >;
-pub type EnclaveSealHandler = SealHandler<
-	EnclaveShieldingKeyRepository,
-	EnclaveStateKeyRepository,
-	EnclaveStateHandler,
-	EnclaveStf,
->;
+pub type EnclaveSealHandler =
+	SealHandler<EnclaveShieldingKeyRepository, EnclaveStateKeyRepository, EnclaveStateHandler>;
 pub type EnclaveOffchainWorkerExecutor = itc_offchain_worker_executor::executor::Executor<
 	ParentchainBlock,
 	EnclaveTopPoolAuthor,

enclave-runtime/src/initialization.rs

@@ -73,9 +73,7 @@ use itp_settings::{
 	files::STATE_SNAPSHOTS_CACHE_SIZE,
 	worker_mode::{ProvideWorkerMode, WorkerMode},
 };
-use itp_sgx_crypto::{
-	aes, ed25519, ed25519_derivation::DeriveEd25519, rsa3072, AesSeal, Ed25519Seal, Rsa3072Seal,
-};
+use itp_sgx_crypto::{aes, ed25519, rsa3072, AesSeal, Ed25519Seal, Rsa3072Seal};
 use itp_sgx_io::StaticSealedIO;
 use itp_stf_state_handler::{
 	handle_state::HandleState, query_shard_state::QueryShardState,
@@ -115,11 +113,8 @@ pub(crate) fn init_enclave(mu_ra_url: String, untrusted_worker_url: String) -> E
 		Arc::new(EnclaveStateKeyRepository::new(state_key, Arc::new(AesSeal)));
 	GLOBAL_STATE_KEY_REPOSITORY_COMPONENT.initialize(state_key_repository.clone());
 
-	let enclave_call_signer_key = shielding_key.derive_ed25519()?;
-	let state_file_io = Arc::new(EnclaveStateFileIo::new(
-		state_key_repository,
-		enclave_call_signer_key.public().into(),
-	));
+	let state_file_io =
+		Arc::new(EnclaveStateFileIo::new(state_key_repository, shielding_key_repository.clone()));
 	let state_snapshot_repository_loader =
 		StateSnapshotRepositoryLoader::<EnclaveStateFileIo>::new(state_file_io);
 	let state_snapshot_repository =

enclave-runtime/src/test/tests_main.rs

@@ -146,6 +146,7 @@ pub extern "C" fn test_main_entrance() -> size_t {
 		tls_ra::seal_handler::test::seal_state_fails_for_invalid_state,
 		tls_ra::seal_handler::test::unseal_seal_state_works,
 		tls_ra::tests::test_tls_ra_server_client_networking,
+		tls_ra::tests::test_state_and_key_provisioning,
 		// RPC tests
 		direct_rpc_tests::get_state_request_works,
 

enclave-runtime/src/tls_ra/mod.rs

@@ -57,7 +57,7 @@ impl From<u8> for Opcode {
 			0 => Opcode::ShieldingKey,
 			1 => Opcode::StateKey,
 			2 => Opcode::State,
-			_ => unimplemented!(),
+			_ => unimplemented!("Unsupported/unknown Opcode for MU-RA exchange"),
 		}
 	}
 }