Wait for Session disconnect before returning

crates/lazaret-client/src/lib.rs

@@ -1,6 +1,6 @@
 use std::fmt;
 use std::fs::File;
-use std::io::{self, Write};
+use std::io::{self, Read, Write};
 use std::net::Shutdown;
 use std::path::PathBuf;
 use std::sync::{Mutex, mpsc};
@@ -13,9 +13,9 @@ use lazaret_core::identity::{ClientId, SessionAnchorId, SessionId};
 use lazaret_core::records::{
     ClientEnrollmentResponse, ControlRequest, ControlResponse, DeveloperDecision,
     DeveloperDecisionKind, InteractiveSessionBatchV2, InteractiveSessionEventV2,
-    InteractiveSessionRequestV2, NonInteractiveSessionMessage, NonInteractiveSessionRequest,
-    ProtocolPublicationIdentity, ProtocolVersion, SessionBinding, SessionExitOutcome,
-    SessionRegistrationRequest, TlsSubjectId,
+    InteractiveSessionMessage, InteractiveSessionRequestV2, NonInteractiveSessionMessage,
+    NonInteractiveSessionRequest, ProtocolPublicationIdentity, ProtocolVersion, SessionBinding,
+    SessionExitOutcome, SessionRegistrationRequest, SessionTerminalMessage, TlsSubjectId,
 };
 use sha2::{Digest, Sha256};
 
@@ -386,7 +386,7 @@ pub fn execute_over_control_face_as<E>(
             }
             NonInteractiveSessionMessage::Exit(exit) => {
                 require_session(exit.session_id(), session_id)?;
-                return match exit.outcome() {
+                let result = match exit.outcome() {
                     SessionExitOutcome::Exited => exit
                         .exit_code()
                         .ok_or(ClientError::Protocol("Exited omitted its exit code")),
@@ -394,6 +394,8 @@ pub fn execute_over_control_face_as<E>(
                         Err(ClientError::LostConnection)
                     }
                 };
+                await_control_close(&mut stream)?;
+                return result;
             }
         }
     }
@@ -576,7 +578,27 @@ impl Drop for ControlFaceInteractiveTransport {
 fn receive_interactive_batch(
     stream: &mut impl std::io::Read,
 ) -> Result<InteractiveSessionBatchV2, ClientError<std::convert::Infallible>> {
-    ControlClient::receive(stream).map_err(map_control_connection)
+    let batch: InteractiveSessionBatchV2 =
+        ControlClient::receive(stream).map_err(map_control_connection)?;
+    if batch.messages.iter().any(|message| {
+        matches!(
+            message,
+            InteractiveSessionMessage::Terminal(SessionTerminalMessage::Exit(_))
+        )
+    }) {
+        await_control_close(stream)?;
+    }
+    Ok(batch)
+}
+
+fn await_control_close<E>(stream: &mut impl Read) -> Result<(), ClientError<E>> {
+    let mut trailing = [0_u8; 1];
+    match stream.read(&mut trailing) {
+        Ok(0) => Ok(()),
+        Ok(_) => Err(ClientError::Protocol("control data followed Session exit")),
+        Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => Ok(()),
+        Err(_) => Err(ClientError::LostConnection),
+    }
 }
 
 fn map_control_connection(error: ControlClientError) -> ClientError<std::convert::Infallible> {
@@ -623,6 +645,33 @@ impl<E> From<io::Error> for ClientError<E> {
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    struct UnexpectedEof;
+
+    impl Read for UnexpectedEof {
+        fn read(&mut self, _buffer: &mut [u8]) -> io::Result<usize> {
+            Err(io::Error::from(io::ErrorKind::UnexpectedEof))
+        }
+    }
+
+    #[test]
+    fn clean_control_close_is_required_after_session_exit() {
+        let mut closed = std::io::Cursor::new(Vec::<u8>::new());
+        assert!(await_control_close::<std::convert::Infallible>(&mut closed).is_ok());
+
+        assert!(await_control_close::<std::convert::Infallible>(&mut UnexpectedEof).is_ok());
+
+        let mut trailing = std::io::Cursor::new(vec![1_u8]);
+        assert!(matches!(
+            await_control_close::<std::convert::Infallible>(&mut trailing),
+            Err(ClientError::Protocol("control data followed Session exit"))
+        ));
+    }
+}
+
 impl<E> From<ControlClientError> for ClientError<E> {
     fn from(error: ControlClientError) -> Self {
         Self::Control(error)