Merge pull request #6 from AxiomChat/node-system

Node system

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
-name = "voxa-server"
-version = "0.1.0-beta2"
+name = "axiom-server"
+version = "0.1.1-beta1"
 edition = "2024"
 
 [dependencies]

README.md

@@ -1,10 +1,10 @@
-# voxa-server
+# axiom-server
 
-Self host your voxa community server with customizable plugins
+Self host your axiom community server with customizable plugins
 
 # Protocol
 
-1. When the client is running in the background or windowed, it's connected to the [Voxa Cloud server](#voxa-cloud)
+1. When the client is running in the background or windowed, it's connected to the [Axiom Cloud server](#axiom-cloud)
 
 ## Handshake
 
@@ -28,6 +28,6 @@ Server -> Cloud: `POST /message { content: <Message>, author: <User-Id> }`
 
 Cloud -> Client(s): `{ Message: { content: <Message>, author: <User-Id> } }`
 
-# Voxa Cloud
+# Axiom Cloud
 
-The voxa cloud server is the main auth and notification handler.
+The axiom cloud server is the main auth and notification handler.

src/main.rs

@@ -1,5 +1,6 @@
 mod cli;
 mod macros;
+mod node_requests;
 mod plugin;
 mod requests;
 mod server;
@@ -22,6 +23,17 @@ fn main() -> Result<()> {
     } else {
         vfs::read_config(&root.join("config.json"))?
     };
+
+    if let Ok(a) = std::env::var("AXIOM_NODE") {
+        if a == "true" {
+            config
+                .build_req(&root, crate::server::Server::call_node_request)
+                .run()?;
+
+            return Ok(());
+        }
+    }
+
     config.build(&root).run()?;
     Ok(())
 }

src/node_requests/chunk.rs

@@ -0,0 +1,18 @@
+use crate::{server::Server, types::message::ServerMessage, utils::client::Client};
+use std::sync::Arc;
+
+crate::logger!(LOGGER "Chunk Loader");
+
+pub fn load_chunk(
+    server: &Arc<Server>,
+    client: &Client,
+    channel_id: &str,
+    chunk_id: usize,
+) -> crate::Result<()> {
+    let mut chunk = server
+        .db
+        .get_chunk_node(&client.get_uuid()?, channel_id, chunk_id)?;
+    chunk.reverse();
+    client.send(ServerMessage::Chunk(chunk))?;
+    Ok(())
+}

src/node_requests/indicator.rs

@@ -0,0 +1,19 @@
+use crate::{requests::indicator::IndicatorContext, server::Server, types, utils::client::Client};
+use std::sync::Arc;
+
+pub fn start_typing(server: &Arc<Server>, client: &Client, channel_id: &str) -> crate::Result<()> {
+    let user_id = client.get_uuid()?;
+
+    server.broadcast_to(
+        &[&user_id.clone(), &channel_id.to_string()],
+        types::message::ServerMessage::Indicator(IndicatorContext {
+            indicator: crate::requests::indicator::Indicator::Typing {
+                user_id,
+                channel_id: channel_id.to_string(),
+            },
+            expires: 2, // 2 secs
+        }),
+    )?;
+
+    Ok(())
+}

src/node_requests/message.rs

