Route production control through lifecycle ownership

crates/lazaret-host-coordinator/src/control.rs

@@ -69,6 +69,7 @@ pub enum ControlError {
     InvalidRequest,
     OversizeRequest { length: usize, maximum: usize },
     LockPoisoned,
+    Lifecycle(String),
 }
 
 impl fmt::Display for ControlError {
@@ -100,6 +101,7 @@ impl fmt::Display for ControlError {
                 )
             }
             Self::LockPoisoned => formatter.write_str("control registry lock is poisoned"),
+            Self::Lifecycle(detail) => write!(formatter, "control lifecycle failed: {detail}"),
         }
     }
 }
@@ -1007,12 +1009,42 @@ impl PublishedControlFace {
         &self.paths.identity
     }
 
+    /// Withdraws exact discovery and connection paths while retaining live accepted streams.
+    pub fn withdraw(&self) {
+        remove_if_identity_matches(&self.paths.socket, self.socket_identity);
+        remove_if_identity_matches(&self.paths.identity, self.identity_file_identity);
+    }
+
     /// Accepts and serves one TLS connection. Registered connections remain live until EOF.
     ///
     /// # Errors
     /// Refuses peer UID, publication, TLS, ALPN, certificate, framing, tuple, duplicate, and
     /// persistence failures.
     pub fn serve_next(&self) -> Result<(), ControlError> {
+        self.serve_next_coordinated(
+            |admission| {
+                self.registry
+                    .register(admission.request, &admission.credential)
+                    .map(|_| ())
+            },
+            |session_id| self.registry.disconnect(session_id),
+        )
+    }
+
+    /// Serves one connection while routing authenticated session state through its lifecycle.
+    ///
+    /// # Errors
+    /// Refuses the same conditions as [`Self::serve_next`] plus lifecycle admission or disconnect
+    /// failure.
+    pub fn serve_next_coordinated<Admit, Disconnect>(
+        &self,
+        mut admit: Admit,
+        mut disconnect: Disconnect,
+    ) -> Result<(), ControlError>
+    where
+        Admit: FnMut(ControlSessionAdmission) -> Result<(), ControlError>,
+        Disconnect: FnMut(SessionId) -> Result<(), ControlError>,
+    {
         let (socket, _) = self.listener.accept()?;
         let peer = socket_peercred(&socket).map_err(io::Error::from)?;
         require_peer_uid(peer.uid.as_raw(), self.uid)?;
@@ -1034,7 +1066,9 @@ impl PublishedControlFace {
         let mut stream = rustls::StreamOwned::new(connection, socket);
         match peer_certificate {
             None => self.serve_enrollment(&mut stream),
-            Some(certificates) => self.serve_registration(&mut stream, &certificates),
+            Some(certificates) => {
+                self.serve_registration(&mut stream, &certificates, &mut admit, &mut disconnect)
+            }
         }
     }
 
@@ -1057,11 +1091,17 @@ impl PublishedControlFace {
         )
     }
 
-    fn serve_registration(
+    fn serve_registration<Admit, Disconnect>(
         &self,
         stream: &mut rustls::StreamOwned<ServerConnection, std::os::unix::net::UnixStream>,
         certificates: &[Vec<u8>],
-    ) -> Result<(), ControlError> {
+        admit: &mut Admit,
+        disconnect: &mut Disconnect,
+    ) -> Result<(), ControlError>
+    where
+        Admit: FnMut(ControlSessionAdmission) -> Result<(), ControlError>,
+        Disconnect: FnMut(SessionId) -> Result<(), ControlError>,
+    {
         let [leaf] = certificates else {
             return Err(ControlError::InvalidCertificateProfile);
         };
@@ -1076,7 +1116,11 @@ impl PublishedControlFace {
             self.anchor.sandbox_protocol_credential.credential_id,
             &self.anchor.coordinator_protocol_key,
         )?;
-        let registered = self.registry.register(request, &credential)?;
+        let session_id = request.session_binding.session_id;
+        admit(ControlSessionAdmission {
+            request,
+            credential,
+        })?;
         let result = (|| {
             write_frame(
                 stream,
@@ -1093,8 +1137,8 @@ impl PublishedControlFace {
                 Err(error) => Err(error.into()),
             }
         })();
-        let disconnect = self.registry.disconnect(registered.binding.session_id);
-        result.and(disconnect)
+        let disconnected = disconnect(session_id);
+        result.and(disconnected)
     }
 
     fn validate_publication(&self) -> Result<(), ControlError> {
@@ -1216,8 +1260,7 @@ impl PublishedControlFace {
 
 impl Drop for PublishedControlFace {
     fn drop(&mut self) {
-        remove_if_identity_matches(&self.paths.socket, self.socket_identity);
-        remove_if_identity_matches(&self.paths.identity, self.identity_file_identity);
+        self.withdraw();
     }
 }