fix(sandbox): check every backend confirmation against the launch

crates/openshell-sandbox-backend/src/runtime.rs

@@ -272,10 +272,6 @@ impl OpenShellRuntimeBackend {
         };
         validate_runtime_descriptor(&runtime_descriptor, &sandbox)?;
         let host_gateway_ip = runtime_descriptor.host_gateway_ip;
-        let resource_claims = runtime_descriptor.resource_claims.clone();
-        let generation = runtime_descriptor.generation.clone();
-        let session_id = runtime_descriptor.session_id;
-        let provisioning_evidence = runtime_descriptor.provisioning_evidence.clone();
         let mut client = BoundaryClient::with_backend(
             runtime_descriptor,
             self.sandbox_bearer.clone(),
@@ -290,14 +286,7 @@ impl OpenShellRuntimeBackend {
                 policy: Box::new(SandboxPolicyWire::from(sandbox.policy.clone())),
             })
             .await?;
-        let Response::Attached { snapshot } = response else {
-            return Err(unexpected_response("attached", &response));
-        };
-        if snapshot.generation != generation {
-            return Err(BackendError::Confirm(
-                "sandbox session snapshot generation does not match runtime descriptor".to_string(),
-            ));
-        }
+        client.check_attached(response)?;
         Ok(Box::new(BoundBoundary {
             client: client.clone(),
             agent: sandbox.agent,
@@ -307,10 +296,6 @@ impl OpenShellRuntimeBackend {
             ca_file_paths: self.ca_file_paths.clone(),
             provider_credentials: self.provider_credentials.clone(),
             identity: sandbox.identity,
-            generation,
-            session_id,
-            resource_claims,
-            provisioning_evidence,
         }))
     }
 }
