From 3ac90e59b07d0ad1c33bd7203c153ec78fd22309 Mon Sep 17 00:00:00 2001
From: Tangent Wantwight <tangent128@gmail.com>
Date: Fri, 22 May 2020 22:50:28 -0400
Subject: [PATCH] Move websocket handler into dedicated function

---
 src/main.rs | 61 ++++++++++++++++++++++++++++++-----------------------
 1 file changed, 35 insertions(+), 26 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 8c7e4ba..6453754 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,11 +1,15 @@
-#[macro_use] extern crate log;
+#[macro_use]
+extern crate log;
 
-use anyhow::{Context, Result};
-use futures::stream::{FuturesUnordered, StreamExt};
+use anyhow::{Context, Error, Result};
+use futures::prelude::*;
+use net::{ClientMessage, ServerMessage};
+use serde_json::{from_str, Value};
 use std::net::ToSocketAddrs;
+use stream::{iter, FuturesUnordered};
 use structopt::StructOpt;
 use warp::{serve, ws, ws::Ws, Filter};
-use ws::Message;
+use ws::{Message, WebSocket};
 
 pub mod net;
 
@@ -23,28 +27,8 @@ async fn main() -> Result<()> {
     env_logger::init();
     let args = Args::from_args();
 
-    let socket_handler = ws().map(|upgrade: Ws| {
-        upgrade.on_upgrade(|websocket| {
-            let (mut sink, mut source) = websocket.split();
-            async move {
-                loop {
-                    match source.next().await {
-                        Some(Ok(msg)) => {
-                            if let Ok(text) = msg.to_str() {
-                                debug!("Text message: {}", text);
-                            } else {
-                                debug!("Other message");
-                            }
-                        },
-                        Some(Err(err)) => {
-                            warn!("Websocket connection lost: {}", err);
-                        },
-                        None => break,
-                    }
-                }
-            }
-        })
-    });
+    // dispatch websockets
+    let socket_handler = ws().map(|upgrade: Ws| upgrade.on_upgrade(handle_socket));
 
     let addrs = args
         .listen
@@ -60,3 +44,28 @@ async fn main() -> Result<()> {
 
     Ok(())
 }
+
+async fn handle_socket(websocket: WebSocket) {
+    let (_sink, mut source) = websocket.split();
+
+    let error: Option<Error> = loop {
+        match source.next().await {
+            Some(Ok(msg)) => match msg.to_str() {
+                Ok(json) => {
+                    match from_str::<ClientMessage<Value, Value>>(json).context("Parsing JSON") {
+                        Ok(msg) => {
+                            debug!("Client message: {:?}", &msg);
+                        }
+                        Err(error) => break Some(error),
+                    }
+                }
+                Err(()) => debug!("Non-text message"),
+            },
+            Some(Err(error)) => break Some(error.into()),
+            None => break None,
+        }
+    };
+    if let Some(error) = error {
+        warn!("Websocket connection lost: {:#}", error);
+    }
+}