alioth/src/mem/mapped.rs - 02504ad

mapped.rs67.21%

// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
 
use std::cell::UnsafeCell;
#[cfg(target_os = "linux")]
use std::ffi::CStr;
use std::fmt::Debug;
use std::fs::File;
use std::io::{IoSlice, IoSliceMut, Read, Write};
use std::mem::{align_of, size_of};
#[cfg(target_os = "linux")]
use std::os::fd::FromRawFd;
use std::os::fd::{AsFd, AsRawFd, BorrowedFd};
use std::ptr::{NonNull, null_mut};
use std::sync::Arc;
 
#[cfg(target_os = "linux")]
use libc::{MADV_HUGEPAGE, MFD_CLOEXEC};
use libc::{
    MAP_ANONYMOUS, MAP_FAILED, MAP_PRIVATE, MAP_SHARED, MS_ASYNC, PROT_READ, PROT_WRITE, c_void,
    madvise, mmap, msync, munmap,
};
use parking_lot::{RwLock, RwLockReadGuard};
use snafu::ResultExt;
use zerocopy::{FromBytes, Immutable, IntoBytes};
 
use crate::ffi;
use crate::mem::addressable::{Addressable, SlotBackend};
use crate::mem::{Error, Result, error};
 
#[derive(Debug)]
struct MemPages {
    addr: NonNull<c_void>,
    len: usize,
    fd: Option<(File, u64)>,
}
 
unsafe impl Send for MemPages {}
unsafe impl Sync for MemPages {}
 
impl Drop for MemPages {
    fn drop(&mut self) {78x
        let ret = unsafe { munmap(self.addr.as_ptr(), self.len) };78x
        if ret != 0 {78x
            log::error!("munmap({:p}, {:x}) = {:x}", self.addr, self.len, ret);
        } else {
            log::info!("munmap({:p}, {:x}) = {:x}, done", self.addr, self.len, ret);78x
        }
    }78x
}
// ArcMemPages uses Arc to manage the underlying memory and caches
// the address and size on the stack. Compared with using Arc<MemPages>,
// it avoids a memory load when a caller tries to read/write the pages.
// TODO: is it really necessary?
#[derive(Debug, Clone)]
pub struct ArcMemPages {
    addr: usize,
    size: usize,
    _inner: Arc<MemPages>,
}
 
impl SlotBackend for ArcMemPages {
    fn size(&self) -> u64 {2085x
        self.size as u642085x
    }2085x
}
 
impl ArcMemPages {
    pub fn addr(&self) -> usize {
        self.addr
    }
 
    pub fn size(&self) -> u64 {
        self.size as u64
    }
 
    pub fn fd(&self) -> Option<(BorrowedFd<'_>, u64)> {
        self._inner
            .fd
            .as_ref()
            .map(|(f, offset)| (f.as_fd(), *offset))
    }
 
    pub fn sync(&self) -> Result<()> {
        ffi!(unsafe { msync(self.addr as *mut _, self.size, MS_ASYNC) })?;
        Ok(())
    }
 
    #[cfg(target_os = "linux")]
    pub fn madvise_hugepage(&self) -> Result<()> {
        ffi!(unsafe { madvise(self.addr as *mut _, self.size, MADV_HUGEPAGE) })?;
        Ok(())
    }
 
    fn from_raw(addr: *mut c_void, len: usize, fd: Option<(File, u64)>) -> Self {78x
        let addr = NonNull::new(addr).expect("address from mmap() should not be null");78x
        ArcMemPages {78x
            addr: addr.as_ptr() as usize,78x
            size: len,78x
            _inner: Arc::new(MemPages { addr, len, fd }),78x
        }78x
    }78x
 
    pub fn from_file(file: File, offset: i64, len: usize, prot: i32) -> Result<Self> {
        let addr = ffi!(
            unsafe { mmap(null_mut(), len, prot, MAP_SHARED, file.as_raw_fd(), offset) },
            MAP_FAILED
        )?;
        Ok(Self::from_raw(addr, len, Some((file, offset as u64))))
    }
 
    #[cfg(target_os = "linux")]
    pub fn from_memfd(name: &CStr, size: usize, prot: Option<i32>) -> Result<Self> {
        let fd = ffi!(unsafe { libc::memfd_create(name.as_ptr(), MFD_CLOEXEC) })?;
        let prot = prot.unwrap_or(PROT_WRITE | PROT_READ);
        let addr = ffi!(
            unsafe { mmap(null_mut(), size, prot, MAP_SHARED, fd, 0) },
            MAP_FAILED
        )?;
        let file = unsafe { File::from_raw_fd(fd) };
        file.set_len(size as _)?;
        Ok(Self::from_raw(addr, size, Some((file, 0))))
    }
 