@@ -0,0 +1,83 @@
+use std::sync::Arc;
+
+use anyhow::anyhow;
+
+use crate::{server::Server, types, utils::client::Client};
+
+crate::logger!(LOGGER "Message Manager");
+
+pub fn send(
+    server: &Arc<Server>,
+    client: &Client,
+    channel_id: &str,
+    contents: &str,
+) -> crate::Result<()> {
+    LOGGER.info(format!("SendMessage to {channel_id}: {contents}"));
+
+    if contents.is_empty() {
+        client.send(types::message::ResponseError::InvalidRequest(format!(
+            "Invalid message: empty message"
+        )))?;
+
+        return Ok(());
+    }
+
+    let msg = server.db.insert_message(
+        &channel_id,
+        &client.get_uuid()?,
+        &contents,
+        chrono::Utc::now().timestamp(),
+    )?;
+
+    server.broadcast_to(
+        &[&msg.channel_id, &msg.from],
+        types::message::ServerMessage::MessageCreate(msg.clone()),
+    )?;
+
+    Ok(())
+}
+
+pub fn edit(
+    server: &Arc<Server>,
+    client: &Client,
+    message_id: i64,
+    new_contents: &str,
+) -> crate::Result<()> {
+    LOGGER.info(format!("EditMessage {message_id}: {new_contents}"));
+    let Some(msg) = server.db.get_message_by_id(message_id)? else {
+        return Err(anyhow!("Message does not exist"));
+    };
+
+    if msg.from != client.get_uuid()? {
+        return Err(anyhow!("You are not the author of this message"));
+    }
+
+    server.db.edit_message(message_id, new_contents)?;
+
+    server.broadcast_to(
+        &[&msg.channel_id, &msg.from],
+        types::message::ServerMessage::MessageCreate(msg.clone()),
+    )?;
+
+    Ok(())
+}
+
+pub fn delete(server: &Arc<Server>, client: &Client, message_id: i64) -> crate::Result<()> {
+    LOGGER.info(format!("DeleteMessage {message_id}"));
+    let Some(msg) = server.db.get_message_by_id(message_id)? else {
+        return Err(anyhow!("Message does not exist"));
+    };
+
+    if msg.from != client.get_uuid()? {
+        return Err(anyhow!("You are not the author of this message"));
+    }
+
+    server.db.delete_message(message_id)?;
+
+    server.broadcast_to(
+        &[&msg.channel_id, &msg.from],
+        types::message::ServerMessage::MessageDelete { message_id },
+    )?;
+
+    Ok(())
+}

src/node_requests/mod.rs

@@ -0,0 +1,86 @@
+pub mod chunk;
+pub mod indicator;
+pub mod message;
+pub mod voice;
+
+use std::sync::Arc;
+
+use crate::{
+    server::Server,
+    types::message::{ClientMessage, ServerMessage, WsMessage},
+    utils::client::Client,
+};
+
+impl Server {
+    pub fn call_node_request(
+        self: &Arc<Self>,
+        req: &WsMessage<ClientMessage>,
+        client: &Client,
+    ) -> crate::Result<()> {
+        match req {
+            WsMessage::Message(req) => match req {
+                ClientMessage::SendMessage {
+                    channel_id,
+                    contents,
+                } => {
+                    message::send(self, client, channel_id, contents)?;
+                }
+
+                ClientMessage::EditMessage {
+                    message_id,
+                    new_contents,
+                } => message::edit(self, client, *message_id, new_contents)?,
+
+                ClientMessage::DeleteMessage { message_id } => {
+                    message::delete(self, client, *message_id)?
+                }
+
+                ClientMessage::LoadChunk {
+                    chunk_id,
+                    channel_id,
+                } => chunk::load_chunk(self, client, channel_id, *chunk_id)?,
+
+                ClientMessage::Typing { channel_id } => {
+                    indicator::start_typing(self, client, channel_id)?
+                }
+
+                ClientMessage::JoinVoice { channel_id } => voice::join(self, client, channel_id)?,
+                ClientMessage::LeaveVoice { channel_id } => voice::leave(self, client, channel_id)?,
+            },
+
+            WsMessage::Binary(data) => {
+                // Self::LOGGER.info(format!("Binary message: {data:?}"));
+                voice::voice(self, client, data)?;
+            }
+
+            WsMessage::String(s) => {
+                Self::LOGGER.info(format!("String message: {s}"));
+            }
+        }
+
+        Ok(())
+    }
+
+    pub fn broadcast_to(
+        self: &Arc<Self>,
+        targets: &[&String],
+        message: ServerMessage,
+    ) -> crate::Result<()> {
+        for c in self.clients.lock().unwrap().iter() {
+            if !targets.contains(&&c.get_uuid()?) {
+                continue;
+            }
+
+            let c = c.clone();
+            let server = self.clone();
+            let message = message.clone();
+            std::thread::spawn(move || {
+                server
+                    .wrap_err(&c, c.send(&message))
+                    .expect("Failed to broadcast");
+            });
+        }
+
+        Ok(())
+    }
+}