@@ -417,10 +402,6 @@ pub struct BoundBoundary {
     ca_file_paths: Arc<std::sync::Mutex<Option<(PathBuf, PathBuf)>>>,
     provider_credentials: openshell_core::provider_credentials::ProviderCredentialState,
     identity: ResolvedWorkloadIdentity,
-    generation: String,
-    session_id: openshell_core::SandboxSessionId,
-    resource_claims: std::collections::BTreeMap<String, String>,
-    provisioning_evidence: openshell_isolation_interface::contract::ProvisioningEvidence,
 }
 
 impl BoundBoundary {
@@ -434,20 +415,7 @@ impl BoundBoundary {
 
     pub async fn confirm(self: Box<Self>) -> Result<ConfirmedBoundary, BackendError> {
         let response = self.client.call_idempotent(Request::Confirm).await?;
-        let Response::Confirmed { confirmation } = response else {
-            return Err(unexpected_response("confirmed", &response));
-        };
-        if confirmation.generation != self.generation
-            || confirmation.session_id != self.session_id
-            || confirmation.resource_claims != self.resource_claims
-            || confirmation.provisioning_evidence_digest
-                != self.provisioning_evidence.evidence_digest
-        {
-            return Err(BackendError::Confirm(
-                "sandbox confirmation generation, session, resource claims, or provisioning evidence do not match runtime descriptor"
-                    .to_string(),
-            ));
-        }
+        let confirmation = self.client.launch_confirmation(response)?;
         if self.client.backend_name == crate::BACKEND_NAME {
             crate::audit::validate_native_linux_audit(&confirmation.backend_audit)?;
         }
@@ -1349,6 +1317,51 @@ impl BoundaryClient {
         }
     }
 
+    /// Check an Attach reply, including replays on reconnect, against the launch.
+    fn check_attached(&self, response: Response) -> Result<(), BackendError> {
+        let Response::Attached { snapshot } = response else {
+            return Err(unexpected_response("attached", &response));
+        };
+        if snapshot.generation != self.runtime_descriptor.generation {
+            return Err(BackendError::Confirm(
+                "sandbox session snapshot generation does not match runtime descriptor".to_string(),
+            ));
+        }
+        Ok(())
+    }
+
+    /// Check a Confirm reply against the launch. Reconnect, credential renewal
+    /// and liveness probes replay Confirm, so every reply gets the same check.
+    fn launch_confirmation(
+        &self,
+        response: Response,
+    ) -> Result<Box<BoundaryConfirmation>, BackendError> {
+        let Response::Confirmed { confirmation } = response else {
+            return Err(unexpected_response("confirmed", &response));
+        };
+        let launch = &self.runtime_descriptor;
+        let mismatch = if confirmation.identity != launch.workload_identity {
+            "identity"
+        } else if confirmation.generation.is_empty() {
+            "generation (empty)"
+        } else if confirmation.generation != launch.generation {
+            "generation"
+        } else if confirmation.session_id != launch.session_id {
+            "session"
+        } else if confirmation.resource_claims != launch.resource_claims {
+            "resource claims"
+        } else if confirmation.provisioning_evidence_digest
+            != launch.provisioning_evidence.evidence_digest
+        {
+            "provisioning evidence"
+        } else {
+            return Ok(confirmation);
+        };
+        Err(BackendError::Confirm(format!(
+            "sandbox confirmation {mismatch} does not match runtime descriptor"
+        )))
+    }
+
     async fn acquire_network_stream(&self) -> tokio::sync::OwnedSemaphorePermit {
         let acquire = self.network_streams.clone().acquire_owned();
         tokio::pin!(acquire);
@@ -1641,7 +1654,12 @@ impl BoundaryClient {
             .map_err(|_| {
                 BackendError::Unavailable("boundary credential renewal timed out".to_string())
             })??;
-            expect_response(response, "confirmed")?;
+            if let Err(error) = self.launch_confirmation(response) {
+                // A backend answering for another launch keeps no live streams.
+                cached.transport.abort();
+                *self.mediation.lock().await = None;
+                return Err(error);
+            }
             if let Some(current) = self.grpc_channel.lock().await.as_mut()
                 && current.generation == cached.generation
             {
@@ -1763,14 +1781,15 @@ impl BoundaryClient {
         let Ok(credential) = BoundaryCredential::capture(&self.sandbox_bearer) else {
             return false;
         };
-        matches!(
-            tokio::time::timeout(
-                CONNECTION_PROBE_TIMEOUT,
-                self.exchange_on_channel(cached.channel.clone(), confirm, &credential),
-            )
-            .await,
-            Ok(Ok(Response::Confirmed { .. }))
+        let Ok(Ok(response)) = tokio::time::timeout(
+            CONNECTION_PROBE_TIMEOUT,
+            self.exchange_on_channel(cached.channel.clone(), confirm, &credential),
         )
+        .await
+        else {
+            return false;
+        };
+        self.launch_confirmation(response).is_ok()
     }
 
     /// Open a new physical connection and replay the authenticated lifecycle,
@@ -1783,11 +1802,15 @@ impl BoundaryClient {
     ) -> Result<(tonic::transport::Channel, Arc<TransportAbort>), BackendError> {
         let (channel, transport) = self.build_grpc_channel(credential).await?;
         let result = async {
-            self.exchange_on_channel(channel.clone(), attach, credential)
+            let attached = self
+                .exchange_on_channel(channel.clone(), attach, credential)
                 .await?;
+            self.check_attached(attached)?;
             if let Some(confirm) = confirm {
-                self.exchange_on_channel(channel.clone(), confirm, credential)
+                let confirmed = self
+                    .exchange_on_channel(channel.clone(), confirm, credential)
                     .await?;
+                self.launch_confirmation(confirmed)?;
             }
             Ok(())
         }
@@ -2354,9 +2377,7 @@ fn connect_host_vsock(
 fn expect_response(response: Response, expected: &str) -> Result<(), BackendError> {
     let matches = matches!(
         (&response, expected),
-        (Response::Attached { .. }, "attached")
-            | (Response::Confirmed { .. }, "confirmed")
-            | (Response::Signaled, "signaled")
+        (Response::Signaled, "signaled")
             | (Response::Terminated, "terminated")
             | (Response::BoundaryTerminated, "boundary_terminated")
     );
@@ -2689,6 +2710,7 @@ mod tests {
 
     #[derive(Clone)]
     struct TestGrpcBoundary {
+        attach_generation: String,
         wait_for_half_close: bool,
         expected_token: String,
         requests: Arc<std::sync::atomic::AtomicUsize>,
@@ -2722,6 +2744,7 @@ mod tests {
             let mediation_ready = self.mediation_ready;
             let provider_environment_generation = self.provider_environment_generation;
             let confirmation = self.confirmation.clone();
+            let attach_generation = self.attach_generation.clone();
             let (outbound, outbound_rx) = tokio::sync::mpsc::channel(1);
             tokio::spawn(async move {
                 let mut frame = Vec::new();
@@ -2764,7 +2787,7 @@ mod tests {
                             },
                             Request::Attach { .. } => Response::Attached {
                                 snapshot: crate::boundary_protocol::SessionSnapshotWire {
-                                    generation: "test-generation".to_string(),
+                                    generation: attach_generation,
                                     processes: Vec::new(),
                                 },
                             },
@@ -2891,6 +2914,7 @@ mod tests {
         let address = listener.local_addr().unwrap();
         let requests = Arc::new(std::sync::atomic::AtomicUsize::new(0));
         let service = TestGrpcBoundary {
+            attach_generation: "test-generation".to_string(),
             wait_for_half_close: false,
             expected_token: "a".repeat(32),
             requests: requests.clone(),
@@ -2965,6 +2989,7 @@ mod tests {
                     .unwrap();
                 server_accepted.fetch_add(1, Ordering::AcqRel);
                 let service = TestGrpcBoundary {
+                    attach_generation: "test-generation".to_string(),
                     wait_for_half_close: false,
                     expected_token: "a".repeat(32),
                     requests: server_requests.clone(),
@@ -3076,6 +3101,7 @@ mod tests {
                     .await
                     .unwrap();
                 let service = TestGrpcBoundary {
+                    attach_generation: "test-generation".to_string(),
                     wait_for_half_close: false,
                     expected_token: "a".repeat(32),
                     requests: server_requests.clone(),
@@ -3131,6 +3157,7 @@ mod tests {
         let address = listener.local_addr().unwrap();
         let requests = Arc::new(std::sync::atomic::AtomicUsize::new(0));
         let service = TestGrpcBoundary {
+            attach_generation: "test-generation".to_string(),
             wait_for_half_close: true,
             expected_token: "a".repeat(32),
             requests,
@@ -3330,6 +3357,7 @@ mod tests {
         confirmation: BoundaryConfirmation,
     ) {
         let service = TestGrpcBoundary {
+            attach_generation: "test-generation".to_string(),
             wait_for_half_close: false,
             expected_token,
             requests: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
@@ -3438,7 +3466,6 @@ mod tests {
         });
 
         let descriptor = tls_runtime_descriptor(address, certificate.client_tls);
-        let provisioning_evidence = descriptor.provisioning_evidence.clone();
         let context = sandbox();
         let client = Arc::new(BoundaryClient::new(
             descriptor,
@@ -3461,10 +3488,6 @@ mod tests {
                     HashMap::new(),
                 ),
             identity: context.identity,
-            generation: "test-generation".to_string(),
-            session_id: test_session_id(),
-            resource_claims: std::collections::BTreeMap::new(),
-            provisioning_evidence,
         };
 
         assert!(matches!(
@@ -3848,6 +3871,7 @@ mod tests {
         let accepted_by_server = accepted.clone();
         let handled = Arc::new(std::sync::atomic::AtomicUsize::new(0));
         let service = TestGrpcBoundary {
+            attach_generation: "test-generation".to_string(),
             wait_for_half_close: false,
             expected_token: "a".repeat(32),
             requests: handled.clone(),
@@ -3921,6 +3945,7 @@ mod tests {
         let address = listener.local_addr().unwrap();
         let requests = Arc::new(std::sync::atomic::AtomicUsize::new(0));
         let service = TestGrpcBoundary {
+            attach_generation: "test-generation".to_string(),
             wait_for_half_close: false,
             expected_token: "a".repeat(32),
             requests: requests.clone(),