    pub fn from_anonymous(size: usize, prot: Option<i32>, flags: Option<i32>) -> Result<Self> {78x
        let prot = prot.unwrap_or(PROT_WRITE | PROT_READ);78x
        let flags = flags.unwrap_or(MAP_PRIVATE) | MAP_ANONYMOUS;78x
        let addr = ffi!(78x
            unsafe { mmap(null_mut(), size, prot, flags, -1, 0) },78x
            MAP_FAILED
        )?;
        Ok(Self::from_raw(addr, size, None))78x
    }78x
 
    /// Given offset and len, return the host virtual address and len;
    /// len might be truncated.
    fn get_valid_range(&self, offset: usize, len: usize) -> Result<(usize, usize)> {2388x
        let end = offset.wrapping_add(len).wrapping_sub(1);2388x
        if offset >= self.size || end < offset {2388x
            return error::ExceedsLimit {
                addr: offset as u64,
                size: len as u64,
            }
            .fail();
        }2388x
        let valid_len = std::cmp::min(self.size - offset, len);2388x
        Ok((self.addr + offset, valid_len))2388x
    }2388x
 
    pub fn as_slice_mut(&mut self) -> &mut [u8] {
        unsafe { std::slice::from_raw_parts_mut(self.addr as *mut u8, self.size) }
    }
 
    pub fn as_slice(&self) -> &[u8] {
        unsafe { std::slice::from_raw_parts(self.addr as *const u8, self.size) }
    }
 
    /// Given offset and len, return a slice, len might be truncated.
    fn get_partial_slice(&self, offset: usize, len: usize) -> Result<&[u8], Error> {1002x
        let (addr, len) = self.get_valid_range(offset, len)?;1002x
        Ok(unsafe { std::slice::from_raw_parts(addr as *const u8, len) })1002x
    }1002x
 
    /// Given offset and len, return a mutable slice, len might be truncated.
    #[allow(clippy::mut_from_ref)]
    fn get_partial_slice_mut(&self, offset: usize, len: usize) -> Result<&mut [u8], Error> {1386x
        let (addr, len) = self.get_valid_range(offset, len)?;1386x
        Ok(unsafe { std::slice::from_raw_parts_mut(addr as *mut u8, len) })1386x
    }1386x
}
 
#[derive(Debug)]
pub struct Ram {
    inner: Addressable<ArcMemPages>,
}
 
#[derive(Debug)]
pub struct RamBus {
    ram: RwLock<Ram>,
}
 
struct Iter<'m> {
    ram: &'m Ram,
    gpa: u64,
    remain: u64,
}
 
impl<'m> Iterator for Iter<'m> {
    type Item = Result<&'m [u8]>;
 
    fn next(&mut self) -> Option<Self::Item> {1224x
        if self.remain == 0 {1224x
            return None;384x
        }840x
        let r = self.ram.get_partial_slice(self.gpa, self.remain);840x
        if let Ok(s) = r {840x
            self.gpa += s.len() as u64;747x
            self.remain -= s.len() as u64;747x
        }747x
        Some(r)840x
    }1224x
}
 
struct IterMut<'m> {
    ram: &'m Ram,
    gpa: u64,
    remain: u64,
}
 
impl<'m> Iterator for IterMut<'m> {
    type Item = Result<&'m mut [u8]>;
 
    fn next(&mut self) -> Option<Self::Item> {1158x
        if self.remain == 0 {1158x
            return None;345x
        }813x
        let r = self.ram.get_partial_slice_mut(self.gpa, self.remain);813x
        if let Ok(ref s) = r {813x
            self.gpa += s.len() as u64;720x
            self.remain -= s.len() as u64;720x
        }720x
        Some(r)813x
    }1158x
}
 
