diff --git a/Cargo.lock b/Cargo.lock index f25b64e58..5e5a11480 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -852,7 +852,7 @@ dependencies = [ "getrandom 0.2.16", "getrandom 0.3.3", "hex", - "indexmap 2.11.1", + "indexmap 2.11.3", "js-sys", "once_cell", "rand 0.9.2", @@ -1092,7 +1092,7 @@ name = "config" version = "1.19.5" dependencies = [ "colored", - "indexmap 2.11.1", + "indexmap 2.11.3", "regex", "serde", "serde_json", @@ -1751,9 +1751,9 @@ dependencies = [ [[package]] name = "fs-err" -version = "3.1.1" +version = "3.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88d7be93788013f265201256d58f04936a8079ad5dc898743aa20525f503b683" +checksum = "44f150ffc8782f35521cec2b23727707cb4045706ba3c854e86bef66b3a8cdbd" dependencies = [ "autocfg", "tokio", @@ -1894,7 +1894,7 @@ dependencies = [ "js-sys", "libc", "r-efi", - "wasi 0.14.5+wasi-0.2.4", + "wasi 0.14.7+wasi-0.2.4", "wasm-bindgen", ] @@ -1947,7 +1947,7 @@ dependencies = [ "futures-sink", "futures-util", "http 0.2.12", - "indexmap 2.11.1", + "indexmap 2.11.3", "slab", "tokio", "tokio-util", @@ -1966,7 +1966,7 @@ dependencies = [ "futures-core", "futures-sink", "http 1.3.1", - "indexmap 2.11.1", + "indexmap 2.11.3", "slab", "tokio", "tokio-util", @@ -2284,9 +2284,9 @@ dependencies = [ [[package]] name = "hyper-util" -version = "0.1.16" +version = "0.1.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8d9b05277c7e8da2c93a568989bb6207bef0112e8d17df7a6eda4a3cf143bc5e" +checksum = "3c6995591a8f1380fcb4ba966a252a4b29188d51d2b89e3a252f5305be65aea8" dependencies = [ "base64 0.22.1", "bytes", @@ -2473,13 +2473,14 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.11.1" +version = "2.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "206a8042aec68fa4a62e8d3f7aa4ceb508177d9324faf261e1959e495b7a1921" +checksum = "92119844f513ffa41556430369ab02c295a3578af21cf945caa3e9e0c2481ac3" dependencies = [ "equivalent", "hashbrown 0.15.5", "serde", + "serde_core", ] [[package]] @@ -2611,9 +2612,9 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.78" +version = "0.3.79" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c0b063578492ceec17683ef2f8c5e89121fbd0b172cbc280635ab7567db2738" +checksum = "6247da8b8658ad4e73a186e747fcc5fc2a29f979d6fe6269127fdb5fd08298d0" dependencies = [ "once_cell", "wasm-bindgen", @@ -2673,7 +2674,7 @@ dependencies = [ "derive_variants", "envy", "futures", - "indexmap 2.11.1", + "indexmap 2.11.3", "ipnetwork", "mongo_indexed", "partial_derive2", @@ -2709,6 +2710,7 @@ dependencies = [ "axum-server", "base64 0.22.1", "bcrypt", + "bytes", "cache", "chrono", "chrono-tz", @@ -2728,7 +2730,7 @@ dependencies = [ "git", "hex", "hmac", - "indexmap 2.11.1", + "indexmap 2.11.3", "interpolate", "jsonwebtoken", "komodo_client", @@ -2758,6 +2760,7 @@ dependencies = [ "toml_pretty", "tower-http", "tracing", + "transport", "typeshare", "urlencoding", "uuid", @@ -2810,6 +2813,7 @@ dependencies = [ "tokio-util", "tracing", "transport", + "uuid", ] [[package]] @@ -4315,7 +4319,7 @@ dependencies = [ "once_cell", "ring", "rustls-pki-types", - "rustls-webpki 0.103.5", + "rustls-webpki 0.103.6", "subtle", "zeroize", ] @@ -4384,9 +4388,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.103.5" +version = "0.103.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5a37813727b78798e53c2bec3f5e8fe12a6d6f8389bf9ca7802add4c9905ad8" +checksum = "8572f3c2cb9934231157b45499fc41e1f58c589fdfb81a844ba873265e80f8eb" dependencies = [ "aws-lc-rs", "ring", @@ -4541,9 +4545,9 @@ checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" [[package]] name = "serde" -version = "1.0.223" +version = "1.0.225" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a505d71960adde88e293da5cb5eda57093379f64e61cf77bf0e6a63af07a7bac" +checksum = "fd6c24dee235d0da097043389623fb913daddf92c76e9f5a1db88607a0bcbd1d" dependencies = [ "serde_core", "serde_derive", @@ -4561,9 +4565,9 @@ dependencies = [ [[package]] name = "serde_bytes" -version = "0.11.18" +version = "0.11.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe07b5d88710e3b807c16a06ccbc9dfecd5fff6a4d2745c59e3e26774f10de6a" +checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8" dependencies = [ "serde", "serde_core", @@ -4571,18 +4575,18 @@ dependencies = [ [[package]] name = "serde_core" -version = "1.0.223" +version = "1.0.225" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "20f57cbd357666aa7b3ac84a90b4ea328f1d4ddb6772b430caa5d9e1309bb9e9" +checksum = "659356f9a0cb1e529b24c01e43ad2bdf520ec4ceaf83047b83ddcc2251f96383" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.223" +version = "1.0.225" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d428d07faf17e306e699ec1e91996e5a165ba5d6bce5b5155173e91a8a01a56" +checksum = "0ea936adf78b1f766949a4977b91d2f5595825bd6ec079aa9543ad2685fc4516" dependencies = [ "proc-macro2", "quote", @@ -4606,7 +4610,7 @@ version = "1.0.145" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "402a6f66d8c709116cf22f558eab210f5a50187f702eb4d7e5ef38d9a7f1c79c" dependencies = [ - "indexmap 2.11.1", + "indexmap 2.11.3", "itoa", "memchr", "ryu", @@ -4616,9 +4620,9 @@ dependencies = [ [[package]] name = "serde_path_to_error" -version = "0.1.19" +version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a30a8abed938137c7183c173848e3c9b3517f5e038226849a4ecc9b21a4b4e2a" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" dependencies = [ "itoa", "serde", @@ -4687,7 +4691,7 @@ dependencies = [ "chrono", "hex", "indexmap 1.9.3", - "indexmap 2.11.1", + "indexmap 2.11.3", "schemars 0.9.0", "schemars 1.0.4", "serde", @@ -4715,7 +4719,7 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7b4db627b98b36d4203a7b458cf3573730f2bb591b28871d916dfa9efabfd41f" dependencies = [ - "indexmap 2.11.1", + "indexmap 2.11.3", "itoa", "ryu", "serde", @@ -5284,7 +5288,7 @@ version = "0.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ae2a4cf385da23d1d53bc15cdfa5c2109e93d8d362393c801e87da2f72f0e201" dependencies = [ - "indexmap 2.11.1", + "indexmap 2.11.3", "serde_core", "serde_spanned", "toml_datetime", @@ -5365,7 +5369,7 @@ checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9" dependencies = [ "futures-core", "futures-util", - "indexmap 2.11.1", + "indexmap 2.11.3", "pin-project-lite", "slab", "sync_wrapper", @@ -5512,7 +5516,6 @@ dependencies = [ "anyhow", "axum", "bytes", - "cache", "futures-util", "rustls 0.23.31", "serror", @@ -5760,27 +5763,27 @@ checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" [[package]] name = "wasi" -version = "0.14.5+wasi-0.2.4" +version = "0.14.7+wasi-0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4494f6290a82f5fe584817a676a34b9d6763e8d9d18204009fb31dceca98fd4" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" dependencies = [ "wasip2", ] [[package]] name = "wasip2" -version = "1.0.0+wasi-0.2.4" +version = "1.0.1+wasi-0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03fa2761397e5bd52002cd7e73110c71af2109aca4e521a9f40473fe685b0a24" +checksum = "0562428422c63773dad2c345a1882263bbf4d65cf3f42e90921f787ef5ad58e7" dependencies = [ "wit-bindgen", ] [[package]] name = "wasm-bindgen" -version = "0.2.101" +version = "0.2.102" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e14915cadd45b529bb8d1f343c4ed0ac1de926144b746e2710f9cd05df6603b" +checksum = "4ad224d2776649cfb4f4471124f8176e54c1cca67a88108e30a0cd98b90e7ad3" dependencies = [ "cfg-if", "once_cell", @@ -5791,9 +5794,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-backend" -version = "0.2.101" +version = "0.2.102" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e28d1ba982ca7923fd01448d5c30c6864d0a14109560296a162f80f305fb93bb" +checksum = "3a1364104bdcd3c03f22b16a3b1c9620891469f5e9f09bc38b2db121e593e732" dependencies = [ "bumpalo", "log", @@ -5805,9 +5808,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.51" +version = "0.4.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ca85039a9b469b38336411d6d6ced91f3fc87109a2a27b0c197663f5144dffe" +checksum = "9c0a08ecf5d99d5604a6666a70b3cde6ab7cc6142f5e641a8ef48fc744ce8854" dependencies = [ "cfg-if", "js-sys", @@ -5818,9 +5821,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.101" +version = "0.2.102" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c3d463ae3eff775b0c45df9da45d68837702ac35af998361e2c84e7c5ec1b0d" +checksum = "0d7ab4ca3e367bb1ed84ddbd83cc6e41e115f8337ed047239578210214e36c76" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -5828,9 +5831,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.101" +version = "0.2.102" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7bb4ce89b08211f923caf51d527662b75bdc9c9c7aab40f86dcb9fb85ac552aa" +checksum = "4a518014843a19e2dbbd0ed5dfb6b99b23fb886b14e6192a00803a3e14c552b0" dependencies = [ "proc-macro2", "quote", @@ -5841,9 +5844,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.101" +version = "0.2.102" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f143854a3b13752c6950862c906306adb27c7e839f7414cec8fea35beab624c1" +checksum = "255eb0aa4cc2eea3662a00c2bbd66e93911b7361d5e0fcd62385acfd7e15dcee" dependencies = [ "unicode-ident", ] @@ -5878,9 +5881,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.78" +version = "0.3.79" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77e4b637749ff0d92b8fad63aa1f7cff3cbe125fd49c175cd6345e7272638b12" +checksum = "50462a022f46851b81d5441d1a6f5bac0b21a1d72d64bd4906fbdd4bf7230ec7" dependencies = [ "js-sys", "wasm-bindgen", @@ -6368,9 +6371,9 @@ dependencies = [ [[package]] name = "wit-bindgen" -version = "0.45.1" +version = "0.46.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c573471f125075647d03df72e026074b7203790d41351cd6edc96f46bcccd36" +checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" [[package]] name = "writeable" diff --git a/bin/core/Cargo.toml b/bin/core/Cargo.toml index 2f960ccf0..bdc42dc73 100644 --- a/bin/core/Cargo.toml +++ b/bin/core/Cargo.toml @@ -20,6 +20,7 @@ periphery_client.workspace = true environment_file.workspace = true interpolate.workspace = true formatting.workspace = true +transport.workspace = true database.workspace = true response.workspace = true command.workspace = true @@ -70,6 +71,7 @@ chrono.workspace = true bcrypt.workspace = true base64.workspace = true rustls.workspace = true +bytes.workspace = true tokio.workspace = true serde.workspace = true regex.workspace = true diff --git a/bin/core/src/monitor/mod.rs b/bin/core/src/monitor/mod.rs index cb9251102..2df06a444 100644 --- a/bin/core/src/monitor/mod.rs +++ b/bin/core/src/monitor/mod.rs @@ -112,7 +112,8 @@ async fn refresh_server_cache(ts: i64) { return; } }; - periphery_client::connection::manage_connections(&servers).await; + periphery_client::connection::manage_outbound_connections(&servers) + .await; let futures = servers.into_iter().map(|server| async move { update_cache_for_server(&server, false).await; }); diff --git a/bin/core/src/ws/container.rs b/bin/core/src/ws/container.rs index 38ab5524a..8d2bd1389 100644 --- a/bin/core/src/ws/container.rs +++ b/bin/core/src/ws/container.rs @@ -21,7 +21,7 @@ pub async fn terminal( ) -> impl IntoResponse { ws.on_upgrade(|socket| async move { let Some((mut client_socket, user)) = - super::ws_login(socket).await + super::user_ws_login(socket).await else { return; }; diff --git a/bin/core/src/ws/deployment.rs b/bin/core/src/ws/deployment.rs index 11462593e..5577cacc3 100644 --- a/bin/core/src/ws/deployment.rs +++ b/bin/core/src/ws/deployment.rs @@ -22,7 +22,7 @@ pub async fn terminal( ) -> impl IntoResponse { ws.on_upgrade(|socket| async move { let Some((mut client_socket, user)) = - super::ws_login(socket).await + super::user_ws_login(socket).await else { return; }; diff --git a/bin/core/src/ws/mod.rs b/bin/core/src/ws/mod.rs index 3458b8612..a10155613 100644 --- a/bin/core/src/ws/mod.rs +++ b/bin/core/src/ws/mod.rs @@ -8,19 +8,29 @@ use axum::{ extract::ws::{CloseFrame, Message, Utf8Bytes, WebSocket}, routing::get, }; +use bytes::Bytes; use futures::{SinkExt, StreamExt}; use komodo_client::{ entities::{server::Server, user::User}, ws::WsLoginMessage, }; -use tokio::net::TcpStream; +use tokio::{ + net::TcpStream, + sync::mpsc::{Receiver, Sender}, +}; use tokio_tungstenite::{ MaybeTlsStream, WebSocketStream, tungstenite, }; use tokio_util::sync::CancellationToken; +use transport::{ + MessageState, + bytes::{data_from_transport_bytes, to_transport_bytes}, +}; +use uuid::Uuid; mod container; mod deployment; +mod periphery; mod stack; mod terminal; mod update; @@ -35,7 +45,7 @@ pub fn router() -> Router { } #[instrument(level = "debug")] -async fn ws_login( +async fn user_ws_login( mut socket: WebSocket, ) -> Option<(WebSocket, User)> { let login_msg = match socket.recv().await { @@ -249,6 +259,107 @@ async fn core_periphery_forward_ws( tokio::join!(core_to_periphery, periphery_to_core); } +async fn core_periphery_forward_ws_channel( + client_socket: axum::extract::ws::WebSocket, + periphery_connection_id: Uuid, + periphery_send: Sender, + mut periphery_receive: Receiver, +) { + let (mut core_send, mut core_receive) = client_socket.split(); + let cancel = CancellationToken::new(); + + trace!("starting ws exchange"); + + let core_to_periphery = async { + loop { + let res = tokio::select! { + res = core_receive.next() => res, + _ = cancel.cancelled() => { + trace!("core to periphery read: cancelled from inside"); + break; + } + }; + match res { + Some(Ok(Message::Binary(data))) => { + if let Err(e) = periphery_send + .send(to_transport_bytes( + data.into(), + periphery_connection_id, + MessageState::Terminal, + )) + .await + { + debug!("Failed to send terminal message | {e:?}",); + cancel.cancel(); + break; + }; + } + Some(Ok(Message::Text(data))) => { + let data: Bytes = data.into(); + if let Err(e) = periphery_send + .send(to_transport_bytes( + data.into(), + periphery_connection_id, + MessageState::Terminal, + )) + .await + { + debug!("Failed to send terminal message | {e:?}",); + cancel.cancel(); + break; + }; + } + // TODO: Disconnect from periphery when client disconnects + Some(Ok(Message::Close(_frame))) => { + cancel.cancel(); + break; + } + // Ignore + Some(Ok(_)) => {} + Some(Err(_e)) => { + cancel.cancel(); + break; + } + None => { + cancel.cancel(); + break; + } + } + } + }; + + let periphery_to_core = async { + loop { + let res = tokio::select! { + res = periphery_receive.recv() => res.map(data_from_transport_bytes), + _ = cancel.cancelled() => { + trace!("periphery to core read: cancelled from inside"); + break; + } + }; + match res { + Some(Ok(bytes)) => { + if let Err(e) = core_send.send(Message::Binary(bytes)).await + { + debug!("{e:?}"); + cancel.cancel(); + break; + }; + } + // No data, ignore + Some(Err(_e)) => {} + None => { + let _ = core_send.send(Message::text("STREAM EOF")).await; + cancel.cancel(); + break; + } + } + } + }; + + tokio::join!(core_to_periphery, periphery_to_core); +} + fn axum_to_tungstenite(msg: Message) -> tungstenite::Message { match msg { Message::Text(text) => tungstenite::Message::Text( diff --git a/bin/core/src/ws/periphery.rs b/bin/core/src/ws/periphery.rs new file mode 100644 index 000000000..e69de29bb diff --git a/bin/core/src/ws/stack.rs b/bin/core/src/ws/stack.rs index c260fb2d3..9e7008cd5 100644 --- a/bin/core/src/ws/stack.rs +++ b/bin/core/src/ws/stack.rs @@ -26,7 +26,7 @@ pub async fn terminal( ) -> impl IntoResponse { ws.on_upgrade(|socket| async move { let Some((mut client_socket, user)) = - super::ws_login(socket).await + super::user_ws_login(socket).await else { return; }; diff --git a/bin/core/src/ws/terminal.rs b/bin/core/src/ws/terminal.rs index 3ec6b6aa1..8270f3155 100644 --- a/bin/core/src/ws/terminal.rs +++ b/bin/core/src/ws/terminal.rs @@ -10,7 +10,7 @@ use komodo_client::{ use crate::{ helpers::periphery_client, permission::get_check_permissions, - ws::core_periphery_forward_ws, + ws::core_periphery_forward_ws_channel, }; #[instrument(name = "ConnectTerminal", skip(ws))] @@ -22,7 +22,7 @@ pub async fn handler( ) -> impl IntoResponse { ws.on_upgrade(|socket| async move { let Some((mut client_socket, user)) = - super::ws_login(socket).await + super::user_ws_login(socket).await else { return; }; @@ -59,21 +59,30 @@ pub async fn handler( trace!("connecting to periphery terminal websocket"); - let periphery_socket = - match periphery.connect_terminal(terminal).await { - Ok(ws) => ws, - Err(e) => { - debug!("Failed connect to periphery terminal | {e:#}"); - let _ = client_socket - .send(Message::text(format!("ERROR: {e:#}"))) - .await; - let _ = client_socket.close().await; - return; - } - }; + let ( + periphery_connection_id, + periphery_sender, + periphery_receiver, + ) = match periphery.connect_terminal(terminal).await { + Ok(ws) => ws, + Err(e) => { + debug!("Failed connect to periphery terminal | {e:#}"); + let _ = client_socket + .send(Message::text(format!("ERROR: {e:#}"))) + .await; + let _ = client_socket.close().await; + return; + } + }; trace!("connected to periphery terminal websocket"); - core_periphery_forward_ws(client_socket, periphery_socket).await + core_periphery_forward_ws_channel( + client_socket, + periphery_connection_id, + periphery_sender, + periphery_receiver, + ) + .await }) } diff --git a/bin/core/src/ws/update.rs b/bin/core/src/ws/update.rs index 705172044..a77a39e06 100644 --- a/bin/core/src/ws/update.rs +++ b/bin/core/src/ws/update.rs @@ -23,7 +23,7 @@ pub async fn handler(ws: WebSocketUpgrade) -> impl IntoResponse { // handle http -> ws updgrade ws.on_upgrade(|socket| async move { - let Some((socket, user)) = super::ws_login(socket).await else { + let Some((socket, user)) = super::user_ws_login(socket).await else { return }; diff --git a/bin/periphery/Cargo.toml b/bin/periphery/Cargo.toml index 3f55f5cda..112d23bc9 100644 --- a/bin/periphery/Cargo.toml +++ b/bin/periphery/Cargo.toml @@ -36,11 +36,12 @@ resolver_api.workspace = true run_command.workspace = true # external pin-project-lite.workspace = true +serde_yaml_ng.workspace = true tokio-stream.workspace = true portable-pty.workspace = true +shell-escape.workspace = true axum-server.workspace = true serde_json.workspace = true -serde_yaml_ng.workspace = true tokio-util.workspace = true arc-swap.workspace = true colored.workspace = true @@ -58,4 +59,4 @@ axum.workspace = true clap.workspace = true envy.workspace = true rand.workspace = true -shell-escape.workspace = true \ No newline at end of file +uuid.workspace = true \ No newline at end of file diff --git a/bin/periphery/src/api/mod.rs b/bin/periphery/src/api/mod.rs index 2761490e7..9d4a7cedf 100644 --- a/bin/periphery/src/api/mod.rs +++ b/bin/periphery/src/api/mod.rs @@ -142,6 +142,8 @@ pub enum PeripheryRequest { // Terminal ListTerminals(ListTerminals), CreateTerminal(CreateTerminal), + ConnectTerminal(ConnectTerminal), + DisconnectTerminal(DisconnectTerminal), DeleteTerminal(DeleteTerminal), DeleteAllTerminals(DeleteAllTerminals), CreateTerminalAuthToken(CreateTerminalAuthToken), diff --git a/bin/periphery/src/api/terminal.rs b/bin/periphery/src/api/terminal.rs index dd15803ea..5f7c061c9 100644 --- a/bin/periphery/src/api/terminal.rs +++ b/bin/periphery/src/api/terminal.rs @@ -16,9 +16,16 @@ use komodo_client::{ use periphery_client::api::terminal::*; use resolver_api::Resolve; use serror::{AddStatusCodeError, Json}; +use tokio::sync::mpsc::channel; use tokio_util::sync::CancellationToken; +use transport::{MessageState, bytes::to_transport_bytes}; +use uuid::Uuid; -use crate::{config::periphery_config, terminal::*}; +use crate::{ + config::periphery_config, + connection::{terminal_channels, ws_sender}, + terminal::*, +}; impl Resolve for ListTerminals { #[instrument(name = "ListTerminals", level = "debug")] @@ -63,6 +70,189 @@ impl Resolve for DeleteAllTerminals { } } +impl Resolve for ConnectTerminal { + #[instrument(name = "ConnectTerminal", level = "debug")] + async fn resolve(self, _: &super::Args) -> serror::Result { + let id = Uuid::new_v4(); + let ws_sender = ws_sender(); + let (sender, mut ws_receiver) = channel(1000); + let cancel = CancellationToken::new(); + + terminal_channels() + .insert(id, (sender, cancel.clone())) + .await; + + clean_up_terminals().await; + let terminal = get_terminal(&self.terminal).await?; + + tokio::spawn(async move { + let init_res = async { + let (a, b) = terminal.history.bytes_parts(); + if !a.is_empty() { + ws_sender + .send(to_transport_bytes( + a.into(), + id, + MessageState::Terminal, + )) + .await + .context("Failed to send history part a")?; + } + if !b.is_empty() { + ws_sender + .send(to_transport_bytes( + b.into(), + id, + MessageState::Terminal, + )) + .await + .context("Failed to send history part b")?; + } + anyhow::Ok(()) + } + .await; + + if let Err(e) = init_res { + // TODO: Handle error + warn!("Failed to init terminal | {e:#}"); + return; + } + + let ws_read = async { + loop { + let res = tokio::select! { + res = ws_receiver.recv() => res, + _ = terminal.cancel.cancelled() => { + trace!("ws read: cancelled from outside"); + break + }, + _ = cancel.cancelled() => { + trace!("ws read: cancelled from inside"); + break; + } + }; + match res { + Some(bytes) if bytes.first() == Some(&0x00) => { + // println!("Got ws read bytes - for stdin"); + if let Err(e) = terminal + .stdin + .send(StdinMsg::Bytes(Bytes::copy_from_slice( + &bytes[1..], + ))) + .await + { + debug!("WS -> PTY channel send error: {e:}"); + terminal.cancel(); + break; + }; + } + Some(bytes) if bytes.first() == Some(&0xFF) => { + // println!("Got ws read bytes - for resize"); + if let Ok(dimensions) = serde_json::from_slice::< + ResizeDimensions, + >(&bytes[1..]) + && let Err(e) = terminal + .stdin + .send(StdinMsg::Resize(dimensions)) + .await + { + debug!("WS -> PTY channel send error: {e:}"); + terminal.cancel(); + break; + } + } + Some(bytes) => { + trace!("Got ws read text"); + if let Err(e) = + terminal.stdin.send(StdinMsg::Bytes(bytes)).await + { + debug!("WS -> PTY channel send error: {e:?}"); + terminal.cancel(); + break; + }; + } + None => { + debug!("Got ws read none"); + cancel.cancel(); + break; + } + } + } + }; + + let ws_write = async { + let mut stdout = terminal.stdout.resubscribe(); + loop { + let res = tokio::select! { + res = stdout.recv() => res.context("Failed to get message over stdout receiver"), + _ = terminal.cancel.cancelled() => { + trace!("ws write: cancelled from outside"); + // let _ = ws_sender.send("PTY KILLED")).await; + // if let Err(e) = ws_write.close().await { + // debug!("Failed to close ws: {e:?}"); + // }; + break + }, + _ = cancel.cancelled() => { + // let _ = ws_write.send(Message::Text(Utf8Bytes::from_static("WS KILLED"))).await; + // if let Err(e) = ws_write.close().await { + // debug!("Failed to close ws: {e:?}"); + // }; + break + } + }; + match res { + Ok(bytes) => { + if let Err(e) = ws_sender + .send(to_transport_bytes( + bytes.into(), + id, + MessageState::Terminal, + )) + .await + { + debug!("Failed to send to WS: {e:?}"); + cancel.cancel(); + break; + } + } + Err(e) => { + debug!("PTY -> WS channel read error: {e:?}"); + let _ = ws_sender + .send(to_transport_bytes( + format!("ERROR: {e:#}").into(), + id, + MessageState::Terminal, + )) + .await; + terminal.cancel(); + break; + } + } + } + }; + + tokio::join!(ws_read, ws_write); + + clean_up_terminals().await; + }); + + Ok(id) + } +} + +impl Resolve for DisconnectTerminal { + #[instrument(name = "DisconnectTerminal", level = "debug")] + async fn resolve(self, _: &super::Args) -> serror::Result { + if let Some((_, cancel)) = + terminal_channels().remove(&self.uuid).await + { + cancel.cancel(); + } + Ok(NoData {}) + } +} + impl Resolve for CreateTerminalAuthToken { #[instrument(name = "CreateTerminalAuthToken", level = "debug")] async fn resolve( @@ -129,6 +319,8 @@ pub async fn connect_container_exec( .await } +// async fn hande_t + async fn handle_terminal_websocket( ConnectTerminalQuery { token, terminal }: ConnectTerminalQuery, ws: WebSocketUpgrade, diff --git a/bin/periphery/src/connection.rs b/bin/periphery/src/connection.rs index 1345e4528..5c756f4a9 100644 --- a/bin/periphery/src/connection.rs +++ b/bin/periphery/src/connection.rs @@ -2,29 +2,35 @@ use std::{sync::OnceLock, time::Duration}; use axum::{extract::WebSocketUpgrade, response::Response}; use bytes::Bytes; +use cache::CloneCache; use resolver_api::Resolve; use response::JsonBytes; use serror::serialize_error_bytes; use tokio::sync::{Mutex, mpsc::Sender}; +use tokio_util::sync::CancellationToken; use transport::{ - MessageState, - bytes::{from_transport_bytes, to_transport_bytes}, + MessageState, TransportHandler, + bytes::{ + data_from_transport_bytes, id_state_from_transport_bytes, + to_transport_bytes, + }, channel::{BufferedReceiver, buffered_channel}, }; +use uuid::Uuid; use crate::api::PeripheryRequest; -static RESPONSE_SENDER: OnceLock> = OnceLock::new(); -fn response_sender() -> &'static Sender { - RESPONSE_SENDER +static WS_SENDER: OnceLock> = OnceLock::new(); +pub fn ws_sender() -> &'static Sender { + WS_SENDER .get() .expect("response_sender accessed before initialized") } -static RESPONSE_RECEIVER: OnceLock>> = +static WS_RECEIVER: OnceLock>> = OnceLock::new(); -fn response_receiver() -> &'static Mutex> { - RESPONSE_RECEIVER +fn ws_receiver() -> &'static Mutex> { + WS_RECEIVER .get() .expect("response_receiver accessed before initialized") } @@ -34,10 +40,10 @@ const RESPONSE_BUFFER_MAX_LEN: usize = 1_024; /// Must call in startup sequence pub fn init_response_channel() { let (sender, receiver) = buffered_channel(RESPONSE_BUFFER_MAX_LEN); - RESPONSE_SENDER + WS_SENDER .set(sender) .expect("response_sender initialized more than once"); - RESPONSE_RECEIVER + WS_RECEIVER .set(Mutex::new(receiver)) .expect("response_receiver initialized more than once"); } @@ -47,26 +53,45 @@ pub async fn inbound_connection( ) -> serror::Result { transport::server::inbound_connection( ws, - handle_msg, - response_receiver(), + PeripheryTransportHandler, + ws_receiver(), ) } -// Creates an execution thread to process the request. -fn handle_msg(msg: Bytes) { - tokio::spawn(async move { - let (id, request) = match from_transport_bytes(&msg) { - Ok((id, _, Some(request))) => (id, request), - Ok(_) => { - // No data, ignore - return; - } +struct PeripheryTransportHandler; + +impl TransportHandler for PeripheryTransportHandler { + async fn handle_incoming_bytes(&self, bytes: Bytes) { + // Maybe wrap all of this on tokio spawn + let (id, state) = match id_state_from_transport_bytes(&bytes) { + Ok(res) => res, Err(e) => { // TODO: handle: warn!("Failed to parse transport bytes | {e:#}"); return; } }; + match state { + MessageState::Request => handle_request(id, bytes), + MessageState::Terminal => { + handle_terminal_message(id, bytes).await + } + // Shouldn't be received by Periphery + MessageState::InProgress => {} + MessageState::Successful => {} + MessageState::Failed => {} + } + } +} + +fn handle_request(id: Uuid, bytes: Bytes) { + tokio::spawn(async move { + let request = match data_from_transport_bytes(bytes) { + Ok(req) if !req.is_empty() => req, + _ => { + return; + } + }; let request = match serde_json::from_slice::(&request) { @@ -93,9 +118,8 @@ fn handle_msg(msg: Bytes) { (MessageState::Failed, serialize_error_bytes(&e.error)) } }; - if let Err(e) = response_sender() - .send(to_transport_bytes(id, state, &data)) - .await + if let Err(e) = + ws_sender().send(to_transport_bytes(data, id, state)).await { error!("Failed to send response over channel | {e:?}"); } @@ -104,8 +128,12 @@ fn handle_msg(msg: Bytes) { let ping_in_progress = async { loop { tokio::time::sleep(Duration::from_secs(5)).await; - if let Err(e) = response_sender() - .send(to_transport_bytes(id, MessageState::InProgress, &[])) + if let Err(e) = ws_sender() + .send(to_transport_bytes( + Vec::new(), + id, + MessageState::InProgress, + )) .await { error!("Failed to ping in progress over channel | {e:?}"); @@ -119,3 +147,25 @@ fn handle_msg(msg: Bytes) { } }); } + +pub type TerminalChannels = + CloneCache, CancellationToken)>; +pub fn terminal_channels() -> &'static TerminalChannels { + static TERMINAL_CHANNELS: OnceLock = + OnceLock::new(); + TERMINAL_CHANNELS.get_or_init(Default::default) +} + +async fn handle_terminal_message(id: Uuid, bytes: Bytes) { + let Some((channel, _)) = terminal_channels().get(&id).await else { + warn!("No terminal channel for {id}"); + return; + }; + let Ok(data) = data_from_transport_bytes(bytes) else { + warn!("Got terminal message with no data for {id}"); + return; + }; + if let Err(e) = channel.send(data).await { + warn!("No receiver for {id} | {e:?}"); + }; +} diff --git a/client/periphery/rs/src/api/terminal.rs b/client/periphery/rs/src/api/terminal.rs index 88d61e0e0..525a235db 100644 --- a/client/periphery/rs/src/api/terminal.rs +++ b/client/periphery/rs/src/api/terminal.rs @@ -4,6 +4,7 @@ use komodo_client::{ }; use resolver_api::Resolve; use serde::{Deserialize, Serialize}; +use uuid::Uuid; #[derive(Serialize, Deserialize, Debug, Clone, Resolve)] #[response(Vec)] @@ -33,6 +34,26 @@ fn default_command() -> String { // +#[derive(Serialize, Deserialize, Debug, Clone, Resolve)] +#[response(Uuid)] +#[error(serror::Error)] +pub struct ConnectTerminal { + /// The name of the terminal to connect to + pub terminal: String, +} + +// + +#[derive(Serialize, Deserialize, Debug, Clone, Resolve)] +#[response(NoData)] +#[error(serror::Error)] +pub struct DisconnectTerminal { + /// The connection of the terminal to disconnect from + pub uuid: Uuid, +} + +// + #[derive(Serialize, Deserialize, Debug, Clone, Resolve)] #[response(NoData)] #[error(serror::Error)] diff --git a/client/periphery/rs/src/connection.rs b/client/periphery/rs/src/connection.rs index 95a41dbf8..04f209971 100644 --- a/client/periphery/rs/src/connection.rs +++ b/client/periphery/rs/src/connection.rs @@ -6,23 +6,17 @@ use std::{ use bytes::Bytes; use cache::CloneCache; use komodo_client::entities::server::Server; -use tokio::sync::mpsc; +use tokio::sync::mpsc::{self, Sender}; use tracing::warn; use transport::{ + TransportHandler, + bytes::id_from_transport_bytes, channel::{BufferedReceiver, buffered_channel}, - client::ClientConnection, + client::{ClientConnection, fix_ws_address}, }; use uuid::Uuid; -/// server id => connection -pub type PeripheryConnections = - CloneCache>; -pub fn periphery_connections() -> &'static PeripheryConnections { - static CONNECTIONS: OnceLock = - OnceLock::new(); - CONNECTIONS.get_or_init(Default::default) -} - +// Server id => Channel sender map pub type ResponseChannels = CloneCache>>>; pub fn periphery_response_channels() -> &'static ResponseChannels { @@ -31,13 +25,50 @@ pub fn periphery_response_channels() -> &'static ResponseChannels { RESPONSE_CHANNELS.get_or_init(Default::default) } +pub struct CoreTransportHandler { + response_channels: Arc>>, +} + +impl CoreTransportHandler { + pub async fn new(server_id: &String) -> CoreTransportHandler { + CoreTransportHandler { + response_channels: periphery_response_channels() + .get_or_insert_default(server_id) + .await, + } + } +} + +impl TransportHandler for CoreTransportHandler { + async fn handle_incoming_bytes(&self, bytes: Bytes) { + let id = match id_from_transport_bytes(&bytes) { + Ok(res) => res, + Err(e) => { + // TODO: handle better + warn!("Failed to read id | {e:#}"); + return; + } + }; + let Some(channel) = self.response_channels.get(&id).await else { + // TODO: handle better + warn!("Failed to send response | No response channel found"); + return; + }; + if let Err(e) = channel.send(bytes).await { + // TODO: handle better + warn!("Failed to send response | Channel failure | {e:#}"); + } + } +} + /// Managed connections to exactly those specified by specs (ServerId -> Address) -pub async fn manage_connections(servers: &[Server]) { +pub async fn manage_outbound_connections(servers: &[Server]) { let periphery_connections = periphery_connections(); let periphery_response_channels = periphery_response_channels(); let specs = servers .iter() + .filter(|s| !s.config.address.is_empty()) .map(|s| (&s.id, &s.config.address)) .collect::>(); @@ -54,14 +85,14 @@ pub async fn manage_connections(servers: &[Server]) { // Apply latest connection specs for (server_id, address) in specs { - let address = fix_address(address); + let address = fix_ws_address(address); match periphery_connections.get(server_id).await { // Existing connection good to go, nothing to do Some(existing) if existing.address == address => {} // All other cases re-spawn connection _ => { if let Err(e) = - spawn_connection(server_id.clone(), address).await + spawn_outbound_connection(server_id.clone(), address).await { warn!( "Failed to spawn new connnection for {server_id} | {e:#}" @@ -72,33 +103,14 @@ pub async fn manage_connections(servers: &[Server]) { } } -/// Fixes server addresses: -/// server.domain => wss://server.domain -/// http://server.domain => ws://server.domain -/// https://server.domain => wss://server.domain -fn fix_address(address: &str) -> String { - if address.starts_with("ws://") || address.starts_with("wss://") { - return address.to_string(); - } - if address.starts_with("http://") { - return address.replace("http://", "ws://"); - } - if address.starts_with("https://") { - return address.replace("https://", "wss://"); - } - format!("wss://{address}") -} - // Assumes address already wss formatted -async fn spawn_connection( +async fn spawn_outbound_connection( server_id: String, address: String, ) -> anyhow::Result<()> { - let response_channels = periphery_response_channels() - .get_or_insert_default(&server_id) - .await; + let transport = CoreTransportHandler::new(&server_id).await; - let (connection, request_receiver) = + let (connection, mut request_receiver) = PeripheryConnection::new(address.clone()); if let Some(existing_connection) = periphery_connections() .insert(server_id, connection.clone()) @@ -107,21 +119,33 @@ async fn spawn_connection( existing_connection.client.cancel(); } - transport::client::spawn_reconnecting_websocket( - address, - connection.client.clone(), - request_receiver, - response_channels, - ); + tokio::spawn(async move { + transport::client::handle_reconnecting_websocket( + &address, + &connection.client, + &transport, + &mut request_receiver, + ) + .await + }); Ok(()) } +/// server id => connection +pub type PeripheryConnections = + CloneCache>; +pub fn periphery_connections() -> &'static PeripheryConnections { + static CONNECTIONS: OnceLock = + OnceLock::new(); + CONNECTIONS.get_or_init(Default::default) +} + #[derive(Debug)] pub struct PeripheryConnection { address: String, - request_sender: mpsc::Sender, - pub client: Arc, + pub request_sender: mpsc::Sender, + pub client: ClientConnection, } impl PeripheryConnection { diff --git a/client/periphery/rs/src/request.rs b/client/periphery/rs/src/request.rs index 6dfde781c..f39d174ed 100644 --- a/client/periphery/rs/src/request.rs +++ b/client/periphery/rs/src/request.rs @@ -4,7 +4,7 @@ use serde::{Serialize, de::DeserializeOwned}; use serde_json::json; use serror::{deserialize_error_bytes, serror_into_anyhow_error}; use tokio::sync::mpsc; -use tracing::warn; +use tracing::{info, warn}; use transport::{ MessageState, bytes::{from_transport_bytes, to_transport_bytes}, @@ -25,11 +25,12 @@ where T::Response: DeserializeOwned, { let connection = - &periphery_connections().get(server_id).await.with_context( + periphery_connections().get(server_id).await.with_context( || format!("No connection found for server {server_id}"), )?; if !connection.connected() { + info!("got server not connected"); if let Some(e) = connection.error().await { return Err(serror_into_anyhow_error(e)); } @@ -57,7 +58,7 @@ where .context("Failed to serialize request to bytes")?; if let Err(e) = connection - .send(to_transport_bytes(id, MessageState::InProgress, &data)) + .send(to_transport_bytes(data, id, MessageState::Request)) .await .context("Failed to send request over channel") { @@ -76,8 +77,8 @@ where )); } }; - let (state, data) = match from_transport_bytes(&bytes) { - Ok((_, state, Some(data))) => (state, data), + let (state, data) = match from_transport_bytes(bytes) { + Ok((_, state, data)) if !data.is_empty() => (state, data), // TODO: Handle no data cases Ok(_) => continue, Err(e) => { @@ -90,16 +91,24 @@ where MessageState::Successful => { // cleanup response_channels.remove(&id).await; - return serde_json::from_slice(data) + return serde_json::from_slice(&data) .context("Failed to parse successful response"); } MessageState::Failed => { // cleanup response_channels.remove(&id).await; - return Err(deserialize_error_bytes(data)); + return Err(deserialize_error_bytes(&data)); } // TODO: bail if this isn't received frequent enough MessageState::InProgress => continue, + // Shouldn't be received by this receiver + other => { + // TODO: delete log + warn!( + "Got other message over over response channel: {other:?}" + ); + continue; + } } } } diff --git a/client/periphery/rs/src/terminal.rs b/client/periphery/rs/src/terminal.rs index e3c2e9de3..c8c22b8db 100644 --- a/client/periphery/rs/src/terminal.rs +++ b/client/periphery/rs/src/terminal.rs @@ -1,10 +1,19 @@ use anyhow::Context; +use bytes::Bytes; use komodo_client::terminal::TerminalStreamResponse; use reqwest::RequestBuilder; -use tokio::net::TcpStream; +use tokio::{ + net::TcpStream, + sync::mpsc::{Receiver, Sender, channel}, +}; use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; +use uuid::Uuid; -use crate::{PeripheryClient, api::terminal::*}; +use crate::{ + PeripheryClient, + api::terminal::*, + connection::{periphery_connections, periphery_response_channels}, +}; impl PeripheryClient { /// Handles ws connect and login. @@ -12,28 +21,28 @@ impl PeripheryClient { pub async fn connect_terminal( &self, terminal: String, - ) -> anyhow::Result>> { + ) -> anyhow::Result<(Uuid, Sender, Receiver)> { tracing::trace!( "request | type: ConnectTerminal | terminal name: {terminal}", ); - let token = self - .request(CreateTerminalAuthToken {}) + let connection = + periphery_connections().get(&self.id).await.with_context( + || format!("No connection found for server {}", self.id), + )?; + + let id = self + .request(ConnectTerminal { terminal }) .await - .context("Failed to create terminal auth token")?; + .context("Failed to create terminal connectionn")?; - let query_str = serde_qs::to_string(&ConnectTerminalQuery { - token: token.token, - terminal, - }) - .context("Failed to serialize query string")?; + let response_channels = periphery_response_channels() + .get_or_insert_default(&self.id) + .await; + let (response_sender, response_receiever) = channel(1000); + response_channels.insert(id, response_sender).await; - let url = format!( - "{}/terminal?{query_str}", - self.address.replacen("http", "ws", 1) - ); - - transport::client::connect_websocket(&url).await + Ok((id, connection.request_sender.clone(), response_receiever)) } /// Executes command on specified terminal, diff --git a/lib/transport/Cargo.toml b/lib/transport/Cargo.toml index b47d31432..1c6c17423 100644 --- a/lib/transport/Cargo.toml +++ b/lib/transport/Cargo.toml @@ -8,8 +8,6 @@ repository.workspace = true homepage.workspace = true [dependencies] -cache.workspace = true -# serror.workspace = true # tokio-tungstenite.workspace = true diff --git a/lib/transport/src/bytes.rs b/lib/transport/src/bytes.rs index 9c18e3f02..cded56504 100644 --- a/lib/transport/src/bytes.rs +++ b/lib/transport/src/bytes.rs @@ -4,65 +4,93 @@ use uuid::Uuid; use crate::MessageState; -/// Serializes channel id + data to byte vec. -/// The first 16 bytes are the Uuid, followed by the json serialized data bytes. +/// Serializes data + channel id + state to byte vec. +/// The last byte is the State, and the 16 before that is the Uuid. pub fn to_transport_bytes( + mut data: Vec, id: Uuid, state: MessageState, - data: &[u8], ) -> Bytes { - // Index 0..15 - let mut res = id.into_bytes().to_vec(); - // Index 16 - res.push(state.as_byte()); - // Index 17..end - res.extend_from_slice(data); - - Bytes::from(res) + data.extend(id.into_bytes()); + data.push(state.as_byte()); + data.into() } /// Deserializes channel id from /// incoming transport bytes. pub fn id_from_transport_bytes(bytes: &[u8]) -> anyhow::Result { - if bytes.len() < 16 { - return Err(anyhow!("Transport bytes too short to include uuid")); + let len = bytes.len(); + if len < 17 { + return Err(anyhow!( + "Transport bytes too short to include uuid + state" + )); } - Uuid::from_slice(&bytes[..16]).context("Invalid Uuid bytes") + Uuid::from_slice(&bytes[(len - 17)..(len - 1)]) + .context("Invalid Uuid bytes") +} + +/// Deserializes channel id from +/// incoming transport bytes. +pub fn id_state_from_transport_bytes( + bytes: &[u8], +) -> anyhow::Result<(Uuid, MessageState)> { + let len = bytes.len(); + if len < 17 { + return Err(anyhow!( + "Transport bytes too short to include uuid + state" + )); + } + let uuid = Uuid::from_slice(&bytes[(len - 17)..(len - 1)]) + .context("Invalid Uuid bytes")?; + let state = MessageState::from_byte(bytes[len - 1]); + Ok((uuid, state)) +} + +/// extracts data from incoming transport bytes, +/// consuming bytes in the process. +pub fn data_from_transport_bytes( + bytes: Bytes, +) -> anyhow::Result { + let len = bytes.len(); + if len < 17 { + return Err(anyhow!( + "Transport bytes too short to include uuid + state + data" + )); + } + let mut res: Vec = bytes.into(); + res.drain((len - 17)..); + Ok(res.into()) } /// Deserializes channel id + data from /// incoming transport bytes. pub fn from_transport_bytes( - bytes: &[u8], -) -> anyhow::Result<(Uuid, MessageState, Option<&[u8]>)> { - if bytes.len() < 17 { - return Err(anyhow!( - "Transport bytes too short to include uuid and state" - )); - } - let (id, state, data) = (&bytes[..16], bytes[16], bytes.get(17..)); - let id = Uuid::from_slice(id).context("Invalid Uuid bytes")?; - let state = MessageState::from_byte(state); - Ok((id, state, data)) + bytes: Bytes, +) -> anyhow::Result<(Uuid, MessageState, Bytes)> { + let (id, state) = id_state_from_transport_bytes(&bytes)?; + let mut res: Vec = bytes.into(); + res.drain((res.len() - 17)..); + Ok((id, state, res.into())) } impl MessageState { - /// - 0 => Successful - /// - 1 => Failed - /// - other => InProgress pub fn from_byte(byte: u8) -> MessageState { match byte { - 0 => MessageState::Successful, - 1 => MessageState::Failed, + 0 => MessageState::Request, + 1 => MessageState::Terminal, + 2 => MessageState::Successful, + 3 => MessageState::Failed, _ => MessageState::InProgress, } } pub fn as_byte(&self) -> u8 { match self { - MessageState::Successful => 0, - MessageState::Failed => 1, - MessageState::InProgress => 2, + MessageState::Request => 0, + MessageState::Terminal => 1, + MessageState::Successful => 2, + MessageState::Failed => 3, + MessageState::InProgress => 4, } } } diff --git a/lib/transport/src/client.rs b/lib/transport/src/client.rs index cbda3f00f..cc5497022 100644 --- a/lib/transport/src/client.rs +++ b/lib/transport/src/client.rs @@ -8,136 +8,128 @@ use std::{ use anyhow::Context; use bytes::Bytes; -use cache::CloneCache; use futures_util::{SinkExt, StreamExt}; use rustls::{ClientConfig, client::danger::ServerCertVerifier}; -use tokio::{ - net::TcpStream, - sync::{RwLock, mpsc::Sender}, -}; +use tokio::{net::TcpStream, sync::RwLock}; use tokio_tungstenite::{ Connector, MaybeTlsStream, WebSocketStream, tungstenite::Message, }; use tokio_util::sync::CancellationToken; use tracing::{info, warn}; -use uuid::Uuid; -use crate::{ - bytes::id_from_transport_bytes, channel::BufferedReceiver, -}; +use crate::{TransportHandler, channel::BufferedReceiver}; -pub fn spawn_reconnecting_websocket( - address: String, - connection: Arc, - mut request_receiver: BufferedReceiver, - response_channels: Arc>>, +/// Fixes server addresses: +/// server.domain => wss://server.domain +/// http://server.domain => ws://server.domain +/// https://server.domain => wss://server.domain +pub fn fix_ws_address(address: &str) -> String { + if address.starts_with("ws://") || address.starts_with("wss://") { + return address.to_string(); + } + if address.starts_with("http://") { + return address.replace("http://", "ws://"); + } + if address.starts_with("https://") { + return address.replace("https://", "wss://"); + } + format!("wss://{address}") +} + +pub async fn handle_reconnecting_websocket< + T: TransportHandler + Send + Sync + 'static, +>( + address: &str, + connection: &ClientConnection, + transport: &T, + write_receiver: &mut BufferedReceiver, ) { - tokio::spawn(async move { - loop { - let socket = match connect_websocket(&address).await { - Ok(socket) => socket, - Err(e) => { - connection.set_error(e).await; - tokio::time::sleep(Duration::from_secs(5)).await; - continue; - } - }; + loop { + let socket = match connect_websocket(address).await { + Ok(socket) => socket, + Err(e) => { + connection.set_error(e).await; + tokio::time::sleep(Duration::from_secs(5)).await; + continue; + } + }; - info!("Connected to {address}"); - connection.connected.store(true, atomic::Ordering::Relaxed); - connection.clear_error().await; + info!("Connected to {address}"); + connection.connected.store(true, atomic::Ordering::Relaxed); + connection.clear_error().await; - let (mut ws_write, mut ws_read) = socket.split(); + let (mut ws_write, mut ws_read) = socket.split(); - let forward_requests = async { - loop { - let next = tokio::select! { - next = request_receiver.recv() => next, - _ = connection.cancel.cancelled() => { - let _ = ws_write.close().await; - break; - } - }; + let forward_writes = async { + loop { + let next = tokio::select! { + next = write_receiver.recv() => next, + _ = connection.cancel.cancelled() => break, + }; - let message = match next { - None => { - info!("Got None over request reciever for {address}"); - break; - } - Some(request) => { - Message::Binary(Bytes::copy_from_slice(request)) - } - }; + let message = match next { + None => { + info!("Got None over request reciever for {address}"); + break; + } + Some(request) => { + Message::Binary(Bytes::copy_from_slice(request)) + } + }; - match ws_write.send(message).await { - Ok(_) => request_receiver.clear_buffer(), - Err(e) => { - warn!("Failed to send request to {address} | {e:#}"); - break; - } + match ws_write.send(message).await { + Ok(_) => write_receiver.clear_buffer(), + Err(e) => { + warn!("Failed to send request to {address} | {e:#}"); + break; } } - }; + } + // Cancel again if not already + let _ = ws_write.close().await; + connection.cancel(); + }; - let read_responses = async { - loop { - let next = tokio::select! { - next = ws_read.next() => next, - _ = connection.cancel.cancelled() => break, - }; + let handle_reads = async { + loop { + let next = tokio::select! { + next = ws_read.next() => next, + _ = connection.cancel.cancelled() => break, + }; - let bytes = match next { - Some(Ok(Message::Binary(bytes))) => bytes, - Some(Ok(Message::Close(frame))) => { - warn!( - "Connection to {address} broken with frame: {frame:?}" - ); - break; - } - Some(Err(e)) => { - warn!( - "Connection to {address} broken with error: {e:?}" - ); - break; - } - None => { - warn!("Connection to {address} closed"); - break; - } - // Can ignore other message types - Some(Ok(_)) => { - continue; - } - }; - - let id = match id_from_transport_bytes(&bytes) { - Ok(res) => res, - Err(e) => { - warn!("Failed to read id from {address} | {e:#}"); - continue; - } - }; - - let Some(channel) = response_channels.get(&id).await else { + match next { + Some(Ok(Message::Binary(bytes))) => { + transport.handle_incoming_bytes(bytes).await + } + Some(Ok(Message::Close(frame))) => { warn!( - "Failed to send response for {address} | No response channel found" + "Connection to {address} broken with frame: {frame:?}" ); + break; + } + Some(Err(e)) => { + warn!("Connection to {address} broken with error: {e:?}"); + break; + } + None => { + warn!("Connection to {address} closed"); + break; + } + // Can ignore other message types + Some(Ok(_)) => { continue; - }; - if let Err(e) = channel.send(bytes).await { - warn!( - "Failed to send response for {address} | Channel failure | {e:#}" - ); } - } - }; + }; + } + // Cancel again if not already + connection.cancel(); + }; - tokio::join!(forward_requests, read_responses); + tokio::join!(forward_writes, handle_reads); - warn!("Disconnnected from {address}"); - connection.connected.store(false, atomic::Ordering::Relaxed); - } - }); + warn!("Disconnnected from {address}"); + connection.connected.store(false, atomic::Ordering::Relaxed); + } } #[derive(Debug)] diff --git a/lib/transport/src/lib.rs b/lib/transport/src/lib.rs index bb8fbebe5..51715ad5a 100644 --- a/lib/transport/src/lib.rs +++ b/lib/transport/src/lib.rs @@ -1,3 +1,5 @@ +use ::bytes::Bytes; + pub mod bytes; pub mod channel; pub mod client; @@ -7,5 +9,14 @@ pub mod server; pub enum MessageState { Successful, Failed, + Request, InProgress, + Terminal, +} + +pub trait TransportHandler { + fn handle_incoming_bytes( + &self, + bytes: Bytes, + ) -> impl Future + Send; } diff --git a/lib/transport/src/server.rs b/lib/transport/src/server.rs index e4373c99a..d07903dda 100644 --- a/lib/transport/src/server.rs +++ b/lib/transport/src/server.rs @@ -9,15 +9,17 @@ use serror::AddStatusCode; use tokio::sync::Mutex; use tracing::{error, warn}; -use crate::channel::BufferedReceiver; +use crate::{TransportHandler, channel::BufferedReceiver}; -pub fn inbound_connection( +pub fn inbound_connection< + T: TransportHandler + Send + Sync + 'static, +>( ws: WebSocketUpgrade, - handle_read: impl Fn(Bytes) + Send + Sync + 'static, + transport: T, write_receiver: &'static Mutex>, ) -> serror::Result { // Limits to only one active websocket connection. - let mut response_receiver = write_receiver + let mut write_receiver = write_receiver .try_lock() .status_code(StatusCode::FORBIDDEN)?; @@ -26,12 +28,34 @@ pub fn inbound_connection( let (mut ws_write, mut ws_read) = socket.split(); - // Handle incoming messages - let ws_read = async { + let forward_writes = async { + loop { + let msg = match write_receiver.recv().await { + // Sender Dropped (shouldn't happen, it is static). + None => break, + // This has to copy the bytes to follow ownership rules. + Some(msg) => Message::Binary(Bytes::copy_from_slice(msg)), + }; + match ws_write.send(msg).await { + // Clears the stored message from receiver buffer. + // TODO: Move after response ack. + Ok(_) => write_receiver.clear_buffer(), + Err(e) => { + warn!("Failed to send response | {e:?}"); + let _ = ws_write.close().await; + break; + } + } + } + }; + + let handle_reads = async { loop { match ws_read.next().await { // Incoming core msg - Some(Ok(Message::Binary(msg))) => handle_read(msg), + Some(Ok(Message::Binary(bytes))) => { + transport.handle_incoming_bytes(bytes).await + } // Disconnection cases. Some(Ok(Message::Close(frame))) => { warn!("Connection closed with frame: {frame:?}"); @@ -50,31 +74,9 @@ pub fn inbound_connection( } }; - // Handle outgoing messages - let ws_write = async { - loop { - let msg = match response_receiver.recv().await { - // Sender Dropped (shouldn't happen, it is static). - None => break, - // This has to copy the bytes to follow ownership rules. - Some(msg) => Message::Binary(Bytes::copy_from_slice(msg)), - }; - match ws_write.send(msg).await { - // Clears the stored message from receiver buffer. - // TODO: Move after response ack from Core. - Ok(_) => response_receiver.clear_buffer(), - Err(e) => { - warn!("Failed to send response | {e:?}"); - let _ = ws_write.close().await; - break; - } - } - } - }; - tokio::select! { - _ = ws_read => {}, - _ = ws_write => {} + _ = forward_writes => {}, + _ = handle_reads => {}, }; })) }