prog

Cargo.lock

@@ -4738,9 +4738,9 @@ dependencies = [
 
 [[package]]
 name = "serror"
-version = "0.5.1"
+version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bee3f9a0b6dee3c0e8c581dddc560f66f20a1c5c20f6609e827167265ac84eb6"
+checksum = "24c14f04906667a9c590474a470a92fa2c76189d4afe186017329e9cb22f1d6a"
 dependencies = [
  "anyhow",
  "axum",

Cargo.toml

@@ -37,7 +37,7 @@ git = { path = "lib/git" }
 
 # MOGH
 run_command = { version = "0.0.6", features = ["async_tokio"] }
-serror = { version = "0.5.1", default-features = false }
+serror = { version = "0.5.2", default-features = false }
 slack = { version = "0.4.0", package = "slack_client_rs", default-features = false, features = ["rustls"] }
 derive_default_builder = "0.1.8"
 derive_empty_traits = "0.1.0"

bin/periphery/src/connection.rs

@@ -8,13 +8,13 @@ use axum::{
 use bytes::Bytes;
 use futures::{SinkExt, StreamExt};
 use periphery_client::message::{
-  MessageState, from_transport_bytes, to_transport_bytes,
+  MessageState, parse_from_transport_bytes, to_transport_bytes,
 };
 use resolver_api::Resolve;
 use response::JsonBytes;
 use serror::{AddStatusCode, serialize_error_bytes};
 use tokio::sync::{Mutex, mpsc::Sender};
-use transport::{BufferedReceiver, buffered_channel};
+use transport::channel::{BufferedReceiver, buffered_channel};
 
 use crate::api::PeripheryRequest;
 
@@ -90,7 +90,7 @@ pub async fn inbound_connection(
 fn handle_msg(msg: Bytes) {
   tokio::spawn(async move {
     let (id, request) =
-      match from_transport_bytes::<PeripheryRequest>(&msg) {
+      match parse_from_transport_bytes::<PeripheryRequest>(&msg) {
         Ok((id, _, Some(request))) => (id, request),
         Ok(_) => {
           // No data, ignore
@@ -151,9 +151,9 @@ fn response_sender() -> &'static Sender<Vec<u8>> {
     .expect("response_sender accessed before initialized")
 }
 
-static RESPONSE_RECEIVER: OnceLock<Mutex<BufferedReceiver>> =
+static RESPONSE_RECEIVER: OnceLock<Mutex<BufferedReceiver<Vec<u8>>>> =
   OnceLock::new();
-fn response_receiver() -> &'static Mutex<BufferedReceiver> {
+fn response_receiver() -> &'static Mutex<BufferedReceiver<Vec<u8>>> {
   RESPONSE_RECEIVER
     .get()
     .expect("response_receiver accessed before initialized")

client/periphery/rs/src/connection.rs

@@ -0,0 +1,227 @@
+use std::{sync::OnceLock, time::Duration};
+
+use anyhow::{Context, anyhow};
+use cache::CloneCache;
+use futures_util::{SinkExt, StreamExt};
+use serror::deserialize_error_bytes;
+use tokio::sync::mpsc;
+use tokio_tungstenite::tungstenite::Message;
+use tokio_util::{bytes::Bytes, sync::CancellationToken};
+use tracing::warn;
+use transport::channel::{BufferedReceiver, buffered_channel};
+use uuid::Uuid;
+
+use crate::message::{
+  MessageState, from_transport_bytes, id_from_transport_bytes,
+  to_transport_bytes,
+};
+
+#[derive(Debug)]
+struct Connection {
+  cancel: CancellationToken,
+  request_sender: mpsc::Sender<Vec<u8>>,
+}
+
+impl Connection {
+  fn new()
+  -> (Connection, BufferedReceiver<Vec<u8>>, CancellationToken) {
+    let (request_sender, request_receiver) = buffered_channel(1000);
+    let cancel = CancellationToken::new();
+    (
+      Connection {
+        cancel: cancel.clone(),
+        request_sender,
+      },
+      request_receiver,
+      cancel,
+    )
+  }
+}
+
+impl Clone for Connection {
+  fn clone(&self) -> Self {
+    Self {
+      cancel: self.cancel.clone(),
+      request_sender: self.request_sender.clone(),
+    }
+  }
+}
+
+/// server id => connection
+type Connections = CloneCache<String, Connection>;
+fn connections() -> &'static Connections {
+  static CONNECTIONS: OnceLock<Connections> = OnceLock::new();
+  CONNECTIONS.get_or_init(Default::default)
+}
+
+type ResponseChannels = CloneCache<Uuid, mpsc::Sender<Vec<u8>>>;
+fn response_channels() -> &'static ResponseChannels {
+  static RESPONSE_CHANNELS: OnceLock<ResponseChannels> =
+    OnceLock::new();
+  RESPONSE_CHANNELS.get_or_init(Default::default)
+}
+
+pub async fn send_request(
+  server_id: &String,
+  data: &[u8],
+) -> anyhow::Result<Vec<u8>> {
+  let request_sender = &connections()
+    .get(server_id)
+    .await
+    .with_context(|| {
+      format!("No connection found for server {server_id}")
+    })?
+    .request_sender;
+  let id = Uuid::new_v4();
+  let (response_sender, mut response_receiever) = mpsc::channel(1000);
+  response_channels().insert(id, response_sender).await;
+  if let Err(e) = request_sender
+    .send(to_transport_bytes(id, MessageState::InProgress, data))
+    .await
+    .context("Failed to send request over channel")
+  {
+    // cleanup
+    response_channels().remove(&id).await;
+    return Err(e);
+  }
+
+  // Poll for the response
+  loop {
+    let bytes = match response_receiever.recv().await {
+      Some(bytes) => bytes,
+      None => {
+        return Err(anyhow!(
+          "Sender dropped before response was recieved"
+        ));
+      }
+    };
+    let (state, data) = match from_transport_bytes(&bytes) {
+      Ok((_, state, Some(data))) => (state, data),
+      // TODO: Handle no data cases
+      Ok(_) => continue,
+      Err(e) => {
+        warn!("Received invalid message | {e:#}");
+        continue;
+      }
+    };
+    match state {
+      // TODO: improve the allocation in .to_vec
+      MessageState::Successful => return Ok(data.to_vec()),
+      MessageState::Failed => {
+        return Err(deserialize_error_bytes(&data));
+      }
+      // TODO: bail if this isn't received frequent enough
+      MessageState::InProgress => continue,
+    }
+  }
+}
+
+// Assumes address already wss formatted
+pub async fn spawn_connection(
+  server_id: String,
+  address: String,
+) -> anyhow::Result<()> {
+  let response_channels = response_channels();
+  let (channel, mut request_receiver, cancel) = Connection::new();
+  if let Some(existing) =
+    connections().insert(server_id, channel).await
+  {
+    existing.cancel.cancel();
+  }
+
+  tokio::spawn(async move {
+    let connection = async {
+      // Outer connection loop
+      loop {
+        let socket =
+          match crate::ws::connect_websocket(&address).await {
+            Ok(socket) => socket,
+            Err(e) => {
+              warn!("{e:#}");
+              tokio::time::sleep(Duration::from_secs(5)).await;
+              continue;
+            }
+          };
+
+        let (mut ws_write, mut ws_read) = socket.split();
+
+        let forward_requests = async {
+          loop {
+            let message = match request_receiver.recv().await {
+              None => 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;
+              }
+            }
+          }
+        };
+
+        let read_responses = async {
+          loop {
+            let bytes = match ws_read.next().await {
+              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 {
+              warn!(
+                "Failed to send response for {address} | No response channel found"
+              );
+              continue;
+            };
+            if let Err(e) = channel.send(bytes.into()).await {
+              warn!(
+                "Failed to send response for {address} | Channel failure | {e:#}"
+              );
+            }
+          }
+        };
+
+        tokio::select! {
+          _ = forward_requests => {},
+          _ = read_responses => {}
+        };
+      }
+    };
+
+    tokio::select! {
+      _ = connection => {},
+      _ = cancel.cancelled() => {}
+    }
+  });
+
+  Ok(())
+}

client/periphery/rs/src/connections.rs

@@ -1,91 +0,0 @@
-use std::sync::OnceLock;
-
-use cache::CloneCache;
-use futures_util::StreamExt;
-use tokio::sync::{broadcast, mpsc};
-use tokio_util::sync::CancellationToken;
-use transport::{BufferedReceiver, buffered_channel};
-
-#[derive(Debug)]
-struct ConnectionChannel {
-  request_sender: mpsc::Sender<Vec<u8>>,
-  response_receiver: broadcast::Receiver<Vec<u8>>,
-  cancel: CancellationToken,
-}
-
-impl ConnectionChannel {
-  fn new() -> (
-    ConnectionChannel,
-    BufferedReceiver,
-    broadcast::Sender<Vec<u8>>,
-    CancellationToken,
-  ) {
-    let (request_sender, request_receiver) = buffered_channel(1000);
-    let (response_sender, response_receiver) =
-      broadcast::channel(1000);
-    let cancel = CancellationToken::new();
-    (
-      ConnectionChannel {
-        request_sender,
-        response_receiver,
-        cancel: cancel.clone(),
-      },
-      request_receiver,
-      response_sender,
-      cancel,
-    )
-  }
-}
-
-impl Clone for ConnectionChannel {
-  fn clone(&self) -> Self {
-    Self {
-      request_sender: self.request_sender.clone(),
-      response_receiver: self.response_receiver.resubscribe(),
-      cancel: self.cancel.clone(),
-    }
-  }
-}
-
-type ConnectionCache = CloneCache<String, ConnectionChannel>;
-fn connections() -> &'static ConnectionCache {
-  static CONNECTIONS: OnceLock<ConnectionCache> = OnceLock::new();
-  CONNECTIONS.get_or_init(Default::default)
-}
-
-// Assumes address already wss formatted
-async fn spawn_connection(address: String) -> anyhow::Result<()> {
-  let (channel, request_receiver, response_sender, cancel) =
-    ConnectionChannel::new();
-  if let Some(existing) =
-    connections().insert(address.clone(), channel).await
-  {
-    existing.cancel.cancel();
-  }
-
-  tokio::spawn(async move {
-    // Outer connection loop
-    loop {
-      let socket = match crate::ws::connect_websocket(&address).await
-      {
-        Ok(socket) => socket,
-        Err(e) => {
-          // TODO: handle connect error
-          return;
-        }
-      };
-      let (mut ws_read, mut ws_write) = socket.split();
-
-      let forward_requests = async {
-        loop {
-          // match request_receiver.recv().await {
-          //   None => break,
-
-          // }
-        }
-      };
-    }
-  });
-
-  Ok(())
-}

client/periphery/rs/src/lib.rs

@@ -7,7 +7,7 @@ use serde::{Serialize, de::DeserializeOwned};
 use serde_json::json;
 
 pub mod api;
-pub mod connections;
+pub mod connection;
 pub mod message;
 
 mod terminal;

client/periphery/rs/src/message.rs

@@ -19,19 +19,37 @@ pub fn to_transport_bytes(
   res
 }
 
+/// Deserializes channel id from
+/// incoming transport bytes.
+pub fn id_from_transport_bytes(bytes: &[u8]) -> anyhow::Result<Uuid> {
+  if bytes.len() < 16 {
+    return Err(anyhow!("Transport bytes too short to include uuid"));
+  }
+  Uuid::from_slice(&bytes[..16]).context("Invalid Uuid bytes")
+}
+
 /// Deserializes channel id + data from
 /// incoming transport bytes.
-pub fn from_transport_bytes<T: DeserializeOwned>(
+pub fn from_transport_bytes(
   bytes: &[u8],
-) -> anyhow::Result<(Uuid, MessageState, Option<T>)> {
+) -> anyhow::Result<(Uuid, MessageState, Option<&[u8]>)> {
   if bytes.len() < 17 {
     return Err(anyhow!(
-      "Transport bytes too short to include uuid at state"
+      "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))
+}
+
+/// Deserializes channel id + data from
+/// incoming transport bytes.
+pub fn parse_from_transport_bytes<T: DeserializeOwned>(
+  bytes: &[u8],
+) -> anyhow::Result<(Uuid, MessageState, Option<T>)> {
+  let (id, state, data) = from_transport_bytes(bytes)?;
   let data = data
     .map(|data| {
       serde_json::from_slice::<T>(data)

lib/cache/src/lib.rs

@@ -59,6 +59,7 @@ fn clone_anyhow_error(e: &anyhow::Error) -> anyhow::Error {
   e
 }
 
+#[derive(Debug)]
 pub struct CloneCache<K: PartialEq + Eq + Hash, T: Clone> {
   cache: RwLock<HashMap<K, T>>,
 }

lib/transport/src/channel.rs

@@ -0,0 +1,48 @@
+use std::ops::Deref;
+
+use tokio::sync::mpsc;
+
+/// Create a buffered channel
+pub fn buffered_channel<T: Deref>(
+  buffer: usize,
+) -> (mpsc::Sender<T>, BufferedReceiver<T>) {
+  let (sender, receiver) = mpsc::channel(buffer);
+  (sender, BufferedReceiver::new(receiver))
+}
+
+/// Wrapper around channel receiver to control when
+/// the latest message is dropped,
+/// in case it must be re-transmitted.
+#[derive(Debug)]
+pub struct BufferedReceiver<T> {
+  receiver: mpsc::Receiver<T>,
+  buffer: Option<T>,
+}
+
+impl<T: Deref> BufferedReceiver<T> {
+  pub fn new(receiver: mpsc::Receiver<T>) -> BufferedReceiver<T> {
+    BufferedReceiver {
+      receiver,
+      buffer: None,
+    }
+  }
+
+  /// If 'next: Some(bytes)':
+  ///   - Immediately returns borrow of next.
+  /// Else:
+  ///   - Wait for next item
+  ///   - store in 'next'
+  ///   - return borrow of next.
+  pub async fn recv(&mut self) -> Option<&<T as Deref>::Target> {
+    if self.buffer.is_none() {
+      self.buffer = Some(self.receiver.recv().await?);
+    }
+    self.buffer.as_deref()
+  }
+
+  /// Clears buffer.
+  /// Should be called after transmission confirmed.
+  pub fn clear_buffer(&mut self) {
+    self.buffer = None;
+  }
+}

lib/transport/src/lib.rs

@@ -1,46 +1 @@
-use tokio::sync::mpsc;
-
-/// Create a buffered channel
-pub fn buffered_channel(
-  buffer: usize,
-) -> (mpsc::Sender<Vec<u8>>, BufferedReceiver) {
-  let (sender, receiver) = mpsc::channel(buffer);
-  (sender, BufferedReceiver::new(receiver))
-}
-
-/// Wrapper around channel receiver to control when
-/// the latest message is dropped,
-/// in case it must be re-transmitted.
-#[derive(Debug)]
-pub struct BufferedReceiver {
-  receiver: mpsc::Receiver<Vec<u8>>,
-  buffer: Option<Vec<u8>>,
-}
-
-impl BufferedReceiver {
-  pub fn new(receiver: mpsc::Receiver<Vec<u8>>) -> BufferedReceiver {
-    BufferedReceiver {
-      receiver,
-      buffer: None,
-    }
-  }
-
-  /// If 'next: Some(bytes)':
-  ///   - Immediately returns borrow of next.
-  /// Else:
-  ///   - Wait for next item
-  ///   - store in 'next'
-  ///   - return borrow of next.
-  pub async fn recv(&mut self) -> Option<&[u8]> {
-    if self.buffer.is_none() {
-      self.buffer = Some(self.receiver.recv().await?);
-    }
-    self.buffer.as_deref()
-  }
-
-  /// Clears buffer.
-  /// Should be called after transmission confirmed.
-  pub fn clear_buffer(&mut self) {
-    self.buffer = None;
-  }
-}
+pub mod channel;