include "session is done" in the session transcript

Including the process exit status.

And cleaner Role.User shutdown sequence.

1 files changed, 40 insertions, 28 deletions

diff --git a/Role/User.hs b/Role/User.hs
index fdf4e53..94625c2 100644
--- a/Role/User.hs
+++ b/Role/User.hs
@@ -26,25 +26,33 @@ import System.IO
 import System.Environment
 
 run :: UserOpts -> IO ExitCode
-run os = do
-	(cmd, cmdparams) <- shellCommand os
-	fromMaybe (ExitFailure 101) <$> go cmd cmdparams startSession
+run os = fromMaybe (ExitFailure 101) <$> connect
   where
-	go cmd cmdparams startmsg = do
+	connect = do
 		putStr "Connecting to debug-me server..."
 		hFlush stdout
-		runClientApp $ clientApp (InitMode mempty) User developerMessages $ \ichan ochan sid -> do
+		runClientApp $ clientApp (InitMode mempty) User developerMessages $ \ochan ichan sid -> do
 			let url = sessionIDUrl sid "localhost" 8081
 			putStrLn ""
 			putStrLn "Others can connect to this session and help you debug by running:"
 			putStrLn $ "  debug-me --debug " ++ url
 			hFlush stdout
-			runWithPty cmd cmdparams $ \(p, ph) -> do
-				uthread <- async (user startmsg p ichan ochan)
-				exitstatus <- waitForProcess ph
-				cancel uthread
-				sessionDone
-				return exitstatus
+			withLogger "debug-me.log" $ go ochan ichan
+	go ochan ichan logger = do
+		(cmd, cmdparams) <- shellCommand os
+		runWithPty cmd cmdparams $ \(p, ph) -> do
+			us <- startProtocol startSession ochan logger
+			p1 <- async $ sendPtyOutput p ochan us logger
+			p2 <- async $ sendPtyInput ichan ochan p us logger
+				`race` forwardTtyInputToPty p
+			exitstatus <- waitForProcess ph
+			displayOutput ochan us logger $
+				rawLine "" <>
+				rawLine (endSession exitstatus)
+			atomically $ closeTMChan ichan
+			cancel p1
+			_ <- waitCatch p2
+			return exitstatus
 
 developerMessages :: LogMessage -> Maybe (Message Entered)
 developerMessages (Developer m) = Just m
@@ -64,9 +72,9 @@ data UserState = UserState
 	, userSigVerifier :: SigVerifier
 	}
 
-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.
+-- | Start by establishing our session key, and displaying the starttxt.
+startProtocol :: B.ByteString -> TMChan (Message Seen) -> Logger -> IO (TVar UserState)
+startProtocol starttxt ochan logger = do
 	let initialmessage msg = do
 		atomically $ writeTMChan ochan msg
 		logger $ User msg
@@ -74,22 +82,18 @@ user starttxt p ochan ichan = withLogger "debug-me.log" $ \logger -> do
 	pk <- myPublicKey sk
 	let c = mkSigned sk $ Control (SessionKey pk)
 	initialmessage $ ControlMessage c
-	let starttxt' = starttxt <> "\r\n"
+	let starttxt' = rawLine starttxt
 	let act = mkSigned sk $ Activity (Seen (Val starttxt')) Nothing
 	let startmsg = ActivityMessage act
 	B.hPut stdout starttxt'
 	hFlush stdout
 	initialmessage startmsg
 	l <- mkLog (User startmsg) <$> getPOSIXTime
-	us <- newTVarIO $ UserState
+	newTVarIO $ UserState
 		{ backLog = l :| []
 		, userSessionKey = sk
 		, userSigVerifier = mempty
 		}
-	_ <- sendPtyOutput p ochan us logger
-		`concurrently` sendPtyInput ichan ochan p us logger
-			`concurrently` forwardTtyInputToPty p
-	return ()
 
 -- | Forward things the user types to the Pty.
 forwardTtyInputToPty :: Pty -> IO ()
@@ -107,16 +111,24 @@ sendPtyOutput :: Pty -> TMChan (Message Seen) -> TVar UserState -> Logger -> IO
 sendPtyOutput p ochan us logger = go
   where
 	go = do
-		b <- readPty p
-		B.hPut stdout b
-		hFlush stdout
-		now <- getPOSIXTime
-		l <- atomically $ do
-			let seen = Seen (Val b)
-			sendDeveloper ochan us seen now
-		logger $ User l
+		displayOutput ochan us logger =<< readPty p
 		go
 
+-- | Display to Tty and send out the TMChan.
+displayOutput :: TMChan (Message Seen) -> TVar UserState -> Logger -> B.ByteString -> IO ()
+displayOutput ochan us logger b = do
+	B.hPut stdout b
+	hFlush stdout
+	now <- getPOSIXTime
+	l <- atomically $ do
+		let seen = Seen (Val b)
+		sendDeveloper ochan us seen now
+	logger $ User l
+
+-- | Since the Tty is in raw mode, need \r before \n
+rawLine :: B.ByteString -> B.ByteString
+rawLine b = b <> "\r\n"
+
 class SendableToDeveloper t where
 	sendDeveloper :: TMChan (Message Seen) -> TVar UserState -> t -> POSIXTime -> STM (Message Seen)