src/node_requests/voice.rs

@@ -0,0 +1,61 @@
+use crate::{server::Server, utils::client::Client};
+use std::sync::Arc;
+
+crate::logger!(LOGGER "Voice chat");
+
+pub fn join(server: &Arc<Server>, client: &Client, channel_id: &str) -> crate::Result<()> {
+    let user_id = client.get_uuid()?;
+
+    let voice_id = server
+        .voice
+        .lock()
+        .unwrap()
+        .set(user_id.clone(), channel_id.to_string());
+
+    server.broadcast(crate::types::message::ServerMessage::VoiceJoin {
+        user_id,
+        channel_id: channel_id.to_string(),
+        voice_id: voice_id,
+    });
+
+    Ok(())
+}
+
+pub fn leave(server: &Arc<Server>, client: &Client, channel_id: &str) -> crate::Result<()> {
+    let user_id = client.get_uuid()?;
+
+    let Some(voice_id) = server.voice.lock().unwrap().remove(channel_id, &user_id) else {
+        return Ok(());
+    };
+
+    server.broadcast(crate::types::message::ServerMessage::VoiceLeave {
+        user_id,
+        channel_id: channel_id.to_string(),
+        voice_id,
+    });
+
+    Ok(())
+}
+
+pub fn voice(server: &Arc<Server>, client: &Client, data: &[u8]) -> crate::Result<()> {
+    let v = server.voice.lock().unwrap();
+    let user_id = client.get_uuid()?;
+    let Some((channel_id, voice_id)) = v.find_user(&user_id) else {
+        return Ok(());
+    };
+
+    let mut targets = v.get(&channel_id);
+
+    if let Some(pos) = targets.iter().position(|x| *x == &user_id) {
+        targets.remove(pos);
+    }
+
+    let prefix = voice_id.to_le_bytes();
+    let mut payload = Vec::with_capacity(prefix.len() + data.len());
+    payload.extend_from_slice(&prefix);
+    payload.extend_from_slice(data);
+
+    server.broadcast_bin_to(&targets, payload)?;
+
+    Ok(())
+}

src/requests/mod.rs

@@ -1,6 +1,7 @@
 pub mod chunk;
-pub mod message;
 pub mod indicator;
+pub mod message;
+pub mod voice;
 
 use std::sync::Arc;
 
