openworm/src/bin/server/main.rs

mod db;
mod net;

use crate::db::{Db, Msg, ServerPerms, User};
use clap::Parser;
use net::{ClientSender, ConAccepter, listen};
use openworm::{
    net::{
        ClientMsg, ClientRequestMsg, CreateAccount, CreateAccountResp, DisconnectReason, LoadMsg,
        Login, LoginResp, RecvHandler, RequestUsersResp, ServerError, ServerMsg, ServerUser,
        install_crypto_provider,
    },
    rsc::DataDir,
};
use rand::distr::{Alphanumeric, SampleString};
use scrypt::{
    Scrypt,
    password_hash::{PasswordHash, PasswordHasher, PasswordVerifier, SaltString, rand_core::OsRng},
};
use std::{
    collections::HashMap,
    sync::{
        Arc,
        atomic::{AtomicU64, Ordering},
    },
};
use tokio::{signal::ctrl_c, sync::RwLock};

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// port to listen on
    #[arg(short, long)]
    port: u16,
}

fn main() {
    let args = Args::parse();
    install_crypto_provider();
    run_server(args.port);
}

#[tokio::main]
pub async fn run_server(port: u16) {
    let dir = DataDir::new(Some("server"));
    let db = Db::open(dir.path.join("db"));
    let handler = ServerListener {
        senders: Default::default(),
        count: 0.into(),
        db: db.clone(),
    };
    if db.users.is_empty() {
        let token = account_token(&db, ServerPerms::ALL);
        println!("no users found, token for admin: {token}");
    }
    let (endpoint, handle) = listen(port, &dir.path, handler);
    let _ = ctrl_c().await;
    println!("stopping server");
    println!("closing connections...");
    endpoint.close(0u32.into(), &[]);
    let _ = handle.await;
    endpoint.wait_idle().await;
}

pub fn account_token(db: &Db, perms: ServerPerms) -> String {
    let token = Alphanumeric.sample_string(&mut rand::rng(), 16);
    db.account_tokens.insert(&token, &perms);
    token
}

type ClientId = u64;

struct ServerListener {
    db: Db,
    senders: Arc<RwLock<HashMap<ClientId, ClientSender>>>,
    count: AtomicU64,
}

#[derive(PartialEq, Eq, Clone, Copy)]
pub enum ClientState {
    Login,
    Authed(u64),
}

impl ConAccepter for ServerListener {
    async fn accept(&self, send: ClientSender) -> impl RecvHandler<ClientRequestMsg> {
        let id = self.count.fetch_add(1, Ordering::Release);
        self.senders.write().await.insert(id, send.clone());
        ClientHandler {
            db: self.db.clone(),
            senders: self.senders.clone(),
            state: Arc::new(RwLock::new(ClientState::Login)),
            send,
            id,
        }
    }
}

struct ClientHandler {
    db: Db,
    send: ClientSender,
    senders: Arc<RwLock<HashMap<ClientId, ClientSender>>>,
    id: ClientId,
    state: Arc<RwLock<ClientState>>,
}

