Improved structure
This commit is contained in:
@@ -0,0 +1,183 @@
|
||||
use std::{
|
||||
hash::{Hash, Hasher},
|
||||
io::{Read, Write},
|
||||
net::TcpStream,
|
||||
};
|
||||
|
||||
use anyhow::anyhow;
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::types::{ClientMessage, WsMessage};
|
||||
|
||||
pub mod handshake {
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::STANDARD as Base64;
|
||||
use sha1::{Digest, Sha1};
|
||||
use std::io::{Read, Write};
|
||||
use std::net::TcpStream;
|
||||
|
||||
const WS_GUID: &str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
|
||||
|
||||
pub fn handle_websocket_handshake(stream: &mut TcpStream) -> std::io::Result<()> {
|
||||
let mut buffer = [0; 1024];
|
||||
let size = stream.read(&mut buffer)?;
|
||||
let request = String::from_utf8_lossy(&buffer[..size]);
|
||||
|
||||
let key_line = request
|
||||
.lines()
|
||||
.find(|line| line.to_lowercase().starts_with("sec-websocket-key"))
|
||||
.ok_or_else(|| {
|
||||
std::io::Error::new(std::io::ErrorKind::InvalidData, "Missing Sec-WebSocket-Key")
|
||||
})?;
|
||||
|
||||
let key = key_line.splitn(2, ':').nth(1).unwrap().trim();
|
||||
|
||||
let mut hasher = Sha1::new();
|
||||
hasher.update(key.as_bytes());
|
||||
hasher.update(WS_GUID.as_bytes());
|
||||
let hash = hasher.finalize();
|
||||
|
||||
let accept_key = Base64.encode(hash);
|
||||
|
||||
let response = format!(
|
||||
"HTTP/1.1 101 Switching Protocols\r\n\
|
||||
Upgrade: websocket\r\n\
|
||||
Connection: Upgrade\r\n\
|
||||
Sec-WebSocket-Accept: {}\r\n\r\n",
|
||||
accept_key
|
||||
);
|
||||
|
||||
stream.write_all(response.as_bytes())?;
|
||||
stream.flush()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Client(TcpStream);
|
||||
|
||||
impl Client {
|
||||
pub fn new(mut stream: TcpStream) -> crate::Result<Self> {
|
||||
handshake::handle_websocket_handshake(&mut stream)?;
|
||||
Ok(Client(stream))
|
||||
}
|
||||
}
|
||||
|
||||
impl Clone for Client {
|
||||
fn clone(&self) -> Self {
|
||||
Client(self.0.try_clone().expect("failed to clone TcpStream"))
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for Client {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.0.peer_addr().unwrap() == other.0.peer_addr().unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for Client {}
|
||||
|
||||
impl Hash for Client {
|
||||
fn hash<H: Hasher>(&self, state: &mut H) {
|
||||
self.0.peer_addr().unwrap().hash(state);
|
||||
}
|
||||
}
|
||||
|
||||
impl Client {
|
||||
/// Read a full WebSocket message, handling fragmentation (FIN)
|
||||
pub fn read(&self) -> crate::Result<Option<WsMessage<ClientMessage>>> {
|
||||
let mut stream = &self.0;
|
||||
let mut message_payload = Vec::new();
|
||||
let mut final_frame = false;
|
||||
|
||||
while !final_frame {
|
||||
let mut header = [0u8; 2];
|
||||
if stream.read_exact(&mut header).is_err() {
|
||||
return Ok(None); // connection closed
|
||||
}
|
||||
|
||||
let fin = header[0] & 0x80 != 0;
|
||||
let opcode = header[0] & 0x0F;
|
||||
let masked = header[1] & 0x80 != 0;
|
||||
let mut payload_len = (header[1] & 0x7F) as u64;
|
||||
|
||||
// Extended payload lengths
|
||||
if payload_len == 126 {
|
||||
let mut ext_len = [0u8; 2];
|
||||
stream.read_exact(&mut ext_len)?;
|
||||
payload_len = u16::from_be_bytes(ext_len) as u64;
|
||||
} else if payload_len == 127 {
|
||||
let mut ext_len = [0u8; 8];
|
||||
stream.read_exact(&mut ext_len)?;
|
||||
payload_len = u64::from_be_bytes(ext_len);
|
||||
}
|
||||
|
||||
// Mask key (client → server)
|
||||
let mut mask = [0u8; 4];
|
||||
if masked {
|
||||
stream.read_exact(&mut mask)?;
|
||||
}
|
||||
|
||||
// Read payload
|
||||
let mut payload = vec![0u8; payload_len as usize];
|
||||
stream.read_exact(&mut payload)?;
|
||||
|
||||
if masked {
|
||||
for i in 0..payload.len() {
|
||||
payload[i] ^= mask[i % 4];
|
||||
}
|
||||
}
|
||||
|
||||
match opcode {
|
||||
0x0 | 0x1 | 0x2 => {
|
||||
// Continuation / Text / Binary
|
||||
message_payload.extend(payload);
|
||||
}
|
||||
0x8 => return Ok(None), // Close
|
||||
0x9 => continue, // Ping → ignore
|
||||
0xA => continue, // Pong → ignore
|
||||
_ => return Err(anyhow!("Unsupported WebSocket opcode: {}", opcode).into()),
|
||||
}
|
||||
|
||||
final_frame = fin;
|
||||
}
|
||||
|
||||
// Try parsing JSON into ClientMessage
|
||||
let message = match String::from_utf8(message_payload.clone()) {
|
||||
Ok(text) => match serde_json::from_str(&text) {
|
||||
Ok(msg) => WsMessage::Message(msg),
|
||||
Err(_) => WsMessage::String(text),
|
||||
},
|
||||
Err(_) => WsMessage::Binary(message_payload),
|
||||
};
|
||||
|
||||
Ok(Some(message))
|
||||
}
|
||||
|
||||
/// Send a JSON-serializable object as a WebSocket text frame
|
||||
pub fn send<T: Serialize>(&self, m: T) -> crate::Result<()> {
|
||||
let payload = serde_json::to_string(&m)?;
|
||||
let payload_bytes = payload.as_bytes();
|
||||
|
||||
let mut stream = self.0.try_clone()?;
|
||||
let mut header = Vec::new();
|
||||
header.push(0x81); // FIN=1, opcode=0x1 (text)
|
||||
|
||||
let len = payload_bytes.len();
|
||||
if len < 126 {
|
||||
header.push(len as u8);
|
||||
} else if len <= 65535 {
|
||||
header.push(126);
|
||||
header.extend_from_slice(&(len as u16).to_be_bytes());
|
||||
} else {
|
||||
header.push(127);
|
||||
header.extend_from_slice(&(len as u64).to_be_bytes());
|
||||
}
|
||||
|
||||
stream.write_all(&header)?;
|
||||
stream.write_all(payload_bytes)?;
|
||||
stream.flush()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
use crate::{ServerConfig, types::data::Message};
|
||||
use rusqlite::{Connection, Result, params};
|
||||
|
||||
pub struct Database {
|
||||
pub messages_db: MessagesDb,
|
||||
}
|
||||
|
||||
impl Database {
|
||||
pub fn new(config: &ServerConfig) -> Option<Self> {
|
||||
let conn = Connection::open("main.db").ok()?;
|
||||
let messages_db = MessagesDb(conn);
|
||||
messages_db.init(config)?;
|
||||
Some(Self { messages_db })
|
||||
}
|
||||
}
|
||||
|
||||
unsafe impl Send for Database {}
|
||||
unsafe impl Sync for Database {}
|
||||
|
||||
pub struct MessagesDb(pub Connection);
|
||||
|
||||
impl MessagesDb {
|
||||
pub fn init(&self, _config: &ServerConfig) -> Option<usize> {
|
||||
self.0
|
||||
.execute(
|
||||
"CREATE TABLE IF NOT EXISTS chat (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
channel_id TEXT NOT NULL,
|
||||
user_id TEXT NOT NULL,
|
||||
contents TEXT NOT NULL,
|
||||
timestamp INTEGER NOT NULL
|
||||
)",
|
||||
[],
|
||||
)
|
||||
.ok()
|
||||
}
|
||||
|
||||
/// Insert a message into the DB
|
||||
pub fn insert(
|
||||
&self,
|
||||
channel_id: &str,
|
||||
user_id: &str,
|
||||
contents: &str,
|
||||
timestamp: i64,
|
||||
) -> Result<Message> {
|
||||
self.0.execute(
|
||||
"INSERT INTO chat (channel_id, user_id, contents, timestamp)
|
||||
VALUES (?1, ?2, ?3, ?4)",
|
||||
params![channel_id, user_id, contents, timestamp],
|
||||
)?;
|
||||
|
||||
let id = self.0.last_insert_rowid();
|
||||
|
||||
Ok(Message {
|
||||
id,
|
||||
channel_id: channel_id.to_string(),
|
||||
from: user_id.to_string(),
|
||||
contents: contents.to_string(),
|
||||
timestamp,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get a message by its ID
|
||||
pub fn get_by_id(&self, message_id: usize) -> Result<Option<Message>> {
|
||||
let mut stmt = self.0.prepare(
|
||||
"SELECT id, channel_id, user_id, contents, timestamp
|
||||
FROM chat
|
||||
WHERE id = ?1",
|
||||
)?;
|
||||
|
||||
let mut rows = stmt.query_map(params![message_id], |row| {
|
||||
Ok((
|
||||
row.get::<_, i64>(0)?, // id
|
||||
row.get::<_, String>(1)?, // channel_id
|
||||
row.get::<_, String>(2)?, // user_id
|
||||
row.get::<_, String>(3)?, // contents
|
||||
row.get::<_, i64>(4)?, // timestamp
|
||||
))
|
||||
})?;
|
||||
|
||||
if let Some(row) = rows.next() {
|
||||
let (id, channel_id, user_id, contents, timestamp) = row?;
|
||||
return Ok(Some(Message {
|
||||
id,
|
||||
channel_id,
|
||||
from: user_id,
|
||||
contents,
|
||||
timestamp,
|
||||
}));
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
use std::path::Path;
|
||||
|
||||
use libloading::{Library, Symbol};
|
||||
|
||||
use crate::{logger, utils::plugin::DynPlugin, utils::vfs};
|
||||
|
||||
logger! {
|
||||
const LOGGER "Loader"
|
||||
}
|
||||
|
||||
pub fn load_plugin(path: &Path) -> anyhow::Result<DynPlugin> {
|
||||
unsafe {
|
||||
LOGGER.info(format!("Loading plugin: {:?}", path));
|
||||
let lib = Library::new(path)?;
|
||||
let func: Symbol<extern "C" fn() -> DynPlugin> = lib.get(b"load_plugin").unwrap();
|
||||
Ok(func())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn load_plugins(arr: &mut Vec<DynPlugin>, path: &Path) -> crate::Result<()> {
|
||||
LOGGER.info("Loading plugins");
|
||||
vfs::dir(path)?;
|
||||
if path.is_dir() {
|
||||
for entry in std::fs::read_dir(path)? {
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|s| s.to_str()) == Some("dylib") {
|
||||
match load_plugin(&path) {
|
||||
Ok(plugin) => {
|
||||
arr.push(plugin);
|
||||
}
|
||||
Err(e) => LOGGER.error(format!("Failed to load plugin {:?}: {}", path, e)),
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
LOGGER.error(format!("{:?} is not a directory", path));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
use std::fmt::Display;
|
||||
|
||||
pub struct Logger {
|
||||
name: String,
|
||||
}
|
||||
|
||||
impl Logger {
|
||||
pub fn new<T: Display>(name: T) -> Self {
|
||||
Logger {
|
||||
name: name.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn info<T: Display>(&self, message: T) {
|
||||
println!("\x1b[32mINFO\x1b[0m ({}) › {}", self.name, message);
|
||||
}
|
||||
|
||||
pub fn warn<T: Display>(&self, message: T) {
|
||||
println!("\x1b[33mWARN\x1b[0m ({}) › {}", self.name, message);
|
||||
}
|
||||
|
||||
pub fn error<T: Display>(&self, message: T) {
|
||||
println!("\x1b[31mERROR\x1b[0m ({}) › {}", self.name, message);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
pub mod client;
|
||||
pub mod database;
|
||||
#[cfg(feature = "loader")]
|
||||
pub mod loader;
|
||||
pub mod logger;
|
||||
pub mod plugin;
|
||||
pub mod vfs;
|
||||
@@ -0,0 +1,22 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::{
|
||||
Server,
|
||||
types::{ClientMessage, WsMessage},
|
||||
utils::client::Client,
|
||||
};
|
||||
|
||||
pub type DynPlugin = Box<dyn Plugin + Send + Sync>;
|
||||
|
||||
pub trait Plugin {
|
||||
fn init(&mut self, server: &Arc<Server>);
|
||||
#[allow(unused_variables)]
|
||||
fn on_request(
|
||||
&mut self,
|
||||
req: &WsMessage<ClientMessage>,
|
||||
client: &Client,
|
||||
server: &Arc<Server>,
|
||||
) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
use std::{fs, path::Path};
|
||||
|
||||
logger! {
|
||||
const LOGGER "File"
|
||||
}
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::logger;
|
||||
|
||||
pub fn dir(path: &Path) -> crate::Result<()> {
|
||||
if !path.exists() {
|
||||
LOGGER.info(format!("Directory {path:?} does not exist, creating it"));
|
||||
fs::create_dir_all(path)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn read(path: &Path, default_content: &str) -> crate::Result<String> {
|
||||
if !path.exists() {
|
||||
LOGGER.info(format!(
|
||||
"File {path:?} does not exist, creating it with default contents"
|
||||
));
|
||||
write(path, default_content)?;
|
||||
return Ok(default_content.to_string());
|
||||
}
|
||||
|
||||
Ok(fs::read_to_string(path)?)
|
||||
}
|
||||
|
||||
pub fn read_bytes<'a>(path: &Path, default_content: Vec<u8>) -> crate::Result<Vec<u8>> {
|
||||
if !path.exists() {
|
||||
LOGGER.info(format!(
|
||||
"File {path:?} does not exist, creating it with default contents"
|
||||
));
|
||||
write_bytes(path, &default_content)?;
|
||||
return Ok(default_content);
|
||||
}
|
||||
|
||||
Ok(fs::read(path)?)
|
||||
}
|
||||
|
||||
pub fn read_config<T: Default + Serialize + for<'de> Deserialize<'de>>(
|
||||
path: &Path,
|
||||
) -> crate::Result<T> {
|
||||
if !path.exists() {
|
||||
LOGGER.info(format!(
|
||||
"File {path:?} does not exist, creating it with default contents"
|
||||
));
|
||||
let default = T::default();
|
||||
write_config(path, &default)?;
|
||||
return Ok(default);
|
||||
}
|
||||
|
||||
let read = fs::read_to_string(path)?;
|
||||
Ok(serde_json::from_str::<T>(&read)?)
|
||||
}
|
||||
|
||||
pub fn write(path: &Path, content: &str) -> crate::Result<()> {
|
||||
dir(path.parent().ok_or(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidFilename,
|
||||
"File doesn't have a parent assigned, example: `config/config.json`",
|
||||
))?)?;
|
||||
fs::write(path, content)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn write_bytes(path: &Path, content: &[u8]) -> crate::Result<()> {
|
||||
dir(path.parent().ok_or(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidFilename,
|
||||
"File doesn't have a parent assigned, example: `config/config.json`",
|
||||
))?)?;
|
||||
fs::write(path, content)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn write_config<T: Serialize>(path: &Path, content: &T) -> crate::Result<()> {
|
||||
dir(path.parent().ok_or(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidFilename,
|
||||
"File doesn't have a parent assigned, example: `config/config.json`",
|
||||
))?)?;
|
||||
fs::write(path, &serde_json::to_string_pretty(content)?)?;
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user