button/src/main.rs

use axum::{
	extract::ws::{Message, WebSocket, WebSocketUpgrade},
	response::{Html, IntoResponse},
	routing::get,
	Router,
};
use std::collections::HashMap;
use axum::extract::State;
use rand::random_bool;
use futures::stream::StreamExt;
use serde_json;
use serde_json::json;
use serde::{Deserialize, Serialize};
use std::fs;
use std::io::Write;
use std::sync::mpsc;
use std::sync::{Mutex, Arc};
use tokio::time::{sleep,Duration};
use regex::Regex;

#[derive(Deserialize,Serialize,Debug,Ord,Eq,PartialEq,PartialOrd,Clone)]
struct Entry
{
	score: u32,
	person: String,
}
enum Leaderboard
{
	Hiscores,
	Loscores,
	Pingscores,
}

#[derive(Clone)]
struct AppState {
	tx: mpsc::Sender<(Entry,Leaderboard)>,
	hiscores: Arc<Mutex<Vec<Entry>>>,
	loscores: Arc<Mutex<Vec<Entry>>>,
	pingscores: Arc<Mutex<HashMap<String,(u64,u32)>>>, // u64 is reset count and u32 is PB
}

static CHANCE: f64 = 1.0/3.0;

#[tokio::main]
async fn main() {
	let file_contents: String = fs::read_to_string("hiscores.json").unwrap();
	let hiscores: Arc<Mutex<Vec<Entry>>> = Arc::new(Mutex::new(serde_json::from_str(&file_contents).unwrap()));
	let file_contents: String = fs::read_to_string("loscores.json").unwrap();
	let loscores: Arc<Mutex<Vec<Entry>>> = Arc::new(Mutex::new(serde_json::from_str(&file_contents).unwrap()));
	let file_contents: String = fs::read_to_string("pingscores.json").unwrap();
	let pingscores: Arc<Mutex<HashMap<String,(u64,u32)>>> = Arc::new(Mutex::new(serde_json::from_str(&file_contents).unwrap()));
	let hiscore_clone1 = Arc::clone(&hiscores);
	let loscore_clone1 = Arc::clone(&loscores);
	let pingscore_clone1 = Arc::clone(&pingscores);

	let (tx, rx) = mpsc::channel::<(Entry,Leaderboard)>();

	tokio::spawn(
		async move {
			handle_hiscores(rx, hiscore_clone1, loscore_clone1, pingscore_clone1);
		}
	);
	let pingscore_clone2 = Arc::clone(&pingscores);
	tokio::spawn(
		async move {
			loop // write pingscores every 30s
			{
				sleep(Duration::from_millis(30000)).await;
				let pingscores = pingscore_clone2.lock().unwrap();
				let file_contents: String = serde_json::to_string(&pingscores.clone()).unwrap();
				drop(pingscores);
				let mut file = fs::OpenOptions::new().write(true).truncate(true).open("pingscores.json").unwrap();
				file.write_all(file_contents.as_bytes()).unwrap();
				file.flush().unwrap();
			}
		}
	);

	let state = AppState { tx,hiscores: Arc::clone(&hiscores),loscores: Arc::clone(&loscores), pingscores: Arc::clone(&pingscores)};
	let app = Router::new()
		.route("/", get(index))
		.route("/ws", get(ws_handler))
         .route("/leaderboard", get(leaderboard))
		.with_state(state);

	let listener = tokio::net::TcpListener::bind("0.0.0.0:8084")
		.await
		.unwrap();

	println!("http://0.0.0.0:8084");

	axum::serve(listener, app).await.unwrap();
}

async fn index() -> Html<&'static str> {
	Html(include_str!("../index.html"))
}

