use anyhow::{Result, bail};
use ax_std::os::arceos::modules::ax_task;
use axvm::AxVMRef;
fn console_reader_isolation_cpu(
host_cpu_count: usize,
vcpu_masks: impl IntoIterator<Item = Option<usize>>,
) -> Option<usize> {
let tracked_cpu_count = host_cpu_count.min(usize::BITS as usize);
if tracked_cpu_count == 0 {
return None;
}
let online_bits = if tracked_cpu_count == usize::BITS as usize {
usize::MAX
} else {
(1usize << tracked_cpu_count) - 1
};
let guest_bits = vcpu_masks.into_iter().fold(0usize, |used, mask| {
let requested = mask.unwrap_or(online_bits) & online_bits;
used | if requested == 0 {
online_bits
} else {
requested
}
});
let host_only_bits = online_bits & !guest_bits;
(host_only_bits != 0).then(|| host_only_bits.trailing_zeros() as usize)
}
pub(crate) fn configure_host_console_reader(vms: &[AxVMRef]) -> Result<()> {
axvm::host::console::set_input_irq_enabled(false);
let isolation_cpu = console_reader_isolation_cpu(
axvm::host::cpu::count(),
vms.iter()
.flat_map(|vm| vm.vcpu_snapshots())
.map(|vcpu| vcpu.phys_cpu_set),
);
let Some(owner_cpu) = isolation_cpu else {
return Ok(());
};
let owner_affinity = ax_task::AxCpuMask::one_shot(owner_cpu);
if !ax_task::set_current_affinity(owner_affinity) {
bail!("failed to pin the host console reader to CPU {owner_cpu}");
}
let actual_owner_cpu = axvm::host::cpu::current_id();
if actual_owner_cpu != owner_cpu {
bail!(
"host console reader affinity selected CPU {owner_cpu}, but migration ended on CPU \
{actual_owner_cpu}"
);
}
Ok(())
}
pub(crate) fn read_host_byte() -> Option<u8> {
let mut byte = [0u8; 1];
(axvm::host::console::read_bytes(&mut byte) == 1).then_some(byte[0])
}
pub(crate) fn wait_for_host_input() {
std::thread::yield_now();
}
pub(super) fn write_host_bytes(bytes: &[u8]) {
axvm::host::console::write_bytes(bytes);
}
#[cfg(test)]
mod tests {
use super::console_reader_isolation_cpu;
#[test]
fn isolation_requires_a_cpu_excluded_by_every_vcpu() {
assert_eq!(
console_reader_isolation_cpu(4, [Some(0b0010), Some(0b1000)]),
Some(0)
);
assert_eq!(console_reader_isolation_cpu(4, [None]), None);
assert_eq!(console_reader_isolation_cpu(4, [Some(0)]), None);
assert_eq!(console_reader_isolation_cpu(4, [Some(usize::MAX)]), None);
}
}