Re-register enclave after state provisioning (#2578)

* remove deprecated request-state subcommand

* re-register itself to update the key

* remove request-state subcommand

tee-worker/docker/multiworker-docker-compose.yml

@@ -208,7 +208,7 @@ services:
     entrypoint:
       "/usr/local/bin/litentry-worker --clean-reset --ws-external -M litentry-worker-2 -T wss://litentry-worker-2
       -u ws://litentry-node -U ws://litentry-worker-2 -P 2011 -w 2101 -p 9912 -h 4645 --enable-mock-server
-      run --dev --skip-ra --request-state"
+      run --dev --skip-ra"
     restart: "no"
   litentry-worker-3:
     image: litentry/litentry-worker:latest
@@ -264,7 +264,7 @@ services:
     entrypoint:
       "/usr/local/bin/litentry-worker --clean-reset --ws-external -M litentry-worker-3 -T wss://litentry-worker-3
       -u ws://litentry-node -U ws://litentry-worker-3 -P 2011 -w 2101 -p 9912 -h 4645 --enable-mock-server
-      run --dev --skip-ra --request-state"
+      run --dev --skip-ra"
     restart: "no"
 volumes:
   ? relaychain-alice

tee-worker/scripts/launch_local_worker.sh

@@ -36,7 +36,6 @@ UNTRUSTED_WORKER_PORT="3000"
 
 F_CLEAN=""
 FSUBCMD_DEV=""
-FSUBCMD_REQ_STATE=""
 
 WAIT_INTERVAL_SECONDS=10
 WAIT_ROUNDS=20
@@ -87,10 +86,6 @@ for ((i = 0; i < ${WORKER_NUM}; i++)); do
 	echo ""
 	echo "--------------------setup worker(${worker_name})----------------------------------------"
 
-	if ((i > 0)); then
-		FSUBCMD_REQ_STATE="--request-state"
-	fi
-
 	if ((i == 0)); then
 		MOCK_SERVER="--enable-mock-server"
 	fi
@@ -128,7 +123,7 @@ for ((i = 0; i < ${WORKER_NUM}; i++)); do
 --untrusted-external-address ws://${WORKER_ENDPOINT} \
 --untrusted-http-port ${untrusted_http_port} \
 --untrusted-worker-port ${untrusted_worker_port} \
-run --skip-ra ${FSUBCMD_DEV} ${FSUBCMD_REQ_STATE}"
+run --skip-ra ${FSUBCMD_DEV}"
 
 	echo "${worker_name} command: ${launch_command}"
 	eval "${launch_command}" > "${ROOTDIR}"/log/${worker_name}.log 2>&1 &

tee-worker/service/src/cli.yml

@@ -151,17 +151,14 @@ subcommands:
                     long: dev
                     short: d
                     help: Set this flag if running in development mode to bootstrap enclave account on parentchain via //Alice.
-              - request-state:
-                    long: request-state
-                    short: r
-                    help: Run the worker and request key and state provisioning from another worker.
-              - marblerun-url:
+              - shielding-target:
                     required: false
-                    long: marblerun-url
-                    help: Set the URL for prometheus metrics. Probably put some meaningful description
+                    long: shielding-target
+                    short: s
+                    help: set parentchain target for shielding / unshielding. only relevant for primary worker upon first start for shard. can't be changed later for a shard
                     takes_value: true
     - request-state:
-          about: (Deprecated - TODO) join a shard by requesting key provisioning from another worker
+          about: (DEPRECATED) join a shard by requesting key provisioning from another worker
           args:
               - shard:
                     long: shard

tee-worker/service/src/main_impl.rs

@@ -505,27 +505,19 @@ fn start_worker<E, T, D, InitializationHandler, WorkerModeProvider>(
 	#[cfg(feature = "dcap")]
 	enclave.set_sgx_qpl_logging().expect("QPL logging setup failed");
 
-	let enclave2 = enclave.clone();
-	let node_api2 = litentry_rpc_api.clone();
-	let tee_accountid2 = tee_accountid.clone();
-	let trusted_url2 = trusted_url.clone();
-
-	#[cfg(not(feature = "dcap"))]
-	let register_xt = move || enclave2.generate_ias_ra_extrinsic(&trusted_url2, skip_ra).unwrap();
-	#[cfg(feature = "dcap")]
-	let register_xt = move || enclave2.generate_dcap_ra_extrinsic(&trusted_url2, skip_ra).unwrap();
-
-	let send_register_xt = move || {
-		println!("[+] Send register enclave extrinsic");
-		send_litentry_extrinsic(register_xt(), &node_api2, &tee_accountid2, is_development_mode)
-	};
-
 	// Litentry: send the registration extrinsic regardless of being registered or not,
 	//           the reason is the mrenclave could change in between, so we rely on the
 	//           on-chain logic to handle everything.
 	//           this is the same behavior as upstream
-	let register_enclave_block_hash =
-		send_register_xt().expect("enclave RA registration must be successful to continue");
+	let register_enclave_block_hash = register_enclave(
+		enclave.clone(),
+		&litentry_rpc_api,
+		&tee_accountid,
+		&trusted_url,
+		skip_ra,
+		is_development_mode,
+	)
+	.expect("enclave RA registration must be successful to continue");
 
 	let api_register_enclave_xt_header =
 		litentry_rpc_api.get_header(Some(register_enclave_block_hash)).unwrap().unwrap();
@@ -598,6 +590,17 @@ fn start_worker<E, T, D, InitializationHandler, WorkerModeProvider>(
 							enclave.as_ref(),
 							skip_ra,
 						);
+
+						info!("re-register the enclave to update the keys after provisioning");
+						register_enclave(
+							enclave.clone(),
+							&litentry_rpc_api,
+							&tee_accountid,
+							&trusted_url,
+							skip_ra,
+							is_development_mode,
+						)
+						.expect("enclave RA registration must be successful to continue");
 					}
 					(false, true)
 				}
