fix(pegboard-runner): clear terminal tunnel routes

Author: NathanFlurry

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

@@ -860,6 +860,9 @@ 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 = should_clear_tunnel_route_mk2(&msg.message_kind);
+
 	// Extract inner data length before consuming msg
 	let inner_data_len = tunnel_message_inner_data_len_mk2(&msg.message_kind);
 
@@ -868,10 +871,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 +899,10 @@ async fn handle_tunnel_message_mk2(
 			)
 		})?;
 
+	if clear_route {
+		authorized_tunnel_routes.remove_async(&route).await;
+	}
+
 	Ok(())
 }
 
@@ -909,6 +913,9 @@ 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 = should_clear_tunnel_route_mk1(&msg.message_kind);
+
 	// Ignore DeprecatedTunnelAck messages (used only for backwards compatibility)
 	if matches!(
 		msg.message_kind,
@@ -925,10 +932,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,9 +954,35 @@ async fn handle_tunnel_message_mk1(
 			)
 		})?;
 
+	if clear_route {
+		authorized_tunnel_routes.remove_async(&route).await;
+	}
+
 	Ok(())
 }
 
+fn should_clear_tunnel_route_mk2(msg_kind: &protocol::mk2::ToServerTunnelMessageKind) -> bool {
+	match msg_kind {
+		protocol::mk2::ToServerTunnelMessageKind::ToServerResponseStart(response) => {
+			!response.stream
+		}
+		protocol::mk2::ToServerTunnelMessageKind::ToServerResponseChunk(chunk) => chunk.finish,
+		protocol::mk2::ToServerTunnelMessageKind::ToServerResponseAbort
+		| protocol::mk2::ToServerTunnelMessageKind::ToServerWebSocketClose(_) => true,
+		_ => false,
+	}
+}
+
+fn should_clear_tunnel_route_mk1(msg_kind: &protocol::ToServerTunnelMessageKind) -> bool {
+	match msg_kind {
+		protocol::ToServerTunnelMessageKind::ToServerResponseStart(response) => !response.stream,
+		protocol::ToServerTunnelMessageKind::ToServerResponseChunk(chunk) => chunk.finish,
+		protocol::ToServerTunnelMessageKind::ToServerResponseAbort
+		| protocol::ToServerTunnelMessageKind::ToServerWebSocketClose(_) => true,
+		_ => false,
+	}
+}
+
 /// Returns the length of the inner data payload for a tunnel message kind.
 fn tunnel_message_inner_data_len_mk2(kind: &protocol::mk2::ToServerTunnelMessageKind) -> usize {
 	use protocol::mk2::ToServerTunnelMessageKind;