diff --git a/bin/core/src/connection/mod.rs b/bin/core/src/connection/mod.rs index 60eda6102..d5fe2c2f9 100644 --- a/bin/core/src/connection/mod.rs +++ b/bin/core/src/connection/mod.rs @@ -19,7 +19,7 @@ use komodo_client::entities::{ server::Server, }; use periphery_client::transport::{ - EncodedTransportMessage, TransportMessage, + EncodedTransportMessage, ResponseMessage, TransportMessage, }; use serror::serror_into_anyhow_error; use tokio::sync::RwLock; @@ -415,28 +415,30 @@ impl PeripheryConnection { } }; match message { - TransportMessage::Response(data) => match data.decode() { - Ok(WithChannel { - channel: channel_id, - data, - }) => { - let Some(channel) = self.responses.get(&channel_id).await - else { - warn!( - "Failed to forward Response message | No response channel found at {channel_id}" - ); - return; - }; - if let Err(e) = channel.send(data).await { - warn!( - "Failed to send response | Channel failure at {channel_id} | {e:#}" - ); + TransportMessage::Response(data) => { + match data.decode().map(ResponseMessage::extract) { + Ok(WithChannel { + channel: channel_id, + data, + }) => { + let Some(channel) = self.responses.get(&channel_id).await + else { + warn!( + "Failed to forward Response message | No response channel found at {channel_id}" + ); + return; + }; + if let Err(e) = channel.send(data).await { + warn!( + "Failed to send response | Channel failure at {channel_id} | {e:#}" + ); + } + } + Err(e) => { + warn!("Failed to read Response message | {e:#}"); } } - Err(e) => { - warn!("Failed to read Response message | {e:#}"); - } - }, + } TransportMessage::Terminal(data) => match data.decode() { Ok(WithChannel { channel: channel_id, diff --git a/bin/periphery/src/connection/mod.rs b/bin/periphery/src/connection/mod.rs index 653080299..cc4ed256b 100644 --- a/bin/periphery/src/connection/mod.rs +++ b/bin/periphery/src/connection/mod.rs @@ -11,7 +11,8 @@ use encoding::{ }; use noise::key::SpkiPublicKey; use periphery_client::transport::{ - EncodedRequestMessage, EncodedTransportMessage, TransportMessage, + EncodedRequestMessage, EncodedTransportMessage, RequestMessage, + TransportMessage, }; use resolver_api::Resolve; use transport::{ @@ -208,14 +209,15 @@ fn handle_request( let WithChannel { channel, data: request, - }: WithChannel = match message.decode() { - Ok(res) => res, - Err(e) => { - // TODO: handle: - warn!("Failed to parse Request bytes | {e:#}"); - return; - } - }; + }: WithChannel = + match message.decode().and_then(RequestMessage::map_decode) { + Ok(res) => res, + Err(e) => { + // TODO: handle: + warn!("Failed to parse Request bytes | {e:#}"); + return; + } + }; let resolve_response = async { let response = match request.resolve(&args).await { diff --git a/client/core/rs/src/api/read/build.rs b/client/core/rs/src/api/read/build.rs index feaa6985d..23c3b1aec 100644 --- a/client/core/rs/src/api/read/build.rs +++ b/client/core/rs/src/api/read/build.rs @@ -236,4 +236,3 @@ pub struct ListCommonBuildExtraArgs { #[typeshare] pub type ListCommonBuildExtraArgsResponse = Vec; - diff --git a/client/periphery/rs/src/transport/login.rs b/client/periphery/rs/src/transport/login.rs index 9d943d48e..5a966ad27 100644 --- a/client/periphery/rs/src/transport/login.rs +++ b/client/periphery/rs/src/transport/login.rs @@ -1,51 +1,26 @@ use anyhow::{Context, anyhow}; use derive_variants::{EnumVariants, ExtractVariant}; -use encoding::{CastBytes, Decode, Encode, EncodedResult}; +use encoding::{ + CastBytes, Decode, Encode, EncodedResult, impl_cast_bytes_vec, +}; use noise::key::SpkiPublicKey; -use crate::transport::{ - DecodedTransportMessage, EncodedTransportMessage, -}; - -impl Encode for LoginMessage { - fn encode(self) -> EncodedTransportMessage { - DecodedTransportMessage::Login(Ok(self)).encode() - } -} +use crate::transport::{EncodedTransportMessage, TransportMessage}; #[derive(Debug)] pub struct EncodedLoginMessage( EncodedResult, ); -impl CastBytes for EncodedLoginMessage { - fn from_vec(bytes: Vec) -> Self { - Self(EncodedResult::from_vec(bytes)) - } - fn into_vec(self) -> Vec { - self.0.into_vec() - } -} - -impl Encode for anyhow::Result { - fn encode(self) -> EncodedLoginMessage { - EncodedLoginMessage(encoding::Result::from(self).map_encode()) - } -} - -impl Encode - for EncodedResult +impl From> + for EncodedLoginMessage { - fn encode(self) -> EncodedLoginMessage { - EncodedLoginMessage(self) + fn from(value: EncodedResult) -> Self { + Self(value) } } -impl Decode for EncodedLoginMessage { - fn decode(self) -> anyhow::Result { - self.0.decode_into() - } -} +impl_cast_bytes_vec!(EncodedLoginMessage, EncodedResult); /// ```markdown /// | -- u8[] -- | --------- u8 ------------ | @@ -54,14 +29,7 @@ impl Decode for EncodedLoginMessage { #[derive(Clone, Debug)] pub struct InnerEncodedLoginMessage(Vec); -impl CastBytes for InnerEncodedLoginMessage { - fn from_vec(vec: Vec) -> Self { - Self(vec) - } - fn into_vec(self) -> Vec { - self.0 - } -} +impl_cast_bytes_vec!(InnerEncodedLoginMessage, Vec); #[derive(EnumVariants, Clone)] #[variant_derive(Debug, Clone, Copy)] @@ -93,8 +61,8 @@ pub enum LoginMessage { V1Passkey(Vec), } -impl Encode for LoginMessage { - fn encode(self) -> InnerEncodedLoginMessage { +impl Encode for LoginMessage { + fn encode(self) -> EncodedTransportMessage { let variant_byte = self.extract_variant().as_byte(); let mut bytes = match self { LoginMessage::Success => Vec::new(), @@ -114,29 +82,33 @@ impl Encode for LoginMessage { LoginMessage::V1Passkey(bytes) => bytes, }; bytes.push(variant_byte); - InnerEncodedLoginMessage(bytes) + let inner = InnerEncodedLoginMessage(bytes); + let res = Ok(inner).encode(); + TransportMessage::Login(EncodedLoginMessage(res)).encode() } } -impl Decode for InnerEncodedLoginMessage { - /// Parses login messages, performing various validations. +impl Decode for EncodedLoginMessage { fn decode(self) -> anyhow::Result { - let mut bytes = self.0; + let mut bytes = self.0.decode()?.into_vec(); + let variant_byte = bytes .pop() .context("Failed to parse login message | Bytes are empty")?; - let variant = LoginMessageVariant::from_byte(variant_byte)?; - use LoginMessageVariant::*; - match variant { - Success => Ok(LoginMessage::Success), - Nonce => Ok(LoginMessage::Nonce( + let variant = LoginMessageVariant::from_byte(variant_byte)?; + + use LoginMessageVariant::*; + let message = match variant { + Success => LoginMessage::Success, + + Nonce => LoginMessage::Nonce( bytes .try_into() .map_err(|_| anyhow!("Invalid connection nonce"))?, - )), + ), - Handshake => Ok(LoginMessage::Handshake(bytes)), + Handshake => LoginMessage::Handshake(bytes), OnboardingFlow => { let onboarding_flow = match bytes.as_slice() { @@ -148,7 +120,7 @@ impl Decode for InnerEncodedLoginMessage { )); } }; - Ok(LoginMessage::OnboardingFlow(onboarding_flow)) + LoginMessage::OnboardingFlow(onboarding_flow) } PublicKey => { @@ -159,7 +131,7 @@ impl Decode for InnerEncodedLoginMessage { } let public_key = String::from_utf8(bytes) .context("Public key is not valid utf-8")?; - Ok(LoginMessage::PublicKey(SpkiPublicKey::from(public_key))) + LoginMessage::PublicKey(SpkiPublicKey::from(public_key)) } // V1 @@ -173,7 +145,7 @@ impl Decode for InnerEncodedLoginMessage { )); } }; - Ok(LoginMessage::V1PasskeyFlow(passkey_login)) + LoginMessage::V1PasskeyFlow(passkey_login) } V1Passkey => { @@ -182,9 +154,11 @@ impl Decode for InnerEncodedLoginMessage { "Got empty LoginMessage V1Passkey bytes" )); } - Ok(LoginMessage::V1Passkey(bytes)) + LoginMessage::V1Passkey(bytes) } - } + }; + + Ok(message) } } diff --git a/client/periphery/rs/src/transport/mod.rs b/client/periphery/rs/src/transport/mod.rs index 3af2c1db9..b8cff26a1 100644 --- a/client/periphery/rs/src/transport/mod.rs +++ b/client/periphery/rs/src/transport/mod.rs @@ -2,29 +2,23 @@ use anyhow::{Context as _, anyhow}; use derive_variants::{EnumVariants, ExtractVariant as _}; use encoding::{ CastBytes, Decode, Encode, EncodedChannel, EncodedJsonMessage, - EncodedOption, EncodedResult, WithChannel, + EncodedOption, EncodedResult, WithChannel, impl_cast_bytes_vec, }; mod login; pub use login::*; use serde::de::DeserializeOwned; - -#[derive(Debug, Clone)] -pub struct EncodedTransportMessage(Vec); - -impl CastBytes for EncodedTransportMessage { - fn from_vec(vec: Vec) -> Self { - Self(vec) - } - fn into_vec(self) -> Vec { - self.0 - } -} +use uuid::Uuid; // =================== // TRANSPORT MESSAGE // =================== +#[derive(Debug, Clone)] +pub struct EncodedTransportMessage(Vec); + +impl_cast_bytes_vec!(EncodedTransportMessage, Vec); + /// When an EncodedTransportMessage is received, /// it is decoded into this type. /// @@ -76,7 +70,7 @@ impl Decode for EncodedTransportMessage { EncodedTerminalMessage(EncodedChannel::from_vec(bytes)), ), }; - Ok(message.into()) + Ok(message) } } @@ -108,72 +102,6 @@ impl TransportMessageVariant { } } -// =================== -// DECODED TRANSPORT -// =================== - -pub enum DecodedTransportMessage { - Login(anyhow::Result), - Request(WithChannel), - Response( - WithChannel>>, // EncodedChannel>>, - ), - Terminal(WithChannel>), -} - -impl Encode for DecodedTransportMessage { - fn encode(self) -> TransportMessage { - use DecodedTransportMessage::*; - match self { - Login(res) => TransportMessage::Login(res.encode()), - Request(data) => TransportMessage::Request( - EncodedRequestMessage(data.encode()), - ), - Response(data) => { - TransportMessage::Response(EncodedResponseMessage( - data - .map(|data| encoding::Option::from(data).map_encode()) - .encode(), - )) - } - Terminal(data) => TransportMessage::Terminal( - EncodedTerminalMessage(data.encode()), - ), - } - } -} - -impl Encode for DecodedTransportMessage { - fn encode(self) -> EncodedTransportMessage { - let res: TransportMessage = self.encode(); - res.encode() - } -} - -impl> Decode - for TransportMessage -{ - fn decode(self) -> anyhow::Result { - let message = match self { - TransportMessage::Login(encoded_result) => { - DecodedTransportMessage::Login(encoded_result.decode()) - } - TransportMessage::Request(encoded_channel) => { - DecodedTransportMessage::Request(encoded_channel.0.decode()?) - } - TransportMessage::Response(encoded_channel) => { - DecodedTransportMessage::Response( - encoded_channel.0.decode()?.map(|data| data.decode_map()), - ) - } - TransportMessage::Terminal(encoded_channel) => { - DecodedTransportMessage::Terminal(encoded_channel.0.decode()?) - } - }; - Ok(message.into()) - } -} - // ================= // REQUEST MESSAGE // ================= @@ -181,30 +109,37 @@ impl> Decode #[derive(Debug)] pub struct EncodedRequestMessage(EncodedChannel); -impl CastBytes for EncodedRequestMessage { - fn from_vec(bytes: Vec) -> Self { - Self(EncodedChannel::from_vec(bytes)) - } - fn into_vec(self) -> Vec { - self.0.into_vec() +impl_cast_bytes_vec!(EncodedRequestMessage, EncodedChannel); + +pub struct RequestMessage(WithChannel); + +impl RequestMessage { + pub fn map_decode( + self, + ) -> anyhow::Result> { + self.0.map_decode() } } -impl Encode - for WithChannel -{ - fn encode(self) -> EncodedRequestMessage { - EncodedRequestMessage(self.encode()) +impl RequestMessage { + pub fn new(channel: Uuid, json: EncodedJsonMessage) -> Self { + Self(WithChannel { + channel, + data: json, + }) } } -impl Decode> - for EncodedRequestMessage -{ - fn decode(self) -> anyhow::Result> { - let WithChannel { channel, data } = self.0.decode()?; - let data = data.decode()?; - Ok(WithChannel { channel, data }) +impl Encode for RequestMessage { + fn encode(self) -> EncodedTransportMessage { + TransportMessage::Request(EncodedRequestMessage(self.0.encode())) + .encode() + } +} + +impl Decode for EncodedRequestMessage { + fn decode(self) -> anyhow::Result { + self.0.decode().map(RequestMessage) } } @@ -214,74 +149,63 @@ impl Decode> #[derive(Debug)] pub struct EncodedResponseMessage( - EncodedChannel>>, + EncodedChannel, ); -impl CastBytes for EncodedResponseMessage { - fn from_vec(bytes: Vec) -> Self { - Self(EncodedChannel::from_vec(bytes)) - } - fn into_vec(self) -> Vec { - self.0.into_vec() - } -} - -impl Encode - for WithChannel< - encoding::Option>, - > -{ - fn encode(self) -> EncodedResponseMessage { - EncodedResponseMessage( - self.map(|data| data.map_encode()).encode(), - ) - } -} - -impl Encode - for WithChannel> -{ - fn encode(self) -> EncodedResponseMessage { - EncodedResponseMessage(self.map(Some).map_encode()) - } -} - -impl - Decode>>> +impl From> for EncodedResponseMessage { - fn decode( - self, - ) -> anyhow::Result>>> { - let WithChannel { channel, data } = self.0.decode()?; - let data = data.decode()?.into_std(); - // Does it this way to make sure the inner result is preserved - // for later handling. - let Some(data) = data else { - return Ok(WithChannel { - channel, - data: None, - }); - }; - let data = data.decode().and_then(|data| data.decode()); - Ok(WithChannel { - channel, - data: Some(data), - }) + fn from( + value: EncodedChannel, + ) -> Self { + Self(value) } } -impl - Decode< - WithChannel>>, - > for EncodedResponseMessage -{ - fn decode( +impl_cast_bytes_vec!(EncodedResponseMessage, EncodedChannel); + +/// This is handled as one bundle in main handler, +/// and passed to response handler for parsing. +#[derive(Debug)] +pub struct InnerEncodedResponseMessage( + EncodedOption>, +); + +impl_cast_bytes_vec!(InnerEncodedResponseMessage, EncodedOption); + +pub struct ResponseMessage(WithChannel); + +impl ResponseMessage { + pub fn new( + channel: Uuid, + response: Option>, + ) -> Self { + Self(WithChannel { + channel, + data: InnerEncodedResponseMessage(response.encode()), + }) + } + + pub fn extract( self, - ) -> anyhow::Result< - WithChannel>>, - > { - self.0.decode() + ) -> WithChannel>> + { + self.0.map(|data| data.0) + } +} + +impl Encode for ResponseMessage { + fn encode(self) -> EncodedTransportMessage { + TransportMessage::Response(EncodedResponseMessage( + self.0.encode(), + )) + .encode() + } +} + +impl Decode for EncodedResponseMessage { + fn decode(self) -> anyhow::Result { + self.0.decode().map(ResponseMessage) } } @@ -292,18 +216,25 @@ impl #[derive(Debug)] pub struct EncodedTerminalMessage(EncodedChannel>); -impl CastBytes for EncodedTerminalMessage { - fn from_vec(bytes: Vec) -> Self { - Self(EncodedChannel::from_vec(bytes)) - } - fn into_vec(self) -> Vec { - self.0.into_vec() +impl TerminalMessage { + pub fn new(channel: Uuid, bytes: Vec) -> Self { + Self(WithChannel { + channel, + data: bytes, + }) } } -impl Encode for WithChannel> { - fn encode(self) -> EncodedTerminalMessage { - EncodedTerminalMessage(self.encode()) +impl_cast_bytes_vec!(EncodedTerminalMessage, EncodedChannel); + +pub struct TerminalMessage(WithChannel>); + +impl Encode for TerminalMessage { + fn encode(self) -> EncodedTransportMessage { + TransportMessage::Terminal(EncodedTerminalMessage( + self.0.encode(), + )) + .encode() } } diff --git a/lib/encoding/src/lib.rs b/lib/encoding/src/lib.rs index d3048e0a5..256b2a6f1 100644 --- a/lib/encoding/src/lib.rs +++ b/lib/encoding/src/lib.rs @@ -35,9 +35,6 @@ pub trait Decode: Sized { impl_identity!(Bytes); impl_identity!(Vec); -impl_identity_into!(Bytes, Vec); -impl_identity_into!(Vec, Bytes); - /// Helps cast between the top level message types. /// Implement whichever ones are most convenient for the source type. pub trait CastBytes: Sized { @@ -73,8 +70,6 @@ impl CastBytes for Vec { } } -// Encode is basically 'Into', -// while decode is 'TryInto'. #[macro_export] macro_rules! impl_identity { ($typ:ty) => { @@ -91,23 +86,6 @@ macro_rules! impl_identity { }; } -/// impl where a: Into -#[macro_export] -macro_rules! impl_identity_into { - ($a:ty, $b:ty) => { - impl Encode<$b> for $a { - fn encode(self) -> $b { - self.into() - } - } - impl Decode<$b> for $a { - fn decode(self) -> anyhow::Result<$b> { - Ok(self.into()) - } - } - }; -} - #[macro_export] macro_rules! impl_wrapper { ($struct:ident) => { @@ -116,7 +94,6 @@ macro_rules! impl_wrapper { Self(value) } } - impl CastBytes for $struct { fn from_bytes(bytes: Bytes) -> Self { Self(T::from_bytes(bytes)) @@ -133,3 +110,17 @@ macro_rules! impl_wrapper { } }; } + +#[macro_export] +macro_rules! impl_cast_bytes_vec { + ($typ:ty, $through:ident) => { + impl CastBytes for $typ { + fn from_vec(bytes: Vec) -> Self { + Self($through::from_vec(bytes)) + } + fn into_vec(self) -> Vec { + self.0.into_vec() + } + } + }; +} diff --git a/lib/transport/src/channel.rs b/lib/transport/src/channel.rs index 7ce650872..1980e7bab 100644 --- a/lib/transport/src/channel.rs +++ b/lib/transport/src/channel.rs @@ -1,10 +1,11 @@ use anyhow::{Context, anyhow}; use encoding::{ - Encode, EncodedJsonMessage, EncodedResult, JsonMessage, WithChannel, + Encode, EncodedJsonMessage, EncodedResult, JsonMessage, }; use futures_util::FutureExt; use periphery_client::transport::{ - DecodedTransportMessage, EncodedTransportMessage, TransportMessage, + EncodedTransportMessage, RequestMessage, ResponseMessage, + TerminalMessage, }; use serde::Serialize; use tokio::sync::{Mutex, MutexGuard, mpsc}; @@ -86,21 +87,15 @@ impl Sender { where &'a T: Send, { - let data = JsonMessage(request).encode().into_anyhow()?; - let message = - DecodedTransportMessage::Request(WithChannel { channel, data }); - self.send_message(message).await + let json = JsonMessage(request).encode().into_anyhow()?; + self.send_message(RequestMessage::new(channel, json)).await } pub async fn send_in_progress( &self, channel: Uuid, ) -> anyhow::Result<()> { - let message = DecodedTransportMessage::Response(WithChannel { - channel, - data: None, - }); - self.send_message(message).await + self.send_message(ResponseMessage::new(channel, None)).await } pub async fn send_response( @@ -108,14 +103,9 @@ impl Sender { channel: Uuid, response: EncodedResult, ) -> anyhow::Result<()> { - let message = TransportMessage::Response( - WithChannel { - channel, - data: response, - } - .encode(), - ); - self.send_message(message).await + self + .send_message(ResponseMessage::new(channel, Some(response))) + .await } pub async fn send_terminal( @@ -123,11 +113,9 @@ impl Sender { channel: Uuid, data: impl Into>, ) -> anyhow::Result<()> { - let message = DecodedTransportMessage::Terminal(WithChannel { - channel, - data: data.into(), - }); - self.send_message(message).await + self + .send_message(TerminalMessage::new(channel, data.into())) + .await } } diff --git a/lib/transport/src/websocket/login.rs b/lib/transport/src/websocket/login.rs index b107382d4..e973ef891 100644 --- a/lib/transport/src/websocket/login.rs +++ b/lib/transport/src/websocket/login.rs @@ -1,7 +1,9 @@ use anyhow::anyhow; use encoding::{Decode as _, Encode as _}; use noise::key::SpkiPublicKey; -use periphery_client::transport::{LoginMessage, TransportMessage}; +use periphery_client::transport::{ + EncodedLoginMessage, LoginMessage, TransportMessage, +}; use crate::{ auth::AUTH_TIMEOUT, @@ -13,7 +15,8 @@ pub trait LoginWebsocketExt: WebsocketExt { &mut self, e: &anyhow::Error, ) -> impl Future> + Send { - let message = TransportMessage::Login(e.encode_into()); + let message = + TransportMessage::Login(EncodedLoginMessage::from(e.encode())); self.send(message) } diff --git a/lib/transport/src/websocket/mod.rs b/lib/transport/src/websocket/mod.rs index 627a24f44..7d527806f 100644 --- a/lib/transport/src/websocket/mod.rs +++ b/lib/transport/src/websocket/mod.rs @@ -4,10 +4,11 @@ use anyhow::{Context, anyhow}; use bytes::Bytes; use encoding::{ CastBytes as _, Decode as _, Encode, EncodedJsonMessage, - EncodedResult, JsonMessage, WithChannel, + EncodedResult, JsonMessage, }; use periphery_client::transport::{ - DecodedTransportMessage, EncodedTransportMessage, TransportMessage, + EncodedTransportMessage, RequestMessage, ResponseMessage, + TerminalMessage, TransportMessage, }; use serde::Serialize; use tokio_util::sync::CancellationToken; @@ -103,7 +104,7 @@ pub trait WebsocketSender { } pub trait WebsocketSenderExt: WebsocketSender + Send { - fn send( + fn send_message( &mut self, message: impl Encode, ) -> impl Future> + Send { @@ -119,12 +120,8 @@ pub trait WebsocketSenderExt: WebsocketSender + Send { &'a T: Send, { async move { - let data = JsonMessage(request).encode().into_anyhow()?; - let message = DecodedTransportMessage::Request(WithChannel { - channel, - data, - }); - self.send(message).await + let json = JsonMessage(request).encode().into_anyhow()?; + self.send_message(RequestMessage::new(channel, json)).await } } @@ -132,11 +129,7 @@ pub trait WebsocketSenderExt: WebsocketSender + Send { &mut self, channel: Uuid, ) -> impl Future> + Send { - let message = DecodedTransportMessage::Response(WithChannel { - channel, - data: None, - }); - self.send(message) + self.send_message(ResponseMessage::new(channel, None)) } fn send_response( @@ -144,14 +137,7 @@ pub trait WebsocketSenderExt: WebsocketSender + Send { channel: Uuid, response: EncodedResult, ) -> impl Future> + Send { - let message = TransportMessage::Response( - WithChannel { - channel, - data: response, - } - .encode(), - ); - self.send(message) + self.send_message(ResponseMessage::new(channel, Some(response))) } fn send_terminal( @@ -159,11 +145,7 @@ pub trait WebsocketSenderExt: WebsocketSender + Send { channel: Uuid, data: impl Into>, ) -> impl Future> + Send { - let message = DecodedTransportMessage::Terminal(WithChannel { - channel, - data: data.into(), - }); - self.send(message) + self.send_message(TerminalMessage::new(channel, data.into())) } }