From 12dd72d88f100697ec074ee420134ad580d2380d Mon Sep 17 00:00:00 2001
From: Rob Bradford <rbradford@meta.com>
Date: Fri, 17 Apr 2026 09:33:10 +0100
Subject: [PATCH] virtio-devices: balloon: Enable use with confidential VMs

Following the pattern used by the existing virtio devices make the
balloon device work with confidential VMs (e.g. SEV-SNP). This requires
advertising the VIRTIO_F_ACCESS_PLATFORM feature. Do not expose this to
the user as a controllable option and instead only enable in on the
"force" case.

Signed-off-by: Rob Bradford <rbradford@meta.com>
---
 fuzz/fuzz_targets/balloon.rs  |  1 +
 virtio-devices/src/balloon.rs | 33 +++++++++++++++++++++++++++++----
 vmm/src/device_manager.rs     |  1 +
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/fuzz/fuzz_targets/balloon.rs b/fuzz/fuzz_targets/balloon.rs
index 69f0c07e8..20260c546 100644
--- a/fuzz/fuzz_targets/balloon.rs
+++ b/fuzz/fuzz_targets/balloon.rs
@@ -49,6 +49,7 @@ fuzz_target!(|bytes: &[u8]| -> Corpus {
         BALLOON_SIZE,
         true,
         true,
+        false,
         SeccompAction::Allow,
         EventFd::new(EFD_NONBLOCK).unwrap(),
         None,
diff --git a/virtio-devices/src/balloon.rs b/virtio-devices/src/balloon.rs
index bb9c46cbc..f9db09bd3 100644
--- a/virtio-devices/src/balloon.rs
+++ b/virtio-devices/src/balloon.rs
@@ -34,14 +34,15 @@ use vm_memory::{
     GuestMemoryError, GuestMemoryRegion,
 };
 use vm_migration::{Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable};
+use vm_virtio::{AccessPlatform, Translatable};
 use vmm_sys_util::eventfd::EventFd;
 
 use crate::seccomp_filters::Thread;
 use crate::thread_helper::spawn_virtio_thread;
 use crate::{
     ActivateResult, EPOLL_HELPER_EVENT_LAST, EpollHelper, EpollHelperError, EpollHelperHandler,
-    GuestMemoryMmap, VIRTIO_F_VERSION_1, VirtioCommon, VirtioDevice, VirtioDeviceType,
-    VirtioInterrupt, VirtioInterruptType,
+    GuestMemoryMmap, VIRTIO_F_ACCESS_PLATFORM, VIRTIO_F_VERSION_1, VirtioCommon, VirtioDevice,
+    VirtioDeviceType, VirtioInterrupt, VirtioInterruptType,
 };
 
 const QUEUE_SIZE: u16 = 128;
@@ -160,6 +161,7 @@ struct BalloonEpollHandler {
     kill_evt: EventFd,
     pause_evt: EventFd,
     pbp: Option<PartiallyBalloonedPage>,
+    access_platform: Option<Arc<dyn AccessPlatform>>,
 }
 
 impl BalloonEpollHandler {
@@ -277,7 +279,12 @@ impl BalloonEpollHandler {
 
             let mut offset = 0u64;
             while offset < desc.len() as u64 {
-                let addr = desc.addr().checked_add(offset).unwrap();
+                let addr = desc
+                    .addr()
+                    .checked_add(offset)
+                    .unwrap()
+                    .translate_gva(self.access_platform.as_deref(), data_chunk_size)
+                    .map_err(|e| Error::GuestMemory(GuestMemoryError::IOError(e)))?;
                 let pfn: u32 = desc_chain
                     .memory()
                     .read_obj(addr)
@@ -324,7 +331,11 @@ impl BalloonEpollHandler {
             let mut descs_len = 0;
             while let Some(desc) = desc_chain.next() {
                 descs_len += desc.len();
-                Self::release_memory_range(desc_chain.memory(), desc.addr(), desc.len() as usize)?;
+                let addr = desc
+                    .addr()
+                    .translate_gva(self.access_platform.as_deref(), desc.len() as usize)
+                    .map_err(|e| Error::GuestMemory(GuestMemoryError::IOError(e)))?;
+                Self::release_memory_range(desc_chain.memory(), addr, desc.len() as usize)?;
             }
 
             self.queues[queue_index]
@@ -437,11 +448,13 @@ pub struct Balloon {
 
 impl Balloon {
     // Create a new virtio-balloon.
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         id: String,
         size: u64,
         deflate_on_oom: bool,
         free_page_reporting: bool,
+        access_platform_enabled: bool,
         seccomp_action: SeccompAction,
         exit_evt: EventFd,
         state: Option<BalloonState>,
@@ -464,6 +477,9 @@ impl Balloon {
             if free_page_reporting {
                 avail_features |= 1u64 << VIRTIO_BALLOON_F_REPORTING;
             }
+            if access_platform_enabled {
+                avail_features |= 1u64 << VIRTIO_F_ACCESS_PLATFORM;
+            }
 
             let config = VirtioBalloonConfig {
                 num_pages: (size >> VIRTIO_BALLOON_PFN_SHIFT) as u32,
@@ -628,6 +644,7 @@ impl VirtioDevice for Balloon {
             kill_evt,
             pause_evt,
             pbp: None,
+            access_platform: self.common.access_platform(),
         };
 
         let paused = self.common.paused.clone();
@@ -648,6 +665,14 @@ impl VirtioDevice for Balloon {
         Ok(())
     }
 
+    fn set_access_platform(&mut self, access_platform: Arc<dyn AccessPlatform>) {
+        self.common.set_access_platform(access_platform);
+    }
+
+    fn access_platform(&self) -> Option<Arc<dyn AccessPlatform>> {
+        self.common.access_platform()
+    }
+
     fn reset(&mut self) -> Option<Arc<dyn VirtioInterrupt>> {
         let result = self.common.reset();
         event!("virtio-device", "reset", "id", &self.id);
diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs
index 2c31efe33..948afdfae 100644
--- a/vmm/src/device_manager.rs
+++ b/vmm/src/device_manager.rs
@@ -3704,6 +3704,7 @@ impl DeviceManager {
                     balloon_config.size,
                     balloon_config.deflate_on_oom,
                     balloon_config.free_page_reporting,
+                    self.force_access_platform,
                     self.seccomp_action.clone(),
                     self.exit_evt
                         .try_clone()