region/

query.rs

use crate::{os, util, Error, Region, Result};

/// An iterator over the [`Region`]s that encompass an address range.
///
/// This `struct` is created by [`query_range`]. See its documentation for more.
pub struct QueryIter {
  iterator: Option<os::QueryIter>,
  origin: *const (),
}

impl QueryIter {
  pub(crate) fn new<T>(origin: *const T, size: usize) -> Result<Self> {
    let origin = origin.cast();

    os::QueryIter::new(origin, size).map(|iterator| Self {
      iterator: Some(iterator),
      origin,
    })
  }
}

impl Iterator for QueryIter {
  type Item = Result<Region>;

  /// Advances the iterator and returns the next region.
  ///
  /// If the iterator has been exhausted (i.e. all [`Region`]s have been
  /// queried), or if an error is encountered during iteration, all further
  /// invocations will return [`None`] (in the case of an error, the error will
  /// be the last item that is yielded before the iterator is fused).
  #[allow(clippy::missing_inline_in_public_items)]
  fn next(&mut self) -> Option<Self::Item> {
    let regions = self.iterator.as_mut()?;

    while let Some(result) = regions.next() {
      match result {
        Ok(region) => {
          let range = region.as_range();

          // Skip the region if it precedes the queried range
          if range.end <= self.origin as usize {
            continue;
          }

          // Stop iteration if the region is past the queried range
          if range.start >= regions.upper_bound() {
            break;
          }

          return Some(Ok(region));
        }
        Err(error) => {
          self.iterator.take();
          return Some(Err(error));
        }
      }
    }

    self.iterator.take();
    None
  }
}

impl std::iter::FusedIterator for QueryIter {}

unsafe impl Send for QueryIter {}
unsafe impl Sync for QueryIter {}

/// Queries the OS with an address, returning the region it resides within.
///
/// If the queried address does not reside within any mapped region, or if it's
/// outside the process' address space, the function will error with
/// [`Error::UnmappedRegion`].
///
/// # Parameters
///
/// - The enclosing region can be of multiple page sizes.
/// - The address is rounded down to the closest page boundary.
///
/// # Errors
///
/// - If an interaction with the underlying operating system fails, an error
/// will be returned.
///
/// # Examples
///
/// ```
/// # fn main() -> region::Result<()> {
/// use region::Protection;
///
/// let data = [0; 100];
/// let region = region::query(data.as_ptr())?;
///
/// assert_eq!(region.protection(), Protection::READ_WRITE);
/// # Ok(())
/// # }
/// ```
#[inline]
pub fn query<T>(address: *const T) -> Result<Region> {
  // For UNIX systems, the address must be aligned to the closest page boundary
  let (address, size) = util::round_to_page_boundaries(address, 1)?;

  QueryIter::new(address, size)?
    .next()
    .ok_or(Error::UnmappedRegion)?
}

/// Queries the OS for mapped regions that overlap with the specified range.
///
/// The implementation clamps any input that exceeds the boundaries of a
/// process' address space. Therefore it's safe to, e.g., pass in
/// [`std::ptr::null`] and [`usize::max_value`] to iterate the mapped memory
/// pages of an entire process.
///
/// If an error is encountered during iteration, the error will be the last item
/// that is yielded. Thereafter the iterator becomes fused.
///
/// A 2-byte range straddling a page boundary, will return both pages (or one
/// region, if the pages share the same properties).
///
/// This function only returns mapped regions. If required, unmapped regions can
/// be manually identified by inspecting the potential gaps between two
/// neighboring regions.
///
/// # Parameters
///
/// - The range is `[address, address + size)`
/// - The address is rounded down to the closest page boundary.
/// - The size may not be zero.
/// - The size is rounded up to the closest page boundary, relative to the
///   address.
///
/// # Errors
///
/// - If an interaction with the underlying operating system fails, an error
/// will be returned.
/// - If size is zero, [`Error::InvalidParameter`] will be returned.
///
/// # Examples
///
/// ```
/// # use region::Result;
/// # fn main() -> Result<()> {
/// let data = [0; 100];
/// let region = region::query_range(data.as_ptr(), data.len())?
///   .collect::<Result<Vec<_>>>()?;
///
/// assert_eq!(region.len(), 1);
/// assert_eq!(region[0].protection(), region::Protection::READ_WRITE);
/// # Ok(())
/// # }
/// ```
#[inline]
pub fn query_range<T>(address: *const T, size: usize) -> Result<QueryIter> {
  let (address, size) = util::round_to_page_boundaries(address, size)?;
  QueryIter::new(address, size)
}

#[cfg(test)]
mod tests {
  use super::*;
  use crate::tests::util::alloc_pages;
  use crate::{page, Protection};

  const TEXT_SEGMENT_PROT: Protection = if cfg!(target_os = "openbsd") {
    Protection::EXECUTE
  } else {
    Protection::READ_EXECUTE
  };

