#![feature(try_blocks)] use std::collections::HashMap; use std::str::FromStr; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; use anyhow::{Context, Result}; use futures_util::{SinkExt, StreamExt}; use lodestone_scraper::lodestone_parser::ffxiv_types::World; use log::{debug, error, info, Level, LevelFilter, warn}; use sha3::Digest; use sqlx::{ConnectOptions, Executor, Pool, Sqlite}; use sqlx::migrate::Migrator; use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; use tokio::{ net::{TcpListener, TcpStream}, }; use tokio::sync::mpsc::{Sender, UnboundedSender}; use tokio::sync::RwLock; use tokio_tungstenite::{ tungstenite::Message as WsMessage, WebSocketStream, }; use uuid::Uuid; use crate::{ types::{ protocol::{ MessageRequest, MessageResponse, RegisterRequest, RegisterResponse, RequestContainer, RequestKind, ResponseContainer, }, user::User, }, }; use crate::handlers::SecretsRequestInfo; use crate::types::config::Config; use crate::types::protocol::{AnnounceResponse, AuthenticateRequest, AuthenticateResponse, ErrorResponse, ResponseKind}; use crate::types::protocol::channel::Rank; pub mod types; pub mod handlers; pub mod util; pub mod updater; pub mod logging; pub mod influx; pub type WsStream = WebSocketStream; pub struct State { pub db: Pool, pub clients: HashMap>>, pub ids: HashMap<(String, u16), u64>, pub secrets_requests: HashMap, pub messages_sent: AtomicU64, pub updater_tx: UnboundedSender, } impl State { pub async fn announce(&self, msg: impl Into) { let msg = msg.into(); for client in self.clients.values() { client.read().await.tx.send(ResponseContainer { number: 0, kind: ResponseKind::Announce(AnnounceResponse::new(&msg)), }).await.ok(); } } pub async fn get_id(&self, state: &RwLock, name: &str, world: u16) -> Option { // if they're logged in, grab the id the easy way if let Some(id) = self.ids.get(&(name.to_string(), world)).copied() { return Some(id); } let world_name = util::world_from_id(world)?.as_str(); let id = sqlx::query!( // language=sqlite "select lodestone_id from users where name = ? and world = ?", name, world_name, ) .fetch_optional(&state.read().await.db) .await .ok()?; id.map(|id| id.lodestone_id as u64) } } static MIGRATOR: Migrator = sqlx::migrate!(); #[tokio::main] async fn main() -> Result<()> { logging::setup()?; // get config let config_path = std::env::args().nth(1).unwrap_or_else(|| "config.toml".to_string()); let config_toml = std::fs::read_to_string(config_path) .context("couldn't read config file")?; let config: Config = toml::from_str(&config_toml) .context("couldn't parse config file")?; // set up database pool let mut options = SqliteConnectOptions::new(); options.log_statements(LevelFilter::Debug); let pool = SqlitePoolOptions::new() .after_connect(|conn, _| Box::pin(async move { conn.execute( // language=sqlite "PRAGMA foreign_keys = ON;" ).await?; Ok(()) })) .connect_with(options.filename(&config.database.path)) .await .context("could not connect to database")?; MIGRATOR.run(&pool) .await .context("could not run database migrations")?; // set up updater channel let (updater_tx, updater_rx) = tokio::sync::mpsc::unbounded_channel(); // set up server let server = TcpListener::bind(&config.server.address).await?; let state = Arc::new(RwLock::new(State { db: pool, clients: Default::default(), ids: Default::default(), secrets_requests: Default::default(), messages_sent: AtomicU64::default(), updater_tx, })); info!("Listening on ws://{}/", server.local_addr()?); let (quit_tx, mut quit_rx) = tokio::sync::mpsc::channel(1); let (announce_tx, mut announce_rx) = tokio::sync::mpsc::channel(1); std::thread::spawn(move || { let mut editor = match rustyline::Editor::<()>::new() { Ok(e) => e, Err(e) => { error!("error creating line editor: {:#?}", e); return; } }; for line in editor.iter("> ") { let line = match line { Ok(l) => l, Err(rustyline::error::ReadlineError::Interrupted) => { quit_tx.blocking_send(()).ok(); return; } Err(e) => { error!("error reading input: {:#?}", e); continue; } }; let command: Vec<_> = line.splitn(2, ' ').collect(); match command[0] { "exit" | "quit" => { quit_tx.blocking_send(()).ok(); return; } "announce" => { if command.len() == 2 { let msg = command[1].to_string(); announce_tx.blocking_send(msg).ok(); } else { info!("usage: announce "); } } "log" | "level" => { if command.len() == 2 { match Level::from_str(command[1]) { Ok(level) => *logging::LOG_LEVEL.write() = level, Err(_) => warn!("invalid log level"), } } else { info!("usage: log "); } } "" => {} x => warn!("unknown command: {}", x), } } }); { let state = Arc::clone(&state); tokio::task::spawn(async move { let mut last_messages = 0; loop { let messages = state.read().await.messages_sent.load(Ordering::SeqCst); let diff = messages - last_messages; last_messages = messages; let clients = state.read().await.clients.len(); info!( "Clients: {}, messages sent: {} (+{})", clients, messages, diff, ); tokio::time::sleep(Duration::from_secs(60)).await; } }); } influx::spawn(&config, Arc::clone(&state)); updater::spawn(Arc::clone(&state), updater_rx); loop { let res: Result<()> = try { tokio::select! { accept = server.accept() => { let (sock, _addr) = accept?; let state = Arc::clone(&state); tokio::task::spawn(async move { let conn = match tokio_tungstenite::accept_async(sock).await { Ok(c) => c, Err(e) => { error!("client error: {:?}", e); return; } }; if let Err(e) = client_loop(state, conn).await { error!("client error: {}", e); } }); } _ = quit_rx.recv() => { break; } msg = announce_rx.recv() => { if let Some(msg) = msg { state.read().await.announce(msg).await; } } } }; if let Err(e) = res { error!("server error: {}", e); } } info!("quitting"); Ok(()) } pub struct ClientState { user: Option, tx: Sender, shutdown_tx: Sender<()>, pk: Vec, allow_invites: bool, } impl ClientState { pub fn lodestone_id(&self) -> Option { self.user.as_ref().map(|u| u.lodestone_id) } pub async fn in_channel(&self, channel_id: Uuid, state: &RwLock) -> Result { let user = match &self.user { Some(user) => user, None => return Ok(false), }; let channel_id_str = channel_id.as_simple().to_string(); let id = user.lodestone_id as i64; let members = sqlx::query!( // language=sqlite "select count(*) as count from user_channels where channel_id = ? and lodestone_id = ?", channel_id_str, id, ) .fetch_one(&state.read().await.db) .await .context("could not get count")?; Ok(members.count > 0) } pub async fn get_rank(&self, channel_id: Uuid, state: &RwLock) -> Result> { let user = match &self.user { Some(user) => user, None => return Ok(None), }; let channel_id_str = channel_id.as_simple().to_string(); let id = user.lodestone_id as i64; let rank = sqlx::query!( // language=sqlite "select rank from user_channels where channel_id = ? and lodestone_id = ?", channel_id_str, id, ) .fetch_optional(&state.read().await.db) .await .context("could not get rank")?; Ok(rank.map(|rank| Rank::from_u8(rank.rank as u8))) } pub async fn get_rank_invite(&self, channel_id: Uuid, state: &RwLock) -> Result> { if let Some(rank) = self.get_rank(channel_id, state).await? { return Ok(Some(rank)); } let user = match &self.user { Some(user) => user, None => return Ok(None), }; let channel_id_str = channel_id.as_simple().to_string(); let id = user.lodestone_id as i64; let count = sqlx::query!( // language=sqlite "select count(*) as count from channel_invites where channel_id = ? and invited = ?", channel_id_str, id, ) .fetch_one(&state.read().await.db) .await .context("could not get count")? .count; if count > 0 { Ok(Some(Rank::Invited)) } else { Ok(None) } } } async fn client_loop(state: Arc>, mut conn: WsStream) -> Result<()> { let (tx, mut rx) = tokio::sync::mpsc::channel(10); let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); let client_state = Arc::new(RwLock::new(ClientState { user: None, tx, shutdown_tx, pk: Default::default(), allow_invites: false, })); loop { let res: Result<()> = try { tokio::select! { _ = shutdown_rx.recv() => { debug!("break due to new login"); break; } msg = rx.recv() => { if let Some(msg) = msg { let encoded = rmp_serde::to_vec(&msg)?; conn.send(WsMessage::Binary(encoded)).await?; } } msg = conn.next() => { // match &msg { // Some(Ok(WsMessage::Pong(_))) => {}, // _ => debug!("{:?}", msg), // } match msg { Some(Ok(WsMessage::Binary(msg))) => { let msg: RequestContainer = rmp_serde::from_slice(&msg)?; debug!("{:#?}", msg); let logged_in = client_state.read().await.user.is_some(); match msg.kind { RequestKind::Ping(_) => { crate::handlers::ping(&mut conn, msg.number).await?; } RequestKind::Version(req) => { if !crate::handlers::version(&mut conn, msg.number, req).await? { break; } } RequestKind::Register(req) => { crate::handlers::register(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::Authenticate(req) => { crate::handlers::authenticate(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::Create(req) if logged_in => { crate::handlers::create(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::PublicKey(req) if logged_in => { crate::handlers::public_key(Arc::clone(&state), &mut conn, msg.number, req).await?; } RequestKind::Invite(req) if logged_in => { crate::handlers::invite(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::Join(req) if logged_in => { crate::handlers::join(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::Message(req) if logged_in => { crate::handlers::message(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::List(req) if logged_in => { crate::handlers::list(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::Leave(req) if logged_in => { crate::handlers::leave(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::Promote(req) if logged_in => { crate::handlers::promote(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::Kick(req) if logged_in => { crate::handlers::kick(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::Disband(req) if logged_in => { crate::handlers::disband(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::Update(req) if logged_in => { crate::handlers::update(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::Secrets(req) if logged_in => { crate::handlers::secrets(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::SendSecrets(req) if logged_in => { crate::handlers::send_secrets(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::AllowInvites(req) if logged_in => { crate::handlers::allow_invites(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } RequestKind::DeleteAccount(req) if logged_in => { crate::handlers::delete_account(Arc::clone(&state), Arc::clone(&client_state), &mut conn, msg.number, req).await?; } _ if !logged_in => { util::send(&mut conn, msg.number, ErrorResponse::new(None, "not logged in")).await?; } _ => { util::send(&mut conn, msg.number, ErrorResponse::new(None, "not yet implemented")).await?; } } } None | Some(Ok(WsMessage::Close(_))) | Some(Err(_)) => { debug!("break"); break; } _ => {} } } } }; if let Err(e) = res { error!("error in client loop: {:#?}", e); break; } } debug!("ending client thread"); if let Some(user) = &client_state.read().await.user { state.write().await.clients.remove(&user.lodestone_id); state.write().await.ids.remove(&(user.name.clone(), util::id_from_world(user.world))); } debug!("client thread ended"); Ok(()) }