[zombinet] initial implementation of zombienet backchannel (#4377)

polkadot/node/zombienet-backchannel/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "zombienet-backchannel"
+description = "Zombienet backchannel to notify test runner and coordinate with malus actors."
+license = "GPL-3.0-only"
+version = "0.9.13"
+authors = ["Parity Technologies <admin@parity.io>"]
+edition = "2021"
+readme = "README.md"
+publish = false
+
+[dependencies]
+tokio = { version = "1.0.0", default-features = false, features = ["macros", "net", "rt-multi-thread", "sync"] }
+url = "2.0.0"
+tokio-tungstenite = "0.16"
+futures-util = "0.3.18"
+lazy_static = "1.4.0"
+parity-scale-codec = { version = "2.3.1", features = ["derive"] }
+reqwest = "0.11"
+thiserror = "1.0.30"
+tracing = "0.1.26"
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1"

polkadot/node/zombienet-backchannel/src/errors.rs

@@ -0,0 +1,39 @@
+// Copyright 2021 Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Polkadot is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Polkadot Zombienet Backchannel error definitions.
+
+#[derive(Debug, thiserror::Error)]
+#[allow(missing_docs)]
+pub enum BackchannelError {
+	#[error("Error connecting websocket server")]
+	CantConnectToWS,
+
+	#[error("Backchannel not initialized yet")]
+	Uninitialized,
+
+	#[error("Backchannel already initialized")]
+	AlreadyInitialized,
+
+	#[error("Error sending new value to backchannel")]
+	SendItemFail,
+
+	#[error("Invalid host for connection backchannel")]
+	InvalidHost,
+
+	#[error("Invalid port for connection backchannel")]
+	InvalidPort,
+}

polkadot/node/zombienet-backchannel/src/lib.rs

@@ -0,0 +1,159 @@
+// Copyright 2017-2021 Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Polkadot is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Provides the possibility to coordination between malicious actors and
+//! the zombienet test-runner, allowing to reference runtime's generated
+//! values in the test specifications, through a bidirectional message passing
+//! implemented as a `backchannel`.
+
+use futures_util::{SinkExt, StreamExt};
+use lazy_static::lazy_static;
+use parity_scale_codec as codec;
+use serde::{Deserialize, Serialize};
+use std::{env, sync::Mutex};
+use tokio::sync::broadcast;
+use tokio_tungstenite::{connect_async, tungstenite::protocol::Message};
+
+mod errors;
+use errors::BackchannelError;
+
+lazy_static! {
+	pub static ref ZOMBIENET_BACKCHANNEL: Mutex<Option<ZombienetBackchannel>> = Mutex::new(None);
+}
+
+#[derive(Debug)]
+pub struct ZombienetBackchannel {
+	broadcast_tx: broadcast::Sender<BackchannelItem>,
+	ws_tx: broadcast::Sender<BackchannelItem>,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct BackchannelItem {
+	key: String,
+	value: String,
+}
+
+pub struct Broadcaster;
+
+pub const ZOMBIENET: &str = "🧟ZOMBIENET🧟";
+
+impl Broadcaster {
+	/// Return a subscriber that will receive all message broadcasted by the zombienet backchannel
+	/// websocket server.
+	pub fn subscribe(&self) -> Result<broadcast::Receiver<BackchannelItem>, BackchannelError> {
+		let mut zombienet_bkc = ZOMBIENET_BACKCHANNEL.lock().unwrap();
+		let zombienet_bkc = zombienet_bkc.as_mut().ok_or(BackchannelError::Uninitialized)?;
+		let sender = zombienet_bkc.broadcast_tx.clone();
+		Ok(sender.subscribe())
+	}
+
+	/// Provides a simple API to send a key/value to the zombienet websocket server.
+	pub async fn send(
+		&mut self,
+		key: &'static str,
+		val: impl codec::Encode,
+	) -> Result<(), BackchannelError> {
+		let mut zombienet_bkc = ZOMBIENET_BACKCHANNEL.lock().unwrap();
+		let zombienet_bkc = zombienet_bkc.as_mut().ok_or(BackchannelError::Uninitialized)?;
+
+		let encoded = val.encode();
+		let backchannel_item = BackchannelItem {
+			key: key.to_string(),
+			value: String::from_utf8_lossy(&encoded).to_string(),
+		};
+
+		let sender = zombienet_bkc.ws_tx.clone();
+		sender.send(backchannel_item).map_err(|e| {
+			tracing::error!(target = ZOMBIENET, "Error sending new item: {}", e);
+			BackchannelError::SendItemFail
+		})?;
+
+		Ok(())
+	}
+}
+
+impl ZombienetBackchannel {
+	pub async fn init() -> Result<(), BackchannelError> {
+		let mut zombienet_bkc = ZOMBIENET_BACKCHANNEL.lock().unwrap();
+		if zombienet_bkc.is_none() {
+			let backchannel_host =
+				env::var("BACKCHANNEL_HOST").unwrap_or_else(|_| "backchannel".to_string());
+			let backchannel_port =
+				env::var("BACKCHANNEL_PORT").unwrap_or_else(|_| "3000".to_string());
+
+			// validate port
+			backchannel_port.parse::<u16>().map_err(|_| BackchannelError::InvalidPort)?;
+			// validate non empty string for host
+			if backchannel_host.trim().is_empty() {
+				return Err(BackchannelError::InvalidHost)
+			};
+
+			let ws_url = format!("ws://{}:{}/ws", backchannel_host, backchannel_port);
+			tracing::debug!(target = ZOMBIENET, "Connecting to : {}", &ws_url);
+			let (ws_stream, _) =
+				connect_async(ws_url).await.map_err(|_| BackchannelError::CantConnectToWS)?;
+			let (mut write, mut read) = ws_stream.split();
+
+			let (tx, _rx) = broadcast::channel(256);
+			let (tx_relay, mut rx_relay) = broadcast::channel::<BackchannelItem>(256);
+
+			// receive from the ws and send to all subcribers
+			let tx1 = tx.clone();
+			tokio::spawn(async move {
+				while let Some(Ok(Message::Text(text))) = read.next().await {
+					match serde_json::from_str::<BackchannelItem>(&text) {
+						Ok(backchannel_item) =>
+							if tx1.send(backchannel_item).is_err() {
+								tracing::error!("Error sending through the channel");
+								return
+							},
+						Err(_) => {
+							tracing::error!("Invalid payload received");
+						},
+					}
+				}
+			});
+
+			// receive from subscribers and relay to ws
+			tokio::spawn(async move {
+				while let Ok(item) = rx_relay.recv().await {
+					if write
+						.send(Message::Text(serde_json::to_string(&item).unwrap()))
+						.await
+						.is_err()
+					{
+						tracing::error!("Error sending through ws");
+					}
+				}
+			});
+
+			*zombienet_bkc = Some(ZombienetBackchannel { broadcast_tx: tx, ws_tx: tx_relay });
+			return Ok(())
+		}
+
+		Err(BackchannelError::AlreadyInitialized)
+	}
+
+	/// Ensure that the backchannel is initialized and return a broadcaster instance
+	/// allowing to subscribe or send new items.
+	pub fn broadcaster() -> Result<Broadcaster, BackchannelError> {
+		if ZOMBIENET_BACKCHANNEL.lock().unwrap().is_some() {
+			Ok(Broadcaster {})
+		} else {
+			Err(BackchannelError::Uninitialized)
+		}
+	}
+}