fix(pegboard): convert ws manager socket to unix socket

Author: MasterPtato

packages/edge/infra/client/config/src/manager.rs

@@ -85,13 +85,6 @@ pub struct Runner {
 	/// Whether or not to use a mount for actor file systems.
 	pub use_mounts: Option<bool>,
 
-	/// Address of the WebSocket server for runners. Should exist on a network interface that both the host
-	/// and containers can access.
-	pub ip: Option<IpAddr>,
-
-	/// WebSocket port for runners on this machine to connect to.
-	pub port: Option<u16>,
-
 	pub container_runner_binary_path: Option<PathBuf>,
 }
 
@@ -100,14 +93,6 @@ impl Runner {
 		self.use_mounts.unwrap_or(true)
 	}
 
-	pub fn ip(&self) -> IpAddr {
-		self.ip.unwrap_or(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)))
-	}
-
-	pub fn port(&self) -> u16 {
-		self.port.unwrap_or(6080)
-	}
-
 	pub fn container_runner_binary_path(&self) -> PathBuf {
 		self.container_runner_binary_path
 			.clone()

packages/edge/infra/client/config/src/runner_protocol.rs

@@ -38,3 +38,24 @@ pub enum ActorState {
 	Running,
 	Exited { exit_code: Option<i32> },
 }
+
+pub fn codec() -> LengthDelimitedCodec {
+	LengthDelimitedCodec::builder()
+		.length_field_type::<u32>()
+		.length_field_length(4)
+		// No offset
+		.length_field_offset(0)
+		// Header length is not included in the length calculation
+		.length_adjustment(4)
+		// header is included in the returned bytes
+		.num_skip(0)
+		.new_codec()
+}
+
+pub fn encode_frame() {
+
+}
+
+pub fn decode_frame() {
+
+}

packages/edge/infra/client/manager/src/ctx.rs

@@ -220,44 +220,6 @@ impl Ctx {
 
 		self.receive_init(&mut rx).await?;
 
-		// Start runner socket and attaches incoming connections to their corresponding runner
-		let self2 = self.clone();
-		let runner_socket: tokio::task::JoinHandle<Result<()>> = tokio::spawn(async move {
-			tracing::info!(port=%self2.config().runner.port(), "listening for runner sockets");
-
-			let listener =
-				TcpListener::bind((self2.config().runner.ip(), self2.config().runner.port()))
-					.await
-					.map_err(RuntimeError::RunnerSocketListenFailed)?;
-
-			loop {
-				match listener.accept().await {
-					Ok((stream, _)) => {
-						let mut ws_stream = Some(tokio_tungstenite::accept_async(stream).await?);
-
-						tracing::info!("received new socket");
-
-						if let Err(err) = self2.receive_runner_init_message(&mut ws_stream).await {
-							tracing::error!(
-								?err,
-								"failed to receive init message from runner socket"
-							);
-						}
-
-						// Close stream
-						if let Some(mut ws_stream) = ws_stream {
-							let close_frame = CloseFrame {
-								code: CloseCode::Error,
-								reason: "init failed".into(),
-							};
-							ws_stream.send(Message::Close(Some(close_frame))).await?;
-						}
-					}
-					Err(err) => tracing::error!(?err, "failed to connect websocket"),
-				}
-			}
-		});
-
 		// Start ping thread after init packet is received because ping denotes this client as "ready"
 		let self2 = self.clone();
 		let ping_thread: tokio::task::JoinHandle<Result<()>> = tokio::spawn(async move {
@@ -549,63 +511,6 @@ impl Ctx {
 		Ok(())
 	}
 
-	async fn receive_runner_init_message(
-		self: &Arc<Self>,
-		ws_stream: &mut Option<WebSocketStream<TcpStream>>,
-	) -> Result<()> {
-		match tokio::time::timeout(
-			RUNNER_INIT_TIMEOUT,
-			ws_stream.as_mut().context("ws_stream should exist")?.next(),
-		)
-		.await
-		{
-			Ok(msg) => match msg {
-				Some(Ok(Message::Close(_))) | None => {
-					tracing::debug!("runner socket closed");
-					Ok(())
-				}
-				Some(Ok(Message::Binary(buf))) => {
-					match serde_json::from_slice::<runner_protocol::ToManager>(&buf) {
-						Ok(runner_protocol::ToManager::Init { access_token }) => {
-							let claims = Claims::decode(&access_token, self.secret())
-								.context("could not decode access token from init packet")?;
-
-							// Read runner id from claims
-							let runner_id = 'ent: {
-								for ent in claims.ent() {
-									#[allow(irrefutable_let_patterns)]
-									if let Entitlement::Runner { runner_id } = ent {
-										break 'ent runner_id;
-									}
-								}
-
-								bail!("no runner entitlement");
-							};
-
-							if let Some(runner) = self.runners.read().await.get(&runner_id) {
-								runner
-									.attach_socket(
-										self,
-										ws_stream.take().context("ws_stream should exist")?,
-									)
-									.await?;
-
-								Ok(())
-							} else {
-								bail!("killing unknown runner: {runner_id}");
-							}
-						}
-						Ok(packet) => bail!("invalid init packet from runner: {packet:?}"),
-						Err(err) => Err(err).context("invalid init packet from runner"),
-					}
-				}
-				Some(Ok(packet)) => bail!("invalid init packet from runner: {packet:?}"),
-				Some(Err(err)) => Err(err).context("runner socket error"),
-			},
-			Err(_) => bail!("runner socket init timed out"),
-		}
-	}
-
 	/// Returns None if the runner could not be found in the runners map on time.
 	async fn get_or_create_runner(
 		&self,
@@ -892,6 +797,10 @@ impl Ctx {
 			let runner = runner.clone();
 			let self2 = self.clone();
 			tokio::spawn(async move {
+				// The socket file already exists, this will reconnect and spawn the appropriate task to
+				// handle the connection
+				runner.start_socket(&self2).await?;
+
 				if let Err(err) = runner.observe(&self2, false).await {
 					tracing::error!(?runner_id, ?err, "observe failed");
 				}

packages/edge/infra/client/manager/src/runner/mod.rs

@@ -111,7 +111,7 @@ impl Runner {
 	pub async fn attach_socket(
 		self: &Arc<Self>,
 		ctx: &Arc<Ctx>,
-		ws_stream: WebSocketStream<TcpStream>,
+		stream: WebSocketStream<TcpStream>,
 	) -> Result<()> {
 		match &self.comms {
 			Comms::Basic => bail!("attempt to attach socket to basic runner"),
@@ -120,33 +120,36 @@ impl Runner {
 
 				let mut guard = tx.lock().await;
 
-				if let Some(existing_ws_tx) = &mut *guard {
+				if let Some(existing_tx) = &mut *guard {
 					tracing::info!(runner_id=?self.runner_id, "runner received another socket, closing old one");
 
 					// Close the old socket
 					let close_frame = CloseFrame {
 						code: CloseCode::Error,
 						reason: "replacing with new socket".into(),
 					};
-					existing_ws_tx
-						.send(Message::Close(Some(close_frame)))
-						.await?;
+
+					if let Err(err) = existing_tx.send(Message::Close(Some(close_frame))).await {
+						tracing::error!(runner_id=?self.runner_id, ?err, "failed to close old socket");
+					};
 
 					tracing::info!(runner_id=?self.runner_id, "socket replaced");
 				} else {
 					tracing::info!(runner_id=?self.runner_id, "socket attached");
 				}
 
-				let (ws_tx, ws_rx) = ws_stream.split();
+				// Wrap the stream in a framed transport
+				let mut framed = Framed::new(stream, runner_protocol::codec());
+				let (tx, rx) = stream.split();
 
-				*guard = Some(ws_tx);
+				*guard = Some(tx);
 				self.bump();
 
 				// Spawn a new thread to handle incoming messages
 				let self2 = self.clone();
 				let ctx2 = ctx.clone();
 				tokio::task::spawn(async move {
-					if let Err(err) = self2.receive_messages(&ctx2, ws_rx).await {
+					if let Err(err) = self2.receive_frames(&ctx2, rx).await {
 						tracing::error!(runner_id=?self2.runner_id, ?err, "socket error, killing runner");
 
 						if let Err(err) = self2.signal(Signal::SIGKILL).await {
@@ -160,52 +163,46 @@ impl Runner {
 		Ok(())
 	}
 
-	async fn receive_messages(
-		&self,
-		ctx: &Ctx,
-		mut ws_rx: SplitStream<WebSocketStream<TcpStream>>,
-	) -> Result<()> {
+	async fn receive_frames(&self, ctx: &Ctx, mut ws_rx: SplitStream<UnixStream>) -> Result<()> {
 		loop {
-			match tokio::time::timeout(PING_TIMEOUT, ws_rx.next()).await {
-				Ok(msg) => match msg {
-					Some(Ok(Message::Ping(_))) => {
-						// Pongs are sent automatically
-					}
-					Some(Ok(Message::Close(_))) | None => {
-						tracing::debug!(runner_id=?self.runner_id, "runner socket closed");
-						break Ok(());
-					}
-					Some(Ok(Message::Binary(buf))) => {
-						let packet = serde_json::from_slice::<runner_protocol::ToManager>(&buf)?;
+			let Ok(frame) = tokio::time::timeout(PING_TIMEOUT, ws_rx.next()).await else {
+				bail!("runner socket ping timed out");
+			};
 
-						self.process_packet(ctx, packet).await?;
-					}
-					Some(Ok(packet)) => bail!("runner socket unexpected packet: {packet:?}"),
-					Some(Err(err)) => break Err(err).context("runner socket error"),
-				},
-				Err(_) => bail!("socket ping timed out"),
-			}
-		}
-	}
+			let Some(buf) = frame else {
+				tracing::debug!(runner_id=?self.runner_id, "runner socket closed");
+				break Ok(());
+			};
+
+			let (_, packet) = runner_protocol::decode_frame::<runner_protocol::ToManager>(&buf)?;
+
+			tracing::debug!(?packet, "runner received packet");
 
-	async fn process_packet(&self, ctx: &Ctx, packet: runner_protocol::ToManager) -> Result<()> {
-		tracing::debug!(?packet, "runner received packet");
-
-		match packet {
-			runner_protocol::ToManager::Init { .. } => bail!("unexpected second init packet"),
-			runner_protocol::ToManager::ActorStateUpdate {
-				actor_id,
-				generation,
-				state,
-			} => {
-				// NOTE: We don't have to verify if the actor id given here is valid because only valid actors
-				// are listening to this runner's `actor_observer_tx`. This means invalid messages are ignored.
-				// NOTE: No receivers is not an error
-				let _ = self.actor_observer_tx.send((actor_id, generation, state));
+			match packet {
+				runner_protocol::ToManager::Ping { .. } => {
+					let socket = socket
+						.lock()
+						.await
+						.context("must have socket if receiving ping")?;
+
+					let buf = runner_protocol::encode_frame(runner_protocol::ToRunner::Pong)?;
+					socket
+						.send(buf)
+						.await
+						.context("failed to send packet to socket")?;
+				}
+				runner_protocol::ToManager::ActorStateUpdate {
+					actor_id,
+					generation,
+					state,
+				} => {
+					// NOTE: We don't have to verify if the actor id given here is valid because only valid actors
+					// are listening to this runner's `actor_observer_tx`. This means invalid messages are ignored.
+					// NOTE: No receivers is not an error
+					let _ = self.actor_observer_tx.send((actor_id, generation, state));
+				}
 			}
 		}
-
-		Ok(())
 	}
 
 	pub async fn send(&self, packet: &runner_protocol::ToRunner) -> Result<()> {
@@ -618,7 +615,7 @@ impl Runner {
 
 pub enum Comms {
 	Basic,
-	Socket(Mutex<Option<SplitSink<WebSocketStream<TcpStream>, Message>>>),
+	Socket(Mutex<Option<SplitSink<UnixStream, u32>>>),
 }
 
 impl Comms {

packages/edge/infra/client/manager/src/runner/oci_config.rs

@@ -14,6 +14,7 @@ pub struct ConfigOpts<'a> {
 	pub cpu: u64,
 	pub memory: u64,
 	pub memory_max: u64,
+	pub mount_socket: bool,
 }
 
 /// Generates base config.json for an OCI bundle.
@@ -160,7 +161,7 @@ fn capabilities() -> Vec<&'static str> {
 }
 
 fn mounts(opts: &ConfigOpts) -> Result<serde_json::Value> {
-	Ok(json!([
+	let mut mounts = json!([
 		{
 			"destination": "/proc",
 			"type": "proc",
@@ -245,7 +246,18 @@ fn mounts(opts: &ConfigOpts) -> Result<serde_json::Value> {
 			"source": opts.runner_path.join("resolv.conf").to_str().context("resolv.conf path")?,
 			"options": ["rbind", "rprivate"]
 		},
-	]))
+	]);
+
+	if let Some(socket_path) = opts.mount_socket {
+		mounts.as_array().unwrap().push(json!({
+			"destination": socket_mount_dest_path().to_str().context("manager.sock dest path")?,
+			"type": "bind",
+			"source": opts.runner_path.join("manager.sock").to_str().context("manager.sock source path")?,
+			"options": ["rbind", "ro"]
+		},));
+	}
+
+	Ok()
 }
 
 fn linux_resources_devices() -> serde_json::Value {
@@ -322,3 +334,8 @@ fn linux_resources_devices() -> serde_json::Value {
 		}
 	])
 }
+
+/// Mounting path of the manager socket inside of the container.
+pub fn socket_mount_dest_path() -> Path {
+	Path::new("/srv/sockets/manager.sock")
+}