button/src/main.rs

// TODO rename pingscores userscores
use std::
{
	collections::HashMap,
	fs,
	io::Write,
	ops::{Deref, DerefMut},
	sync::Arc,
};

use axum::
{
	Router,
	extract::
	{
		State,
		ws::{Message, WebSocket, WebSocketUpgrade},
	},
	response::IntoResponse,
	routing::get,
};
use tokio::
{
	sync::{Mutex, mpsc},
	time::{Duration, sleep},
};
use tower_http::services::ServeDir;
use futures::StreamExt as _;
use rand::random_bool;
use serde::{Deserialize, Serialize, de};
use serde_json::json;
use regex::Regex;

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

#[derive(Clone)]
struct AppState
{
	tx: mpsc::Sender<LeaderboardUpdate>,
	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
}

struct LeaderboardUpdate
{
	name: Arc<str>,
	update: LeaderboardUpdateType,
}

enum LeaderboardUpdateType
{
	Reset { hiscore_pingscore: u32 },
	Increment { loscore: u32 },
}

const CHANCE: f64 = 1.0 / 3.0;

const PATH_HISCORES: &str = "hiscores.json";
const PATH_LOSCORES: &str = "loscores.json";
const PATH_PINGSCORES: &str = "pingscores.json";

const MAX_LEADERBOARD: usize = 20;

#[tokio::main]
async fn main() -> anyhow::Result<()>
{
	fn read_file<T: for<'de> de::Deserialize<'de>>
	(
		file_path: &str,
	)
	 -> anyhow::Result<Arc<Mutex<T>>>
	{
		let file_contents: String = fs::read_to_string(file_path)?;
		Ok(Arc::new(Mutex::new(serde_json::from_str(&file_contents)?)))
	}

	/// Makes the vector at `vec` one with a capacity of exactly [`MAX_LEADERBOARD`] if `vec` is
	/// smaller or equal.
	fn exact_leaderboard<T>(mut vec: impl DerefMut<Target = Vec<T>>)
	{
		let old_vec = std::mem::replace(&mut *vec, Vec::with_capacity(MAX_LEADERBOARD));
		for e in old_vec
		{
			vec.push(e);
		}
	}

	let hiscores: Arc<Mutex<Vec<Entry>>> = read_file(PATH_HISCORES)?;
	exact_leaderboard(hiscores.lock().await);
	let loscores: Arc<Mutex<Vec<Entry>>> = read_file(PATH_LOSCORES)?;
	exact_leaderboard(loscores.lock().await);
	let pingscores: Arc<Mutex<HashMap<String, (u64, u32)>>> = read_file(PATH_PINGSCORES)?;

	let (tx, rx) = mpsc::channel::<LeaderboardUpdate>(1024);
	{
		let (hiscores, loscores, pingscores) =
			(hiscores.clone(), loscores.clone(), pingscores.clone());
		tokio::spawn
		(async move {
			handle_hiscores(rx, &hiscores, &loscores, &pingscores).await;
		});
	}

	{
		let pingscores = pingscores.clone();
		tokio::spawn
		(async move {
			// write pingscores every 30s
			loop
			{
				sleep(Duration::from_secs(30)).await;
				let pingscores = pingscores.lock().await;
				let file_contents: String = serde_json::to_string(&pingscores.clone())
					.expect("failed to serialize pingscores");
				drop(pingscores);
				let mut file = fs::OpenOptions::new()
					.write(true)
					.truncate(true)
					.open(PATH_PINGSCORES)
					.expect("failed to open pingscores file");
				file.write_all(file_contents.as_bytes())
					.expect("failed to write pingscores");
				drop(file);
			}
		});
	}

	let static_files = ServeDir::new("./static");
	let app = Router::new()
		.fallback_service(static_files)
		.route("/ws", get(ws_handler))
		.route("/ws-leaderboard", get(leaderboard_handler))
		.with_state
		(AppState {
			tx,
			hiscores: Arc::clone(&hiscores),
			loscores: Arc::clone(&loscores),
			pingscores: Arc::clone(&pingscores),
		});

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

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

	axum::serve(listener, app).await?;

	Ok(())
}
// receiver: 0 for hiscore, 1 for loscore, 2 for pingscore
async fn handle_hiscores
(
	mut rx: mpsc::Receiver<LeaderboardUpdate>,
	hiscores: &Mutex<Vec<Entry>>,
	loscores: &Mutex<Vec<Entry>>,
	pingscores: &Mutex<HashMap<String, (u64, u32)>>,
)
{
	fn update_scoretable<G: Deref<Target = Vec<Entry>> + DerefMut>
	(
		score_name: &str,
		mut scoretable_lock: G,
		name: &str,
		score: u32,
		file_path: &str,
	)
	 -> anyhow::Result<()>
	{
		let scoretable = &mut *scoretable_lock;
		if let Some(index_to_insert_at) = scoretable.iter().position(|e| score > e.score)
		{
			println!("New {score_name} {score} by {name}");
			scoretable[index_to_insert_at..].rotate_right(1);
			let push_out = std::mem::replace
			(
				&mut scoretable[index_to_insert_at],
				Entry
				{
					score,
					person: name.to_string(),
				},
			);
			if scoretable.len() < MAX_LEADERBOARD
			{
				scoretable.push(push_out);
			}

			let file_contents: String = serde_json::to_string(&*scoretable_lock)?;
			drop(scoretable_lock);
			let mut file = fs::OpenOptions::new()
				.write(true)
				.truncate(true)
				.open(file_path)?;
			file.write_all(file_contents.as_bytes())?;
			file.flush()?;
		}
		else if scoretable.len() < MAX_LEADERBOARD
		{
			println!("New {score_name} {score} by {name}");
			scoretable.push
			(Entry {
				score,
				person: name.to_string(),
			});
		}

		Ok(())
	}

	// Panic galore
	let mut hiscores_lock = hiscores.lock().await;
	hiscores_lock.sort();
	hiscores_lock.reverse();
	let file_contents: String =
		serde_json::to_string(&hiscores_lock.clone()).expect("failed to serialize hiscores");
	drop(hiscores_lock);
	let mut file = fs::OpenOptions::new()
		.write(true)
		.truncate(true)
		.open(PATH_HISCORES)
		.expect("failed to open hiscores");
	file.write_all(file_contents.as_bytes())
		.expect("failed to write hiscores");
	drop(file);

	loop
	{
		let LeaderboardUpdate { name, update } = rx.recv().await.expect("channel error");

		match update
		{
			LeaderboardUpdateType::Reset { hiscore_pingscore } =>
			{
				// Hiscore
				update_scoretable
				(
					"hiscore",
					hiscores.lock().await,
					&name,
					hiscore_pingscore,
					PATH_HISCORES,
				)
				.expect("failed to update hiscores");

				// Pingscore
				let mut pingscores = pingscores.lock().await;
				// pb
				if hiscore_pingscore > pingscores.get(&*name).unwrap_or(&(0, 0)).1
				{
					pingscores.entry(name.to_string()).or_insert((0, 0)).1 = hiscore_pingscore;
					println!("{name} new PB: {hiscore_pingscore}");
				}
				pingscores.entry(name.to_string()).or_insert((0, 0)).0 += 1; // reset count
				drop(pingscores);
			}

			LeaderboardUpdateType::Increment { loscore } =>
			{
				update_scoretable
				(
					"loscore",
					loscores.lock().await,
					&name,
					loscore,
					PATH_LOSCORES,
				)
				.expect("failed to update loscores");
			}
		}
	}
}