@@ -11,7 +12,7 @@ use crate::{
 };
 
 impl Server {
-    pub fn call_request(
+    pub fn call_server_request(
         self: &Arc<Self>,
         req: &WsMessage<ClientMessage>,
         client: &Client,
@@ -42,10 +43,14 @@ impl Server {
                 ClientMessage::Typing { channel_id } => {
                     indicator::start_typing(self, client, channel_id)?
                 }
+
+                ClientMessage::JoinVoice { channel_id } => voice::join(self, client, channel_id)?,
+                ClientMessage::LeaveVoice { channel_id } => voice::leave(self, client, channel_id)?,
             },
 
-            WsMessage::Binary(b) => {
-                Self::LOGGER.info(format!("Binary message: {b:?}"));
+            WsMessage::Binary(data) => {
+                // Self::LOGGER.info(format!("Binary message: {data:?}"));
+                voice::voice(self, client, data)?;
             }
 
             WsMessage::String(s) => {
@@ -68,4 +73,27 @@ impl Server {
             });
         }
     }
+
+    pub fn broadcast_bin_to(
+        self: &Arc<Self>,
+        targets: &[&String],
+        bytes: Vec<u8>,
+    ) -> crate::Result<()> {
+        for c in self.clients.lock().unwrap().iter() {
+            if !targets.contains(&&c.get_uuid()?) {
+                continue;
+            }
+
+            let c = c.clone();
+            let server = self.clone();
+            let bytes = bytes.clone();
+            std::thread::spawn(move || {
+                server
+                    .wrap_err(&c, c.send_bin(&bytes))
+                    .expect("Failed to broadcast");
+            });
+        }
+
+        Ok(())
+    }
 }

src/requests/voice.rs

@@ -0,0 +1,61 @@
+use crate::{server::Server, utils::client::Client};
+use std::sync::Arc;
+
+crate::logger!(LOGGER "Voice chat");
+
+pub fn join(server: &Arc<Server>, client: &Client, channel_id: &str) -> crate::Result<()> {
+    let user_id = client.get_uuid()?;
+
+    let voice_id = server
+        .voice
+        .lock()
+        .unwrap()
+        .set(user_id.clone(), channel_id.to_string());
+
+    server.broadcast(crate::types::message::ServerMessage::VoiceJoin {
+        user_id,
+        channel_id: channel_id.to_string(),
+        voice_id: voice_id,
+    });
+
+    Ok(())
+}
+
+pub fn leave(server: &Arc<Server>, client: &Client, channel_id: &str) -> crate::Result<()> {
+    let user_id = client.get_uuid()?;
+
+    let Some(voice_id) = server.voice.lock().unwrap().remove(channel_id, &user_id) else {
+        return Ok(());
+    };
+
+    server.broadcast(crate::types::message::ServerMessage::VoiceLeave {
+        user_id,
+        channel_id: channel_id.to_string(),
+        voice_id,
+    });
+
+    Ok(())
+}
+
+pub fn voice(server: &Arc<Server>, client: &Client, data: &[u8]) -> crate::Result<()> {
+    let v = server.voice.lock().unwrap();
+    let user_id = client.get_uuid()?;
+    let Some((channel_id, voice_id)) = v.find_user(&user_id) else {
+        return Ok(());
+    };
+
+    let mut targets = v.get(&channel_id);
+
+    if let Some(pos) = targets.iter().position(|x| *x == &user_id) {
+        targets.remove(pos);
+    }
+
+    let prefix = voice_id.to_le_bytes();
+    let mut payload = Vec::with_capacity(prefix.len() + data.len());
+    payload.extend_from_slice(&prefix);
+    payload.extend_from_slice(data);
+
+    server.broadcast_bin_to(&targets, payload)?;
+
+    Ok(())
+}

src/server.rs

@@ -8,24 +8,23 @@ use std::{
     },
 };
 
