button/src/main.rs

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

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

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

#[derive(Clone)]
struct AppState
{
	tx: mpsc::UnboundedSender<LeaderboardUpdate>,
	hiscores: Arc<RwLock<Vec<Entry>>>,
	loscores: Arc<RwLock<Vec<Entry>>>,
	pingscores: Arc<RwLock<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;

async fn write_file(file_path: &str, file_contents: &str) -> anyhow::Result<()>
{
	let mut file = tokio::fs::OpenOptions::new()
		.write(true)
		.truncate(true)
		.open(file_path)
		.await?;
	file.write_all(file_contents.as_bytes()).await?;
	file.flush().await?;
	Ok(())
}

#[tokio::main]
async fn main() -> anyhow::Result<()>
{
	fn read_file<T: for<'de> de::Deserialize<'de>>(
		file_path: &str,
	) -> anyhow::Result<Arc<RwLock<T>>>
	{
		let file_contents: String = fs::read_to_string(file_path)?;
		Ok(Arc::new(RwLock::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<RwLock<Vec<Entry>>> = read_file(PATH_HISCORES)?;
	exact_leaderboard(hiscores.write().await);
	let loscores: Arc<RwLock<Vec<Entry>>> = read_file(PATH_LOSCORES)?;
	exact_leaderboard(loscores.write().await);
	let pingscores: Arc<RwLock<HashMap<String, (u64, u32)>>> = read_file(PATH_PINGSCORES)?;

	let (tx, rx) = mpsc::unbounded_channel::<LeaderboardUpdate>();
	{
		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.read().await.clone();
				let file_contents: String =
					serde_json::to_string(&pingscores).expect("failed to serialize pingscores");
				write_file(PATH_PINGSCORES, &file_contents)
					.await
					.expect("failed to write pingscores");
			}
		});
	}

	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::UnboundedReceiver<LeaderboardUpdate>,
	hiscores: &RwLock<Vec<Entry>>,
	loscores: &RwLock<Vec<Entry>>,
	pingscores: &RwLock<HashMap<String, (u64, u32)>>,
)
{
	async 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 file_contents = {
			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);
				}

				Some(serde_json::to_string(&*scoretable_lock)?)
			}
			else if scoretable.len() < MAX_LEADERBOARD
			{
				println!("New {score_name} {score} by {name}");
				scoretable.push(Entry {
					score,
					person: name.to_string(),
				});
				None
			}
			else
			{
				None
			}
		};
		drop(scoretable_lock);
		if let Some(file_contents) = file_contents
		{
			write_file(file_path, &file_contents).await?;
		}

		Ok(())
	}

	// Panic galore
	let mut hiscores_lock = hiscores.write().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);
	write_file(PATH_HISCORES, &file_contents)
		.await
		.expect("failed to write hiscores");

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

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

				// Pingscore
				let mut pingscores = pingscores.write().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.write().await,
					&name,
					loscore,
					PATH_LOSCORES,
				)
				.await
				.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: &RwLock<Vec<Entry>>,
	loscores: &RwLock<Vec<Entry>>,
	pingscores: &RwLock<HashMap<String, (u64, u32)>>,
)
{
	match socket.next().await
	{
		Some(Ok(Message::Text(selection))) =>
		{
			let msg = match selection.as_str()
			{
				// all the leaderboards
				"0" =>
				{
					let hiscores = hiscores.read().await.clone();
					let loscores = loscores.read().await.clone();
					let pingscores = pingscores.read().await.clone();
					json!
					({
						"hiscores": hiscores,
						"loscores": loscores,
						"pingscores": pingscores
					})
					.to_string()
				}
				// just the hiscores table
				"1" => json! ({ "hiscores": hiscores.read().await.clone() }).to_string(),
				// just the loscores table
				"2" => json! ({ "loscores": hiscores.read().await.clone() }).to_string(),
				// just the pingscores table
				"3" => json! ({ "pingscores": hiscores.read().await.clone() }).to_string(),
				_ => "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::UnboundedSender<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,
						},
					});
					resets += 1;
					value = 0;
				}
				else
				{
					value += 1;
					if prev == 0
					{
						let _ = tx.send(LeaderboardUpdate {
							name: name.clone(),
							update: LeaderboardUpdateType::Increment { loscore: resets },
						});
						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";
	}

	if input
		.bytes()
		.all(|byte| byte.is_ascii_alphanumeric() || byte == b'_' || byte == b'-')
	{
		input
	}
	else
	{
		"anon"
	}
}