impl Ram {
    fn slice_iter(&self, gpa: u64, len: u64) -> Iter<'_> {477x
        Iter {477x
            ram: self,477x
            gpa,477x
            remain: len,477x
        }477x
    }477x
 
    fn slice_iter_mut(&self, gpa: u64, len: u64) -> IterMut<'_> {438x
        IterMut {438x
            ram: self,438x
            gpa,438x
            remain: len,438x
        }438x
    }438x
 
    fn get_partial_slice(&self, gpa: u64, len: u64) -> Result<&[u8]> {1095x
        let Some((start, user_mem)) = self.inner.search(gpa) else {1095x
            return error::NotMapped { addr: gpa }.fail();93x
        };
        user_mem.get_partial_slice((gpa - start) as usize, len as usize)1002x
    }1095x
 
    fn get_partial_slice_mut(&self, gpa: u64, len: u64) -> Result<&mut [u8]> {1479x
        let Some((start, user_mem)) = self.inner.search(gpa) else {1479x
            return error::NotMapped { addr: gpa }.fail();93x
        };
        user_mem.get_partial_slice_mut((gpa - start) as usize, len as usize)1386x
    }1479x
 
    pub fn get_slice<T>(&self, gpa: u64, len: u64) -> Result<&[UnsafeCell<T>], Error> {
        let total_len = len * size_of::<T>() as u64;
        let host_ref = self.get_partial_slice(gpa, total_len)?;
        let ptr = host_ref.as_ptr() as *const UnsafeCell<T>;
        if host_ref.len() as u64 != total_len {
            error::NotContinuous {
                addr: gpa,
                size: total_len,
            }
            .fail()
        } else if !ptr.is_aligned() {
            error::NotAligned {
                addr: ptr as u64,
                align: align_of::<T>(),
            }
            .fail()
        } else {
            Ok(unsafe { &*core::ptr::slice_from_raw_parts(ptr, len as usize) })
        }
    }
 
    pub fn get_ptr<T>(&self, gpa: u64) -> Result<*mut T, Error> {390x
        let host_ref = self.get_partial_slice_mut(gpa, size_of::<T>() as u64)?;390x
        let ptr = host_ref.as_mut_ptr();390x
        if host_ref.len() != size_of::<T>() {390x
            error::NotContinuous {
                addr: gpa,
                size: size_of::<T>() as u64,
            }
            .fail()
        } else if !ptr.is_aligned() {390x
            error::NotAligned {
                addr: ptr as u64,
                align: align_of::<T>(),
            }
            .fail()
        } else {
            Ok(ptr as *mut T)390x
        }
    }390x
 
    pub fn read(&self, gpa: u64, buf: &mut [u8]) -> Result<()> {255x
        let host_ref = self.get_partial_slice(gpa, buf.len() as u64)?;255x
        if host_ref.len() == buf.len() {255x
            buf.copy_from_slice(host_ref);69x
        } else {69x
            let mut cur = 0;186x
            for r in self.slice_iter(gpa, buf.len() as u64) {372x
                let s = r?;372x
                let s_len = s.len();279x
                buf[cur..(cur + s_len)].copy_from_slice(s);279x
                cur += s_len;279x
            }
        }
        Ok(())162x
    }255x
 
    pub fn read_t<T>(&self, gpa: u64) -> Result<T>225x
    where225x
        T: FromBytes + IntoBytes,225x
    {
        let mut v = T::new_zeroed();225x
        self.read(gpa, v.as_mut_bytes())?;225x
        Ok(v)132x
    }225x
 
    pub fn write(&self, gpa: u64, buf: &[u8]) -> Result<()> {276x
        let len = buf.len() as u64;276x
        let host_ref = self.get_partial_slice_mut(gpa, len)?;276x
        if host_ref.len() == buf.len() {276x
            host_ref.copy_from_slice(buf);90x
            Ok(())90x
        } else {
            let mut cur = 0;186x
            for r in self.slice_iter_mut(gpa, len) {372x
                let s = r?;372x
                let s_len = s.len();279x
                s.copy_from_slice(&buf[cur..(cur + s_len)]);279x
                cur += s_len;279x
            }
            Ok(())93x
        }
    }276x
 
    pub fn write_t<T>(&self, gpa: u64, val: &T) -> Result<(), Error>225x
    where225x
        T: IntoBytes + Immutable,225x
    {
        self.write(gpa, val.as_bytes())225x
    }225x
 
    pub fn translate(&self, gpa: u64) -> Result<*const u8> {
        let s = self.get_partial_slice(gpa, 1)?;
        Ok(s.as_ptr())
    }
 
    pub fn translate_iov<'a>(&'a self, iov: &[(u64, u64)]) -> Result<Vec<IoSlice<'a>>> {90x
        let mut slices = vec![];90x
        for (gpa, len) in iov {90x
            for r in self.slice_iter(*gpa, *len) {87x
                slices.push(IoSlice::new(r?));72x
            }
        }
        Ok(slices)90x
    }90x
 
    pub fn translate_iov_mut<'a>(&'a self, iov: &[(u64, u64)]) -> Result<Vec<IoSliceMut<'a>>> {87x
        let mut slices = vec![];87x
        for (gpa, len) in iov {87x
            for r in self.slice_iter_mut(*gpa, *len) {48x
                slices.push(IoSliceMut::new(r?));45x
            }
        }
        Ok(slices)87x
    }87x
 
    pub fn iter(&self) -> impl DoubleEndedIterator<Item = (u64, &ArcMemPages)> {
        self.inner.iter()
    }
 
    pub fn madvise(&self, gpa: u64, size: u64, advice: i32) -> Result<()> {
        for r in self.slice_iter_mut(gpa, size) {
            let s = r?;
            ffi!(unsafe { madvise(s.as_mut_ptr() as _, s.len(), advice) })?;
        }
        Ok(())
    }
}
 
