From 0bbe0a7621500788dd841be7e4d869cfba5f0676 Mon Sep 17 00:00:00 2001
From: Marcin S <marcin@realemail.net>
Date: Mon, 21 Aug 2023 19:31:02 +0200
Subject: [PATCH] PVF worker: random fixes (#7649)

* PVF worker: random fixes

- Fixes possible panic due to non-UTF-8 env vars
  (https://github.com/paritytech/polkadot/pull/7330#discussion_r1300101716)
- Very small refactor of some duplicated code

* Don't need `to_str()` for comparison between OsString and str

* Check edge cases that can cause env::remove_var to panic

In case of a key or value that would cause env::remove_var to panic, we first
log a warning and then proceed to attempt to remove the env var.

* Make warning message clearer for end users

* Backslash was unescaped, but can just remove it from error messages
---
 node/core/pvf/common/src/worker/mod.rs | 64 ++++++++++++++++++++------
 1 file changed, 51 insertions(+), 13 deletions(-)

diff --git a/node/core/pvf/common/src/worker/mod.rs b/node/core/pvf/common/src/worker/mod.rs
index d249007ec36e..8b41cb82f73b 100644
--- a/node/core/pvf/common/src/worker/mod.rs
+++ b/node/core/pvf/common/src/worker/mod.rs
@@ -121,22 +121,13 @@ pub fn worker_event_loop<F, Fut>(
 				"Node and worker version mismatch, node needs restarting, forcing shutdown",
 			);
 			kill_parent_node_in_emergency();
-			let err: io::Result<Never> =
-				Err(io::Error::new(io::ErrorKind::Unsupported, "Version mismatch"));
-			gum::debug!(target: LOG_TARGET, %worker_pid, "quitting pvf worker({}): {:?}", debug_id, err);
+			let err = io::Error::new(io::ErrorKind::Unsupported, "Version mismatch");
+			worker_shutdown_message(debug_id, worker_pid, err);
 			return
 		}
 	}
 
-	// Delete all env vars to prevent malicious code from accessing them.
-	for (key, _) in std::env::vars() {
-		// TODO: *theoretically* the value (or mere presence) of `RUST_LOG` can be a source of
-		// randomness for malicious code. In the future we can remove it also and log in the host;
-		// see <https://github.com/paritytech/polkadot/issues/7117>.
-		if key != "RUST_LOG" {
-			std::env::remove_var(key);
-		}
-	}
+	remove_env_vars(debug_id);
 
 	// Run the main worker loop.
 	let rt = Runtime::new().expect("Creates tokio runtime. If this panics the worker will die and the host will detect that and deal with it.");
@@ -152,7 +143,7 @@ pub fn worker_event_loop<F, Fut>(
 		// It's never `Ok` because it's `Ok(Never)`.
 		.unwrap_err();
 
-	gum::debug!(target: LOG_TARGET, %worker_pid, "quitting pvf worker ({}): {:?}", debug_id, err);
+	worker_shutdown_message(debug_id, worker_pid, err);
 
 	// We don't want tokio to wait for the tasks to finish. We want to bring down the worker as fast
 	// as possible and not wait for stalled validation to finish. This isn't strictly necessary now,
@@ -160,6 +151,53 @@ pub fn worker_event_loop<F, Fut>(
 	rt.shutdown_background();
 }
 
+/// Delete all env vars to prevent malicious code from accessing them.
+fn remove_env_vars(debug_id: &'static str) {
+	for (key, value) in std::env::vars_os() {
+		// TODO: *theoretically* the value (or mere presence) of `RUST_LOG` can be a source of
+		// randomness for malicious code. In the future we can remove it also and log in the host;
+		// see <https://github.com/paritytech/polkadot/issues/7117>.
+		if key == "RUST_LOG" {
+			continue
+		}
+
+		// In case of a key or value that would cause [`env::remove_var` to
+		// panic](https://doc.rust-lang.org/std/env/fn.remove_var.html#panics), we first log a
+		// warning and then proceed to attempt to remove the env var.
+		let mut err_reasons = vec![];
+		let (key_str, value_str) = (key.to_str(), value.to_str());
+		if key.is_empty() {
+			err_reasons.push("key is empty");
+		}
+		if key_str.is_some_and(|s| s.contains('=')) {
+			err_reasons.push("key contains '='");
+		}
+		if key_str.is_some_and(|s| s.contains('\0')) {
+			err_reasons.push("key contains null character");
+		}
+		if value_str.is_some_and(|s| s.contains('\0')) {
+			err_reasons.push("value contains null character");
+		}
+		if !err_reasons.is_empty() {
+			gum::warn!(
+				target: LOG_TARGET,
+				%debug_id,
+				?key,
+				?value,
+				"Attempting to remove badly-formatted env var, this may cause the PVF worker to crash. Please remove it yourself. Reasons: {:?}",
+				err_reasons
+			);
+		}
+
+		std::env::remove_var(key);
+	}
+}
+
+/// Provide a consistent message on worker shutdown.
+fn worker_shutdown_message(debug_id: &'static str, worker_pid: u32, err: io::Error) {
+	gum::debug!(target: LOG_TARGET, %worker_pid, "quitting pvf worker ({}): {:?}", debug_id, err);
+}
+
 /// Loop that runs in the CPU time monitor thread on prepare and execute jobs. Continuously wakes up
 /// and then either blocks for the remaining CPU time, or returns if we exceed the CPU timeout.
 ///