From 1b33d301af0940e17e78fc009d34129c551ef762 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?On=C3=A9sime=20BERNARD?= Date: Mon, 7 Apr 2025 14:07:34 +0000 Subject: [PATCH 1/2] Add support for multiple plc connections --- Cargo.lock | 107 +++++++++++++++++++++++++++++++++++++++++++++++++++- Cargo.toml | 1 + src/main.rs | 46 +++++++++++++++------- 3 files changed, 140 insertions(+), 14 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5890b35..a6dcb73 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "addr2line" @@ -41,6 +41,7 @@ dependencies = [ "clap", "ctor", "env_logger", + "futures", "log", "tokio", ] @@ -244,6 +245,95 @@ dependencies = [ "termcolor", ] +[[package]] +name = "futures" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" + +[[package]] +name = "futures-executor" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" + +[[package]] +name = "futures-macro" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-sink" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" + +[[package]] +name = "futures-task" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" + +[[package]] +name = "futures-util" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "pin-utils", + "slab", +] + [[package]] name = "gimli" version = "0.31.1" @@ -387,6 +477,12 @@ version = "0.2.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bda66fc9667c18cb2758a2ac84d1167245054bcf85d5d1aaa6923f45801bdd02" +[[package]] +name = "pin-utils" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" + [[package]] name = "proc-macro2" version = "1.0.87" @@ -464,6 +560,15 @@ dependencies = [ "libc", ] +[[package]] +name = "slab" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f92a496fb766b417c996b9c5e57daf2f7ad3b0bebe1ccfca4856390e3d3bb67" +dependencies = [ + "autocfg", +] + [[package]] name = "smallvec" version = "1.13.2" diff --git a/Cargo.toml b/Cargo.toml index bc75162..9d93cac 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,6 +26,7 @@ bytes = "1.4.0" clap = { version = "4.3.3", features = ["derive"] } ctor = "0.2.2" env_logger = "0.10.0" +futures = "0.3.31" log = "0.4.18" tokio = { version = "1.28.2", features = ["net", "rt", "full"] } diff --git a/src/main.rs b/src/main.rs index 4235941..f21c1da 100644 --- a/src/main.rs +++ b/src/main.rs @@ -11,10 +11,13 @@ use byteorder::{LittleEndian, ReadBytesExt}; use bytes::BytesMut; use clap::Parser; use env_logger::Env; +use futures::future::select_all; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::select; use tokio::sync::{broadcast, mpsc, RwLock}; +use futures::stream::FuturesUnordered; +use tokio::task::JoinHandle; #[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Hash)] pub struct AmsNetId(pub [u8; 6]); @@ -233,7 +236,7 @@ struct Args { pub listen_addr: SocketAddr, /// PLC address, e.g. 172.18.0.10:48898 - pub plc_addr: SocketAddr, + pub plc_addr: Vec, } async fn read_ams_packet(reader: &mut T, buf: &mut BytesMut, packet_size: usize) -> Result @@ -358,9 +361,9 @@ where result } -async fn connect_plc(args: Arc, table: Table) -> Result<()> { +async fn connect_plc(args: Arc, table: Table, plc_addr: SocketAddr) -> Result<()> { // detect plc info - let plc_info = ads::udp::get_info((&args.plc_addr.ip().to_string(), ads::UDP_PORT))?; + let plc_info = ads::udp::get_info((&plc_addr.ip().to_string(), ads::UDP_PORT))?; log::info!("plc net_id={}", plc_info.netid); log::info!("plc hostname={}", plc_info.hostname); log::info!("plc twincat_version={:?}", plc_info.twincat_version); @@ -368,14 +371,14 @@ async fn connect_plc(args: Arc, table: Table) -> Result<()> { log::info!("plc fingerprint={}", plc_info.fingerprint); // connect plc backend - log::info!("connecting plc {}...", args.plc_addr); - let mut plc_client = TcpStream::connect(args.plc_addr).await?; + log::info!("connecting plc {}...", plc_addr); + let mut plc_client = TcpStream::connect(plc_addr).await?; let plc_addr = plc_client.peer_addr()?; let host_addr = plc_client.local_addr()?; // add extra route if let Some(ams_net_id) = args.route_ams_net_id { - let route_host = args.route_host.clone().unwrap_or(host_addr.ip().to_string()); + let route_host = args.route_host.clone().unwrap_or(plc_addr.ip().to_string()); log::info!("add route {} host {} to plc", ams_net_id, route_host); ads::udp::add_route( (&plc_addr.ip().to_string(), ads::UDP_PORT), @@ -389,7 +392,7 @@ async fn connect_plc(args: Arc, table: Table) -> Result<()> { // reconnect plc plc_client.shutdown().await?; - plc_client = TcpStream::connect(args.plc_addr).await?; + plc_client = TcpStream::connect(plc_addr).await?; } // prepare queues @@ -476,12 +479,29 @@ async fn main() -> Result<()> { // global proxy tables and stop event let proxy_table: Table = Arc::new(RwLock::new(HashMap::new())); - let plc_task = connect_plc(args.clone(), proxy_table.clone()); - let accept_task = accept_client(args.clone(), proxy_table.clone()); - - let e = select! { - r = tokio::spawn(plc_task) => r, - r = tokio::spawn(accept_task) => r, + let plc_tasks: Vec>> = args + .plc_addr + .iter() + .cloned() + .map(|addr| { + let args = args.clone(); + let proxy_table = proxy_table.clone(); + tokio::spawn(async move { + connect_plc(args, proxy_table, addr).await + }) + }) + .collect(); + + // spawn the accept_client task + let accept_task = tokio::spawn(accept_client(args.clone(), proxy_table.clone())); + + // run select! to wait for either task group + let e = tokio::select! { + r = accept_task => r, + r = async { + let (res, _idx, _remaining) = select_all(plc_tasks).await; + res + } => r, }; log::error!("ads-proxy error: {:?}", e); log::warn!("ads-proxy stopped"); From 20fdce58d73cac9611b4329319b40c5cb6dbd154 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?On=C3=A9sime=20BERNARD?= Date: Mon, 7 Apr 2025 14:53:10 +0000 Subject: [PATCH 2/2] Fix host_addr plc_addr mixup --- src/main.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/main.rs b/src/main.rs index f21c1da..310f0b5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,7 +16,6 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::select; use tokio::sync::{broadcast, mpsc, RwLock}; -use futures::stream::FuturesUnordered; use tokio::task::JoinHandle; #[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Hash)] @@ -378,7 +377,7 @@ async fn connect_plc(args: Arc, table: Table, plc_addr: SocketAddr) -> Res // add extra route if let Some(ams_net_id) = args.route_ams_net_id { - let route_host = args.route_host.clone().unwrap_or(plc_addr.ip().to_string()); + let route_host = args.route_host.clone().unwrap_or(host_addr.ip().to_string()); log::info!("add route {} host {} to plc", ams_net_id, route_host); ads::udp::add_route( (&plc_addr.ip().to_string(), ads::UDP_PORT), @@ -492,10 +491,8 @@ async fn main() -> Result<()> { }) .collect(); - // spawn the accept_client task let accept_task = tokio::spawn(accept_client(args.clone(), proxy_table.clone())); - // run select! to wait for either task group let e = tokio::select! { r = accept_task => r, r = async {