impl Default for RamBus {
    fn default() -> Self {
        Self::new()
    }
}
 
impl RamBus {
    pub fn lock_layout(&self) -> RwLockReadGuard<'_, Ram> {81x
        self.ram.read()81x
    }81x
 
    pub fn new() -> Self {81x
        Self {81x
            ram: RwLock::new(Ram {81x
                inner: Addressable::default(),81x
            }),81x
        }81x
    }81x
 
    pub(crate) fn add(&self, gpa: u64, user_mem: ArcMemPages) -> Result<(), Error> {78x
        let mut ram = self.ram.write();78x
        ram.inner.add(gpa, user_mem)?;78x
        Ok(())78x
    }78x
 
    pub(crate) fn remove(&self, gpa: u64) -> Result<ArcMemPages, Error> {3x
        let mut ram = self.ram.write();3x
        ram.inner.remove(gpa)3x
    }3x
 
    pub fn read(&self, gpa: u64, buf: &mut [u8]) -> Result<()> {
        let ram = self.ram.read();
        ram.read(gpa, buf)
    }
 
    pub fn write(&self, gpa: u64, buf: &[u8]) -> Result<()> {
        let ram = self.ram.read();
        ram.write(gpa, buf)
    }
 
    pub fn read_t<T>(&self, gpa: u64) -> Result<T, Error>225x
    where225x
        T: FromBytes + IntoBytes,225x
    {
        let ram = self.ram.read();225x
        ram.read_t(gpa)225x
    }225x
 
    pub fn write_t<T>(&self, gpa: u64, val: &T) -> Result<(), Error>225x
    where225x
        T: IntoBytes + Immutable,225x
    {
        let ram = self.ram.read();225x
        ram.write_t(gpa, val)225x
    }225x
 
    pub fn read_range(&self, gpa: u64, len: u64, dst: &mut impl Write) -> Result<()> {195x
        let ram = self.ram.read();195x
        for r in ram.slice_iter(gpa, len) {384x
            dst.write_all(r?).context(error::Write)?;384x
        }
        Ok(())195x
    }195x
 
    pub fn write_range(&self, gpa: u64, len: u64, mut src: impl Read) -> Result<()> {195x
        let ram = self.ram.read();195x
        for r in ram.slice_iter_mut(gpa, len) {384x
            src.read_exact(r?).context(error::Read)?;384x
        }
        Ok(())195x
    }195x
 
    pub fn read_vectored<T, F>(&self, bufs: &[(u64, u64)], callback: F) -> Result<T, Error>3x
    where3x
        F: FnOnce(&[IoSlice<'_>]) -> T,3x
    {
        let ram = self.ram.read();3x
        let mut iov = vec![];3x
        for (gpa, len) in bufs {9x
            for r in ram.slice_iter(*gpa, *len) {12x
                iov.push(IoSlice::new(r?));12x
            }
        }
        Ok(callback(&iov))3x
    }3x
 
    pub fn write_vectored<T, F>(&self, bufs: &[(u64, u64)], callback: F) -> Result<T, Error>3x
    where3x
        F: FnOnce(&mut [IoSliceMut<'_>]) -> T,3x
    {
        let ram = self.ram.read();3x
        let mut iov = vec![];3x
        for (gpa, len) in bufs {9x
            for r in ram.slice_iter_mut(*gpa, *len) {12x
                iov.push(IoSliceMut::new(r?));12x
            }
        }
        Ok(callback(&mut iov))3x
    }3x
}
 
#[cfg(test)]
#[path = "mapped_test.rs"]
mod tests;
 