fix(pegboard-runner): clear terminal tunnel routes

Author: NathanFlurry

engine/packages/pegboard-runner/src/ws_to_tunnel_task.rs

@@ -860,6 +860,13 @@ async fn handle_tunnel_message_mk2(
 	authorized_tunnel_routes: &HashMap<(protocol::mk2::GatewayId, protocol::mk2::RequestId), ()>,
 	msg: protocol::mk2::ToServerTunnelMessage,
 ) -> Result<()> {
+	let route = (msg.message_id.gateway_id, msg.message_id.request_id);
+	let clear_route = matches!(
+		msg.message_kind,
+		protocol::mk2::ToServerTunnelMessageKind::ToServerResponseAbort
+			| protocol::mk2::ToServerTunnelMessageKind::ToServerWebSocketClose(_)
+	);
+
 	// Extract inner data length before consuming msg
 	let inner_data_len = tunnel_message_inner_data_len_mk2(&msg.message_kind);
 
@@ -868,10 +875,7 @@ async fn handle_tunnel_message_mk2(
 		return Err(errors::WsError::InvalidPacket("payload too large".to_string()).build());
 	}
 
-	if !authorized_tunnel_routes
-		.contains_async(&(msg.message_id.gateway_id, msg.message_id.request_id))
-		.await
-	{
+	if !authorized_tunnel_routes.contains_async(&route).await {
 		return Err(
 			errors::WsError::InvalidPacket("unauthorized tunnel message".to_string()).build(),
 		);
@@ -899,6 +903,10 @@ async fn handle_tunnel_message_mk2(
 			)
 		})?;
 
+	if clear_route {
+		authorized_tunnel_routes.remove_async(&route).await;
+	}
+
 	Ok(())
 }
 
@@ -909,6 +917,13 @@ async fn handle_tunnel_message_mk1(
 	authorized_tunnel_routes: &HashMap<(protocol::mk2::GatewayId, protocol::mk2::RequestId), ()>,
 	msg: protocol::ToServerTunnelMessage,
 ) -> Result<()> {
+	let route = (msg.message_id.gateway_id, msg.message_id.request_id);
+	let clear_route = matches!(
+		msg.message_kind,
+		protocol::ToServerTunnelMessageKind::ToServerResponseAbort
+			| protocol::ToServerTunnelMessageKind::ToServerWebSocketClose(_)
+	);
+
 	// Ignore DeprecatedTunnelAck messages (used only for backwards compatibility)
 	if matches!(
 		msg.message_kind,
@@ -925,10 +940,7 @@ async fn handle_tunnel_message_mk1(
 		return Err(errors::WsError::InvalidPacket("payload too large".to_string()).build());
 	}
 
-	if !authorized_tunnel_routes
-		.contains_async(&(msg.message_id.gateway_id, msg.message_id.request_id))
-		.await
-	{
+	if !authorized_tunnel_routes.contains_async(&route).await {
 		return Err(
 			errors::WsError::InvalidPacket("unauthorized tunnel message".to_string()).build(),
 		);
@@ -950,6 +962,10 @@ async fn handle_tunnel_message_mk1(
 			)
 		})?;
 
+	if clear_route {
+		authorized_tunnel_routes.remove_async(&route).await;
+	}
+
 	Ok(())
 }
 

engine/packages/pegboard-runner/tests/support/ws_to_tunnel_task.rs

@@ -92,6 +92,11 @@ async fn republishes_issued_mk2_tunnel_message_pairs() {
 		.unwrap()
 		.unwrap();
 	assert!(matches!(msg, NextOutput::Message(_)));
+	assert!(
+		authorized_tunnel_routes
+			.contains_async(&(gateway_id, request_id))
+			.await
+	);
 }
 
 #[tokio::test]
@@ -147,4 +152,61 @@ async fn republishes_issued_mk1_tunnel_message_pairs() {
 		.unwrap()
 		.unwrap();
 	assert!(matches!(msg, NextOutput::Message(_)));
+	assert!(
+		authorized_tunnel_routes
+			.contains_async(&(gateway_id, request_id))
+			.await
+	);
+}
+
+#[tokio::test]
+async fn removes_terminal_mk2_tunnel_message_pairs() {
+	let pubsub = memory_pubsub("pegboard-runner-ws-to-tunnel-test-remove-mk2");
+	let gateway_id = [33, 34, 35, 36];
+	let request_id = [37, 38, 39, 40];
+	let authorized_tunnel_routes = HashMap::new();
+	let _ = authorized_tunnel_routes
+		.insert_async((gateway_id, request_id), ())
+		.await;
+
+	handle_tunnel_message_mk2(
+		&pubsub,
+		1024,
+		&authorized_tunnel_routes,
+		response_abort_message_mk2(gateway_id, request_id),
+	)
+	.await
+	.unwrap();
+
+	assert!(
+		!authorized_tunnel_routes
+			.contains_async(&(gateway_id, request_id))
+			.await
+	);
+}
+
+#[tokio::test]
+async fn removes_terminal_mk1_tunnel_message_pairs() {
+	let pubsub = memory_pubsub("pegboard-runner-ws-to-tunnel-test-remove-mk1");
+	let gateway_id = [41, 42, 43, 44];
+	let request_id = [45, 46, 47, 48];
+	let authorized_tunnel_routes = HashMap::new();
+	let _ = authorized_tunnel_routes
+		.insert_async((gateway_id, request_id), ())
+		.await;
+
+	handle_tunnel_message_mk1(
+		&pubsub,
+		1024,
+		&authorized_tunnel_routes,
+		response_abort_message_mk1(gateway_id, request_id),
+	)
+	.await
+	.unwrap();
+
+	assert!(
+		!authorized_tunnel_routes
+			.contains_async(&(gateway_id, request_id))
+			.await
+	);
 }