Accept PVF code hashes in validation host

node/core/candidate-validation/src/lib.rs

@@ -154,7 +154,7 @@ where
 					let res = validate_candidate_exhaustive(
 						&mut validation_host,
 						persisted_validation_data,
-						validation_code,
+						Some(validation_code),
 						descriptor,
 						pov,
 						&metrics,
@@ -197,11 +197,29 @@ where
 
 #[derive(Debug)]
 enum AssumptionCheckOutcome {
-	Matches(PersistedValidationData, ValidationCode),
+	Matches(PersistedValidationData),
 	DoesNotMatch,
 	BadRequest,
 }
 
+async fn request_validation_code_by_hash<Context>(
+	ctx: &mut Context,
+	descriptor: &CandidateDescriptor,
+) -> SubsystemResult<Result<Option<ValidationCode>, RuntimeApiError>>
+where
+	Context: SubsystemContext<Message = CandidateValidationMessage>,
+	Context: overseer::SubsystemContext<Message = CandidateValidationMessage>,
+{
+	let (tx, rx) = oneshot::channel();
+	runtime_api_request(
+		ctx,
+		descriptor.relay_parent,
+		RuntimeApiRequest::ValidationCodeByHash(descriptor.validation_code_hash, tx),
+		rx,
+	)
+	.await
+}
+
 async fn check_assumption_validation_data<Context>(
 	ctx: &mut Context,
 	descriptor: &CandidateDescriptor,
@@ -213,37 +231,25 @@ where
 {
 	let validation_data = {
 		let (tx, rx) = oneshot::channel();
-		let d = runtime_api_request(
+		let data = runtime_api_request(
 			ctx,
 			descriptor.relay_parent,
 			RuntimeApiRequest::PersistedValidationData(descriptor.para_id, assumption, tx),
 			rx,
 		)
 		.await?;
 
-		match d {
+		match data {
 			Ok(None) | Err(_) => return Ok(AssumptionCheckOutcome::BadRequest),
-			Ok(Some(d)) => d,
+			Ok(Some(data)) => data,
 		}
 	};
 
 	let persisted_validation_data_hash = validation_data.hash();
 
 	SubsystemResult::Ok(
 		if descriptor.persisted_validation_data_hash == persisted_validation_data_hash {
-			let (code_tx, code_rx) = oneshot::channel();
-			let validation_code = runtime_api_request(
-				ctx,
-				descriptor.relay_parent,
-				RuntimeApiRequest::ValidationCode(descriptor.para_id, assumption, code_tx),
-				code_rx,
-			)
-			.await?;
-
-			match validation_code {
-				Ok(None) | Err(_) => AssumptionCheckOutcome::BadRequest,
-				Ok(Some(v)) => AssumptionCheckOutcome::Matches(validation_data, v),
-			}
+			AssumptionCheckOutcome::Matches(validation_data)
 		} else {
 			AssumptionCheckOutcome::DoesNotMatch
 		},
@@ -276,7 +282,7 @@ where
 		let outcome = check_assumption_validation_data(ctx, descriptor, *assumption).await?;
 
 		match outcome {
-			AssumptionCheckOutcome::Matches(_, _) => return Ok(outcome),
+			AssumptionCheckOutcome::Matches(_) => return Ok(outcome),
 			AssumptionCheckOutcome::BadRequest => return Ok(outcome),
 			AssumptionCheckOutcome::DoesNotMatch => continue,
 		}
@@ -296,30 +302,53 @@ where
 	Context: SubsystemContext<Message = CandidateValidationMessage>,
 	Context: overseer::SubsystemContext<Message = CandidateValidationMessage>,
 {
-	let (validation_data, validation_code) =
-		match find_assumed_validation_data(ctx, &descriptor).await? {
-			AssumptionCheckOutcome::Matches(validation_data, validation_code) =>
-				(validation_data, validation_code),
-			AssumptionCheckOutcome::DoesNotMatch => {
-				// If neither the assumption of the occupied core having the para included or the assumption
-				// of the occupied core timing out are valid, then the persisted_validation_data_hash in the descriptor
-				// is not based on the relay parent and is thus invalid.
-				return Ok(Ok(ValidationResult::Invalid(InvalidCandidate::BadParent)))
-			},
-			AssumptionCheckOutcome::BadRequest =>
-				return Ok(Err(ValidationFailed("Assumption Check: Bad request".into()))),
-		};
+	let validation_data = match find_assumed_validation_data(ctx, &descriptor).await? {
+		AssumptionCheckOutcome::Matches(validation_data) => validation_data,
+		AssumptionCheckOutcome::DoesNotMatch => {
+			// If neither the assumption of the occupied core having the para included or the assumption
+			// of the occupied core timing out are valid, then the persisted_validation_data_hash in the descriptor
+			// is not based on the relay parent and is thus invalid.
+			return Ok(Ok(ValidationResult::Invalid(InvalidCandidate::BadParent)))
+		},
+		AssumptionCheckOutcome::BadRequest =>
+			return Ok(Err(ValidationFailed("Assumption Check: Bad request".into()))),
+	};
 
-	let validation_result = validate_candidate_exhaustive(
-		validation_host,
-		validation_data,
-		validation_code,
+	let mut validation_result = validate_candidate_exhaustive(
+		&mut validation_host.clone(),
+		validation_data.clone(),
+		None,
 		descriptor.clone(),
-		pov,
+		pov.clone(),
 		metrics,
 	)
 	.await;
 
+	// In case preimage for the supplied code hash was not found by the
+	// validation host, request the code from Runtime API and try again.
+	if let Ok(Err(ref err)) = validation_result {
+		if err.0.as_str() == "Code not found" {
+			let validation_code = match request_validation_code_by_hash(ctx, &descriptor).await? {
+				Ok(Some(validation_code)) => validation_code,
+				Ok(None) =>
+					return Ok(Err(ValidationFailed(
+						"Runtime API didn't return validation code".into(),
+					))),
+				Err(err) => return Ok(Err(ValidationFailed(err.to_string()))),
+			};
+
+			validation_result = validate_candidate_exhaustive(
+				validation_host,
+				validation_data,
+				Some(validation_code),
+				descriptor.clone(),
+				pov,
+				metrics,
+			)
+			.await;
+		}
+	}
+
 	if let Ok(Ok(ValidationResult::Valid(ref outputs, _))) = validation_result {
 		let (tx, rx) = oneshot::channel();
 		match runtime_api_request(
@@ -337,20 +366,30 @@ where
 		}
 	}
 
+	if let Ok(Err(ref err)) = validation_result {
+		if err.0.as_str() == "Code not found" {
+			tracing::error!(
+				target: LOG_TARGET,
+				para_id = ?descriptor.para_id,
+				"Failed to validate candidate: validation code not found in cache"
+			);
+		}
+	}
+
 	validation_result
 }
 
 async fn validate_candidate_exhaustive(
 	mut validation_backend: impl ValidationBackend,
 	persisted_validation_data: PersistedValidationData,
-	validation_code: ValidationCode,
+	validation_code: Option<ValidationCode>,
 	descriptor: CandidateDescriptor,
 	pov: Arc<PoV>,
 	metrics: &Metrics,
 ) -> SubsystemResult<Result<ValidationResult, ValidationFailed>> {
 	let _timer = metrics.time_validate_candidate_exhaustive();
 
-	let validation_code_hash = validation_code.hash();
+	let validation_code_hash = validation_code.as_ref().map(ValidationCode::hash);
 	tracing::debug!(
 		target: LOG_TARGET,
 		?validation_code_hash,
@@ -362,22 +401,30 @@ async fn validate_candidate_exhaustive(
 		&descriptor,
 		persisted_validation_data.max_pov_size,
 		&*pov,
-		&validation_code_hash,
+		validation_code_hash.as_ref(),
 	) {
 		return Ok(Ok(ValidationResult::Invalid(e)))
 	}
 
-	let raw_validation_code = match sp_maybe_compressed_blob::decompress(
-		&validation_code.0,
-		VALIDATION_CODE_BOMB_LIMIT,
-	) {
-		Ok(code) => code,
-		Err(e) => {
-			tracing::debug!(target: LOG_TARGET, err=?e, "Invalid validation code");
+	let validation_code = if let Some(validation_code) = validation_code {
+		let raw_code = match sp_maybe_compressed_blob::decompress(
+			&validation_code.0,
+			VALIDATION_CODE_BOMB_LIMIT,
+		) {
+			Ok(code) => code,
+			Err(e) => {
+				tracing::debug!(target: LOG_TARGET, err=?e, "Invalid validation code");
 
-			// If the validation code is invalid, the candidate certainly is.
-			return Ok(Ok(ValidationResult::Invalid(InvalidCandidate::CodeDecompressionFailure)))
-		},
+				// If the validation code is invalid, the candidate certainly is.
+				return Ok(Ok(ValidationResult::Invalid(InvalidCandidate::CodeDecompressionFailure)))
+			},
+		}
+		.to_vec();
+		Pvf::from_code(raw_code)
+	} else {
+		// In case validation code is not provided, ask the backend to obtain
+		// it from the cache using the hash.
+		Pvf::Hash(ValidationCodeHash::from(descriptor.persisted_validation_data_hash))
 	};
 
 	let raw_block_data =
@@ -398,9 +445,7 @@ async fn validate_candidate_exhaustive(
 		relay_parent_storage_root: persisted_validation_data.relay_parent_storage_root,
 	};
 
-	let result = validation_backend
-		.validate_candidate(raw_validation_code.to_vec(), params)
-		.await;
+	let result = validation_backend.validate_candidate(validation_code, params).await;
 
 	if let Err(ref e) = result {
 		tracing::debug!(
@@ -421,6 +466,8 @@ async fn validate_candidate_exhaustive(
 			Ok(ValidationResult::Invalid(InvalidCandidate::ExecutionError(
 				"ambigious worker death".to_string(),
 			))),
+		Err(ValidationError::ArtifactNotFound) =>
+			Err(ValidationFailed("Code not found".to_string())),
 
 		Ok(res) =>
 			if res.head_data.hash() != descriptor.para_head {
@@ -445,7 +492,7 @@ async fn validate_candidate_exhaustive(
 trait ValidationBackend {
 	async fn validate_candidate(
 		&mut self,
-		raw_validation_code: Vec<u8>,
+		validation_code: Pvf,
 		params: ValidationParams,
 	) -> Result<WasmValidationResult, ValidationError>;
 }
@@ -454,13 +501,13 @@ trait ValidationBackend {
 impl ValidationBackend for &'_ mut ValidationHost {
 	async fn validate_candidate(
 		&mut self,
-		raw_validation_code: Vec<u8>,
+		validation_code: Pvf,
 		params: ValidationParams,
 	) -> Result<WasmValidationResult, ValidationError> {
 		let (tx, rx) = oneshot::channel();
 		if let Err(err) = self
 			.execute_pvf(
-				Pvf::from_code(raw_validation_code),
+				validation_code,
 				params.encode(),
 				polkadot_node_core_pvf::Priority::Normal,
 				tx,
@@ -487,7 +534,7 @@ fn perform_basic_checks(
 	candidate: &CandidateDescriptor,
 	max_pov_size: u32,
 	pov: &PoV,
-	validation_code_hash: &ValidationCodeHash,
+	validation_code_hash: Option<&ValidationCodeHash>,
 ) -> Result<(), InvalidCandidate> {
 	let pov_hash = pov.hash();
 
@@ -500,7 +547,7 @@ fn perform_basic_checks(
 		return Err(InvalidCandidate::PoVHashMismatch)
 	}
 
-	if *validation_code_hash != candidate.validation_code_hash {
+	if let Some(false) = validation_code_hash.map(|hash| *hash == candidate.validation_code_hash) {
 		return Err(InvalidCandidate::CodeHashMismatch)
 	}