async fn leaderboard() -> Html<&'static str> {
	Html(include_str!("../leaderboard.html"))
}
// receiver: 0 for hiscore, 1 for loscore, 2 for pingscore
fn handle_hiscores(rx: mpsc::Receiver<(Entry, Leaderboard)>, hiscores_arc: Arc<Mutex<Vec<Entry>>>, loscores_arc: Arc<Mutex<Vec<Entry>>>,pingscores_arc: Arc<Mutex<HashMap<String,(u64,u32)>>>,) 
{
	// Panic galore
	let mut hiscores = hiscores_arc.lock().unwrap();
	hiscores.sort();
	hiscores.reverse();
	let file_contents: String = serde_json::to_string(&hiscores.clone()).unwrap();
	drop(hiscores);
	let mut file = fs::OpenOptions::new().write(true).truncate(true).open("hiscores.json").unwrap();
	file.write_all(file_contents.as_bytes()).unwrap();
	file.flush().unwrap();
	loop
	{
		match rx.recv()
		{
			Ok((new_entry,Leaderboard::Hiscores)) =>
			{
				let mut hiscores = hiscores_arc.lock().unwrap();
				if hiscores.get(19).is_none_or(|hiscore| new_entry.score > hiscore.score) {
					println!("New hiscore {new_entry:?}");
					hiscores.push(new_entry);
					hiscores.sort();
					hiscores.reverse();
					hiscores.truncate(20);
					let file_contents: String = serde_json::to_string(&hiscores.clone()).unwrap();
					drop(hiscores);
					let mut file = fs::OpenOptions::new().write(true).truncate(true).open("hiscores.json").unwrap();
					file.write_all(file_contents.as_bytes()).unwrap();
					file.flush().unwrap();
				}
			},
			Ok((new_entry,Leaderboard::Loscores)) =>
			{
				let mut loscores = loscores_arc.lock().unwrap();
				if loscores.get(19).is_none_or(|loscore| new_entry.score > loscore.score) {
					println!("New loscore {new_entry:?}");
					loscores.push(new_entry);
					loscores.sort();
					loscores.reverse();
					loscores.truncate(20);
					let file_contents: String = serde_json::to_string(&loscores.clone()).unwrap();
					drop(loscores);
					let mut file = fs::OpenOptions::new().write(true).truncate(true).open("loscores.json").unwrap();
					file.write_all(file_contents.as_bytes()).unwrap();
					file.flush().unwrap();
				}
			},
			Ok((new_entry,Leaderboard::Pingscores)) =>
			{
				let name = new_entry.person;
				let mut pingscores = pingscores_arc.lock().unwrap();
				if new_entry.score > pingscores.get(&name).unwrap_or(&(0,0)).1 // pb
				{
					pingscores.entry(name.clone()).or_insert((0,0)).1 = new_entry.score;
					println!("{name} new PB: {}",new_entry.score);
				};
				pingscores.entry(name.clone()).or_insert((0,0)).0 += 1; // reset count
				drop(pingscores);
			}
			Err(error) => println!("{error}"),
		}
	}
}

async fn ws_handler(
	ws: WebSocketUpgrade,
	State(state): State<AppState>,
) -> impl IntoResponse {
	ws.on_upgrade(move |socket| {
		let tx = state.tx.clone();
		let hiscores = Arc::clone(&state.hiscores);
		let loscores = Arc::clone(&state.loscores);
		let pingscores = Arc::clone(&state.pingscores);
		async move {
			handle_socket(socket, tx, hiscores, loscores, pingscores).await;
		}
	})
}

async fn handle_socket
(
	mut socket: WebSocket,
	tx: mpsc::Sender<(Entry,Leaderboard)>,
	hiscores_arc: Arc<Mutex<Vec<Entry>>>,
	loscores_arc: Arc<Mutex<Vec<Entry>>>,
	pingscores_arc: Arc<Mutex<HashMap<String,(u64,u32)>>>,
) {
	let mut value: u32 = 0;
	let msg = 
	{
		let hiscores = hiscores_arc.lock().unwrap();
		let loscores = loscores_arc.lock().unwrap();
		let pingscores = pingscores_arc.lock().unwrap();
		json!({ "hiscores": &*hiscores, "loscores": &*loscores, "pingscores": &*pingscores}).to_string()
	};
	let name_message = socket.next().await.unwrap().unwrap();
	let name: String = match name_message 
	{
		Message::Text(text) => validate_name(text.to_string()),
		_ => "anon".to_string(),
	};
	println!("Client connected: {name}");
	let mut resets: u32 = 0;
	let mut prev: u32 = 0;

	let _ = socket
		.send(Message::Text(msg.into()))
		.await;

	while let Some(msg) = socket.next().await {
		match msg {
			Ok(Message::Text(_)) => {
				if random_bool(CHANCE) { // reset
					let _ = tx.send((Entry{ person: name.clone(), score: value },Leaderboard::Hiscores)); //hiscores
					let _ = tx.send((Entry{ person: name.clone(), score: value },Leaderboard::Pingscores)); //pingscores
					resets += 1;
					value = 0
				} // 1/3 chance of failing
				else {
					value += 1;
					if prev == 0 {
						let _ = tx.send((Entry{ person: name.clone(), score: resets },Leaderboard::Loscores));//loscores
						resets = 0;
					}
				}
				let _ = socket.send(Message::Text(value.to_string().into())).await;
				prev = value;
			}

			Ok(Message::Close(_)) => {
				break;
			}

			_ => {}
		}
	}
}


fn validate_name(input: String) -> String {
	let input = input.trim();
    if input == "null"
    {
        return "anon".to_string();
    }
	// Length check
	if input.is_empty() || input.len() > 32 {
		return "anon".to_string();
	}

	// Allow only letters, numbers, _ and -
	let re = Regex::new(r"^[a-zA-Z0-9_-]+$").unwrap();

	if re.is_match(input) {
		input.to_string()
	} else {
		"anon".to_string()
	}
}