impl RecvHandler<ClientRequestMsg> for ClientHandler {
    async fn connect(&self) -> () {}
    async fn msg(&self, req: ClientRequestMsg) {
        let msg = ClientMsg::from(req.msg);
        let replier = self.send.replier(req.id);
        match msg {
            ClientMsg::SendMsg(msg) => {
                let ClientState::Authed(uid) = &*self.state.read().await else {
                    let _ = replier.send(ServerError::NotLoggedIn).await;
                    return;
                };
                let msg = Msg {
                    author: *uid,
                    content: msg.content,
                };
                // TODO: it is technically possible to send 2 messages at the exact same time...
                // should probably append a number if one already exists at that time,
                // but also I can't see this ever happening...?
                // should be an easy fix later (write tx)
                let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos();
                self.db.msgs.insert(&timestamp, &msg);
                let mut handles = Vec::new();
                let msg = LoadMsg {
                    content: msg.content,
                    author: *uid,
                };
                for (&id, send) in self.senders.read().await.iter() {
                    if id == self.id {
                        continue;
                    }
                    let send = send.clone();
                    let msg = msg.clone();
                    let fut = async move {
                        let _ = send
                            .send(LoadMsg {
                                content: msg.content,
                                author: msg.author,
                            })
                            .await;
                    };
                    handles.push(tokio::spawn(fut));
                }
                for h in handles {
                    h.await.unwrap();
                }
            }
            ClientMsg::RequestMsgs => {
                let ClientState::Authed(_uid) = &*self.state.read().await else {
                    let _ = replier.send(ServerError::NotLoggedIn).await;
                    return;
                };
                let msgs = self
                    .db
                    .msgs
                    .values()
                    .map(|msg| LoadMsg {
                        content: msg.content,
                        author: msg.author,
                    })
                    .collect();
                let _ = replier.send(ServerMsg::LoadMsgs(msgs)).await;
            }
            ClientMsg::CreateAccount(info) => {
                let CreateAccount {
                    token,
                    username,
                    password,
                } = &info;
                let salt = SaltString::generate(&mut OsRng);
                let params = scrypt::Params::new(11, 8, 1, 32).unwrap();
                let hash = Scrypt
                    .hash_password_customized(password.as_bytes(), None, None, params, &salt)
                    .unwrap()
                    .to_string();
                let mut id;
                loop {
                    let mut tx = self.db.write_tx();
                    let Some(perms) = tx.remove(&self.db.account_tokens, token.to_string()) else {
                        let _ = replier.send(CreateAccountResp::InvalidToken).await;
                        println!("invalid token: {:?}", self.send.remote());
                        return;
                    };
                    if tx.has_key(&self.db.usernames, username.clone()) {
                        let _ = replier.send(CreateAccountResp::UsernameExists).await;
                        return;
                    }
                    id = rand::random();
                    while tx.has_key(&self.db.users, id) {
                        id = rand::random();
                    }
                    tx.insert(
                        &self.db.users,
                        &id,
                        &User {
                            username: username.clone(),
                            password_hash: hash.clone(),
                            bio: String::new(),
                            pfp: None,
                            server_perms: perms,
                        },
                    );
                    tx.insert(&self.db.usernames, username, &id);
                    if tx.commit() {
                        break;
                    }
                }
                println!("account created: \"{username}\"");
                *self.state.write().await = ClientState::Authed(id);
                let _ = replier.send(CreateAccountResp::Ok { id }).await;
            }
            ClientMsg::Login(info) => {
                let Login { username, password } = &info;
                let Some(id) = self.db.usernames.get(username) else {
                    let _ = replier.send(LoginResp::UnknownUsername).await;
                    return;
                };
                let Some(user) = self.db.users.get(&id) else {
                    panic!("invalid state! (should be a user)");
                };
                let hash = PasswordHash::new(&user.password_hash).unwrap();
                if Scrypt.verify_password(password.as_bytes(), &hash).is_err() {
                    println!("invalid password: \"{username}\"");
                    let _ = replier.send(LoginResp::InvalidPassword).await;
                    return;
                }
                *self.state.write().await = ClientState::Authed(id);
                let _ = replier.send(LoginResp::Ok { id }).await;
            }
            ClientMsg::RequestUsers(_) => {
                if self
                    .db
                    .users
                    .get(&self.id)
                    .is_some_and(|u| !u.server_perms.contains(ServerPerms::ALL))
                {
                    let _ = replier.send(ServerError::NoPermission).await;
                    return;
                }
                let users: Vec<_> = self
                    .db
                    .users
                    .iter()
                    .map(|(id, u)| ServerUser {
                        id,
                        username: u.username,
                    })
                    .collect();
                let _ = replier.send(RequestUsersResp { users }).await;
            }
        }
    }

    async fn disconnect(&self, reason: DisconnectReason) -> () {
        match reason {
            DisconnectReason::Closed | DisconnectReason::Timeout => (),
            DisconnectReason::Other(e) => println!("connection issue: {e}"),
        }
    }
}

impl ClientState {
    pub fn is_authed(&self) -> bool {
        matches!(self, Self::Authed(_))
    }
}