starred/komodo
1
0
Fork 0
mirror of https://github.com/moghtech/komodo.git synced 2026-10-27 08:00:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

abstract websocket handling implementations on both sides

Browse Source
This commit is contained in:
mbecker20
2025-09-27 14:21:12 -07:00
parent 2d83105500
commit 951ff34a9e
15 changed files with 468 additions and 509 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
-2
View File
@@ -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",
bin/core/Cargo.toml
-1
View File
@@ -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
bin/periphery/src/connection/client.rs
+27 -81
View File
@@ -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<u8> = 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<u8> =
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::<ClientLoginFlow>(
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 => {},
};
}
}
bin/periphery/src/connection/mod.rs
+68
View File
@@ -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<L: LoginFlow>(
mut socket: impl Websocket,
connection_identifiers: ConnectionIdentifiers<'_>,
private_key: &[u8],
write_receiver: &mut BufferedReceiver<Bytes>,
) -> 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) {
bin/periphery/src/connection/server.rs
+15 -69
View File
@@ -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::<ServerLoginFlow>(
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 => {},
};
}
}))
}
client/periphery/rs/Cargo.toml
-1
View File
@@ -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
client/periphery/rs/src/connection/client.rs
+33 -102
View File
@@ -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::<ClientLoginFlow>(
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);
}
});
client/periphery/rs/src/connection/mod.rs
+83
View File
@@ -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<L: LoginFlow>(
mut socket: impl Websocket,
connection_identifiers: ConnectionIdentifiers<'_>,
private_key: &[u8],
write_receiver: &mut BufferedReceiver<Bytes>,
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<CloneCache<Uuid, Sender<Bytes>>>,
}
client/periphery/rs/src/connection/server.rs
+21 -99
View File
@@ -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::<ServerLoginFlow>(
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);
}
}))
}
lib/transport/Cargo.toml
+1 -1
View File
@@ -21,4 +21,4 @@ rand.workspace = true
snow.workspace = true
sha1.workspace = true
sha2.workspace = true
uuid.workspace = true
uuid.workspace = true
lib/transport/src/auth.rs
+196 -111
View File
@@ -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<Output = anyhow::Result<()>>;
}
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<Bytes> {
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")?,
))
}
lib/transport/src/bytes.rs
+1 -17
View File
@@ -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,
}
}
}
lib/transport/src/websocket/axum.rs
+5 -7
View File
@@ -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<WebsocketMessage<Self::CloseFrame>, 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<axum::extract::ws::WebSocket>;
lib/transport/src/websocket/mod.rs
+10 -6
View File
@@ -20,9 +20,12 @@ pub enum WebsocketMessage<CloseFrame> {
/// 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<WebsocketMessage<Self::CloseFrame>, 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<Output = Result<(), Self::Error>>;
) -> impl Future<Output = Result<(), Self::Error>> + Send + Sync;
/// Send close message
fn close(
&mut self,
frame: Option<Self::CloseFrame>,
) -> impl Future<Output = Result<(), Self::Error>>;
) -> impl Future<Output = Result<(), Self::Error>> + Send + Sync;
}
lib/transport/src/websocket/tungstenite.rs
+8 -12
View File
@@ -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<WebsocketMessage<Self::CloseFrame>, 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<WebSocketStream<MaybeTlsStream<TcpStream>>>;