  #[test]
  fn query_returns_unmapped_for_oob_address() {
    let (min, max) = (std::ptr::null::<()>(), usize::max_value() as *const ());
    assert!(matches!(query(min), Err(Error::UnmappedRegion)));
    assert!(matches!(query(max), Err(Error::UnmappedRegion)));
  }

  #[test]
  fn query_returns_correct_descriptor_for_text_segment() -> Result<()> {
    let region = query(query_returns_correct_descriptor_for_text_segment as *const ())?;

    assert_eq!(region.protection(), TEXT_SEGMENT_PROT);
    assert_eq!(region.is_shared(), cfg!(windows));
    assert!(!region.is_guarded());
    Ok(())
  }

  #[test]
  fn query_returns_one_region_for_multiple_page_allocation() -> Result<()> {
    let alloc = crate::alloc(page::size() + 1, Protection::READ_EXECUTE)?;
    let region = query(alloc.as_ptr::<()>())?;

    assert_eq!(region.protection(), Protection::READ_EXECUTE);
    assert_eq!(region.as_ptr::<()>(), alloc.as_ptr());
    assert_eq!(region.len(), alloc.len());
    assert!(!region.is_guarded());
    Ok(())
  }

  #[test]
  #[cfg(not(target_os = "android"))] // TODO: Determine why this fails on Android in QEMU
  fn query_is_not_off_by_one() -> Result<()> {
    let pages = [Protection::READ, Protection::READ_EXECUTE, Protection::READ];
    let map = alloc_pages(&pages);

    let page_mid = unsafe { map.as_ptr().add(page::size()) };
    let region = query(page_mid)?;

    assert_eq!(region.protection(), Protection::READ_EXECUTE);
    assert_eq!(region.len(), page::size());

    let region = query(unsafe { page_mid.offset(-1) })?;

    assert_eq!(region.protection(), Protection::READ);
    assert_eq!(region.len(), page::size());
    Ok(())
  }

  #[test]
  fn query_range_does_not_return_unmapped_regions() -> Result<()> {
    let regions = query_range(std::ptr::null::<()>(), 1)?.collect::<Result<Vec<_>>>()?;
    assert!(regions.is_empty());
    Ok(())
  }

  #[test]
  fn query_range_returns_both_regions_for_straddling_range() -> Result<()> {
    let pages = [Protection::READ_EXECUTE, Protection::READ_WRITE];
    let map = alloc_pages(&pages);

    // Query an area that overlaps both pages
    let address = unsafe { map.as_ptr().offset(page::size() as isize - 1) };
    let regions = query_range(address, 2)?.collect::<Result<Vec<_>>>()?;

    assert_eq!(regions.len(), pages.len());
    for (page, region) in pages.iter().zip(regions.iter()) {
      assert_eq!(*page, region.protection);
    }
    Ok(())
  }

  #[test]
  fn query_range_has_inclusive_lower_and_exclusive_upper_bound() -> Result<()> {
    let pages = [Protection::READ, Protection::READ_WRITE, Protection::READ];
    let map = alloc_pages(&pages);

    let regions = query_range(map.as_ptr(), page::size())?.collect::<Result<Vec<_>>>()?;
    assert_eq!(regions.len(), 1);
    assert_eq!(regions[0].protection(), Protection::READ);

    let regions = query_range(map.as_ptr(), page::size() + 1)?.collect::<Result<Vec<_>>>()?;
    assert_eq!(regions.len(), 2);
    assert_eq!(regions[0].protection(), Protection::READ);
    assert_eq!(regions[1].protection(), Protection::READ_WRITE);
    Ok(())
  }

  #[test]
  fn query_range_can_iterate_over_entire_process() -> Result<()> {
    let regions =
      query_range(std::ptr::null::<()>(), usize::max_value())?.collect::<Result<Vec<_>>>()?;

    // This test is a bit rough around the edges
    assert!(regions
      .iter()
      .any(|region| region.protection() == Protection::READ));
    assert!(regions
      .iter()
      .any(|region| region.protection() == Protection::READ_WRITE));
    assert!(regions
      .iter()
      .any(|region| region.protection() == TEXT_SEGMENT_PROT));
    assert!(regions.len() > 5);
    Ok(())
  }

  #[test]
  fn query_range_iterator_is_fused_after_exhaustion() -> Result<()> {
    let pages = [Protection::READ, Protection::READ_WRITE];
    let map = alloc_pages(&pages);
    let mut iter = query_range(map.as_ptr(), page::size() + 1)?;

    assert_eq!(
      iter.next().transpose()?.map(|r| r.protection()),
      Some(Protection::READ)
    );
    assert_eq!(
      iter.next().transpose()?.map(|r| r.protection()),
      Some(Protection::READ_WRITE)
    );
    assert_eq!(iter.next().transpose()?, None);
    assert_eq!(iter.next().transpose()?, None);
    Ok(())
  }
}