bwc9876.dev
Live location tracking and playback for the game "manhunt"
1use std::{
2 collections::{HashMap, HashSet},
3 net::SocketAddr,
4};
5
6use axum::{Error as AxumError, extract::ws::Message, http::StatusCode};
7use matchbox_protocol::PeerId;
8use matchbox_signaling::{
9 SignalingError, SignalingState,
10 common_logic::{self, StateObj},
11};
12use rand::{rngs::ThreadRng, seq::IndexedRandom};
13use tokio::sync::mpsc::UnboundedSender;
14use tokio_util::sync::CancellationToken;
15
16pub type RoomId = String;
17pub type Sender = UnboundedSender<Result<Message, AxumError>>;
18
19#[derive(Debug, Clone)]
20struct Match {
21 pub open_lobby: bool,
22 cancel: CancellationToken,
23 pub players: HashSet<PeerId>,
24}
25
26#[derive(Debug, Clone)]
27struct Peer {
28 pub room: RoomId,
29 sender: Sender,
30}
31
32impl Match {
33 pub fn new() -> Self {
34 Self {
35 open_lobby: true,
36 cancel: CancellationToken::new(),
37 players: HashSet::with_capacity(10),
38 }
39 }
40}
41
42#[derive(Default, Debug, Clone)]
43pub struct ServerState {
44 waiting_clients: StateObj<HashMap<SocketAddr, (RoomId, bool)>>,
45 queued_clients: StateObj<HashMap<PeerId, (RoomId, bool)>>,
46 matches: StateObj<HashMap<RoomId, Match>>,
47 clients: StateObj<HashMap<PeerId, Peer>>,
48}
49
50impl SignalingState for ServerState {}
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub enum RoomError {
54 /// Room already exists
55 Exists,
56 /// Room was not found
57 NotFound,
58}
59
60impl From<RoomError> for StatusCode {
61 fn from(val: RoomError) -> Self {
62 match val {
63 RoomError::Exists => StatusCode::CONFLICT,
64 RoomError::NotFound => StatusCode::NOT_FOUND,
65 }
66 }
67}
68
69const ROOM_CODE_CHAR_POOL: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890";
70const ROOM_CODE_LEN: usize = 6;
71const MAX_ROOM_TRIES: usize = 25;
72
73#[derive(Debug, Default, Clone, Copy)]
74pub struct NoRoomsError;
75
76impl ServerState {
77 fn random_room_code(rng: &mut ThreadRng) -> RoomId {
78 ROOM_CODE_CHAR_POOL
79 .sample(rng, ROOM_CODE_LEN)
80 .copied()
81 .map(char::from)
82 .collect()
83 }
84
85 fn check_room_taken(&self, code: &RoomId) -> bool {
86 self.matches.lock().unwrap().contains_key(code)
87 }
88
89 pub fn generate_room_code(&self) -> Result<RoomId, NoRoomsError> {
90 let mut rng = rand::rng();
91 for _ in 0..MAX_ROOM_TRIES {
92 let code = Self::random_room_code(&mut rng);
93
94 if !self.check_room_taken(&code) {
95 return Ok(code);
96 }
97 }
98 Err(NoRoomsError)
99 }
100
101 fn add_client(&mut self, origin: SocketAddr, code: RoomId, host: bool) {
102 self.waiting_clients
103 .lock()
104 .unwrap()
105 .insert(origin, (code.clone(), host));
106 }
107
108 pub fn room_is_open(&self, room_id: &str) -> bool {
109 self.matches
110 .lock()
111 .unwrap()
112 .get(room_id)
113 .is_some_and(|m| m.open_lobby)
114 }
115
116 /// Mark a match as started, disallowing others from joining
117 pub fn mark_started(&mut self, room: &RoomId) {
118 if let Some(mat) = self.matches.lock().unwrap().get_mut(room) {
119 mat.open_lobby = false;
120 }
121 }
122
123 /// Create a new room with the given code, should be called when someone wants to host a game.
124 /// Returns false if a room with that code already exists.
125 fn create_room(&mut self, origin: SocketAddr, code: RoomId) -> bool {
126 let mut matches = self.matches.lock().unwrap();
127 if matches.contains_key(&code) {
128 false
129 } else {
130 matches.insert(code.clone(), Match::new());
131 drop(matches);
132 self.add_client(origin, code, true);
133 true
134 }
135 }
136
137 /// Try to join a room by a code, returns `true` if successful
138 fn try_join_room(&mut self, origin: SocketAddr, code: RoomId) -> bool {
139 if self.room_is_open(&code) {
140 self.waiting_clients
141 .lock()
142 .unwrap()
143 .insert(origin, (code, false));
144 true
145 } else {
146 false
147 }
148 }
149
150 /// Try to create / join a room
151 pub fn handle_room(
152 &mut self,
153 create: bool,
154 origin: SocketAddr,
155 code: RoomId,
156 ) -> Result<(), RoomError> {
157 match create {
158 true => match self.create_room(origin, code) {
159 true => Ok(()),
160 false => Err(RoomError::Exists),
161 },
162 false => match self.try_join_room(origin, code) {
163 true => Ok(()),
164 false => Err(RoomError::NotFound),
165 },
166 }
167 }
168
169 /// Assign a peer an id
170 pub fn assign_peer_id(&mut self, origin: SocketAddr, peer_id: PeerId) {
171 let target_room = self
172 .waiting_clients
173 .lock()
174 .unwrap()
175 .remove(&origin)
176 .expect("origin not waiting?");
177
178 self.queued_clients
179 .lock()
180 .unwrap()
181 .insert(peer_id, target_room);
182 }
183
184 /// Add a peer to a room, returns other peers in the match currently
185 pub fn add_peer(
186 &mut self,
187 peer_id: PeerId,
188 sender: Sender,
189 ) -> (bool, CancellationToken, Vec<PeerId>) {
190 let (target_room, host) = self
191 .queued_clients
192 .lock()
193 .unwrap()
194 .remove(&peer_id)
195 .expect("peer not waiting?");
196 let mut matches = self.matches.lock().unwrap();
197 let mat = matches.get_mut(&target_room).expect("Room not found?");
198 let peers = mat.players.iter().copied().collect::<Vec<_>>();
199 mat.players.insert(peer_id);
200 let cancel = mat.cancel.clone();
201 drop(matches);
202 let peer = Peer {
203 room: target_room,
204 sender,
205 };
206 self.clients.lock().unwrap().insert(peer_id, peer);
207 (host, cancel, peers)
208 }
209
210 /// Disconnect a peer from a room. Automatically deletes the room if no peers remain. Returns
211 /// the removed peer and the set of other peers in the room that need to be notified
212 pub fn remove_peer(&mut self, peer_id: PeerId, host: bool) -> Option<Vec<PeerId>> {
213 let removed_peer = self.clients.lock().unwrap().remove(&peer_id)?;
214
215 let mut matches = self.matches.lock().unwrap();
216
217 let other_peers = matches
218 .get_mut(&removed_peer.room)
219 .map(|m| {
220 m.players.remove(&peer_id);
221 m.players.iter().copied().collect::<Vec<_>>()
222 })
223 .unwrap_or_default();
224
225 if host {
226 if let Some(mat) = matches.get_mut(&removed_peer.room).filter(|m| m.open_lobby) {
227 // If we're host, disconnect everyone else
228 mat.open_lobby = false;
229 mat.cancel.cancel();
230 }
231 }
232
233 if other_peers.is_empty() {
234 matches.remove(&removed_peer.room);
235 }
236
237 Some(other_peers)
238 }
239
240 pub fn try_send(&self, peer: PeerId, msg: Message) -> Result<(), SignalingError> {
241 self.clients
242 .lock()
243 .unwrap()
244 .get(&peer)
245 .ok_or(SignalingError::UnknownPeer)
246 .and_then(|peer| common_logic::try_send(&peer.sender, msg))
247 }
248}
249
250#[cfg(test)]
251mod tests {
252 use std::net::{IpAddr, Ipv4Addr};
253
254 use uuid::Uuid;
255
256 use super::*;
257
258 const fn origin(p: u16) -> SocketAddr {
259 SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), p)
260 }
261
262 const fn peer(p: u16) -> PeerId {
263 PeerId(Uuid::from_u128(p as u128))
264 }
265
266 fn dummy_sender() -> Sender {
267 let (s, _) = tokio::sync::mpsc::unbounded_channel();
268 s
269 }
270
271 fn handle_assign_add(state: &mut ServerState, create: bool, code: &str, p: u16) {
272 state
273 .handle_room(create, origin(p), code.to_string())
274 .expect("Failed to handle room");
275 state.assign_peer_id(origin(p), peer(p));
276 state.add_peer(peer(p), dummy_sender());
277 }
278
279 fn quick_create(state: &mut ServerState, code: &str, p: u16) {
280 handle_assign_add(state, true, code, p);
281 }
282
283 fn quick_join(state: &mut ServerState, code: &str, p: u16) {
284 handle_assign_add(state, false, code, p);
285 }
286
287 #[test]
288 fn test_add_waiting_host() {
289 let mut state = ServerState::default();
290
291 let code = "aaa";
292
293 state
294 .handle_room(true, origin(1), code.to_string())
295 .expect("Could not create room");
296 assert_eq!(
297 *state.waiting_clients.lock().unwrap(),
298 HashMap::from_iter([(origin(1), (code.to_string(), true))])
299 );
300 assert!(state.room_is_open(code))
301 }
302
303 #[test]
304 fn test_add_waiting_player() {
305 let mut state = ServerState::default();
306
307 let code = "aaaa";
308
309 quick_create(&mut state, code, 1);
310
311 state
312 .handle_room(false, origin(2), code.to_string())
313 .expect("Failed to join room");
314 assert_eq!(
315 *state.waiting_clients.lock().unwrap(),
316 HashMap::from_iter([(origin(2), (code.to_string(), false))])
317 );
318 }
319
320 #[test]
321 fn test_assign_id() {
322 let mut state = ServerState::default();
323
324 let code = "aaa";
325
326 state
327 .handle_room(true, origin(1), code.to_string())
328 .expect("Could not create room");
329
330 state.assign_peer_id(origin(1), peer(1));
331
332 assert!(state.waiting_clients.lock().unwrap().is_empty());
333 assert_eq!(
334 *state.queued_clients.lock().unwrap(),
335 HashMap::from_iter([(peer(1), (code.to_string(), true))]),
336 )
337 }
338
339 #[test]
340 fn test_add_peer() {
341 let mut state = ServerState::default();
342
343 let code = "aaa";
344
345 state
346 .handle_room(true, origin(1), code.to_string())
347 .expect("Could not create room");
348
349 state.assign_peer_id(origin(1), peer(1));
350
351 let (_, _, others) = state.add_peer(peer(1), dummy_sender());
352
353 assert!(state.waiting_clients.lock().unwrap().is_empty());
354 assert!(state.queued_clients.lock().unwrap().is_empty());
355 assert!(others.is_empty());
356 assert!(
357 state
358 .clients
359 .lock()
360 .unwrap()
361 .get(&peer(1))
362 .is_some_and(|p| p.room == code)
363 );
364 assert!(
365 state
366 .matches
367 .lock()
368 .unwrap()
369 .get(&code.to_string())
370 .is_some_and(|m| m.players.contains(&peer(1)))
371 );
372 }
373
374 #[test]
375 fn test_join_add_peer() {
376 let mut state = ServerState::default();
377
378 let code = "abcd";
379
380 quick_create(&mut state, code, 1);
381
382 state
383 .handle_room(false, origin(2), code.to_string())
384 .expect("Failed to join");
385 state.assign_peer_id(origin(2), peer(2));
386
387 let (_, _, others) = state.add_peer(peer(2), dummy_sender());
388
389 assert_eq!(others, vec![peer(1)]);
390 assert!(
391 state
392 .clients
393 .lock()
394 .unwrap()
395 .get(&peer(2))
396 .is_some_and(|p| p.room == code)
397 );
398 assert!(
399 state
400 .matches
401 .lock()
402 .unwrap()
403 .get(&code.to_string())
404 .is_some_and(|m| m.players.contains(&peer(1)) && m.players.contains(&peer(2)))
405 );
406 }
407
408 #[test]
409 fn test_player_leave() {
410 let mut state = ServerState::default();
411
412 let code = "asdfasdfasdfasdf";
413
414 quick_create(&mut state, code, 1);
415 quick_join(&mut state, code, 2);
416
417 let others = state.remove_peer(peer(2), false);
418
419 assert_eq!(others, Some(vec![peer(1)]));
420 assert!(
421 state
422 .matches
423 .lock()
424 .unwrap()
425 .get(&code.to_string())
426 .is_some_and(|m| m.players.contains(&peer(1)) && !m.players.contains(&peer(2)))
427 );
428 assert!(!state.clients.lock().unwrap().contains_key(&peer(2)));
429 }
430
431 #[test]
432 fn test_player_leave_only_one() {
433 let mut state = ServerState::default();
434
435 let code = "asdfasdfasdfasdf";
436
437 quick_create(&mut state, code, 1);
438
439 let others = state.remove_peer(peer(1), true);
440
441 assert!(others.is_some_and(|v| v.is_empty()));
442 assert!(state.matches.lock().unwrap().is_empty());
443 assert!(state.clients.lock().unwrap().is_empty());
444 }
445
446 #[test]
447 fn test_host_leave_with_players() {
448 let mut state = ServerState::default();
449
450 let code = "asdfasdfasdfasdf";
451
452 quick_create(&mut state, code, 1);
453 quick_join(&mut state, code, 2);
454
455 let others = state.remove_peer(peer(1), true);
456
457 assert_eq!(others, Some(vec![peer(2)]));
458 let matches = state.matches.lock().unwrap();
459 let mat = &matches[&code.to_string()];
460 assert!(mat.cancel.is_cancelled());
461 assert!(!mat.open_lobby);
462 }
463
464 #[test]
465 fn test_host_leave_with_players_but_started() {
466 let mut state = ServerState::default();
467
468 let code = "asdfasdfasdfasdf";
469
470 quick_create(&mut state, code, 1);
471 quick_join(&mut state, code, 2);
472
473 state.mark_started(&code.to_string());
474
475 let others = state.remove_peer(peer(1), true);
476
477 assert_eq!(others, Some(vec![peer(2)]));
478 let matches = state.matches.lock().unwrap();
479 let mat = &matches[&code.to_string()];
480 assert!(!mat.cancel.is_cancelled());
481 assert!(!mat.open_lobby);
482 }
483
484 #[test]
485 fn test_join_no_match() {
486 let mut state = ServerState::default();
487
488 let code = "asdfasdf";
489
490 let res = state.handle_room(false, origin(1), code.to_string());
491 assert_eq!(res, Err(RoomError::NotFound));
492 }
493
494 #[test]
495 fn test_create_exists() {
496 let mut state = ServerState::default();
497
498 let code = "cdf";
499
500 quick_create(&mut state, code, 1);
501
502 let res = state.handle_room(true, origin(2), code.to_string());
503 assert_eq!(res, Err(RoomError::Exists));
504 }
505
506 #[test]
507 fn test_join_started() {
508 let mut state = ServerState::default();
509
510 let code = "qwerty";
511
512 quick_create(&mut state, code, 1);
513 quick_join(&mut state, code, 2);
514
515 state.mark_started(&code.to_string());
516
517 assert!(
518 state
519 .matches
520 .lock()
521 .unwrap()
522 .get(&code.to_string())
523 .is_some_and(|m| !m.open_lobby)
524 );
525
526 let res = state.handle_room(false, origin(3), code.to_string());
527 assert_eq!(res, Err(RoomError::NotFound));
528 }
529}