async fn leaderboard_handler
(
	ws: WebSocketUpgrade,
	State(state): State<AppState>
)
-> impl IntoResponse
{
	ws.on_upgrade
	(|socket| async move {
		handle_leaderboard
		(
			socket,
			&state.hiscores,
			&state.loscores,
			&state.pingscores,
		)
		.await;
	})
}
async fn handle_leaderboard
(
	mut socket: WebSocket,
	hiscores: &Mutex<Vec<Entry>>,
	loscores: &Mutex<Vec<Entry>>,
	pingscores: &Mutex<HashMap<String, (u64, u32)>>,
)
{
	match socket.next().await
	{
		Some(Ok(Message::Text(selection))) =>
		{
			let msg;
			match selection.as_str()
			{
				"0" => // all the leaderboards
				{
					msg =
					{
						json!
						({
							"hiscores": &*hiscores.lock().await,
							"loscores": &*loscores.lock().await,
							"pingscores": &*pingscores.lock().await
						})
						.to_string()
					};
				},
				"1" => // just the hiscores table
				{
					msg = json! ({ "hiscores": &*hiscores.lock().await }).to_string()
				},
				"2" => // just the loscores table
				{
					msg = json! ({ "loscores": &*hiscores.lock().await }).to_string()
				},
				"3" => // just the pingscores table
				{
					msg = json! ({ "pingscores": &*hiscores.lock().await }).to_string()
				},
				_ => { msg = "Invalid leaderboard selection, please use 0,1,2 or 3".to_string() },
			}
			let _ = socket.send(Message::Text(msg.into())).await;
		},
		_ =>
		{
			let _ = socket.send(Message::Text("Invalid leaderboard selection, please use 0,1,2 or 3".into())).await;
		},
	}
}

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

async fn handle_socket
(
	mut socket: WebSocket,
	tx: &mpsc::Sender<LeaderboardUpdate>,
)
{
	let mut value: u32 = 0;

	let Some(name) = socket.next().await else
	{
		eprintln!("No username");
		return;
	};
	let name: Arc<str> = match name.expect("failed to recv socket msg")
	{
		Message::Text(text) =>
		{
			Arc::from(validate_name(&text))
		}
		_ => Arc::from("anon"),
	};

	println!("Client connected: {name}");

	let mut resets: u32 = 0;
	let mut prev: u32 = 0;

	while let Some(msg) = socket.next().await
	{
		match msg
		{
			Ok(Message::Text(_)) =>
			{
				// 1/3 chance of failing
				if random_bool(CHANCE)
				{
					// reset
					let _ = tx
						.send(LeaderboardUpdate
						{
							name: name.clone(),
							update: LeaderboardUpdateType::Reset
							{
								hiscore_pingscore: value,
							},
						})
						.await;
					resets += 1;
					value = 0;
				}
				else
				{
					value += 1;
					if prev == 0
					{
						let _ = tx
							.send
							(LeaderboardUpdate {
								name: name.clone(),
								update: LeaderboardUpdateType::Increment { loscore: resets },
							})
							.await;
						resets = 0;
					}
				}
				let _ = socket.send(Message::Text(value.to_string().into())).await;
				prev = value;
			}

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

			_ => {}
		}
	}
}

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

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

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