hyperv_ic: correlate responses with outstanding requests (#5029)

IC requests currently reuse transaction ID zero and accept replies with
unexpected protocol versions. A delayed response can therefore complete
a later request.

Assign advancing transaction IDs and ignore replies that do not match
the outstanding request or expected protocol versions. Preserve
transaction identity across shutdown save/restore, including older
snapshots, and require shutdown acknowledgments to have the shutdown
message type.

Regression tests cover stale replies, invalid IDs and versions,
wraparound, and save/restore. All 16 package tests pass on Windows and
Linux.

Author: mebersol

vm/devices/hyperv_ic/src/common.rs

@@ -32,6 +32,7 @@ pub(crate) struct IcPipe<M: RingMem = GpadlRingMem> {
     pub pipe: MessagePipe<M>,
     #[inspect(skip)]
     buf: Vec<u8>,
+    pub transaction_id: u8,
 }
 
 #[derive(Inspect, Default)]
@@ -54,7 +55,11 @@ impl<M: RingMem> IcPipe<M> {
     pub fn new(raw: RawAsyncChannel<M>) -> Result<Self, std::io::Error> {
         let pipe = MessagePipe::new(raw)?;
         let buf = vec![0; hyperv_ic_protocol::MAX_MESSAGE_SIZE];
-        Ok(Self { pipe, buf })
+        Ok(Self {
+            pipe,
+            buf,
+            transaction_id: u8::MAX,
+        })
     }
 
     pub async fn negotiate(
@@ -64,6 +69,7 @@ impl<M: RingMem> IcPipe<M> {
     ) -> anyhow::Result<Option<Versions>> {
         match state {
             NegotiateState::SendVersion => {
+                let transaction_id = self.transaction_id.wrapping_add(1);
                 let message = hyperv_ic_protocol::NegotiateMessage {
                     framework_version_count: FRAMEWORK_VERSIONS.len() as u16,
                     message_version_count: message_versions.len() as u16,
@@ -76,7 +82,7 @@ impl<M: RingMem> IcPipe<M> {
                         + size_of_val(FRAMEWORK_VERSIONS)
                         + size_of_val(message_versions)) as u16,
                     status: Status::SUCCESS,
-                    transaction_id: 0,
+                    transaction_id,
                     flags: HeaderFlags::new().with_transaction(true).with_request(true),
                     ..FromZeros::new_zeroed()
                 };
@@ -91,11 +97,14 @@ impl<M: RingMem> IcPipe<M> {
                     .await
                     .context("ring buffer error")?;
 
+                self.transaction_id = transaction_id;
                 *state = NegotiateState::WaitVersion;
                 Ok(None)
             }
             NegotiateState::WaitVersion => {
-                let (message_type, status, buf) = self.read_message().await?;
+                let Some((message_type, status, buf)) = self.read_message(None).await? else {
+                    return Ok(None);
+                };
                 // A response to a request issued before the channel was
                 // reset (e.g. across save/restore) may still be in flight.
                 // There is no one left to receive it, so drop it.
@@ -148,29 +157,40 @@ impl<M: RingMem> IcPipe<M> {
         flags: HeaderFlags,
         message: &[u8],
     ) -> anyhow::Result<()> {
+        let transaction_id = self.transaction_id.wrapping_add(1);
         let header = hyperv_ic_protocol::Header {
             framework_version: versions.framework_version,
             message_type,
             message_size: message.len() as u16,
             message_version: versions.message_version,
             status: Status::SUCCESS,
-            transaction_id: 0,
+            transaction_id,
             flags,
             ..FromZeros::new_zeroed()
         };
 
         self.pipe
             .send_vectored(&[IoSlice::new(header.as_bytes()), IoSlice::new(message)])
             .await
-            .context("ring buffer error")
+            .context("ring buffer error")?;
+        self.transaction_id = transaction_id;
+        Ok(())
     }
 
-    pub async fn read_response(&mut self) -> anyhow::Result<(Status, &[u8])> {
-        let (_message_type, status, buf) = self.read_message().await?;
-        Ok((status, buf))
+    pub async fn read_response(
+        &mut self,
+        versions: &Versions,
+    ) -> anyhow::Result<Option<(Status, &[u8])>> {
+        Ok(self
+            .read_message(Some(versions))
+            .await?
+            .map(|(_, status, buf)| (status, buf)))
     }
 
-    async fn read_message(&mut self) -> anyhow::Result<(MessageType, Status, &[u8])> {
+    pub async fn read_message(
+        &mut self,
+        versions: Option<&Versions>,
+    ) -> anyhow::Result<Option<(MessageType, Status, &[u8])>> {
         let n = self
             .pipe
             .recv(&mut self.buf)
@@ -181,15 +201,42 @@ impl<M: RingMem> IcPipe<M> {
             .ok()
             .context("missing header")?;
 
-        if header.transaction_id != 0 || !header.flags.transaction() || !header.flags.response() {
+        if !header.flags.transaction() || !header.flags.response() {
             anyhow::bail!("invalid transaction response");
         }
 
+        if header.transaction_id != self.transaction_id {
+            tracelimit::warn_ratelimited!(
+                transaction_id = header.transaction_id,
+                expected_transaction_id = self.transaction_id,
+                "dropping response for an unexpected transaction"
+            );
+            return Ok(None);
+        }
+
         let rest = rest
             .get(..header.message_size as usize)
             .context("missing message body")?;
 
-        Ok((header.message_type, header.status, rest))
+        let valid_versions = if let Some(versions) = versions {
+            header.framework_version == versions.framework_version
+                && header.message_version == versions.message_version
+        } else if header.message_type == MessageType::VERSION_NEGOTIATION {
+            header.framework_version == Version::new(0, 0)
+                && header.message_version == Version::new(0, 0)
+        } else {
+            true
+        };
+        if !valid_versions {
+            tracelimit::warn_ratelimited!(
+                framework_version = %header.framework_version,
+                message_version = %header.message_version,
+                "dropping response with incorrect protocol versions"
+            );
+            return Ok(None);
+        }
+
+        Ok(Some((header.message_type, header.status, rest)))
     }
 }
 
@@ -265,6 +312,160 @@ mod tests {
         assert!(header.flags.request());
     }
 
+    #[async_test]
+    async fn transactions_advance_and_wrap_ids() {
+        let (mut pipe, mut guest) = new_pipes();
+        let mut state = NegotiateState::default();
+        send_version_request(&mut pipe, &mut state, &mut guest).await;
+        let negotiation_response = version_response(FRAMEWORK_VERSION_3, TEST_MESSAGE_VERSION);
+        guest.send(&negotiation_response).await.unwrap();
+        let versions = pipe
+            .negotiate(&mut state, TEST_MESSAGE_VERSIONS)
+            .await
+            .unwrap()
+            .unwrap();
+        for expected_id in 1..=257u16 {
+            pipe.write_message(
+                &versions,
+                MessageType::SHUTDOWN,
+                HeaderFlags::new().with_transaction(true).with_request(true),
+                &[],
+            )
+            .await
+            .unwrap();
+            let mut buffer = [0; hyperv_ic_protocol::MAX_MESSAGE_SIZE];
+            let size = guest.recv(&mut buffer).await.unwrap();
+            let (mut header, _) = Header::read_from_prefix(&buffer[..size]).unwrap();
+            assert_eq!(header.transaction_id, expected_id as u8);
+            if header.transaction_id == 0 {
+                guest.send(&negotiation_response).await.unwrap();
+                assert!(pipe.read_response(&versions).await.unwrap().is_none());
+                assert_eq!(pipe.transaction_id, header.transaction_id);
+            }
+            header.flags = HeaderFlags::new()
+                .with_transaction(true)
+                .with_response(true);
+            guest.send(header.as_bytes()).await.unwrap();
+            let (status, body) = pipe.read_response(&versions).await.unwrap().unwrap();
+            assert_eq!(status, Status::SUCCESS);
+            assert!(body.is_empty());
+        }
+    }
+
+    #[async_test]
+    async fn responses_match_transaction_and_versions() {
+        let versions = Versions {
+            framework_version: FRAMEWORK_VERSION_3,
+            message_version: TEST_MESSAGE_VERSION,
+        };
+        for message_type in [
+            MessageType::SHUTDOWN,
+            MessageType::TIME_SYNC,
+            MessageType::KVP_EXCHANGE,
+        ] {
+            let (mut pipe, mut guest) = new_pipes();
+            let mut buffer = [0; hyperv_ic_protocol::MAX_MESSAGE_SIZE];
+            for expected_id in 0..2 {
+                pipe.write_message(
+                    &versions,
+                    message_type,
+                    HeaderFlags::new().with_transaction(true).with_request(true),
+                    &[],
+                )
+                .await
+                .unwrap();
+                let size = guest.recv(&mut buffer).await.unwrap();
+                let (mut response, _) = Header::read_from_prefix(&buffer[..size]).unwrap();
+                assert_eq!(response.transaction_id, expected_id);
+                response.flags = HeaderFlags::new()
+                    .with_transaction(true)
+                    .with_response(true);
+
+                for status in [Status::SUCCESS, Status::FAIL] {
+                    let mut negotiation_response =
+                        version_response(FRAMEWORK_VERSION_3, TEST_MESSAGE_VERSION);
+                    let (header, _) =
+                        Header::mut_from_prefix(negotiation_response.as_mut_slice()).unwrap();
+                    header.transaction_id = expected_id;
+                    header.status = status;
+                    guest.send(&negotiation_response).await.unwrap();
+                    assert!(pipe.read_response(&versions).await.unwrap().is_none());
+                    assert_eq!(pipe.transaction_id, expected_id);
+                }
+
+                if expected_id == 1 {
+                    for (transaction_id, framework_version, message_version) in [
+                        (0, versions.framework_version, versions.message_version),
+                        (2, versions.framework_version, versions.message_version),
+                        (1, Version::new(0, 0), versions.message_version),
+                        (1, FRAMEWORK_VERSION_1, versions.message_version),
+                        (1, Version::new(3, 1), versions.message_version),
+                        (1, versions.framework_version, Version::new(0, 0)),
+                        (1, versions.framework_version, Version::new(5, 1)),
+                    ] {
+                        let invalid = Header {
+                            transaction_id,
+                            framework_version,
+                            message_version,
+                            status: Status::NOT_SUPPORTED,
+                            ..response
+                        };
+                        guest.send(invalid.as_bytes()).await.unwrap();
+                        assert!(pipe.read_response(&versions).await.unwrap().is_none());
+                        assert_eq!(pipe.transaction_id, expected_id);
+                    }
+                    response.status = Status::FAIL;
+                }
+
+                let body = b"ack";
+                response.message_size = body.len() as u16;
+                guest
+                    .send(&[response.as_bytes(), body].concat())
+                    .await
+                    .unwrap();
+                let (status, received_body) = pipe.read_response(&versions).await.unwrap().unwrap();
+                assert_eq!(status, response.status);
+                assert_eq!(received_body, body);
+            }
+        }
+    }
+
+    #[async_test]
+    async fn negotiate_drops_unmatched_ids_and_invalid_header_versions() {
+        let (mut pipe, mut guest) = new_pipes();
+        let mut state = NegotiateState::default();
+        send_version_request(&mut pipe, &mut state, &mut guest).await;
+        for (transaction_id, framework_version, message_version) in [
+            (1, Version::new(0, 0), Version::new(0, 0)),
+            (0, FRAMEWORK_VERSION_3, Version::new(0, 0)),
+            (0, Version::new(0, 0), TEST_MESSAGE_VERSION),
+        ] {
+            let mut response = version_response(FRAMEWORK_VERSION_3, TEST_MESSAGE_VERSION);
+            let (header, _) = Header::mut_from_prefix(response.as_mut_slice()).unwrap();
+            header.transaction_id = transaction_id;
+            header.framework_version = framework_version;
+            header.message_version = message_version;
+            guest.send(&response).await.unwrap();
+            assert!(
+                pipe.negotiate(&mut state, TEST_MESSAGE_VERSIONS)
+                    .await
+                    .unwrap()
+                    .is_none()
+            );
+            assert!(matches!(state, NegotiateState::WaitVersion));
+        }
+        guest
+            .send(&version_response(FRAMEWORK_VERSION_3, TEST_MESSAGE_VERSION))
+            .await
+            .unwrap();
+        assert!(
+            pipe.negotiate(&mut state, TEST_MESSAGE_VERSIONS)
+                .await
+                .unwrap()
+                .is_some()
+        );
+    }
+
     #[async_test]
     async fn negotiate() {
         for &framework_version in FRAMEWORK_VERSIONS {

vm/devices/hyperv_ic/src/kvp.rs

@@ -385,7 +385,9 @@ impl<M: RingMem + Unpin> KvpChannel<M> {
                     *state = ReadyState::WaitingResponse(rpc);
                 }
                 ReadyState::WaitingResponse(_) => {
-                    let (status, response) = self.pipe.read_response().await?;
+                    let Some((status, response)) = self.pipe.read_response(versions).await? else {
+                        return Ok(());
+                    };
                     let r = if status == Status::SUCCESS {
                         Ok(())
                     } else {