Receive and interpret node updates

client/Main.hs

@@ -3,12 +3,9 @@
 module Main where
 
 import           Brick
-import           Brick.Widgets.Border
-import           Brick.Widgets.Border.Style
 import           Control.Concurrent.Chan
 import           Control.Exception
 import           Control.Monad
-import           Data.Maybe
 import qualified Data.Set                   as Set
 import qualified Data.Text                  as T
 import qualified Graphics.Vty               as Vty
@@ -26,19 +23,38 @@ import           Forest.Util
 
 data Event = EventNode Node | EventConnectionClosed T.Text
 
-wsClientApp :: Chan Event -> WS.ClientApp ()
+sendCloseEvent :: Chan Event -> WS.ConnectionException -> IO ()
-wsClientApp eventChan conn = handle handleConnectionException $ forever $ do
+sendCloseEvent eventChan e = writeChan eventChan $ EventConnectionClosed $ T.pack $ show e
+
+receiveUpdates :: Chan Event -> Node -> WS.Connection -> IO ()
+receiveUpdates eventChan node conn = handle (sendCloseEvent eventChan) $ do
   maybePacket <- receivePacket conn
   case maybePacket of
-    Nothing -> pure ()
+    -- Wait for ws exception since 'receivePacket' should have already closed
+    -- the connection.
+    Nothing -> waitForCloseException conn
     Just packet -> case packet of
-      ServerHello _ node  -> writeChan eventChan (EventNode node)
+      ServerUpdate path subnode -> do
-      -- TODO properly insert node into existing tree
+        let newNode = replaceAt subnode path node
-      ServerUpdate _ node -> writeChan eventChan (EventNode node)
+        writeChan eventChan $ EventNode newNode
-  where
+        receiveUpdates eventChan newNode conn
-    handleConnectionException :: WS.ConnectionException -> IO ()
+      _ -> do
-    handleConnectionException e =
+        closeWithErrorMessage conn "Invalid packet: Expected update"
-      writeChan eventChan $ EventConnectionClosed $ T.pack $ show e
+        waitForCloseException conn
+
+wsClientApp :: Chan Event -> WS.ClientApp ()
+wsClientApp eventChan conn = handle (sendCloseEvent eventChan) $ do
+  maybePacket <- receivePacket conn
+  case maybePacket of
+    -- Wait for ws exception since 'receivePacket' should have already closed
+    -- the connection.
+    Nothing -> waitForCloseException conn
+    Just (ServerHello _ node) -> do
+      writeChan eventChan $ EventNode node
+      receiveUpdates eventChan node conn
+    _ -> do
+      closeWithErrorMessage conn "Invalid packet: Expected hello"
+      waitForCloseException conn
 
 {- Brick client application-}
 

src/Forest/Util.hs

@@ -7,9 +7,11 @@ module Forest.Util
   , sendPacket
   , receivePacket
   , closeWithErrorMessage
+  , waitForCloseException
   ) where
 
 import           Control.Concurrent.Async
+import           Control.Monad
 import           Data.Aeson
 import           Data.List
 import qualified Data.Text                as T
@@ -27,9 +29,6 @@ withThread thread main = withAsync thread $ const main
 sendPacket :: ToJSON a => WS.Connection -> a -> IO ()
 sendPacket conn packet = WS.sendTextData conn $ encode packet
 
-closeWithErrorMessage :: WS.Connection -> T.Text -> IO ()
-closeWithErrorMessage conn = WS.sendCloseCode conn 1003
-
 receivePacket :: FromJSON a => WS.Connection -> IO (Maybe a)
 receivePacket conn = do
   dataMessage <- WS.receiveDataMessage conn
@@ -43,3 +42,9 @@ receivePacket conn = do
     closeOnErrorMessage (Right a) = pure $ Just a
     closeOnErrorMessage (Left errorMsg) =
       Nothing <$ closeWithErrorMessage conn errorMsg
+
+closeWithErrorMessage :: WS.Connection -> T.Text -> IO ()
+closeWithErrorMessage conn = WS.sendCloseCode conn 1003
+
+waitForCloseException :: WS.Connection -> IO ()
+waitForCloseException conn = forever $ void $ WS.receiveDataMessage conn