From c940d3561622963ada091651ef1037f0fa670d0d Mon Sep 17 00:00:00 2001
From: ModZero <modzero@modzero.xyz>
Date: Mon, 1 Jun 2020 17:05:42 +0200
Subject: [PATCH] Some error handling

---
 src/main.rs | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index e462cb7..13a8799 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,13 +1,13 @@
 mod simulation;
 mod state;
 
-use std::{convert::Infallible, sync::Arc};
+use std::{sync::Arc};
 
 use futures::SinkExt;
 use serde_json::json;
 use tera::{Context, Tera};
 use tokio::sync::watch;
-use warp::{Filter, ws::{Message, WebSocket}};
+use warp::{Filter, reject::Reject, Rejection, ws::{Message, WebSocket}};
 
 use simulation::Simulation;
 use state::State;
@@ -73,7 +73,7 @@ async fn watch_state(mut ws: WebSocket, mut state_channel: watch::Receiver<State
 
         let result = match state {
             Some(state) => ws.send(Message::text(json!(state).to_string())).await,
-            None => ws.send(Message::text("Oh...")).await,
+            None => break
         };
 
         if let Err(_) = result {
@@ -82,8 +82,16 @@ async fn watch_state(mut ws: WebSocket, mut state_channel: watch::Receiver<State
     };
 }
 
-async fn get_state(mut channel: watch::Receiver<State>) -> Result<State, Infallible> {
-    Ok(channel.recv().await.unwrap())
+#[derive(Debug)]
+struct SimulationUnavailable;
+
+impl Reject for SimulationUnavailable {}
+
+async fn get_state(mut channel: watch::Receiver<State>) -> Result<State, Rejection> {
+    match channel.recv().await {
+        Some(state) => Ok(state),
+        None => Err(warp::reject::custom(SimulationUnavailable))
+    }
 }
 
 fn render(template: WithTemplate, tera: Arc<Tera>) -> impl warp::Reply {