-const WELCOME: &str = "
-
-$$\\    $$\\  $$$$$$\\  $$\\   $$\\  $$$$$$\\  
-$$ |   $$ |$$  __$$\\ $$ |  $$ |$$  __$$\\ 
-$$ |   $$ |$$ /  $$ |\\$$\\ $$  |$$ /  $$ |
-\\$$\\  $$  |$$ |  $$ | \\$$$$  / $$$$$$$$ |
- \\$$\\$$  / $$ |  $$ | $$  $$<  $$  __$$ |
-  \\$$$  /  $$ |  $$ |$$  /\\$$\\ $$ |  $$ |
-   \\$  /    $$$$$$  |$$ /  $$ |$$ |  $$ |
-    \\_/     \\______/ \\__|  \\__|\\__|  \\__|
-
-";
+const WELCOME: &str = "\x1b[38;2;169;86;252m
+    _          _                 
+   / \\   __  _(_) ___  _ __ ___  
+  / _ \\  \\ \\/ / |/ _ \\| '_ ` _ \\ 
+ / ___ \\  >  <| | (_) | | | | | |
+/_/   \\_\\/_/\\_\\_|\\___/|_| |_| |_|
+\x1b[0m";
 
 use crate::{
     cli, logger,
     plugin::{Plugin, loader::PluginLoader, types::LoaderMessage},
-    types::{self, message::WsMessage},
-    utils::{self, auth, client::Client},
+    requests::voice,
+    types::{
+        self,
+        message::{ClientMessage, WsMessage},
+    },
+    utils::{self, auth, client::Client, voice::Voice},
 };
 
 #[derive(serde::Serialize, serde::Deserialize)]
@@ -45,6 +44,8 @@ pub struct Server {
     pub db: utils::database::Database,
     pub shutting_down: AtomicBool,
     pub indicators: Mutex<Vec<crate::requests::indicator::IndicatorContext>>,
+    pub voice: Mutex<crate::utils::voice::Voice>,
+    pub call_request: fn(&Arc<Self>, &WsMessage<ClientMessage>, &Client) -> crate::Result<()>,
 }
 
 impl Default for ServerConfig {
@@ -64,6 +65,14 @@ impl ServerConfig {
         Server::new_config(root, self)
     }
 
+    pub fn build_req(
+        self,
+        root: &Path,
+        req: fn(&Arc<Server>, &WsMessage<ClientMessage>, &Client) -> crate::Result<()>,
+    ) -> Arc<Server> {
+        Server::new_req_config(root, req, self)
+    }
+
     pub fn from_str(s: &str) -> std::result::Result<Self, serde_json::Error> {
         serde_json::from_str(s)
     }
@@ -72,6 +81,24 @@ impl ServerConfig {
 impl Server {
     logger!(LOGGER "Server");
 
+    pub fn new_req_config(
+        root: &Path,
+        call_request: fn(&Arc<Self>, &WsMessage<ClientMessage>, &Client) -> crate::Result<()>,
+        config: ServerConfig,
+    ) -> Arc<Self> {
+        Arc::new(Self {
+            db: utils::database::Database::new(&config).unwrap(),
+            root: root.to_path_buf(),
+            config,
+            clients: Mutex::new(HashSet::new()),
+            plugins: Mutex::new(Vec::new()),
+            shutting_down: AtomicBool::new(false),
+            indicators: Mutex::new(Vec::new()),
+            voice: Mutex::new(Voice::new()),
+            call_request: call_request,
+        })
+    }
+
     pub fn new_config(root: &Path, config: ServerConfig) -> Arc<Self> {
         Arc::new(Self {
             db: utils::database::Database::new(&config).unwrap(),
@@ -81,6 +108,8 @@ impl Server {
             plugins: Mutex::new(Vec::new()),
             shutting_down: AtomicBool::new(false),
             indicators: Mutex::new(Vec::new()),
+            voice: Mutex::new(Voice::new()),
+            call_request: Self::call_server_request,
         })
     }
 
@@ -173,6 +202,7 @@ impl Server {
                     client.send(types::message::ServerMessage::Authenticated {
                         uuid,
                         indicators: self.indicators.lock().unwrap().clone(),
+                        voice_chat: self.voice.lock().unwrap().get_connections(),
                     }),
                 )?;
             }
@@ -210,7 +240,7 @@ impl Server {
                     }
                 }
 
-                self.wrap_err(&client, self.call_request(r, &client))?;
+                self.wrap_err(&client, (self.call_request)(self, r, &client))?;
             }
         }
         Ok(())
@@ -227,7 +257,21 @@ impl Server {
             if client
                 .send(types::message::ResponseError::InternalError(e.to_string()))
                 .is_err()
-            {}
+            {
+                let Ok(user_id) = client.get_uuid() else {
+                    return res;
+                };
+
+                let channel_id = {
+                    let v = self.voice.lock().unwrap();
+                    let Some((channel_id, _)) = v.find_user(&user_id) else {
+                        return res;
+                    };
+                    channel_id.clone()
+                };
+
+                voice::leave(self, client, &channel_id).unwrap();
+            }
         }
 
         res

src/types.rs

@@ -53,6 +53,8 @@ pub mod handshake {
 }
 
 pub mod message {
+    use std::collections::HashMap;
+
     use serde::{Deserialize, Serialize};
 
     use crate::{
@@ -92,6 +94,14 @@ pub mod message {
         Typing {
             channel_id: String,
         },
+
+        JoinVoice {
+            channel_id: String,
+        },
+
+        LeaveVoice {
+            channel_id: String,
+        },
     }
 
     /// Messages sent *from the server* to the client
@@ -102,6 +112,7 @@ pub mod message {
         Authenticated {
             uuid: Author,
             indicators: Vec<IndicatorContext>,
+            voice_chat: HashMap<String, HashMap<String, u16>>,
         },
 
         TempMessage {
@@ -136,6 +147,18 @@ pub mod message {
         },
 
         Chunk(Vec<Message>),
+
+        VoiceJoin {
+            user_id: String,
+            channel_id: String,
+            voice_id: u16,
+        },
+
+        VoiceLeave {
+            user_id: String,
+            channel_id: String,
+            voice_id: u16,
+        },
     }
 
     #[derive(Debug, Clone, Serialize, Deserialize)]

src/utils/database.rs

@@ -129,6 +129,41 @@ impl Database {
 
         Ok(messages)
     }
+
+    pub fn get_chunk_node(
+        &self,
+        author: &str,
+        channel_id: &str,
+        chunk_id: usize,
+    ) -> Result<Vec<Message>> {
+        let mut stmt = self.0.prepare(
+            "SELECT id, channel_id, user_id, contents, timestamp
+            FROM chat
+            WHERE (
+                (channel_id = ?1 AND user_id = ?3)
+             OR (channel_id = ?3 AND user_id = ?1)
+            )
+            ORDER BY id DESC
+            LIMIT 16 OFFSET (?2 * 16)",
+        )?;
+
+        let rows = stmt.query_map(params![channel_id, chunk_id, author], |row| {
+            Ok(Message {
+                id: row.get::<_, i64>(0)?,
+                channel_id: row.get::<_, String>(1)?,
+                from: row.get::<_, String>(2)?,
+                contents: row.get::<_, String>(3)?,
+                timestamp: row.get::<_, i64>(4)?,
+            })
+        })?;
+
+        let mut messages = Vec::new();
+        for row in rows {
+            messages.push(row?);
+        }
+
+        Ok(messages)
+    }
 }
 
 unsafe impl Send for Database {}

src/utils/mod.rs

@@ -3,3 +3,4 @@ pub mod client;
 pub mod database;
 pub mod logger;
 pub mod vfs;
+pub mod voice;

src/utils/voice.rs

@@ -0,0 +1,58 @@
+use std::collections::HashMap;
+
+pub struct Voice {
+    // channel_id -> (user_id -> voice_id)
+    connections: HashMap<String, HashMap<String, u16>>,
+}
+
+impl Voice {
+    pub fn new() -> Self {
+        Self {
+            connections: HashMap::new(),
+        }
+    }
+
+    /// Insert or update a user's voice_id in a channel
+    pub fn set(&mut self, user_id: String, channel_id: String) -> u16 {
+        let voice_id = rand::random::<u16>();
+
+        self.connections
+            .entry(channel_id)
+            .or_insert_with(HashMap::new)
+            .insert(user_id, voice_id);
+
+        voice_id
+    }
+
+    /// Remove a user from a channel
+    /// Returns their voice_id if they existed
+    pub fn remove(&mut self, channel_id: &str, user_id: &str) -> Option<u16> {
+        let channel = self.connections.get_mut(channel_id)?;
+        let voice_id = channel.remove(user_id)?;
+
+        // Clean up empty channels
+        if channel.is_empty() {
+            self.connections.remove(channel_id);
+        }
+
+        Some(voice_id)
+    }
+
+    /// Get all users + voice_ids in a channel
+    pub fn get(&self, channel_id: &str) -> Vec<&String> {
+        self.connections
+            .get(channel_id)
+            .map(|users| users.keys().collect())
+            .unwrap_or_default()
+    }
+
+    pub fn find_user(&self, user_id: &str) -> Option<(&String, u16)> {
+        self.connections
+            .iter()
+            .find_map(|(channel_id, users)| users.get(user_id).map(|v| (channel_id, *v)))
+    }
+
+    pub fn get_connections(&self) -> HashMap<String, HashMap<String, u16>> {
+        self.connections.clone()
+    }
+}