@@ -1118,3 +1121,23 @@ pub fn enclave_account<E: EnclaveBase>(enclave_api: &E) -> AccountId32 {
 	trace!("[+] Got ed25519 account of TEE = {}", tee_public.to_ss58check());
 	AccountId32::from(*tee_public.as_array_ref())
 }
+
+fn register_enclave<E>(
+	enclave: Arc<E>,
+	api: &ParentchainApi,
+	tee_account: &AccountId32,
+	url: &str,
+	skip_ra: bool,
+	is_development_mode: bool,
+) -> ServiceResult<Hash>
+where
+	E: EnclaveBase + DirectRequest + Sidechain + RemoteAttestation + TlsRemoteAttestation + Clone,
+{
+	#[cfg(not(feature = "dcap"))]
+	let register_xt = move || enclave.generate_ias_ra_extrinsic(url, skip_ra).unwrap();
+	#[cfg(feature = "dcap")]
+	let register_xt = move || enclave.generate_dcap_ra_extrinsic(url, skip_ra).unwrap();
+
+	println!("[+] Send register enclave extrinsic");
+	send_litentry_extrinsic(register_xt(), api, tee_account, is_development_mode)
+}

tee-worker/ts-tests/worker/resuming_worker.test.ts

@@ -24,7 +24,6 @@ type WorkerConfig = {
     trustedWorkerPort: number;
     untrustedWorkerPort: number;
     enableMockServer: boolean;
-    requestStateOnLaunch: boolean;
 };
 
 const workerConfig: Record<'worker0' | 'worker1', WorkerConfig> = {
@@ -35,7 +34,6 @@ const workerConfig: Record<'worker0' | 'worker1', WorkerConfig> = {
         trustedWorkerPort: 2000,
         untrustedWorkerPort: 3000,
         enableMockServer: true,
-        requestStateOnLaunch: false,
     },
     worker1: {
         name: 'worker1',
@@ -44,7 +42,6 @@ const workerConfig: Record<'worker0' | 'worker1', WorkerConfig> = {
         trustedWorkerPort: 2001,
         untrustedWorkerPort: 3001,
         enableMockServer: false,
-        requestStateOnLaunch: true,
     },
 } as const;
 
@@ -78,7 +75,6 @@ function generateWorkerCommandArguments(
         `run`,
         `--skip-ra`,
         `--dev`,
-        ...(isLaunch && workerParams.requestStateOnLaunch ? ['--request-state'] : []),
     ].join(' ');
 }