switch to TMChans so they can be closed when a connection is Done

1 files changed, 37 insertions, 26 deletions

diff --git a/Role/Developer.hs b/Role/Developer.hs
index 0b8fdd9..ffba5c4 100644
--- a/Role/Developer.hs
+++ b/Role/Developer.hs
@@ -13,16 +13,18 @@ import Pty
 
 import Control.Concurrent.Async
 import Control.Concurrent.STM
+import Control.Concurrent.STM.TMChan
 import System.IO
 import qualified Data.ByteString as B
 import qualified Data.Text as T
 import Data.List
+import Data.Maybe
 import Control.Monad
 
 run :: DeveloperOpts -> IO ()
 run = run' developer . debugUrl
 
-run' :: (TChan (Message Entered) -> TChan (Message Seen) -> SessionID -> IO ()) -> UrlString -> IO ()
+run' :: (TMChan (Message Entered) -> TMChan (Message Seen) -> SessionID -> IO ()) -> UrlString -> IO ()
 run' runner url = void $ runClientApp app
   where
 	connect = ConnectMode (T.pack url)
@@ -32,7 +34,7 @@ userMessages :: LogMessage -> Maybe (Message Seen)
 userMessages (User m) = Just m
 userMessages (Developer _) = Nothing
 
-developer :: TChan (Message Entered) -> TChan (Message Seen) -> SessionID -> IO ()
+developer :: TMChan (Message Entered) -> TMChan (Message Seen) -> SessionID -> IO ()
 developer ichan ochan _ = withLogger "debug-me-developer.log" $ \logger -> do
 	devstate <- processSessionStart ochan logger
 	ok <- authUser ichan ochan devstate logger
@@ -53,8 +55,8 @@ data DeveloperState = DeveloperState
 	, developerSigVerifier :: SigVerifier
 	}
 
--- | Read things typed by the developer, and forward them to the TChan.
-sendTtyInput :: TChan (Message Entered) -> TVar DeveloperState -> Logger -> IO ()
+-- | Read things typed by the developer, and forward them to the TMChan.
+sendTtyInput :: TMChan (Message Entered) -> TVar DeveloperState -> Logger -> IO ()
 sendTtyInput ichan devstate logger = go
   where
 	go = do
@@ -76,7 +78,7 @@ sendTtyInput ichan devstate logger = go
 				}
 			let act = mkSigned (developerSessionKey ds) $
 				Activity entered (Just $ lastActivity ds)
-			writeTChan ichan (ActivityMessage act)
+			writeTMChan ichan (ActivityMessage act)
 			let acth = hash act
 			let ds' = ds
 				{ sentSince = sentSince ds ++ [b]
@@ -88,31 +90,35 @@ sendTtyInput ichan devstate logger = go
 		logger $ Developer $ ActivityMessage act
 		go
 
--- | Read activity from the TChan and display it to the developer.
-sendTtyOutput :: TChan (Message Seen) -> TVar DeveloperState -> Logger -> IO ()
+-- | Read activity from the TMChan and display it to the developer.
+sendTtyOutput :: TMChan (Message Seen) -> TVar DeveloperState -> Logger -> IO ()
 sendTtyOutput ochan devstate logger = go
   where
 	go = do
-		(o, msg) <- atomically $ getUserMessage ochan devstate
-		logger $ User msg
-		emitOutput o
-		go
+		v <- atomically $ getUserMessage ochan devstate
+		case v of
+			Nothing -> return ()
+			Just (o, msg) -> do
+				logger $ User msg
+				emitOutput o
+				go
 
 -- | Present our session key to the user. 
 -- Wait for them to accept or reject it, while displaying any Seen data
 -- in the meantime.
-authUser :: TChan (Message Entered) -> TChan (Message Seen) -> TVar DeveloperState -> Logger -> IO Bool
+authUser :: TMChan (Message Entered) -> TMChan (Message Seen) -> TVar DeveloperState -> Logger -> IO Bool
 authUser ichan ochan devstate logger = do
 	ds <- atomically $ readTVar devstate
 	pk <- myPublicKey (developerSessionKey ds)
 	let msg = ControlMessage $ mkSigned (developerSessionKey ds) 
 		(Control (SessionKey pk))
-	atomically $ writeTChan ichan msg
+	atomically $ writeTMChan ichan msg
 	logger $ Developer msg
 	waitresp pk
   where
 	waitresp pk = do
-		(o, msg) <- atomically $ getUserMessage ochan devstate
+		(o, msg) <- fromMaybe (error "No response from server to our session key")
+			<$> atomically (getUserMessage ochan devstate)
 		logger $ User msg
 		emitOutput o
 		case o of
@@ -142,16 +148,19 @@ emitOutput (GotControl _) =
 
 -- | Get messages from user, check their signature, and make sure that they
 -- are properly chained from past messages, before returning.
-getUserMessage :: TChan (Message Seen) -> TVar DeveloperState -> STM (Output, Message Seen)
+getUserMessage :: TMChan (Message Seen) -> TVar DeveloperState -> STM (Maybe (Output, Message Seen))
 getUserMessage ochan devstate = do
-	msg <- readTChan ochan
-	ds <- readTVar devstate
-	-- Check signature before doing anything else.
-	if verifySigned (developerSigVerifier ds) msg
-		then do
-			o <- process ds msg
-			return (o, msg)
-		else getUserMessage ochan devstate
+	mmsg <- readTMChan ochan
+	case mmsg of
+		Nothing -> return Nothing
+		Just msg -> do
+			ds <- readTVar devstate
+			-- Check signature before doing anything else.
+			if verifySigned (developerSigVerifier ds) msg
+				then do
+					o <- process ds msg
+					return (Just (o, msg))
+				else getUserMessage ochan devstate
   where
 	process ds (ActivityMessage act@(Activity (Seen (Val b)) _ _)) = do
 		let (legal, ds') = isLegalSeen act ds
@@ -224,9 +233,10 @@ isLegalSeen act@(Activity (Seen (Val b)) (Just hp) _) ds
 
 -- | Start by reading the initial two messages from the user side,
 -- their session key and the startup message.
-processSessionStart :: TChan (Message Seen) -> Logger -> IO (TVar DeveloperState)
+processSessionStart :: TMChan (Message Seen) -> Logger -> IO (TVar DeveloperState)
 processSessionStart ochan logger = do
-	sessionmsg <- atomically $ readTChan ochan
+	sessionmsg <- fromMaybe (error "Did not get session initialization message")
+		<$> atomically (readTMChan ochan)
 	logger $ User sessionmsg
 	sigverifier <- case sessionmsg of
 		ControlMessage c@(Control (SessionKey pk) _) ->
@@ -235,7 +245,8 @@ processSessionStart ochan logger = do
 				then return sv
 				else error "Badly signed session initialization message"
 		_ -> error $ "Unexpected session initialization message: " ++ show sessionmsg
-	startmsg <- atomically $ readTChan ochan
+	startmsg <- fromMaybe (error "Did not get session startup message")
+		<$> atomically (readTMChan ochan)
 	logger $ User startmsg
 	starthash <- case startmsg of
 		ActivityMessage act@(Activity (Seen (Val b)) Nothing _)