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

1 files changed, 24 insertions, 17 deletions

diff --git a/Role/User.hs b/Role/User.hs
index 49c263c..fdf4e53 100644
--- a/Role/User.hs
+++ b/Role/User.hs
@@ -14,6 +14,7 @@ import SessionID
 
 import Control.Concurrent.Async
 import Control.Concurrent.STM
+import Control.Concurrent.STM.TMChan
 import System.Process
 import System.Exit
 import qualified Data.ByteString as B
@@ -63,11 +64,11 @@ data UserState = UserState
 	, userSigVerifier :: SigVerifier
 	}
 
-user :: B.ByteString -> Pty -> TChan (Message Seen) -> TChan (Message Entered) -> IO ()
+user :: B.ByteString -> Pty -> TMChan (Message Seen) -> TMChan (Message Entered) -> IO ()
 user starttxt p ochan ichan = withLogger "debug-me.log" $ \logger -> do
 	-- Start by establishing our session key, and displaying the starttxt.
 	let initialmessage msg = do
-		atomically $ writeTChan ochan msg
+		atomically $ writeTMChan ochan msg
 		logger $ User msg
 	sk <- genMySessionKey
 	pk <- myPublicKey sk
@@ -100,9 +101,9 @@ forwardTtyInputToPty p = do
 			writePty p b
 			forwardTtyInputToPty p
 
--- | Forward things written to the Pty out the TChan, and also display
+-- | Forward things written to the Pty out the TMChan, and also display
 -- it on their Tty.
-sendPtyOutput :: Pty -> TChan (Message Seen) -> TVar UserState -> Logger -> IO ()
+sendPtyOutput :: Pty -> TMChan (Message Seen) -> TVar UserState -> Logger -> IO ()
 sendPtyOutput p ochan us logger = go
   where
 	go = do
@@ -117,7 +118,7 @@ sendPtyOutput p ochan us logger = go
 		go
 
 class SendableToDeveloper t where
-	sendDeveloper :: TChan (Message Seen) -> TVar UserState -> t -> POSIXTime -> STM (Message Seen)
+	sendDeveloper :: TMChan (Message Seen) -> TVar UserState -> t -> POSIXTime -> STM (Message Seen)
 
 instance SendableToDeveloper Seen where
 	sendDeveloper ochan us seen now = do
@@ -127,7 +128,7 @@ instance SendableToDeveloper Seen where
 			mkSigned (userSessionKey st) $ 
 				Activity seen (loggedHash prev)
 		let l = mkLog (User msg) now
-		writeTChan ochan msg
+		writeTMChan ochan msg
 		writeTVar us $ st { backLog = l :| toList bl }
 		return msg
 
@@ -137,23 +138,24 @@ instance SendableToDeveloper ControlAction where
 		let msg = ControlMessage $
 			mkSigned (userSessionKey st) (Control c)
 		-- Control messages are not kept in the backlog.
-		writeTChan ochan msg
+		writeTMChan ochan msg
 		return msg
 
--- | Read things to be entered from the TChan, verify if they're legal,
+-- | Read things to be entered from the TMChan, verify if they're legal,
 -- and send them to the Pty.
-sendPtyInput :: TChan (Message Entered) -> TChan (Message Seen) -> Pty -> TVar UserState -> Logger -> IO ()
+sendPtyInput :: TMChan (Message Entered) -> TMChan (Message Seen) -> Pty -> TVar UserState -> Logger -> IO ()
 sendPtyInput ichan ochan p us logger = go
   where
 	go = do
 		now <- getPOSIXTime
 		v <- atomically $ getDeveloperMessage ichan ochan us now
 		case v of
-			InputMessage msg@(ActivityMessage entered) -> do
+			Nothing -> return ()
+			Just (InputMessage msg@(ActivityMessage entered)) -> do
 				logger $ Developer msg
 				writePty p $ val $ enteredData $ activity entered
 				go
-			InputMessage msg@(ControlMessage (Control c _)) -> do
+			Just (InputMessage msg@(ControlMessage (Control c _))) -> do
 				logger $ Developer msg
 				case c of
 					SessionKey pk -> do
@@ -162,10 +164,10 @@ sendPtyInput ichan ochan p us logger = go
 					Rejected r -> error $ "User side received a Rejected: " ++ show r
 					SessionKeyAccepted _ -> error "User side received a SessionKeyAccepted"
 					SessionKeyRejected _ -> error "User side received a SessionKeyRejected"
-			RejectedMessage rej -> do
+			Just (RejectedMessage rej) -> do
 				logger $ User rej
 				go
-			BadlySignedMessage _ -> go
+			Just (BadlySignedMessage _) -> go
 
 data Input
 	= InputMessage (Message Entered)
@@ -177,9 +179,14 @@ data Input
 -- signature of the message is only verified against the key in it), and
 -- make sure it's legal before returning it. If it's not legal, sends a
 -- Reject message.
-getDeveloperMessage :: TChan (Message Entered) -> TChan (Message Seen) -> TVar UserState -> POSIXTime -> STM Input
-getDeveloperMessage ichan ochan us now = do
-	msg <- readTChan ichan
+getDeveloperMessage :: TMChan (Message Entered) -> TMChan (Message Seen) -> TVar UserState -> POSIXTime -> STM (Maybe Input)
+getDeveloperMessage ichan ochan us now = maybe
+	(return Nothing)
+	(\msg -> Just <$> getDeveloperMessage' msg ochan us now)
+	=<< readTMChan ichan
+
+getDeveloperMessage' :: Message Entered -> TMChan (Message Seen) -> TVar UserState -> POSIXTime -> STM Input
+getDeveloperMessage' msg ochan us now = do
 	st <- readTVar us
 	case msg of
 		ControlMessage (Control (SessionKey pk) _) -> do
@@ -209,7 +216,7 @@ getDeveloperMessage ichan ochan us now = do
 
 -- | Check if the public key a developer presented is one we want to use,
 -- and if so, add it to the userSigVerifier.
-checkDeveloperPublicKey :: TChan (Message Seen) -> TVar UserState -> Logger -> PublicKey -> IO ()
+checkDeveloperPublicKey :: TMChan (Message Seen) -> TVar UserState -> Logger -> PublicKey -> IO ()
 checkDeveloperPublicKey ochan us logger pk = do
 	now <- getPOSIXTime
 	-- TODO check gpg sig..