From 951ff34a9efa154ae86e1c0d5fb942288ea992d4 Mon Sep 17 00:00:00 2001 From: mbecker20 Date: Sat, 20 Sep 2025 22:43:15 -0700 Subject: [PATCH] abstract websocket handling implementations on both sides --- Cargo.lock | 2 - bin/core/Cargo.toml | 1 - bin/periphery/src/connection/client.rs | 108 ++----- bin/periphery/src/connection/mod.rs | 68 ++++ bin/periphery/src/connection/server.rs | 84 +---- client/periphery/rs/Cargo.toml | 1 - client/periphery/rs/src/connection/client.rs | 135 ++------ client/periphery/rs/src/connection/mod.rs | 83 +++++ client/periphery/rs/src/connection/server.rs | 120 ++------ lib/transport/Cargo.toml | 2 +- lib/transport/src/auth.rs | 307 ++++++++++++------- lib/transport/src/bytes.rs | 18 +- lib/transport/src/websocket/axum.rs | 12 +- lib/transport/src/websocket/mod.rs | 16 +- lib/transport/src/websocket/tungstenite.rs | 20 +- 15 files changed, 468 insertions(+), 509 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4823be528..39b981c5d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2873,7 +2873,6 @@ dependencies = [ "serde_json", "serde_yaml_ng", "serror", - "sha1", "sha2", "slack_client_rs", "svi", @@ -3781,7 +3780,6 @@ version = "1.19.5" dependencies = [ "anyhow", "axum", - "base64 0.22.1", "bytes", "cache", "futures-util", diff --git a/bin/core/Cargo.toml b/bin/core/Cargo.toml index fccb942d8..31b4e8052 100644 --- a/bin/core/Cargo.toml +++ b/bin/core/Cargo.toml @@ -80,6 +80,5 @@ uuid.workspace = true envy.workspace = true rand.workspace = true hmac.workspace = true -sha1.workspace = true sha2.workspace = true hex.workspace = true diff --git a/bin/periphery/src/connection/client.rs b/bin/periphery/src/connection/client.rs index 172e31c4e..cdbf64f4a 100644 --- a/bin/periphery/src/connection/client.rs +++ b/bin/periphery/src/connection/client.rs @@ -2,17 +2,14 @@ use std::time::Duration; use anyhow::Context; use axum::http::HeaderValue; -use bytes::Bytes; use transport::{ - auth::{ConnectionIdentifiers, handle_client_side_login}, + auth::{ClientLoginFlow, ConnectionIdentifiers}, fix_ws_address, - websocket::{ - WebsocketMessage, WebsocketReceiver, WebsocketSender, - tungstenite::TungsteniteWebsocket, - }, + websocket::tungstenite::TungsteniteWebsocket, }; +use url::Url; -use crate::connection::{handle_incoming_bytes, ws_receiver}; +use crate::connection::ws_receiver; pub async fn handler( core_host: &str, @@ -30,101 +27,50 @@ pub async fn handler( .context("Websocket handler called more than once.")?; let core_host = fix_ws_address(core_host); - let url = format!( + let core_endpoint = format!( "{core_host}/ws/periphery?server={}", urlencoding::encode(connect_as) ); - let parsed_url = - ::url::Url::parse(&url).context("Failed to parse ws url")?; - let host: Vec = parsed_url - .host() - .context("url has no host")? - .to_string() - .into(); - let query = - parsed_url.query().context("url has no query")?.as_bytes(); + let url = + Url::parse(&core_endpoint).context("Failed to parse ws url")?; + let host: Vec = + url.host().context("url has no host")?.to_string().into(); + let query = url.query().context("url has no query")?.as_bytes(); info!("Initiating outbound connection to {url}"); loop { - let (mut socket, accept) = match connect_websocket(&url).await { - Ok(res) => res, - Err(e) => { - warn!("{e:#}"); - tokio::time::sleep(Duration::from_secs(5)).await; - continue; - } - }; + let (socket, accept) = + match connect_websocket(&core_endpoint).await { + Ok(res) => res, + Err(e) => { + warn!("{e:#}"); + tokio::time::sleep(Duration::from_secs(5)).await; + continue; + } + }; info!("Connected to core connection websocket"); - let id = ConnectionIdentifiers { + let connection_identifiers = ConnectionIdentifiers { host: &host, accept: accept.as_bytes(), query, }; // TODO: source the pk - if let Err(e) = - handle_client_side_login(&mut socket, id, b"RANDOM_PRIVATE_KEY") - .await + if let Err(e) = super::handle_websocket::( + socket, + connection_identifiers, + b"TEST", + &mut write_receiver, + ) + .await { warn!("Failed to login | {e:#}"); tokio::time::sleep(Duration::from_secs(5)).await; continue; }; - - info!("Logged in to core connection websocket"); - - let (mut ws_write, mut ws_read) = socket.split(); - - 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) => 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(None).await; - break; - } - } - } - }; - - let handle_reads = async { - loop { - match ws_read.recv().await { - Ok(WebsocketMessage::Binary(bytes)) => { - handle_incoming_bytes(bytes).await - } - Ok(WebsocketMessage::Close(frame)) => { - warn!("Connection closed with frame: {frame:?}"); - break; - } - Ok(WebsocketMessage::Closed) => { - warn!("Connection already closed"); - break; - } - Err(e) => { - error!("Failed to read websocket message | {e:?}"); - break; - } - }; - } - }; - - tokio::select! { - _ = forward_writes => {}, - _ = handle_reads => {}, - }; } } diff --git a/bin/periphery/src/connection/mod.rs b/bin/periphery/src/connection/mod.rs index f3ece7fe1..bdf694f26 100644 --- a/bin/periphery/src/connection/mod.rs +++ b/bin/periphery/src/connection/mod.rs @@ -9,11 +9,16 @@ use tokio::sync::{Mutex, mpsc::Sender}; use tokio_util::sync::CancellationToken; use transport::{ MessageState, + auth::{ConnectionIdentifiers, LoginFlow}, bytes::{ data_from_transport_bytes, id_state_from_transport_bytes, to_transport_bytes, }, channel::{BufferedReceiver, buffered_channel}, + websocket::{ + Websocket, WebsocketMessage, WebsocketReceiver, + WebsocketSender as _, + }, }; use uuid::Uuid; @@ -50,6 +55,69 @@ pub fn init_response_channel() { .expect("response_receiver initialized more than once"); } +async fn handle_websocket( + mut socket: impl Websocket, + connection_identifiers: ConnectionIdentifiers<'_>, + private_key: &[u8], + write_receiver: &mut BufferedReceiver, +) -> anyhow::Result<()> { + L::login(&mut socket, connection_identifiers, private_key).await?; + + info!("Logged in to core connection websocket"); + + let (mut ws_write, mut ws_read) = socket.split(); + + 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) => 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(None).await; + break; + } + } + } + }; + + let handle_reads = async { + loop { + match ws_read.recv().await { + Ok(WebsocketMessage::Binary(bytes)) => { + handle_incoming_bytes(bytes).await + } + Ok(WebsocketMessage::Close(frame)) => { + warn!("Connection closed with frame: {frame:?}"); + break; + } + Ok(WebsocketMessage::Closed) => { + warn!("Connection already closed"); + break; + } + Err(e) => { + error!("Failed to read websocket message | {e:?}"); + break; + } + }; + } + }; + + tokio::select! { + _ = forward_writes => {}, + _ = handle_reads => {}, + }; + + Ok(()) +} + async fn handle_incoming_bytes(bytes: Bytes) { // Maybe wrap all of this on tokio spawn let (id, state) = match id_state_from_transport_bytes(&bytes) { diff --git a/bin/periphery/src/connection/server.rs b/bin/periphery/src/connection/server.rs index caa7a7933..0322d62b9 100644 --- a/bin/periphery/src/connection/server.rs +++ b/bin/periphery/src/connection/server.rs @@ -14,22 +14,13 @@ use axum::{ routing::get, }; use axum_server::tls_rustls::RustlsConfig; -use bytes::Bytes; use serror::{AddStatusCode, AddStatusCodeError}; use transport::{ - auth::{ - ConnectionIdentifiers, compute_accept, handle_server_side_login, - }, - websocket::{ - WebsocketMessage, WebsocketReceiver, WebsocketSender, - axum::AxumWebsocket, - }, + auth::{ConnectionIdentifiers, ServerLoginFlow, compute_accept}, + websocket::axum::AxumWebsocket, }; -use crate::{ - config::periphery_config, - connection::{handle_incoming_bytes, ws_receiver}, -}; +use crate::{config::periphery_config, connection::ws_receiver}; pub async fn run() -> anyhow::Result<()> { let config = periphery_config(); @@ -90,71 +81,26 @@ async fn handler( let ws_accept = compute_accept(ws_key.as_bytes()); Ok(ws.on_upgrade(|socket| async move { - let mut socket = AxumWebsocket(socket); + let socket = AxumWebsocket(socket); - let id = ConnectionIdentifiers { + let connection_identifiers = ConnectionIdentifiers { host: host.as_bytes(), query: &[], accept: ws_accept.as_bytes(), }; - if let Err(e) = - handle_server_side_login(&mut socket, id, b"TEMP_SERVER_PK") - .await + // TODO: source the pk + if let Err(e) = super::handle_websocket::( + socket, + connection_identifiers, + b"TEST", + &mut write_receiver, + ) + .await { - warn!("Client failed to login | {e:#}"); + warn!("Core failed to login to connection | {e:#}"); return; - }; - - let (mut ws_write, mut ws_read) = socket.split(); - - 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) => 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(None).await; - break; - } - } - } - }; - - let handle_reads = async { - loop { - match ws_read.recv().await { - Ok(WebsocketMessage::Binary(bytes)) => { - handle_incoming_bytes(bytes).await - } - Ok(WebsocketMessage::Close(frame)) => { - warn!("Connection closed with frame: {frame:?}"); - break; - } - Ok(WebsocketMessage::Closed) => { - warn!("Connection already closed"); - break; - } - Err(e) => { - error!("Failed to read websocket message | {e:?}"); - break; - } - }; - } - }; - - tokio::select! { - _ = forward_writes => {}, - _ = handle_reads => {}, - }; + } })) } diff --git a/client/periphery/rs/Cargo.toml b/client/periphery/rs/Cargo.toml index 544ff7c04..bd8340f36 100644 --- a/client/periphery/rs/Cargo.toml +++ b/client/periphery/rs/Cargo.toml @@ -25,7 +25,6 @@ serde_json.workspace = true tracing.workspace = true anyhow.workspace = true rustls.workspace = true -base64.workspace = true bytes.workspace = true tokio.workspace = true serde.workspace = true diff --git a/client/periphery/rs/src/connection/client.rs b/client/periphery/rs/src/connection/client.rs index f07c4f6a7..8b99752ca 100644 --- a/client/periphery/rs/src/connection/client.rs +++ b/client/periphery/rs/src/connection/client.rs @@ -2,18 +2,14 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use anyhow::Context; use axum::http::HeaderValue; -use bytes::Bytes; use komodo_client::entities::server::Server; use rustls::{ClientConfig, client::danger::ServerCertVerifier}; use tokio_tungstenite::Connector; use tracing::{info, warn}; use transport::{ - auth::{ConnectionIdentifiers, handle_client_side_login}, + auth::{ClientLoginFlow, ConnectionIdentifiers}, fix_ws_address, - websocket::{ - WebsocketMessage, WebsocketReceiver as _, WebsocketSender as _, - tungstenite::TungsteniteWebsocket, - }, + websocket::tungstenite::TungsteniteWebsocket, }; use crate::{ @@ -56,28 +52,32 @@ pub async fn manage_client_connections(servers: &[Server]) { fix_ws_address(address) }; match ( - periphery_connections.get(server_id).await, address.is_empty(), + periphery_connections.get(server_id).await, ) { - (Some(existing), true) if existing.address.is_none() => { + // Periphery -> Core connections + (true, Some(existing)) if existing.address.is_none() => { continue; } - (Some(existing), true) => { + (true, Some(existing)) => { existing.cancel(); continue; } - (Some(existing), false) + (true, None) => continue, + // Core -> Periphery connections + (false, Some(existing)) if existing .address .as_ref() .map(|a| a == &address) .unwrap_or_default() => { + // Connection OK continue; } - (Some(_), false) => {} - (None, false) => continue, - (None, true) => {} + // Recreate connection cases + (false, Some(_)) => {} + (false, None) => {} }; // If reaches here, recreate the connection. if let Err(e) = @@ -116,28 +116,29 @@ async fn spawn_client_connection( tokio::spawn(async move { loop { - let (mut socket, accept) = - match connect_websocket(&address).await { - Ok(res) => res, - Err(e) => { - connection.set_error(e).await; - tokio::time::sleep(Duration::from_secs(5)).await; - continue; - } - }; + let (socket, accept) = match connect_websocket(&address).await { + Ok(res) => res, + Err(e) => { + connection.set_error(e).await; + tokio::time::sleep(Duration::from_secs(5)).await; + continue; + } + }; info!("PERIPHERY: Connected to {address}"); - let id = ConnectionIdentifiers { - host: &host, - accept: accept.as_bytes(), - query: &[], - }; - - if let Err(e) = handle_client_side_login( - &mut socket, - id, - b"RANDOM_PRIVATE_KEY", + if let Err(e) = super::handle_websocket::( + socket, + ConnectionIdentifiers { + host: &host, + accept: accept.as_bytes(), + query: &[], + }, + b"TEST", + &mut write_receiver, + &connection, + &handler, + &address, ) .await { @@ -145,76 +146,6 @@ async fn spawn_client_connection( tokio::time::sleep(Duration::from_secs(5)).await; continue; }; - - info!("PERIPHERY: Logged into {address}"); - - connection.set_connected(true); - connection.clear_error().await; - - let (mut ws_write, mut ws_read) = socket.split(); - - 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) => Bytes::copy_from_slice(request), - }; - - 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(None).await; - connection.cancel(); - }; - - let handle_reads = async { - loop { - let next = tokio::select! { - next = ws_read.recv() => next, - _ = connection.cancel.cancelled() => break, - }; - - match next { - Ok(WebsocketMessage::Binary(bytes)) => { - handler.handle_incoming_bytes(bytes).await - } - Ok(WebsocketMessage::Close(frame)) => { - warn!( - "Connection to {address} broken with frame: {frame:?}" - ); - break; - } - Ok(WebsocketMessage::Closed) => {} - Err(e) => { - warn!( - "Connection to {address} broken with error: {e:?}" - ); - break; - } - }; - } - // Cancel again if not already - connection.cancel(); - }; - - tokio::join!(forward_writes, handle_reads); - - warn!("PERIPHERY: Disconnnected from {address}"); - connection.set_connected(false); } }); diff --git a/client/periphery/rs/src/connection/mod.rs b/client/periphery/rs/src/connection/mod.rs index 6566b39e7..f49c73ef1 100644 --- a/client/periphery/rs/src/connection/mod.rs +++ b/client/periphery/rs/src/connection/mod.rs @@ -12,8 +12,13 @@ use tokio::sync::{ use tokio_util::sync::CancellationToken; use tracing::{info, warn}; use transport::{ + auth::{ConnectionIdentifiers, LoginFlow}, bytes::id_from_transport_bytes, channel::{BufferedReceiver, buffered_channel}, + websocket::{ + Websocket, WebsocketMessage, WebsocketReceiver as _, + WebsocketSender as _, + }, }; use uuid::Uuid; @@ -22,6 +27,84 @@ use crate::periphery_channels; pub mod client; pub mod server; +async fn handle_websocket( + mut socket: impl Websocket, + connection_identifiers: ConnectionIdentifiers<'_>, + private_key: &[u8], + write_receiver: &mut BufferedReceiver, + connection: &PeripheryConnection, + handler: &MessageHandler, + name: &str, +) -> anyhow::Result<()> { + L::login(&mut socket, connection_identifiers, private_key).await?; + + info!("PERIPHERY: Logged into {name}"); + + connection.set_connected(true); + connection.clear_error().await; + + let (mut ws_write, mut ws_read) = socket.split(); + + let forward_writes = async { + loop { + let next = tokio::select! { + next = write_receiver.recv() => next, + _ = connection.cancel.cancelled() => break, + }; + + let message = match next { + Some(request) => Bytes::copy_from_slice(request), + // Sender Dropped (shouldn't happen, a reference is held on 'connection'). + None => break, + }; + + match ws_write.send(message).await { + Ok(_) => write_receiver.clear_buffer(), + Err(e) => { + warn!("Failed to send request to {name} | {e:#}"); + break; + } + } + } + // Cancel again if not already + let _ = ws_write.close(None).await; + connection.cancel(); + }; + + let handle_reads = async { + loop { + let next = tokio::select! { + next = ws_read.recv() => next, + _ = connection.cancel.cancelled() => break, + }; + + match next { + Ok(WebsocketMessage::Binary(bytes)) => { + handler.handle_incoming_bytes(bytes).await + } + Ok(WebsocketMessage::Close(frame)) => { + warn!("Connection to {name} broken with frame: {frame:?}"); + break; + } + Ok(WebsocketMessage::Closed) => {} + Err(e) => { + warn!("Connection to {name} broken with error: {e:?}"); + break; + } + }; + } + // Cancel again if not already + connection.cancel(); + }; + + tokio::join!(forward_writes, handle_reads); + + warn!("PERIPHERY: Disconnnected from {name}"); + connection.set_connected(false); + + Ok(()) +} + pub struct MessageHandler { channels: Arc>>, } diff --git a/client/periphery/rs/src/connection/server.rs b/client/periphery/rs/src/connection/server.rs index 005ac550b..d738bd1b4 100644 --- a/client/periphery/rs/src/connection/server.rs +++ b/client/periphery/rs/src/connection/server.rs @@ -4,17 +4,11 @@ use axum::{ http::{HeaderMap, StatusCode}, response::Response, }; -use bytes::Bytes; use serror::AddStatusCode; -use tracing::{error, info, warn}; +use tracing::warn; use transport::{ - auth::{ - ConnectionIdentifiers, compute_accept, handle_server_side_login, - }, - websocket::{ - WebsocketMessage, WebsocketReceiver, WebsocketSender, - axum::AxumWebsocket, - }, + auth::{ConnectionIdentifiers, ServerLoginFlow, compute_accept}, + websocket::axum::AxumWebsocket, }; use crate::connection::{ @@ -51,98 +45,26 @@ pub async fn handler( } Ok(ws.on_upgrade(|socket| async move { - let mut socket = AxumWebsocket(socket); - - let id = ConnectionIdentifiers { - host: host.as_bytes(), - query: query.as_bytes(), - accept: ws_accept.as_bytes(), - }; - - // TODO: use proper private key - if let Err(e) = - handle_server_side_login(&mut socket, id, b"RANDOM_SERVER_PK") - .await + if let Err(e) = super::handle_websocket::( + AxumWebsocket(socket), + ConnectionIdentifiers { + host: host.as_bytes(), + query: query.as_bytes(), + accept: ws_accept.as_bytes(), + }, + b"TEST", + &mut write_receiver, + &connection, + &handler, + &server_id, + ) + .await { - warn!("PERIPHERY: Client failed to login | {e:#}"); + warn!( + "PERIPHERY: [{server_id}] Client failed to login | {e:#}" + ); connection.set_error(e).await; return; - }; - - info!("PERIPHERY: {server_id} logged in"); - - connection.set_connected(true); - connection.clear_error().await; - - let (mut ws_write, mut ws_read) = socket.split(); - - let forward_writes = async { - loop { - let next = tokio::select! { - next = write_receiver.recv() => next, - _ = connection.cancel.cancelled() => { - info!("WRITE: Connection cancelled"); - break - }, - }; - - let msg = match next { - // Sender Dropped (shouldn't happen, a reference is held on 'connection'). - None => break, - // This has to copy the bytes to follow ownership rules. - Some(msg) => 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:?}"); - break; - } - } - } - // Cancel again if not already - let _ = ws_write.close(None).await; - connection.cancel(); - }; - - let handle_reads = async { - loop { - let next = tokio::select! { - next = ws_read.recv() => next, - _ = connection.cancel.cancelled() => { - info!("READ: Connection cancelled"); - break - }, - }; - - match next { - Ok(WebsocketMessage::Binary(bytes)) => { - handler.handle_incoming_bytes(bytes).await - } - Ok(WebsocketMessage::Close(frame)) => { - warn!("Connection closed with frame: {frame:?}"); - break; - } - Ok(WebsocketMessage::Closed) => { - warn!("Connection already closed"); - break; - } - Err(e) => { - error!("Failed to read websocket message | {e:?}"); - break; - } - }; - } - // Cancel again if not already - connection.cancel(); - }; - - tokio::join!(forward_writes, handle_reads); - - warn!("PERIPHERY: {server_id} Disconnnected"); - connection.set_connected(false); + } })) } diff --git a/lib/transport/Cargo.toml b/lib/transport/Cargo.toml index 486543a5d..dab28e477 100644 --- a/lib/transport/Cargo.toml +++ b/lib/transport/Cargo.toml @@ -21,4 +21,4 @@ rand.workspace = true snow.workspace = true sha1.workspace = true sha2.workspace = true -uuid.workspace = true \ No newline at end of file +uuid.workspace = true diff --git a/lib/transport/src/auth.rs b/lib/transport/src/auth.rs index ec7da4c17..300c0725f 100644 --- a/lib/transport/src/auth.rs +++ b/lib/transport/src/auth.rs @@ -6,15 +6,16 @@ //! This is trivial for Periphery -> Core connection, but presents a challenge //! for Core -> Periphery, where untrusted TLS certs are being used. -use anyhow::Context; +use anyhow::{Context, anyhow}; use base64::{Engine, prelude::BASE64_STANDARD}; use bytes::Bytes; use rand::RngCore; use sha2::{Digest, Sha256}; +use tracing::warn; -use crate::websocket::{ - Websocket, axum::AxumWebsocket, tungstenite::TungsteniteWebsocket, -}; +use crate::{MessageState, websocket::Websocket}; + +const NOISE_XX_PARAMS: &str = "Noise_XX_25519_ChaChaPoly_BLAKE2s"; pub struct ConnectionIdentifiers<'a> { /// Server hostname @@ -25,130 +26,178 @@ pub struct ConnectionIdentifiers<'a> { pub accept: &'a [u8], } -#[derive(Debug, Clone, Copy)] -pub enum AuthType { - Passkey = 0, - Noise = 1, +pub trait LoginFlow { + fn login( + socket: &mut impl Websocket, + connection_identifiers: ConnectionIdentifiers<'_>, + private_key: &[u8], + ) -> impl Future>; } -const NOISE_XX_PARAMS: &str = "Noise_XX_25519_ChaChaPoly_BLAKE2s"; +pub struct ServerLoginFlow; -pub async fn handle_server_side_login( - socket: &mut AxumWebsocket, - id: ConnectionIdentifiers<'_>, - private_key: &[u8], -) -> anyhow::Result<()> { - // Server generates random nonce and sends to client - let nonce = nonce(); - socket - .send(Bytes::from_owner(nonce)) - .await - .context("Failed to send connection nonce")?; +impl LoginFlow for ServerLoginFlow { + async fn login( + socket: &mut impl Websocket, + connection_identifiers: ConnectionIdentifiers<'_>, + private_key: &[u8], + ) -> anyhow::Result<()> { + let res = async { + // Server generates random nonce and sends to client + let nonce = nonce(); + socket + .send(Bytes::from_owner(nonce)) + .await + .context("Failed to send connection nonce")?; - // Build the handshake using the unique prologue hash. - // The prologue must be the same on both sides of connection. - let mut handshake = snow::Builder::new(NOISE_XX_PARAMS.parse()?) - .local_private_key(private_key)? - .prologue(&id.hash(&nonce))? - .build_responder()?; + // Build the handshake using the connection-unique prologue hash. + // The prologue must be the same on both sides of connection. + let mut handshake = + snow::Builder::new(NOISE_XX_PARAMS.parse()?) + .local_private_key(private_key)? + .prologue(&connection_identifiers.hash(&nonce))? + .build_responder()?; - // Receive and read handshake_m1 - let handshake_m1 = socket - .recv_bytes() - .await - .context("Failed to get handshake_m1")?; - handshake - .read_message(&handshake_m1, &mut []) - .context("Failed to read handshake_m1")?; + // Receive and read handshake_m1 + let handshake_m1 = socket + .recv_bytes() + .await + .context("Failed to get handshake_m1")?; + handshake + .read_message(&handshake_m1, &mut []) + .context("Failed to read handshake_m1")?; - // Send handshake_m2 - let mut handshake_m2 = [0u8; 1024]; - let written = handshake - .write_message(&[], &mut handshake_m2) - .context("Failed to write handshake_m2")?; - socket - .send(Bytes::copy_from_slice(&handshake_m2[..written])) - .await - .context("Failed to send handshake_m2")?; + // Send handshake_m2 + let mut handshake_m2 = [0u8; 1024]; + let written = handshake + .write_message(&[], &mut handshake_m2) + .context("Failed to write handshake_m2")?; + socket + .send(Bytes::copy_from_slice(&handshake_m2[..written])) + .await + .context("Failed to send handshake_m2")?; - // Receive and read handshake_m3 - let handshake_m3 = socket - .recv_bytes() - .await - .context("Failed to get handshake_m3")?; - handshake - .read_message(&handshake_m3, &mut []) - .context("Failed to read handshake_m3")?; + // Receive and read handshake_m3 + let handshake_m3 = socket + .recv_bytes() + .await + .context("Failed to get handshake_m3")?; + handshake + .read_message(&handshake_m3, &mut []) + .context("Failed to read handshake_m3")?; - // Server now has client public key - let client_public_key = handshake - .get_remote_static() - .context("Failed to get remote public key")?; - println!( - "Server got client public key: {}", - BASE64_STANDARD.encode(client_public_key) - ); + // Server now has client public key + let client_public_key = handshake + .get_remote_static() + .context("Failed to get remote public key")?; + println!( + "Server got client public key: {}", + BASE64_STANDARD.encode(client_public_key) + ); - Ok(()) + anyhow::Ok(()) + } + .await; + + match res { + Ok(_) => { + socket + .send(MessageState::Successful.into()) + .await + .context("Failed to send login successful to client")?; + Ok(()) + } + Err(e) => { + if let Err(e) = socket + .send(MessageState::Successful.into()) + .await + .context("Failed to send login successful to client") + { + // Log additional error + warn!("{e:#}"); + // Close socket + let _ = socket.close(None).await; + } + // Return the original error + Err(e) + } + } + } } -pub async fn handle_client_side_login( - socket: &mut TungsteniteWebsocket, - id: ConnectionIdentifiers<'_>, - private_key: &[u8], -) -> anyhow::Result<()> { - // Receive nonce from server - let nonce = socket - .recv_bytes() - .await - .context("Failed to receive connection nonce")?; +pub struct ClientLoginFlow; - // Build the handshake using the unique prologue hash. - // The prologue must be the same on both sides of connection. - let mut handshake = snow::Builder::new(NOISE_XX_PARAMS.parse()?) - .local_private_key(private_key)? - .prologue(&id.hash(&nonce))? - .build_initiator()?; +impl LoginFlow for ClientLoginFlow { + async fn login( + socket: &mut impl Websocket, + connection_identifiers: ConnectionIdentifiers<'_>, + private_key: &[u8], + ) -> anyhow::Result<()> { + // Receive nonce from server + let nonce = socket + .recv_bytes() + .await + .context("Failed to receive connection nonce")?; - // Send handshake_m1 - let mut handshake_m1 = [0u8; 1024]; - let written = handshake - .write_message(&[], &mut handshake_m1) - .context("Failed to write handshake_m1")?; - socket - .send(Bytes::copy_from_slice(&handshake_m1[..written])) - .await - .context("Failed to send handshake_m1")?; + // Build the handshake using the connection-unique prologue hash. + // The prologue must be the same on both sides of connection. + let mut handshake = snow::Builder::new(NOISE_XX_PARAMS.parse()?) + .local_private_key(private_key)? + .prologue(&connection_identifiers.hash(&nonce))? + .build_initiator()?; - // Receive and read handshake_m2 - let handshake_m2 = socket - .recv_bytes() - .await - .context("Failed to get handshake_m2")?; - handshake - .read_message(&handshake_m2, &mut []) - .context("Failed to read handshake_m2")?; + // Send handshake_m1 + let mut handshake_m1 = [0u8; 1024]; + let written = handshake + .write_message(&[], &mut handshake_m1) + .context("Failed to write handshake_m1")?; + socket + .send(Bytes::copy_from_slice(&handshake_m1[..written])) + .await + .context("Failed to send handshake_m1")?; - // Client now has server public key - let server_public_key = handshake - .get_remote_static() - .context("Failed to get remote public key")?; - println!( - "Client got server public key: {}", - BASE64_STANDARD.encode(server_public_key) - ); + // Receive and read handshake_m2 + let handshake_m2 = socket + .recv_bytes() + .await + .context("Failed to get handshake_m2")?; + handshake + .read_message(&handshake_m2, &mut []) + .context("Failed to read handshake_m2")?; - // Send handshake_m3 - let mut handshake_m3 = [0u8; 1024]; - let written = handshake - .write_message(&[], &mut handshake_m3) - .context("Failed to write handshake_m3")?; - socket - .send(Bytes::copy_from_slice(&handshake_m3[..written])) - .await - .context("Failed to send handshake_m3")?; + // Client now has server public key + let server_public_key = handshake + .get_remote_static() + .context("Failed to get remote public key")?; + println!( + "Client got server public key: {}", + BASE64_STANDARD.encode(server_public_key) + ); - Ok(()) + // Send handshake_m3 + let mut handshake_m3 = [0u8; 1024]; + let written = handshake + .write_message(&[], &mut handshake_m3) + .context("Failed to write handshake_m3")?; + socket + .send(Bytes::copy_from_slice(&handshake_m3[..written])) + .await + .context("Failed to send handshake_m3")?; + + // Receive login state message and return based on value + let state = socket + .recv_bytes() + .await + .context("Failed to receive authentication state message")?; + let state = state.first().context( + "Authentication state message did not contain state byte", + )?; + match MessageState::from_byte(*state) { + MessageState::Successful => Ok(()), + // Todo: More descriptive error? + _ => Err(anyhow!("Authentication failed")), + } + } } fn nonce() -> [u8; 32] { @@ -182,3 +231,39 @@ pub fn compute_accept(sec_websocket_key: &[u8]) -> String { let digest = sha1.finalize(); BASE64_STANDARD.encode(digest) } + +/// This completes an example handshake +/// to produce the correct public key +/// for a given private key. +pub fn generage_public_key( + private_key: &[u8], +) -> anyhow::Result { + let mut ch = snow::Builder::new(NOISE_XX_PARAMS.parse()?) + .local_private_key(b"ANY")? + .build_initiator()?; + // Use the target private key with server handshake, + // since its public key is the first available in the flow. + let mut sh = snow::Builder::new(NOISE_XX_PARAMS.parse()?) + .local_private_key(private_key)? + .build_responder()?; + // write m1 + let mut m1 = [0u8; 1024]; + let written = ch + .write_message(&[], &mut m1) + .context("CLIENT: failed to write m1")?; + // read m1 + sh.read_message(&m1[..written], &mut []) + .context("SERVER: failed to read m1")?; + // write m2 + let mut m2 = [0u8; 1024]; + let written = sh + .write_message(&[], &mut m2) + .context("SERVER: failed to write m2")?; + // read m2 + ch.read_message(&m2[..written], &mut []) + .context("CLIENT: failed to read m2")?; + // client now has server public key + Ok(Bytes::copy_from_slice( + ch.get_remote_static().context("Failed to get public key")?, + )) +} diff --git a/lib/transport/src/bytes.rs b/lib/transport/src/bytes.rs index e0459560d..e0f3dc9d2 100644 --- a/lib/transport/src/bytes.rs +++ b/lib/transport/src/bytes.rs @@ -2,7 +2,7 @@ use anyhow::{Context, anyhow}; use bytes::Bytes; use uuid::Uuid; -use crate::{MessageState, auth::AuthType}; +use crate::MessageState; /// Serializes data + channel id + state to byte vec. /// The last byte is the State, and the 16 before that is the Uuid. @@ -94,19 +94,3 @@ impl MessageState { } } } - -impl AuthType { - pub fn from_byte(byte: u8) -> AuthType { - match byte { - 0 => AuthType::Passkey, - _ => AuthType::Noise, - } - } - - pub fn as_byte(&self) -> u8 { - match self { - AuthType::Passkey => 0, - AuthType::Noise => 1, - } - } -} diff --git a/lib/transport/src/websocket/axum.rs b/lib/transport/src/websocket/axum.rs index 529fcc751..fd637f3d8 100644 --- a/lib/transport/src/websocket/axum.rs +++ b/lib/transport/src/websocket/axum.rs @@ -14,6 +14,11 @@ impl Websocket for AxumWebsocket { type CloseFrame = CloseFrame; type Error = axum::Error; + fn split(self) -> (impl WebsocketSender, impl WebsocketReceiver) { + let (tx, rx) = self.0.split(); + (AxumWebsocketSender(tx), AxumWebsocketReceiver(rx)) + } + async fn recv( &mut self, ) -> Result, Self::Error> { @@ -35,13 +40,6 @@ impl Websocket for AxumWebsocket { } } -impl AxumWebsocket { - pub fn split(self) -> (AxumWebsocketSender, AxumWebsocketReceiver) { - let (tx, rx) = self.0.split(); - (AxumWebsocketSender(tx), AxumWebsocketReceiver(rx)) - } -} - pub type InnerWebsocketReceiver = SplitStream; diff --git a/lib/transport/src/websocket/mod.rs b/lib/transport/src/websocket/mod.rs index 6c6a883fd..b7a2245ff 100644 --- a/lib/transport/src/websocket/mod.rs +++ b/lib/transport/src/websocket/mod.rs @@ -20,9 +20,12 @@ pub enum WebsocketMessage { /// Standard traits for websocket pub trait Websocket { - type CloseFrame: std::fmt::Debug; + type CloseFrame: std::fmt::Debug + Send + Sync + 'static; type Error: std::error::Error + Send + Sync + 'static; + /// Abstraction over websocket splitting + fn split(self) -> (impl WebsocketSender, impl WebsocketReceiver); + /// Looping receiver for websocket messages which only returns /// on significant messages. fn recv( @@ -64,7 +67,7 @@ pub trait Websocket { /// Traits for split websocket receiver pub trait WebsocketReceiver { - type CloseFrame: std::fmt::Debug; + type CloseFrame: std::fmt::Debug + Send + Sync + 'static; type Error: std::error::Error + Send + Sync + 'static; /// Looping receiver for websocket messages which only returns @@ -73,23 +76,24 @@ pub trait WebsocketReceiver { &mut self, ) -> impl Future< Output = Result, Self::Error>, - >; + > + Send + + Sync; } /// Traits for split websocket receiver pub trait WebsocketSender { - type CloseFrame: std::fmt::Debug; + type CloseFrame: std::fmt::Debug + Send + Sync + 'static; type Error: std::error::Error + Send + Sync + 'static; /// Streamlined sending on bytes fn send( &mut self, bytes: Bytes, - ) -> impl Future>; + ) -> impl Future> + Send + Sync; /// Send close message fn close( &mut self, frame: Option, - ) -> impl Future>; + ) -> impl Future> + Send + Sync; } diff --git a/lib/transport/src/websocket/tungstenite.rs b/lib/transport/src/websocket/tungstenite.rs index c7c3f0eca..a028589cb 100644 --- a/lib/transport/src/websocket/tungstenite.rs +++ b/lib/transport/src/websocket/tungstenite.rs @@ -20,6 +20,14 @@ impl Websocket for TungsteniteWebsocket { type CloseFrame = CloseFrame; type Error = tungstenite::Error; + fn split(self) -> (impl WebsocketSender, impl WebsocketReceiver) { + let (tx, rx) = self.0.split(); + ( + TungsteniteWebsocketSender(tx), + TungsteniteWebsocketReceiver(rx), + ) + } + async fn recv( &mut self, ) -> Result, Self::Error> { @@ -41,18 +49,6 @@ impl Websocket for TungsteniteWebsocket { } } -impl TungsteniteWebsocket { - pub fn split( - self, - ) -> (TungsteniteWebsocketSender, TungsteniteWebsocketReceiver) { - let (tx, rx) = self.0.split(); - ( - TungsteniteWebsocketSender(tx), - TungsteniteWebsocketReceiver(rx), - ) - } -} - pub type InnerWebsocketReceiver = SplitStream>>;