improve connection to done session display

1 files changed, 21 insertions, 16 deletions

diff --git a/Role/Developer.hs b/Role/Developer.hs
index 9400ddf..9450e57 100644
--- a/Role/Developer.hs
+++ b/Role/Developer.hs
@@ -51,16 +51,17 @@ developer :: TMVar (TVar DeveloperState) -> TMChan (Message Entered) -> TMChan A
 developer dsv ichan ochan sid = withSessionLogger sid $ \logger -> do
 	sk <- genMySessionKey
 	spk <- myPublicKey sk (GpgSign True)
-	(controlinput, controloutput) <- openControlWindow
 	(devstate, startoutput) <- processSessionStart sk ochan logger dsv
 	emitOutput startoutput
-	ok <- authUser spk ichan ochan devstate logger
-	if ok
-		then inRawMode $ void $ 
+	res <- authUser spk ichan ochan devstate logger
+	case res of
+		Authed -> inRawMode $ void $ do
+			(controlinput, controloutput) <- openControlWindow
 			sendTtyInput ichan devstate logger
 				`race` sendTtyOutput ochan devstate controlinput logger
 				`race` sendControlOutput controloutput ichan devstate logger
-		else hPutStrLn stderr "\nUser did not grant access to their terminal."
+		AuthFailed -> hPutStrLn stderr "\n** User did not grant access to their terminal."
+		SessionEnded -> hPutStrLn stderr "\n** This debug-me session has already ended."
 
 data DeveloperState = DeveloperState
 	{ lastSeen :: Hash
@@ -170,10 +171,12 @@ sendTtyOutput ochan devstate controlinput logger = go
 		_ -> return ()
 	fwd = atomically . writeTMChan controlinput . ControlInputAction . control
 
+data AuthResult = Authed | AuthFailed | SessionEnded
+
 -- | Present our session key to the user. 
 -- Wait for them to accept or reject it, while displaying any Seen data
 -- in the meantime.
-authUser :: PerhapsSigned PublicKey -> TMChan (Message Entered) -> TMChan AnyMessage -> TVar DeveloperState -> Logger -> IO Bool
+authUser :: PerhapsSigned PublicKey -> TMChan (Message Entered) -> TMChan AnyMessage -> TVar DeveloperState -> Logger -> IO AuthResult
 authUser spk ichan ochan devstate logger = do
 	ds <- atomically $ readTVar devstate
 	let msg = ControlMessage $ mkSigned (developerSessionKey ds) 
@@ -186,16 +189,18 @@ authUser spk ichan ochan devstate logger = do
   where
 	waitresp pk = do
 		ts <- getPOSIXTime
-		(o, msg) <- fromMaybe (error "Looks like that debug-me session is over.")
-			<$> atomically (getServerMessage ochan devstate ts)
-		logger msg
-		emitOutput o
-		case o of
-			GotControl (SessionKeyAccepted pk')
-				| pk' == pk -> return True
-			GotControl (SessionKeyRejected pk')
-				| pk' == pk -> return False
-			_ -> waitresp pk
+		v <- atomically (getServerMessage ochan devstate ts)
+		case v of
+			Nothing -> return SessionEnded
+			Just (o, msg) -> do
+				logger msg
+				emitOutput o
+				case o of
+					GotControl (SessionKeyAccepted pk')
+						| pk' == pk -> return Authed
+					GotControl (SessionKeyRejected pk')
+						| pk' == pk -> return AuthFailed
+					_ -> waitresp pk
 
 data Output
 	= TtyOutput B.ByteString