<commit_before>#[cfg(feature = "zstd")]
use zstd::block::{compress, decompress};

use crate::*;

use super::{
    arr_to_u32, gc_blobs, pwrite_all, raw_segment_iter_from, u32_to_arr,
    u64_to_arr, BasedBuf, DiskPtr, LogIter, LogKind, LogOffset, Lsn,
    MessageKind,
};

/// A snapshot of the state required to quickly restart
/// the `PageCache` and `SegmentAccountant`.
#[derive(Debug, Default)]
#[cfg_attr(test, derive(Clone, PartialEq))]
pub struct Snapshot {
    /// The last read message lsn
    pub stable_lsn: Lsn,
    /// The last read message lid
    pub tip_lid: Option<LogOffset>,
    /// the mapping from pages to (lsn, lid)
    pub pt: Vec<PageState>,
}

#[derive(Clone, Debug, PartialEq)]
pub enum PageState {
    Present { base: (Lsn, DiskPtr, u64), frags: Vec<(Lsn, DiskPtr, u64)> },
    Free(Lsn, DiskPtr),
    Uninitialized,
}

impl PageState {
    fn push(&mut self, item: (Lsn, DiskPtr, u64)) {
        match *self {
            PageState::Present { base, ref mut frags } => {
                if frags.last().map_or(base.0, |f| f.0) < item.0 {
                    frags.push(item)
                } else {
                    debug!(
                        "skipping merging item {:?} into \
                        existing PageState::Present({:?})",
                        item, frags
                    );
                }
            }
            _ => panic!("pushed frags to {:?}", self),
        }
    }

    pub(crate) fn is_free(&self) -> bool {
        match *self {
            PageState::Free(_, _) => true,
            _ => false,
        }
    }
}

impl Snapshot {
    fn apply(
        &mut self,
        log_kind: LogKind,
        pid: PageId,
        lsn: Lsn,
        disk_ptr: DiskPtr,
        sz: u64,
    ) {
        // unwrapping this because it's already passed the crc check
        // in the log iterator
        trace!("trying to deserialize buf for ptr {} lsn {}", disk_ptr, lsn);
        let _measure = Measure::new(&M.snapshot_apply);

        if self.pt.len() <= usize::try_from(pid).unwrap() {
            self.pt.resize(
                usize::try_from(pid + 1).unwrap(),
                PageState::Uninitialized,
            );
        }

        match log_kind {
            LogKind::Replace => {
                trace!(
                    "compact of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn,
                );

                self.pt[usize::try_from(pid).unwrap()] = PageState::Present {
                    base: (lsn, disk_ptr, sz),
                    frags: vec![],
                };
            }
            LogKind::Link => {
                // Because we rewrite pages over time, we may have relocated
                // a page's initial Compact to a later segment. We should skip
                // over pages here unless we've encountered a Compact for them.
                if let Some(lids @ PageState::Present { .. }) =
                    self.pt.get_mut(usize::try_from(pid).unwrap())
                {
                    trace!(
                        "append of pid {} at lid {} lsn {}",
                        pid,
                        disk_ptr,
                        lsn,
                    );

                    if lids.is_free() {
                        // this can happen if the allocate or replace
                        // has been moved to a later segment.

                        trace!(
                            "we have not yet encountered an \
                             allocation of this page, skipping push"
                        );

                        return;
                    }

                    lids.push((lsn, disk_ptr, sz));
                } else {
                    trace!(
                        "skipping dangling append of pid {} at lid {} lsn {}",
                        pid,
                        disk_ptr,
                        lsn,
                    );
                }
            }
            LogKind::Free => {
                trace!("free of pid {} at ptr {} lsn {}", pid, disk_ptr, lsn);
                self.pt[usize::try_from(pid).unwrap()] =
                    PageState::Free(lsn, disk_ptr);
            }
            LogKind::Corrupted | LogKind::Skip => panic!(
                "unexppected messagekind in snapshot application: {:?}",
                log_kind
            ),
        }
    }
}

fn advance_snapshot(
    mut iter: LogIter,
    mut snapshot: Snapshot,
    config: &RunningConfig,
) -> Result<Snapshot> {
    let _measure = Measure::new(&M.advance_snapshot);

    trace!("building on top of old snapshot: {:?}", snapshot);

    let old_stable_lsn = snapshot.stable_lsn;

    while let Some((log_kind, pid, lsn, ptr, sz)) = iter.next() {
        trace!(
            "in advance_snapshot looking at item with pid {} lsn {} ptr {}",
            pid,
            lsn,
            ptr
        );

        if lsn < snapshot.stable_lsn {
            // don't process already-processed Lsn's. stable_lsn is for the last
            // item ALREADY INCLUDED lsn in the snapshot.
            trace!(
                "continuing in advance_snapshot, lsn {} ptr {} stable_lsn {:?}",
                lsn,
                ptr,
                snapshot.stable_lsn
            );
            continue;
        }

        snapshot.apply(log_kind, pid, lsn, ptr, sz);
    }

    // `snapshot.tip_lid` can be set based on 4 possibilities for the tip of the
    // log:
    // 1. an empty DB - tip set to None, causing a fresh segment to be
    //    allocated on initialization
    // 2. the recovered tip is at the end of a
    //    segment with less space left than would fit MAX_MSG_HEADER_LEN -
    //    tip set to None, causing a fresh segment to be allocated on
    //    initialization, as in #1 above
    // 3. the recovered tip is in the middle of a segment - both set to the end
    //    of the last valid message, causing the system to be initialized to
    //    that point without allocating a new segment
    // 4. the recovered tip is at the beginning of a new segment, but without
    //    any valid messages in it yet. treat as #3 above, but also take care
    //    in te SA initialization to properly initialize any segment tracking
    //    state despite not having any pages currently residing there.
    //
    // The information we have at this point is iter.cur_lsn and the
    // iter.segment_base.offset. Situation 1 & 2 happen if iter.cur_lsn %
    // segment_size == 0, and result in None for the recovered log tip.
    // Situations 3 & 4 otherwise. We don't need to differentiate between them
    // here.

    let segment_progress: Lsn = iter.cur_lsn % (config.segment_size as Lsn);

    snapshot.stable_lsn = if segment_progress + MAX_MSG_HEADER_LEN as Lsn
        >= config.segment_size as Lsn
    {
        config.normalize(iter.cur_lsn) + config.segment_size as Lsn
    } else if iter.cur_lsn != 0 {
        iter.cur_lsn
    } else {
        snapshot.stable_lsn
    };

    if segment_progress == 0 || iter.segment_base.is_none() {
        // either situation 1 or situation 2
        snapshot.tip_lid = None;
    } else if let Some(BasedBuf { offset, .. }) = iter.segment_base {
        // either situation 3 or situation 4. we need to zero the

        snapshot.tip_lid = Some(offset + segment_progress as LogOffset);

        // zero the tail of the segment after the recovered tip
        let shred_len = config.segment_size - segment_progress as usize;
        let shred_zone = vec![MessageKind::Corrupted.into(); shred_len];
        debug!(
            "zeroing the end of the recovered segment between {} and {}",
            snapshot.tip_lid.unwrap(),
            snapshot.tip_lid.unwrap() + shred_len as LogOffset
        );
        pwrite_all(&config.file, &shred_zone, snapshot.tip_lid.unwrap())?;
        config.file.sync_all()?;
    } else {
        unreachable!()
    }

    trace!("generated snapshot: {:?}", snapshot);

    assert!(snapshot.stable_lsn >= old_stable_lsn);

    if snapshot.stable_lsn > old_stable_lsn {
        write_snapshot(config, &snapshot)?;
    }

    for to_zero in iter.segments.values().copied() {
        debug!("zeroing torn segment at lid {}", to_zero);

        // NB we intentionally corrupt this header to prevent any segment
        // from being allocated which would duplicate its LSN, messing
        // up recovery in the future.
        io_fail!(config, "segment initial free zero");
        pwrite_all(
            &config.file,
            &*vec![MessageKind::Corrupted.into(); SEG_HEADER_LEN],
            to_zero,
        )?;
        if !config.temporary {
            config.file.sync_all()?;
        }
    }

    // remove all blob files larger than our stable offset
    gc_blobs(config, snapshot.stable_lsn)?;

    #[cfg(feature = "event_log")]
    config.event_log.recovered_lsn(snapshot.stable_lsn);

    Ok(snapshot)
}

/// Read a `Snapshot` or generate a default, then advance it to
/// the tip of the data file, if present.
pub fn read_snapshot_or_default(config: &RunningConfig) -> Result<Snapshot> {
    let last_snap = read_snapshot(config)?.unwrap_or_else(Snapshot::default);

    let log_iter = raw_segment_iter_from(last_snap.stable_lsn, config)?;

    let res = advance_snapshot(log_iter, last_snap, config)?;

    Ok(res)
}

/// Read a `Snapshot` from disk.
fn read_snapshot(config: &RunningConfig) -> std::io::Result<Option<Snapshot>> {
    let mut f = loop {
        let mut candidates = config.get_snapshot_files()?;
        if candidates.is_empty() {
            debug!("no previous snapshot found");
            return Ok(None);
        }

        candidates.sort();

        let path = candidates.pop().unwrap();

        match std::fs::OpenOptions::new().read(true).open(&path) {
            Ok(f) => break f,
            Err(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
                // this can happen if there's a race
                continue;
            }
            Err(other) => return Err(other),
        }
    };

    let mut buf = vec![];
    let _read = f.read_to_end(&mut buf)?;
    let len = buf.len();
    if len <= 12 {
        warn!("empty/corrupt snapshot file found");
        return Ok(None);
    }

    let mut len_expected_bytes = [0; 8];
    len_expected_bytes.copy_from_slice(&buf[len - 12..len - 4]);

    let mut crc_expected_bytes = [0; 4];
    crc_expected_bytes.copy_from_slice(&buf[len - 4..]);

    let _ = buf.split_off(len - 12);
    let crc_expected: u32 = arr_to_u32(&crc_expected_bytes);

    let crc_actual = crc32(&buf);

    if crc_expected != crc_actual {
        return Ok(None);
    }

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        use std::convert::TryInto;

        let len_expected: u64 =
            u64::from_le_bytes(len_expected_bytes.as_ref().try_into().unwrap());

        decompress(&*buf, usize::try_from(len_expected).unwrap()).unwrap()
    } else {
        buf
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = buf;

    Ok(Snapshot::deserialize(&mut bytes.as_slice()).ok())
}

fn write_snapshot(config: &RunningConfig, snapshot: &Snapshot) -> Result<()> {
    trace!("writing snapshot {:?}", snapshot);

    let raw_bytes = snapshot.serialize();
    let decompressed_len = raw_bytes.len();

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        compress(&*raw_bytes, config.compression_factor).unwrap()
    } else {
        raw_bytes
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = raw_bytes;

    let crc32: [u8; 4] = u32_to_arr(crc32(&bytes));
    let len_bytes: [u8; 8] = u64_to_arr(decompressed_len as u64);

    let path_1_suffix = format!("snap.{:016X}.generating", snapshot.stable_lsn);

    let mut path_1 = config.get_path();
    path_1.push(path_1_suffix);

    let path_2_suffix = format!("snap.{:016X}", snapshot.stable_lsn);

    let mut path_2 = config.get_path();
    path_2.push(path_2_suffix);

    let parent = path_1.parent().unwrap();
    std::fs::create_dir_all(parent)?;
    let mut f =
        std::fs::OpenOptions::new().write(true).create(true).open(&path_1)?;

    // write the snapshot bytes, followed by a crc64 checksum at the end
    io_fail!(config, "snap write");
    f.write_all(&*bytes)?;
    io_fail!(config, "snap write len");
    f.write_all(&len_bytes)?;
    io_fail!(config, "snap write crc");
    f.write_all(&crc32)?;
    io_fail!(config, "snap write post");
    f.sync_all()?;

    trace!("wrote snapshot to {}", path_1.to_string_lossy());

    io_fail!(config, "snap write mv");
    std::fs::rename(&path_1, &path_2)?;
    io_fail!(config, "snap write mv post");

    trace!("renamed snapshot to {}", path_2.to_string_lossy());

    // clean up any old snapshots
    let candidates = config.get_snapshot_files()?;
    for path in candidates {
        let path_str = path.file_name().unwrap().to_str().unwrap();
        if !path_2.to_string_lossy().ends_with(&*path_str) {
            debug!("removing old snapshot file {:?}", path);

            io_fail!(config, "snap write rm old");

            if let Err(e) = std::fs::remove_file(&path) {
                // TODO should this just be a try return?
                warn!(
                    "failed to remove old snapshot file, maybe snapshot race? {}",
                    e
                );
            }
        }
    }
    Ok(())
}
<commit_msg>when zeroing torn segments on recovery, zero the entire segment to prevent replay attacks, not just the segment header<commit_after>#[cfg(feature = "zstd")]
use zstd::block::{compress, decompress};

use crate::*;

use super::{
    arr_to_u32, gc_blobs, pwrite_all, raw_segment_iter_from, u32_to_arr,
    u64_to_arr, BasedBuf, DiskPtr, LogIter, LogKind, LogOffset, Lsn,
    MessageKind,
};

/// A snapshot of the state required to quickly restart
/// the `PageCache` and `SegmentAccountant`.
#[derive(Debug, Default)]
#[cfg_attr(test, derive(Clone, PartialEq))]
pub struct Snapshot {
    /// The last read message lsn
    pub stable_lsn: Lsn,
    /// The last read message lid
    pub tip_lid: Option<LogOffset>,
    /// the mapping from pages to (lsn, lid)
    pub pt: Vec<PageState>,
}

#[derive(Clone, Debug, PartialEq)]
pub enum PageState {
    Present { base: (Lsn, DiskPtr, u64), frags: Vec<(Lsn, DiskPtr, u64)> },
    Free(Lsn, DiskPtr),
    Uninitialized,
}

impl PageState {
    fn push(&mut self, item: (Lsn, DiskPtr, u64)) {
        match *self {
            PageState::Present { base, ref mut frags } => {
                if frags.last().map_or(base.0, |f| f.0) < item.0 {
                    frags.push(item)
                } else {
                    debug!(
                        "skipping merging item {:?} into \
                        existing PageState::Present({:?})",
                        item, frags
                    );
                }
            }
            _ => panic!("pushed frags to {:?}", self),
        }
    }

    pub(crate) fn is_free(&self) -> bool {
        match *self {
            PageState::Free(_, _) => true,
            _ => false,
        }
    }
}

impl Snapshot {
    fn apply(
        &mut self,
        log_kind: LogKind,
        pid: PageId,
        lsn: Lsn,
        disk_ptr: DiskPtr,
        sz: u64,
    ) {
        // unwrapping this because it's already passed the crc check
        // in the log iterator
        trace!("trying to deserialize buf for ptr {} lsn {}", disk_ptr, lsn);
        let _measure = Measure::new(&M.snapshot_apply);

        if self.pt.len() <= usize::try_from(pid).unwrap() {
            self.pt.resize(
                usize::try_from(pid + 1).unwrap(),
                PageState::Uninitialized,
            );
        }

        match log_kind {
            LogKind::Replace => {
                trace!(
                    "compact of pid {} at ptr {} lsn {}",
                    pid,
                    disk_ptr,
                    lsn,
                );

                self.pt[usize::try_from(pid).unwrap()] = PageState::Present {
                    base: (lsn, disk_ptr, sz),
                    frags: vec![],
                };
            }
            LogKind::Link => {
                // Because we rewrite pages over time, we may have relocated
                // a page's initial Compact to a later segment. We should skip
                // over pages here unless we've encountered a Compact for them.
                if let Some(lids @ PageState::Present { .. }) =
                    self.pt.get_mut(usize::try_from(pid).unwrap())
                {
                    trace!(
                        "append of pid {} at lid {} lsn {}",
                        pid,
                        disk_ptr,
                        lsn,
                    );

                    if lids.is_free() {
                        // this can happen if the allocate or replace
                        // has been moved to a later segment.

                        trace!(
                            "we have not yet encountered an \
                             allocation of this page, skipping push"
                        );

                        return;
                    }

                    lids.push((lsn, disk_ptr, sz));
                } else {
                    trace!(
                        "skipping dangling append of pid {} at lid {} lsn {}",
                        pid,
                        disk_ptr,
                        lsn,
                    );
                }
            }
            LogKind::Free => {
                trace!("free of pid {} at ptr {} lsn {}", pid, disk_ptr, lsn);
                self.pt[usize::try_from(pid).unwrap()] =
                    PageState::Free(lsn, disk_ptr);
            }
            LogKind::Corrupted | LogKind::Skip => panic!(
                "unexppected messagekind in snapshot application: {:?}",
                log_kind
            ),
        }
    }
}

fn advance_snapshot(
    mut iter: LogIter,
    mut snapshot: Snapshot,
    config: &RunningConfig,
) -> Result<Snapshot> {
    let _measure = Measure::new(&M.advance_snapshot);

    trace!("building on top of old snapshot: {:?}", snapshot);

    let old_stable_lsn = snapshot.stable_lsn;

    while let Some((log_kind, pid, lsn, ptr, sz)) = iter.next() {
        trace!(
            "in advance_snapshot looking at item with pid {} lsn {} ptr {}",
            pid,
            lsn,
            ptr
        );

        if lsn < snapshot.stable_lsn {
            // don't process already-processed Lsn's. stable_lsn is for the last
            // item ALREADY INCLUDED lsn in the snapshot.
            trace!(
                "continuing in advance_snapshot, lsn {} ptr {} stable_lsn {:?}",
                lsn,
                ptr,
                snapshot.stable_lsn
            );
            continue;
        }

        snapshot.apply(log_kind, pid, lsn, ptr, sz);
    }

    // `snapshot.tip_lid` can be set based on 4 possibilities for the tip of the
    // log:
    // 1. an empty DB - tip set to None, causing a fresh segment to be
    //    allocated on initialization
    // 2. the recovered tip is at the end of a
    //    segment with less space left than would fit MAX_MSG_HEADER_LEN -
    //    tip set to None, causing a fresh segment to be allocated on
    //    initialization, as in #1 above
    // 3. the recovered tip is in the middle of a segment - both set to the end
    //    of the last valid message, causing the system to be initialized to
    //    that point without allocating a new segment
    // 4. the recovered tip is at the beginning of a new segment, but without
    //    any valid messages in it yet. treat as #3 above, but also take care
    //    in te SA initialization to properly initialize any segment tracking
    //    state despite not having any pages currently residing there.
    //
    // The information we have at this point is iter.cur_lsn and the
    // iter.segment_base.offset. Situation 1 & 2 happen if iter.cur_lsn %
    // segment_size == 0, and result in None for the recovered log tip.
    // Situations 3 & 4 otherwise. We don't need to differentiate between them
    // here.

    let segment_progress: Lsn = iter.cur_lsn % (config.segment_size as Lsn);

    snapshot.stable_lsn = if segment_progress + MAX_MSG_HEADER_LEN as Lsn
        >= config.segment_size as Lsn
    {
        config.normalize(iter.cur_lsn) + config.segment_size as Lsn
    } else if iter.cur_lsn != 0 {
        iter.cur_lsn
    } else {
        snapshot.stable_lsn
    };

    if segment_progress == 0 || iter.segment_base.is_none() {
        // either situation 1 or situation 2
        snapshot.tip_lid = None;
    } else if let Some(BasedBuf { offset, .. }) = iter.segment_base {
        // either situation 3 or situation 4. we need to zero the

        snapshot.tip_lid = Some(offset + segment_progress as LogOffset);

        // zero the tail of the segment after the recovered tip
        let shred_len = config.segment_size - segment_progress as usize;
        let shred_zone = vec![MessageKind::Corrupted.into(); shred_len];
        debug!(
            "zeroing the end of the recovered segment between {} and {}",
            snapshot.tip_lid.unwrap(),
            snapshot.tip_lid.unwrap() + shred_len as LogOffset
        );
        pwrite_all(&config.file, &shred_zone, snapshot.tip_lid.unwrap())?;
        config.file.sync_all()?;
    } else {
        unreachable!()
    }

    trace!("generated snapshot: {:?}", snapshot);

    assert!(snapshot.stable_lsn >= old_stable_lsn);

    if snapshot.stable_lsn > old_stable_lsn {
        write_snapshot(config, &snapshot)?;
    }

    for to_zero in iter.segments.values().copied() {
        debug!("zeroing torn segment at lid {}", to_zero);

        // NB we intentionally corrupt this header to prevent any segment
        // from being allocated which would duplicate its LSN, messing
        // up recovery in the future.
        io_fail!(config, "segment initial free zero");
        pwrite_all(
            &config.file,
            &*vec![MessageKind::Corrupted.into(); config.segment_size],
            to_zero,
        )?;
        if !config.temporary {
            config.file.sync_all()?;
        }
    }

    // remove all blob files larger than our stable offset
    gc_blobs(config, snapshot.stable_lsn)?;

    #[cfg(feature = "event_log")]
    config.event_log.recovered_lsn(snapshot.stable_lsn);

    Ok(snapshot)
}

/// Read a `Snapshot` or generate a default, then advance it to
/// the tip of the data file, if present.
pub fn read_snapshot_or_default(config: &RunningConfig) -> Result<Snapshot> {
    let last_snap = read_snapshot(config)?.unwrap_or_else(Snapshot::default);

    let log_iter = raw_segment_iter_from(last_snap.stable_lsn, config)?;

    let res = advance_snapshot(log_iter, last_snap, config)?;

    Ok(res)
}

/// Read a `Snapshot` from disk.
fn read_snapshot(config: &RunningConfig) -> std::io::Result<Option<Snapshot>> {
    let mut f = loop {
        let mut candidates = config.get_snapshot_files()?;
        if candidates.is_empty() {
            debug!("no previous snapshot found");
            return Ok(None);
        }

        candidates.sort();

        let path = candidates.pop().unwrap();

        match std::fs::OpenOptions::new().read(true).open(&path) {
            Ok(f) => break f,
            Err(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
                // this can happen if there's a race
                continue;
            }
            Err(other) => return Err(other),
        }
    };

    let mut buf = vec![];
    let _read = f.read_to_end(&mut buf)?;
    let len = buf.len();
    if len <= 12 {
        warn!("empty/corrupt snapshot file found");
        return Ok(None);
    }

    let mut len_expected_bytes = [0; 8];
    len_expected_bytes.copy_from_slice(&buf[len - 12..len - 4]);

    let mut crc_expected_bytes = [0; 4];
    crc_expected_bytes.copy_from_slice(&buf[len - 4..]);

    let _ = buf.split_off(len - 12);
    let crc_expected: u32 = arr_to_u32(&crc_expected_bytes);

    let crc_actual = crc32(&buf);

    if crc_expected != crc_actual {
        return Ok(None);
    }

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        use std::convert::TryInto;

        let len_expected: u64 =
            u64::from_le_bytes(len_expected_bytes.as_ref().try_into().unwrap());

        decompress(&*buf, usize::try_from(len_expected).unwrap()).unwrap()
    } else {
        buf
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = buf;

    Ok(Snapshot::deserialize(&mut bytes.as_slice()).ok())
}

fn write_snapshot(config: &RunningConfig, snapshot: &Snapshot) -> Result<()> {
    trace!("writing snapshot {:?}", snapshot);

    let raw_bytes = snapshot.serialize();
    let decompressed_len = raw_bytes.len();

    #[cfg(feature = "zstd")]
    let bytes = if config.use_compression {
        compress(&*raw_bytes, config.compression_factor).unwrap()
    } else {
        raw_bytes
    };

    #[cfg(not(feature = "zstd"))]
    let bytes = raw_bytes;

    let crc32: [u8; 4] = u32_to_arr(crc32(&bytes));
    let len_bytes: [u8; 8] = u64_to_arr(decompressed_len as u64);

    let path_1_suffix = format!("snap.{:016X}.generating", snapshot.stable_lsn);

    let mut path_1 = config.get_path();
    path_1.push(path_1_suffix);

    let path_2_suffix = format!("snap.{:016X}", snapshot.stable_lsn);

    let mut path_2 = config.get_path();
    path_2.push(path_2_suffix);

    let parent = path_1.parent().unwrap();
    std::fs::create_dir_all(parent)?;
    let mut f =
        std::fs::OpenOptions::new().write(true).create(true).open(&path_1)?;

    // write the snapshot bytes, followed by a crc64 checksum at the end
    io_fail!(config, "snap write");
    f.write_all(&*bytes)?;
    io_fail!(config, "snap write len");
    f.write_all(&len_bytes)?;
    io_fail!(config, "snap write crc");
    f.write_all(&crc32)?;
    io_fail!(config, "snap write post");
    f.sync_all()?;

    trace!("wrote snapshot to {}", path_1.to_string_lossy());

    io_fail!(config, "snap write mv");
    std::fs::rename(&path_1, &path_2)?;
    io_fail!(config, "snap write mv post");

    trace!("renamed snapshot to {}", path_2.to_string_lossy());

    // clean up any old snapshots
    let candidates = config.get_snapshot_files()?;
    for path in candidates {
        let path_str = path.file_name().unwrap().to_str().unwrap();
        if !path_2.to_string_lossy().ends_with(&*path_str) {
            debug!("removing old snapshot file {:?}", path);

            io_fail!(config, "snap write rm old");

            if let Err(e) = std::fs::remove_file(&path) {
                // TODO should this just be a try return?
                warn!(
                    "failed to remove old snapshot file, maybe snapshot race? {}",
                    e
                );
            }
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::fs::{File, OpenOptions};
use std::io::{Seek, SeekFrom};
use std::path::PathBuf;
use std::result;
use std::sync::{Arc, Mutex};

use arch::DeviceType;
use device_manager::mmio::MMIO_CFG_SPACE_OFF;
use devices::virtio::{Block, MmioTransport, Net, TYPE_BLOCK, TYPE_NET};
use logger::METRICS;
use resources::VmResources;
use rpc_interface::VmmActionError;
use vmm_config;
use vmm_config::drive::DriveError;
use vmm_config::machine_config::VmConfig;
use vmm_config::net::NetworkInterfaceUpdateConfig;
use Vmm;

/// Shorthand result type for external VMM commands.
pub type ActionResult = std::result::Result<(), VmmActionError>;

/// Enables runtime configuration of a Firecracker VMM.
pub struct VmmController {
    vm_resources: VmResources,
    vmm: Arc<Mutex<Vmm>>,
}

impl VmmController {
    /// Returns the VmConfig.
    pub fn vm_config(&self) -> &VmConfig {
        self.vm_resources.vm_config()
    }

    /// Write the metrics on user demand (flush). We use the word `flush` here to highlight the fact
    /// that the metrics will be written immediately.
    /// Defer to inner Vmm. We'll move to a variant where the Vmm simply exposes functionality like
    /// getting the dirty pages, and then we'll have the metrics flushing logic entirely on the outside.
    pub fn flush_metrics(&mut self) -> ActionResult {
        // FIXME: we're losing the bool saying whether metrics were actually written.
        METRICS
            .write()
            .map(|_| ())
            .map_err(super::Error::Metrics)
            .map_err(VmmActionError::InternalVmm)
    }

    /// Injects CTRL+ALT+DEL keystroke combo to the inner Vmm (if present).
    #[cfg(target_arch = "x86_64")]
    pub fn send_ctrl_alt_del(&mut self) -> ActionResult {
        self.vmm
            .lock()
            .unwrap()
            .send_ctrl_alt_del()
            .map_err(VmmActionError::InternalVmm)
    }

    /// Stops the inner Vmm and exits the process with the provided exit_code.
    pub fn stop(&mut self, exit_code: i32) {
        self.vmm.lock().unwrap().stop(exit_code)
    }

    /// Creates a new `VmmController`.
    pub fn new(vm_resources: VmResources, vmm: Arc<Mutex<Vmm>>) -> Self {
        VmmController { vm_resources, vmm }
    }

    /// Triggers a rescan of the host file backing the emulated block device with id `drive_id`.
    pub fn rescan_block_device(&mut self, drive_id: &str) -> ActionResult {
        for drive_config in self.vm_resources.block.config_list.iter() {
            if drive_config.drive_id != *drive_id {
                continue;
            }

            // Use seek() instead of stat() (std::fs::Metadata) to support block devices.
            let new_size = File::open(&drive_config.path_on_host)
                .and_then(|mut f| f.seek(SeekFrom::End(0)))
                .map_err(|_| DriveError::BlockDeviceUpdateFailed)
                .map_err(VmmActionError::DriveConfig)?;

            return match self
                .vmm
                .lock()
                .unwrap()
                .get_bus_device(DeviceType::Virtio(TYPE_BLOCK), drive_id)
            {
                Some(device) => {
                    let data = devices::virtio::build_config_space(new_size);
                    let mut busdev = device
                        .lock()
                        .map_err(|_| DriveError::BlockDeviceUpdateFailed)
                        .map_err(VmmActionError::DriveConfig)?;

                    busdev.write(MMIO_CFG_SPACE_OFF, &data[..]);
                    let _ = busdev.interrupt(devices::virtio::VIRTIO_MMIO_INT_CONFIG);

                    Ok(())
                }
                None => Err(VmmActionError::DriveConfig(
                    DriveError::BlockDeviceUpdateFailed,
                )),
            };
        }

        Err(VmmActionError::DriveConfig(
            DriveError::InvalidBlockDeviceID,
        ))
    }

    fn update_drive_disk_image(
        &mut self,
        drive_id: &str,
        disk_image: File,
    ) -> result::Result<(), DriveError> {
        if let Some(busdev) = self
            .vmm
            .lock()
            .unwrap()
            .get_bus_device(DeviceType::Virtio(TYPE_BLOCK), drive_id)
        {
            let virtio_device = busdev
                .lock()
                .expect("Poisoned device lock")
                .as_any()
                .downcast_ref::<MmioTransport>()
                // Only MmioTransport implements BusDevice at this point.
                .expect("Unexpected BusDevice type")
                .device();

            // This call wraps the temporary `virtio_device` inside a `MutexGuard`.
            let mut lock = virtio_device.lock().expect("Poisoned device lock");

            // Downcast the inner virtio_device to a Block.
            let block_device: &mut Block = lock
                .as_mut_any()
                .downcast_mut::<Block>()
                .expect("Unexpected Block type");

            block_device
                .update_disk_image(disk_image)
                .map_err(|_| DriveError::BlockDeviceUpdateFailed)
        } else {
            Err(DriveError::InvalidBlockDeviceID)
        }
    }

    /// Updates the path of the host file backing the emulated block device with id `drive_id`.
    pub fn update_block_device_path(
        &mut self,
        drive_id: String,
        path_on_host: String,
    ) -> ActionResult {
        // Get the block device configuration specified by drive_id.
        let block_device_index = self
            .vm_resources
            .block
            .get_index_of_drive_id(&drive_id)
            .ok_or(VmmActionError::DriveConfig(
                DriveError::InvalidBlockDeviceID,
            ))?;

        let file_path = PathBuf::from(path_on_host);
        // Try to open the file specified by path_on_host using the permissions of the block_device.
        let disk_file = OpenOptions::new()
            .read(true)
            .write(!self.vm_resources.block.config_list[block_device_index].is_read_only())
            .open(&file_path)
            .map_err(DriveError::CannotOpenBlockDevice)
            .map_err(VmmActionError::DriveConfig)?;

        // Update the path of the block device with the specified path_on_host.
        self.vm_resources.block.config_list[block_device_index].path_on_host = file_path;

        // When the microvm is running, we also need to update the disk image and send a
        // rescan command to the drive.
        self.update_drive_disk_image(&drive_id, disk_file)
            .map_err(VmmActionError::DriveConfig)?;
        self.rescan_block_device(&drive_id)?;
        Ok(())
    }

    /// Updates configuration for an emulated net device as described in `new_cfg`.
    pub fn update_net_rate_limiters(
        &mut self,
        new_cfg: NetworkInterfaceUpdateConfig,
    ) -> ActionResult {
        if let Some(busdev) = self
            .vmm
            .lock()
            .unwrap()
            .get_bus_device(DeviceType::Virtio(TYPE_NET), &new_cfg.iface_id)
        {
            let virtio_device = busdev
                .lock()
                .expect("Poisoned device lock")
                .as_any()
                .downcast_ref::<MmioTransport>()
                // Only MmioTransport implements BusDevice at this point.
                .expect("Unexpected BusDevice type")
                .device();

            macro_rules! get_handler_arg {
                ($rate_limiter: ident, $metric: ident) => {{
                    new_cfg
                        .$rate_limiter
                        .map(|rl| rl.$metric.map(vmm_config::TokenBucketConfig::into))
                        .unwrap_or(None)
                }};
            }

            virtio_device
                .lock()
                .expect("Poisoned device lock")
                .as_mut_any()
                .downcast_mut::<Net>()
                .unwrap()
                .patch_rate_limiters(
                    get_handler_arg!(rx_rate_limiter, bandwidth),
                    get_handler_arg!(rx_rate_limiter, ops),
                    get_handler_arg!(tx_rate_limiter, bandwidth),
                    get_handler_arg!(tx_rate_limiter, ops),
                );
        } else {
            // TODO: Update this error after all devices have been ported.
        }

        Ok(())
    }
}
<commit_msg>vmm_controller: cleanup<commit_after>// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::fs::{File, OpenOptions};
use std::io::{Seek, SeekFrom};
use std::path::PathBuf;
use std::result;
use std::sync::{Arc, Mutex};

use arch::DeviceType;
use device_manager::mmio::MMIO_CFG_SPACE_OFF;
use devices::virtio::{Block, MmioTransport, Net, TYPE_BLOCK, TYPE_NET};
use logger::METRICS;
use resources::VmResources;
use rpc_interface::VmmActionError;
use vmm_config;
use vmm_config::drive::DriveError;
use vmm_config::machine_config::VmConfig;
use vmm_config::net::{NetworkInterfaceError, NetworkInterfaceUpdateConfig};
use Vmm;

/// Shorthand result type for external VMM commands.
pub type ActionResult = std::result::Result<(), VmmActionError>;

/// Enables runtime configuration of a Firecracker VMM.
pub struct VmmController {
    vm_resources: VmResources,
    vmm: Arc<Mutex<Vmm>>,
}

impl VmmController {
    /// Returns the VmConfig.
    pub fn vm_config(&self) -> &VmConfig {
        self.vm_resources.vm_config()
    }

    /// Write the metrics on user demand (flush). We use the word `flush` here to highlight the fact
    /// that the metrics will be written immediately.
    /// Defer to inner Vmm. We'll move to a variant where the Vmm simply exposes functionality like
    /// getting the dirty pages, and then we'll have the metrics flushing logic entirely on the outside.
    pub fn flush_metrics(&mut self) -> ActionResult {
        // FIXME: we're losing the bool saying whether metrics were actually written.
        METRICS
            .write()
            .map(|_| ())
            .map_err(super::Error::Metrics)
            .map_err(VmmActionError::InternalVmm)
    }

    /// Injects CTRL+ALT+DEL keystroke combo to the inner Vmm (if present).
    #[cfg(target_arch = "x86_64")]
    pub fn send_ctrl_alt_del(&mut self) -> ActionResult {
        self.vmm
            .lock()
            .unwrap()
            .send_ctrl_alt_del()
            .map_err(VmmActionError::InternalVmm)
    }

    /// Creates a new `VmmController`.
    pub fn new(vm_resources: VmResources, vmm: Arc<Mutex<Vmm>>) -> Self {
        VmmController { vm_resources, vmm }
    }

    /// Triggers a rescan of the host file backing the emulated block device with id `drive_id`.
    pub fn rescan_block_device(&mut self, drive_id: &str) -> ActionResult {
        for drive_config in self.vm_resources.block.config_list.iter() {
            if drive_config.drive_id != *drive_id {
                continue;
            }

            // Use seek() instead of stat() (std::fs::Metadata) to support block devices.
            let new_size = File::open(&drive_config.path_on_host)
                .and_then(|mut f| f.seek(SeekFrom::End(0)))
                .map_err(|_| DriveError::BlockDeviceUpdateFailed)
                .map_err(VmmActionError::DriveConfig)?;

            return match self
                .vmm
                .lock()
                .unwrap()
                .get_bus_device(DeviceType::Virtio(TYPE_BLOCK), drive_id)
            {
                Some(device) => {
                    let data = devices::virtio::build_config_space(new_size);
                    let mut busdev = device
                        .lock()
                        .map_err(|_| DriveError::BlockDeviceUpdateFailed)
                        .map_err(VmmActionError::DriveConfig)?;

                    busdev.write(MMIO_CFG_SPACE_OFF, &data[..]);
                    let _ = busdev.interrupt(devices::virtio::VIRTIO_MMIO_INT_CONFIG);

                    Ok(())
                }
                None => Err(VmmActionError::DriveConfig(
                    DriveError::BlockDeviceUpdateFailed,
                )),
            };
        }

        Err(VmmActionError::DriveConfig(
            DriveError::InvalidBlockDeviceID,
        ))
    }

    fn update_drive_disk_image(
        &mut self,
        drive_id: &str,
        disk_image: File,
    ) -> result::Result<(), DriveError> {
        if let Some(busdev) = self
            .vmm
            .lock()
            .unwrap()
            .get_bus_device(DeviceType::Virtio(TYPE_BLOCK), drive_id)
        {
            let virtio_device = busdev
                .lock()
                .expect("Poisoned device lock")
                .as_any()
                .downcast_ref::<MmioTransport>()
                // Only MmioTransport implements BusDevice at this point.
                .expect("Unexpected BusDevice type")
                .device();

            // This call wraps the temporary `virtio_device` inside a `MutexGuard`.
            let mut lock = virtio_device.lock().expect("Poisoned device lock");

            // Downcast the inner virtio_device to a Block.
            let block_device: &mut Block = lock
                .as_mut_any()
                .downcast_mut::<Block>()
                .expect("Unexpected Block type");

            block_device
                .update_disk_image(disk_image)
                .map_err(|_| DriveError::BlockDeviceUpdateFailed)
        } else {
            Err(DriveError::InvalidBlockDeviceID)
        }
    }

    /// Updates the path of the host file backing the emulated block device with id `drive_id`.
    pub fn update_block_device_path(
        &mut self,
        drive_id: String,
        path_on_host: String,
    ) -> ActionResult {
        // Get the block device configuration specified by drive_id.
        let block_device_index = self
            .vm_resources
            .block
            .get_index_of_drive_id(&drive_id)
            .ok_or(VmmActionError::DriveConfig(
                DriveError::InvalidBlockDeviceID,
            ))?;

        let file_path = PathBuf::from(path_on_host);
        // Try to open the file specified by path_on_host using the permissions of the block_device.
        let disk_file = OpenOptions::new()
            .read(true)
            .write(!self.vm_resources.block.config_list[block_device_index].is_read_only())
            .open(&file_path)
            .map_err(DriveError::CannotOpenBlockDevice)
            .map_err(VmmActionError::DriveConfig)?;

        // Update the path of the block device with the specified path_on_host.
        self.vm_resources.block.config_list[block_device_index].path_on_host = file_path;

        // When the microvm is running, we also need to update the disk image and send a
        // rescan command to the drive.
        self.update_drive_disk_image(&drive_id, disk_file)
            .map_err(VmmActionError::DriveConfig)?;
        self.rescan_block_device(&drive_id)?;
        Ok(())
    }

    /// Updates configuration for an emulated net device as described in `new_cfg`.
    pub fn update_net_rate_limiters(
        &mut self,
        new_cfg: NetworkInterfaceUpdateConfig,
    ) -> ActionResult {
        if let Some(busdev) = self
            .vmm
            .lock()
            .unwrap()
            .get_bus_device(DeviceType::Virtio(TYPE_NET), &new_cfg.iface_id)
        {
            let virtio_device = busdev
                .lock()
                .expect("Poisoned device lock")
                .as_any()
                .downcast_ref::<MmioTransport>()
                // Only MmioTransport implements BusDevice at this point.
                .expect("Unexpected BusDevice type")
                .device();

            macro_rules! get_handler_arg {
                ($rate_limiter: ident, $metric: ident) => {{
                    new_cfg
                        .$rate_limiter
                        .map(|rl| rl.$metric.map(vmm_config::TokenBucketConfig::into))
                        .unwrap_or(None)
                }};
            }

            virtio_device
                .lock()
                .expect("Poisoned device lock")
                .as_mut_any()
                .downcast_mut::<Net>()
                .unwrap()
                .patch_rate_limiters(
                    get_handler_arg!(rx_rate_limiter, bandwidth),
                    get_handler_arg!(rx_rate_limiter, ops),
                    get_handler_arg!(tx_rate_limiter, bandwidth),
                    get_handler_arg!(tx_rate_limiter, ops),
                );
        } else {
            return Err(VmmActionError::NetworkConfig(
                NetworkInterfaceError::DeviceIdNotFound,
            ));
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Linux DMABUF protocol
//!
//! This module provides helper to handle the linux-dmabuf protocol, which allows clients to submit their
//! contents as dmabuf file descriptors. These handlers automate the aggregation of the metadata associated
//! with a dma buffer, and do some basic checking of the sanity of what the client sends.
//!
//! ## How to use
//!
//! To setup the dmabuf global, you will need to provide 2 things:
//!
//! - a list of the dmabuf formats you wish to support
//! - a closure to test if a dmabuf buffer can be imported by your renderer
//!
//! The list of supported format is just a `Vec<Format>`, where you will enter all the (code, modifier)
//! couples you support. You can typically receive a list of supported formats for one renderer by calling
//! [`crate::backend::renderer::ImportDma::dmabuf_formats`].
//!
//! ```
//! # extern crate wayland_server;
//! # extern crate smithay;
//! use smithay::{
//!     backend::allocator::dmabuf::Dmabuf,
//!     reexports::{wayland_server::protocol::wl_buffer::WlBuffer},
//!     wayland::dmabuf::init_dmabuf_global,
//! };
//!
//! # let mut display = wayland_server::Display::new();
//! // define your supported formats
//! let formats = vec![
//!     /* ... */
//! ];
//! let dmabuf_global = init_dmabuf_global(
//!     &mut display,
//!     formats,
//!     |buffer, dispatch_data| {
//!         /* validate the dmabuf and import it into your renderer state */
//!         true
//!     },
//!     None // we don't provide a logger in this example
//! );
//! ```

use std::{
    cell::RefCell,
    convert::TryFrom,
    os::unix::io::{IntoRawFd, RawFd},
    rc::Rc,
};

pub use wayland_protocols::unstable::linux_dmabuf::v1::server::zwp_linux_buffer_params_v1::Flags;
use wayland_protocols::unstable::linux_dmabuf::v1::server::{
    zwp_linux_buffer_params_v1::{
        Error as ParamError, Request as ParamsRequest, ZwpLinuxBufferParamsV1 as BufferParams,
    },
    zwp_linux_dmabuf_v1,
};
use wayland_server::{protocol::wl_buffer, DispatchData, Display, Filter, Global, Main};

use crate::backend::allocator::{
    dmabuf::{Dmabuf, DmabufFlags, Plane},
    Format, Fourcc, Modifier,
};

/// Handler trait for dmabuf validation
///
/// You need to provide an implementation of this trait

/// Initialize a dmabuf global.
///
/// You need to provide a vector of the supported formats, as well as a closure,
/// that will validate the parameters provided by the client and tests the import as a dmabuf.
pub fn init_dmabuf_global<F, L>(
    display: &mut Display,
    formats: Vec<Format>,
    handler: F,
    logger: L,
) -> Global<zwp_linux_dmabuf_v1::ZwpLinuxDmabufV1>
where
    L: Into<Option<::slog::Logger>>,
    F: for<'a> FnMut(&Dmabuf, DispatchData<'a>) -> bool + 'static,
{
    let log = crate::slog_or_fallback(logger).new(o!("smithay_module" => "dmabuf_handler"));

    let formats = Rc::<[Format]>::from(formats);
    let handler = Rc::new(RefCell::new(handler));

    trace!(
        log,
        "Initializing DMABUF handler with {} supported formats",
        formats.len()
    );

    display.create_global(
        3,
        Filter::new(
            move |(dmabuf, version): (Main<zwp_linux_dmabuf_v1::ZwpLinuxDmabufV1>, u32), _, _| {
                let dma_formats = formats.clone();
                let dma_handler = handler.clone();
                let dma_log = log.clone();
                dmabuf.quick_assign(move |_, req, _| {
                    if let zwp_linux_dmabuf_v1::Request::CreateParams { params_id } = req {
                        let mut handler = ParamsHandler {
                            pending_planes: Vec::new(),
                            max_planes: 4,
                            used: false,
                            formats: dma_formats.clone(),
                            handler: dma_handler.clone(),
                            log: dma_log.clone(),
                        };
                        params_id.quick_assign(move |params, req, ddata| match req {
                            ParamsRequest::Add {
                                fd,
                                plane_idx,
                                offset,
                                stride,
                                modifier_hi,
                                modifier_lo,
                            } => handler.add(
                                &*params,
                                fd,
                                plane_idx,
                                offset,
                                stride,
                                ((modifier_hi as u64) << 32) + (modifier_lo as u64),
                            ),
                            ParamsRequest::Create {
                                width,
                                height,
                                format,
                                flags,
                            } => handler.create(&*params, width, height, format, flags, ddata),
                            ParamsRequest::CreateImmed {
                                buffer_id,
                                width,
                                height,
                                format,
                                flags,
                            } => {
                                handler.create_immed(&*params, buffer_id, width, height, format, flags, ddata)
                            }
                            _ => {}
                        });
                    }
                });

                // send the supported formats
                for f in &*formats {
                    dmabuf.format(f.code as u32);
                    if version >= 3 {
                        dmabuf.modifier(
                            f.code as u32,
                            (Into::<u64>::into(f.modifier) >> 32) as u32,
                            Into::<u64>::into(f.modifier) as u32,
                        );
                    }
                }
            },
        ),
    )
}

struct ParamsHandler<H: for<'a> FnMut(&Dmabuf, DispatchData<'a>) -> bool + 'static> {
    pending_planes: Vec<Plane>,
    max_planes: u32,
    used: bool,
    formats: Rc<[Format]>,
    handler: Rc<RefCell<H>>,
    log: ::slog::Logger,
}

impl<H> ParamsHandler<H>
where
    H: for<'a> FnMut(&Dmabuf, DispatchData<'a>) -> bool + 'static,
{
    fn add(
        &mut self,
        params: &BufferParams,
        fd: RawFd,
        plane_idx: u32,
        offset: u32,
        stride: u32,
        modifier: u64,
    ) {
        // protocol checks:
        // Cannot reuse a params:
        if self.used {
            params.as_ref().post_error(
                ParamError::AlreadyUsed as u32,
                "This buffer_params has already been used to create a buffer.".into(),
            );
            return;
        }
        // plane_idx is not too large
        if plane_idx >= self.max_planes {
            // plane_idx starts at 0
            params.as_ref().post_error(
                ParamError::PlaneIdx as u32,
                format!("Plane index {} is out of bounds.", plane_idx),
            );
            return;
        }
        // plane_idx has already been set
        if self.pending_planes.iter().any(|d| d.plane_idx == plane_idx) {
            params.as_ref().post_error(
                ParamError::PlaneSet as u32,
                format!("Plane index {} is already set.", plane_idx),
            );
            return;
        }
        // all checks passed, store the plane
        self.pending_planes.push(Plane {
            fd: Some(fd),
            plane_idx,
            offset,
            stride,
            modifier: Modifier::from(modifier),
        });
    }

    fn create<'a>(
        &mut self,
        params: &BufferParams,
        width: i32,
        height: i32,
        format: u32,
        flags: u32,
        ddata: DispatchData<'a>,
    ) {
        // Cannot reuse a params:
        if self.used {
            params.as_ref().post_error(
                ParamError::AlreadyUsed as u32,
                "This buffer_params has already been used to create a buffer.".into(),
            );
            return;
        }
        self.used = true;

        let format = match Fourcc::try_from(format) {
            Ok(format) => format,
            Err(_) => {
                params.as_ref().post_error(
                    ParamError::InvalidFormat as u32,
                    format!("Format {:x} is not supported", format),
                );
                return;
            }
        };

        if !buffer_basic_checks(
            &self.formats,
            &self.pending_planes,
            &params,
            format,
            width,
            height,
        ) {
            trace!(self.log, "Killing client providing bogus dmabuf buffer params.");
            return;
        }

        let mut buf = Dmabuf::builder(
            width as u32,
            height as u32,
            format,
            DmabufFlags::from_bits_truncate(flags),
        );
        let planes = ::std::mem::replace(&mut self.pending_planes, Vec::new());
        for (i, plane) in planes.into_iter().enumerate() {
            let offset = plane.offset;
            let stride = plane.stride;
            let modi = plane.modifier;
            buf.add_plane(plane.into_raw_fd(), i as u32, offset, stride, modi);
        }
        let dmabuf = match buf.build() {
            Some(buf) => buf,
            None => {
                params.as_ref().post_error(
                    ParamError::Incomplete as u32,
                    "Provided buffer is incomplete, it has zero planes".to_string(),
                );
                return;
            }
        };

        let mut handler = self.handler.borrow_mut();
        if handler(&dmabuf, ddata) {
            if let Some(buffer) = params
                .as_ref()
                .client()
                .and_then(|c| c.create_resource::<wl_buffer::WlBuffer>(1))
            {
                buffer.as_ref().user_data().set_threadsafe(|| dmabuf);
                buffer.quick_assign(|_, _, _| {});

                trace!(self.log, "Created a new validated dma wl_buffer.");
                params.created(&buffer);
            } else {
                trace!(self.log, "Failed to create a wl_buffer");
                params.failed();
            }
        } else {
            trace!(self.log, "Refusing creation of an invalid dma wl_buffer.");
            params.failed();
        }
    }

    #[allow(clippy::clippy::too_many_arguments)]
    fn create_immed<'a>(
        &mut self,
        params: &BufferParams,
        buffer: Main<wl_buffer::WlBuffer>,
        width: i32,
        height: i32,
        format: u32,
        flags: u32,
        ddata: DispatchData<'a>,
    ) {
        // Cannot reuse a params:
        if self.used {
            params.as_ref().post_error(
                ParamError::AlreadyUsed as u32,
                "This buffer_params has already been used to create a buffer.".into(),
            );
            return;
        }
        self.used = true;

        let format = match Fourcc::try_from(format) {
            Ok(format) => format,
            Err(_) => {
                params.as_ref().post_error(
                    ParamError::InvalidFormat as u32,
                    format!("Format {:x} is not supported", format),
                );
                return;
            }
        };

        if !buffer_basic_checks(
            &self.formats,
            &self.pending_planes,
            &params,
            format,
            width,
            height,
        ) {
            trace!(self.log, "Killing client providing bogus dmabuf buffer params.");
            return;
        }

        let mut buf = Dmabuf::builder(
            width as u32,
            height as u32,
            format,
            DmabufFlags::from_bits_truncate(flags),
        );
        let planes = ::std::mem::replace(&mut self.pending_planes, Vec::new());
        for (i, plane) in planes.into_iter().enumerate() {
            let offset = plane.offset;
            let stride = plane.stride;
            let modi = plane.modifier;
            buf.add_plane(plane.into_raw_fd(), i as u32, offset, stride, modi);
        }
        let dmabuf = match buf.build() {
            Some(buf) => buf,
            None => {
                params.as_ref().post_error(
                    ParamError::Incomplete as u32,
                    "Provided buffer is incomplete, it has zero planes".to_string(),
                );
                return;
            }
        };

        let mut handler = self.handler.borrow_mut();
        if handler(&dmabuf, ddata) {
            buffer.as_ref().user_data().set_threadsafe(|| dmabuf);
            buffer.quick_assign(|_, _, _| {});
            trace!(self.log, "Created a new validated dma wl_buffer.");
        } else {
            trace!(
                self.log,
                "Refusing creation of an invalid immediate dma wl_buffer, killing client."
            );
            params.as_ref().post_error(
                ParamError::InvalidWlBuffer as u32,
                "create_immed resulted in an invalid buffer.".into(),
            );
        }
    }
}

fn buffer_basic_checks(
    formats: &[Format],
    pending_planes: &[Plane],
    params: &BufferParams,
    format: Fourcc,
    width: i32,
    height: i32,
) -> bool {
    // protocol_checks:
    // This must be a known format
    let _format = match formats.iter().find(|f| f.code == format) {
        Some(f) => f,
        None => {
            params.as_ref().post_error(
                ParamError::InvalidFormat as u32,
                format!("Format {:?}/{:x} is not supported.", format, format as u32),
            );
            return false;
        }
    };
    // Width and height must be positivie
    if width < 1 || height < 1 {
        params.as_ref().post_error(
            ParamError::InvalidDimensions as u32,
            format!("Dimensions ({},{}) are not valid.", width, height),
        );
        return false;
    }
    // check the size of each plane buffer
    for plane in pending_planes {
        // check size for overflow
        let end = match plane
            .stride
            .checked_mul(height as u32)
            .and_then(|o| o.checked_add(plane.offset))
        {
            None => {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!("Size overflow for plane {}.", plane.plane_idx),
                );
                return false;
            }
            Some(e) => e,
        };
        if let Ok(size) = ::nix::unistd::lseek(plane.fd.unwrap(), 0, ::nix::unistd::Whence::SeekEnd) {
            // reset the seek point
            let _ = ::nix::unistd::lseek(plane.fd.unwrap(), 0, ::nix::unistd::Whence::SeekSet);
            if plane.offset as i64 > size {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!("Invalid offset {} for plane {}.", plane.offset, plane.plane_idx),
                );
                return false;
            }
            if (plane.offset + plane.stride) as i64 > size {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!("Invalid stride {} for plane {}.", plane.stride, plane.plane_idx),
                );
                return false;
            }
            // Planes > 0 can be subsampled, in which case 'size' will be smaller
            // than expected.
            if plane.plane_idx == 0 && end as i64 > size {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!(
                        "Invalid stride ({}) or height ({}) for plane {}.",
                        plane.stride, height, plane.plane_idx
                    ),
                );
                return false;
            }
        }
    }
    true
}
<commit_msg>fix clippy unknown lint<commit_after>//! Linux DMABUF protocol
//!
//! This module provides helper to handle the linux-dmabuf protocol, which allows clients to submit their
//! contents as dmabuf file descriptors. These handlers automate the aggregation of the metadata associated
//! with a dma buffer, and do some basic checking of the sanity of what the client sends.
//!
//! ## How to use
//!
//! To setup the dmabuf global, you will need to provide 2 things:
//!
//! - a list of the dmabuf formats you wish to support
//! - a closure to test if a dmabuf buffer can be imported by your renderer
//!
//! The list of supported format is just a `Vec<Format>`, where you will enter all the (code, modifier)
//! couples you support. You can typically receive a list of supported formats for one renderer by calling
//! [`crate::backend::renderer::ImportDma::dmabuf_formats`].
//!
//! ```
//! # extern crate wayland_server;
//! # extern crate smithay;
//! use smithay::{
//!     backend::allocator::dmabuf::Dmabuf,
//!     reexports::{wayland_server::protocol::wl_buffer::WlBuffer},
//!     wayland::dmabuf::init_dmabuf_global,
//! };
//!
//! # let mut display = wayland_server::Display::new();
//! // define your supported formats
//! let formats = vec![
//!     /* ... */
//! ];
//! let dmabuf_global = init_dmabuf_global(
//!     &mut display,
//!     formats,
//!     |buffer, dispatch_data| {
//!         /* validate the dmabuf and import it into your renderer state */
//!         true
//!     },
//!     None // we don't provide a logger in this example
//! );
//! ```

use std::{
    cell::RefCell,
    convert::TryFrom,
    os::unix::io::{IntoRawFd, RawFd},
    rc::Rc,
};

pub use wayland_protocols::unstable::linux_dmabuf::v1::server::zwp_linux_buffer_params_v1::Flags;
use wayland_protocols::unstable::linux_dmabuf::v1::server::{
    zwp_linux_buffer_params_v1::{
        Error as ParamError, Request as ParamsRequest, ZwpLinuxBufferParamsV1 as BufferParams,
    },
    zwp_linux_dmabuf_v1,
};
use wayland_server::{protocol::wl_buffer, DispatchData, Display, Filter, Global, Main};

use crate::backend::allocator::{
    dmabuf::{Dmabuf, DmabufFlags, Plane},
    Format, Fourcc, Modifier,
};

/// Handler trait for dmabuf validation
///
/// You need to provide an implementation of this trait

/// Initialize a dmabuf global.
///
/// You need to provide a vector of the supported formats, as well as a closure,
/// that will validate the parameters provided by the client and tests the import as a dmabuf.
pub fn init_dmabuf_global<F, L>(
    display: &mut Display,
    formats: Vec<Format>,
    handler: F,
    logger: L,
) -> Global<zwp_linux_dmabuf_v1::ZwpLinuxDmabufV1>
where
    L: Into<Option<::slog::Logger>>,
    F: for<'a> FnMut(&Dmabuf, DispatchData<'a>) -> bool + 'static,
{
    let log = crate::slog_or_fallback(logger).new(o!("smithay_module" => "dmabuf_handler"));

    let formats = Rc::<[Format]>::from(formats);
    let handler = Rc::new(RefCell::new(handler));

    trace!(
        log,
        "Initializing DMABUF handler with {} supported formats",
        formats.len()
    );

    display.create_global(
        3,
        Filter::new(
            move |(dmabuf, version): (Main<zwp_linux_dmabuf_v1::ZwpLinuxDmabufV1>, u32), _, _| {
                let dma_formats = formats.clone();
                let dma_handler = handler.clone();
                let dma_log = log.clone();
                dmabuf.quick_assign(move |_, req, _| {
                    if let zwp_linux_dmabuf_v1::Request::CreateParams { params_id } = req {
                        let mut handler = ParamsHandler {
                            pending_planes: Vec::new(),
                            max_planes: 4,
                            used: false,
                            formats: dma_formats.clone(),
                            handler: dma_handler.clone(),
                            log: dma_log.clone(),
                        };
                        params_id.quick_assign(move |params, req, ddata| match req {
                            ParamsRequest::Add {
                                fd,
                                plane_idx,
                                offset,
                                stride,
                                modifier_hi,
                                modifier_lo,
                            } => handler.add(
                                &*params,
                                fd,
                                plane_idx,
                                offset,
                                stride,
                                ((modifier_hi as u64) << 32) + (modifier_lo as u64),
                            ),
                            ParamsRequest::Create {
                                width,
                                height,
                                format,
                                flags,
                            } => handler.create(&*params, width, height, format, flags, ddata),
                            ParamsRequest::CreateImmed {
                                buffer_id,
                                width,
                                height,
                                format,
                                flags,
                            } => {
                                handler.create_immed(&*params, buffer_id, width, height, format, flags, ddata)
                            }
                            _ => {}
                        });
                    }
                });

                // send the supported formats
                for f in &*formats {
                    dmabuf.format(f.code as u32);
                    if version >= 3 {
                        dmabuf.modifier(
                            f.code as u32,
                            (Into::<u64>::into(f.modifier) >> 32) as u32,
                            Into::<u64>::into(f.modifier) as u32,
                        );
                    }
                }
            },
        ),
    )
}

struct ParamsHandler<H: for<'a> FnMut(&Dmabuf, DispatchData<'a>) -> bool + 'static> {
    pending_planes: Vec<Plane>,
    max_planes: u32,
    used: bool,
    formats: Rc<[Format]>,
    handler: Rc<RefCell<H>>,
    log: ::slog::Logger,
}

impl<H> ParamsHandler<H>
where
    H: for<'a> FnMut(&Dmabuf, DispatchData<'a>) -> bool + 'static,
{
    fn add(
        &mut self,
        params: &BufferParams,
        fd: RawFd,
        plane_idx: u32,
        offset: u32,
        stride: u32,
        modifier: u64,
    ) {
        // protocol checks:
        // Cannot reuse a params:
        if self.used {
            params.as_ref().post_error(
                ParamError::AlreadyUsed as u32,
                "This buffer_params has already been used to create a buffer.".into(),
            );
            return;
        }
        // plane_idx is not too large
        if plane_idx >= self.max_planes {
            // plane_idx starts at 0
            params.as_ref().post_error(
                ParamError::PlaneIdx as u32,
                format!("Plane index {} is out of bounds.", plane_idx),
            );
            return;
        }
        // plane_idx has already been set
        if self.pending_planes.iter().any(|d| d.plane_idx == plane_idx) {
            params.as_ref().post_error(
                ParamError::PlaneSet as u32,
                format!("Plane index {} is already set.", plane_idx),
            );
            return;
        }
        // all checks passed, store the plane
        self.pending_planes.push(Plane {
            fd: Some(fd),
            plane_idx,
            offset,
            stride,
            modifier: Modifier::from(modifier),
        });
    }

    fn create<'a>(
        &mut self,
        params: &BufferParams,
        width: i32,
        height: i32,
        format: u32,
        flags: u32,
        ddata: DispatchData<'a>,
    ) {
        // Cannot reuse a params:
        if self.used {
            params.as_ref().post_error(
                ParamError::AlreadyUsed as u32,
                "This buffer_params has already been used to create a buffer.".into(),
            );
            return;
        }
        self.used = true;

        let format = match Fourcc::try_from(format) {
            Ok(format) => format,
            Err(_) => {
                params.as_ref().post_error(
                    ParamError::InvalidFormat as u32,
                    format!("Format {:x} is not supported", format),
                );
                return;
            }
        };

        if !buffer_basic_checks(
            &self.formats,
            &self.pending_planes,
            &params,
            format,
            width,
            height,
        ) {
            trace!(self.log, "Killing client providing bogus dmabuf buffer params.");
            return;
        }

        let mut buf = Dmabuf::builder(
            width as u32,
            height as u32,
            format,
            DmabufFlags::from_bits_truncate(flags),
        );
        let planes = ::std::mem::replace(&mut self.pending_planes, Vec::new());
        for (i, plane) in planes.into_iter().enumerate() {
            let offset = plane.offset;
            let stride = plane.stride;
            let modi = plane.modifier;
            buf.add_plane(plane.into_raw_fd(), i as u32, offset, stride, modi);
        }
        let dmabuf = match buf.build() {
            Some(buf) => buf,
            None => {
                params.as_ref().post_error(
                    ParamError::Incomplete as u32,
                    "Provided buffer is incomplete, it has zero planes".to_string(),
                );
                return;
            }
        };

        let mut handler = self.handler.borrow_mut();
        if handler(&dmabuf, ddata) {
            if let Some(buffer) = params
                .as_ref()
                .client()
                .and_then(|c| c.create_resource::<wl_buffer::WlBuffer>(1))
            {
                buffer.as_ref().user_data().set_threadsafe(|| dmabuf);
                buffer.quick_assign(|_, _, _| {});

                trace!(self.log, "Created a new validated dma wl_buffer.");
                params.created(&buffer);
            } else {
                trace!(self.log, "Failed to create a wl_buffer");
                params.failed();
            }
        } else {
            trace!(self.log, "Refusing creation of an invalid dma wl_buffer.");
            params.failed();
        }
    }

    #[allow(clippy::too_many_arguments)]
    fn create_immed<'a>(
        &mut self,
        params: &BufferParams,
        buffer: Main<wl_buffer::WlBuffer>,
        width: i32,
        height: i32,
        format: u32,
        flags: u32,
        ddata: DispatchData<'a>,
    ) {
        // Cannot reuse a params:
        if self.used {
            params.as_ref().post_error(
                ParamError::AlreadyUsed as u32,
                "This buffer_params has already been used to create a buffer.".into(),
            );
            return;
        }
        self.used = true;

        let format = match Fourcc::try_from(format) {
            Ok(format) => format,
            Err(_) => {
                params.as_ref().post_error(
                    ParamError::InvalidFormat as u32,
                    format!("Format {:x} is not supported", format),
                );
                return;
            }
        };

        if !buffer_basic_checks(
            &self.formats,
            &self.pending_planes,
            &params,
            format,
            width,
            height,
        ) {
            trace!(self.log, "Killing client providing bogus dmabuf buffer params.");
            return;
        }

        let mut buf = Dmabuf::builder(
            width as u32,
            height as u32,
            format,
            DmabufFlags::from_bits_truncate(flags),
        );
        let planes = ::std::mem::replace(&mut self.pending_planes, Vec::new());
        for (i, plane) in planes.into_iter().enumerate() {
            let offset = plane.offset;
            let stride = plane.stride;
            let modi = plane.modifier;
            buf.add_plane(plane.into_raw_fd(), i as u32, offset, stride, modi);
        }
        let dmabuf = match buf.build() {
            Some(buf) => buf,
            None => {
                params.as_ref().post_error(
                    ParamError::Incomplete as u32,
                    "Provided buffer is incomplete, it has zero planes".to_string(),
                );
                return;
            }
        };

        let mut handler = self.handler.borrow_mut();
        if handler(&dmabuf, ddata) {
            buffer.as_ref().user_data().set_threadsafe(|| dmabuf);
            buffer.quick_assign(|_, _, _| {});
            trace!(self.log, "Created a new validated dma wl_buffer.");
        } else {
            trace!(
                self.log,
                "Refusing creation of an invalid immediate dma wl_buffer, killing client."
            );
            params.as_ref().post_error(
                ParamError::InvalidWlBuffer as u32,
                "create_immed resulted in an invalid buffer.".into(),
            );
        }
    }
}

fn buffer_basic_checks(
    formats: &[Format],
    pending_planes: &[Plane],
    params: &BufferParams,
    format: Fourcc,
    width: i32,
    height: i32,
) -> bool {
    // protocol_checks:
    // This must be a known format
    let _format = match formats.iter().find(|f| f.code == format) {
        Some(f) => f,
        None => {
            params.as_ref().post_error(
                ParamError::InvalidFormat as u32,
                format!("Format {:?}/{:x} is not supported.", format, format as u32),
            );
            return false;
        }
    };
    // Width and height must be positivie
    if width < 1 || height < 1 {
        params.as_ref().post_error(
            ParamError::InvalidDimensions as u32,
            format!("Dimensions ({},{}) are not valid.", width, height),
        );
        return false;
    }
    // check the size of each plane buffer
    for plane in pending_planes {
        // check size for overflow
        let end = match plane
            .stride
            .checked_mul(height as u32)
            .and_then(|o| o.checked_add(plane.offset))
        {
            None => {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!("Size overflow for plane {}.", plane.plane_idx),
                );
                return false;
            }
            Some(e) => e,
        };
        if let Ok(size) = ::nix::unistd::lseek(plane.fd.unwrap(), 0, ::nix::unistd::Whence::SeekEnd) {
            // reset the seek point
            let _ = ::nix::unistd::lseek(plane.fd.unwrap(), 0, ::nix::unistd::Whence::SeekSet);
            if plane.offset as i64 > size {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!("Invalid offset {} for plane {}.", plane.offset, plane.plane_idx),
                );
                return false;
            }
            if (plane.offset + plane.stride) as i64 > size {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!("Invalid stride {} for plane {}.", plane.stride, plane.plane_idx),
                );
                return false;
            }
            // Planes > 0 can be subsampled, in which case 'size' will be smaller
            // than expected.
            if plane.plane_idx == 0 && end as i64 > size {
                params.as_ref().post_error(
                    ParamError::OutOfBounds as u32,
                    format!(
                        "Invalid stride ({}) or height ({}) for plane {}.",
                        plane.stride, height, plane.plane_idx
                    ),
                );
                return false;
            }
        }
    }
    true
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
// Copyright 2016-2017 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The binomial distribution.

use Rng;
use distributions::Distribution;

/// The binomial distribution `Binomial(n, p)`.
///
/// This distribution has density function:
/// `f(k) = n!/(k! (n-k)!) p^k (1-p)^(n-k)` for `k >= 0`.
///
/// # Example
///
/// ```
/// use rand::distributions::{Binomial, Distribution};
///
/// let bin = Binomial::new(20, 0.3);
/// let v = bin.sample(&mut rand::thread_rng());
/// println!("{} is from a binomial distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Binomial {
    /// Number of trials.
    n: u64,
    /// Probability of success.
    p: f64,
}

impl Binomial {
    /// Construct a new `Binomial` with the given shape parameters `n` (number
    /// of trials) and `p` (probability of success).
    ///
    /// Panics if `p < 0` or `p > 1`.
    pub fn new(n: u64, p: f64) -> Binomial {
        assert!(p >= 0.0, "Binomial::new called with p < 0");
        assert!(p <= 1.0, "Binomial::new called with p > 1");
        Binomial { n, p }
    }
}

impl Distribution<u64> for Binomial {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
        // Handle these values directly.
        if self.p == 0.0 {
            return 0;
        } else if self.p == 1.0 {
            return self.n;
        }

        // The binomial distribution is symmetrical with respect to p -> 1-p,
        // k -> n-k switch p so that it is less than 0.5 - this allows for lower
        // expected values we will just invert the result at the end
        let p = if self.p <= 0.5 {
            self.p
        } else {
            1.0 - self.p
        };

        let result;
        let q = 1. - p;

        // For small n * min(p, 1 - p), the BINV algorithm based on the inverse
        // transformation of the binomial distribution is efficient. Otherwise,
        // the BTPE algorithm is used.
        //
        // Voratas Kachitvichyanukul and Bruce W. Schmeiser. 1988. Binomial
        // random variate generation. Commun. ACM 31, 2 (February 1988),
        // 216-222. http://dx.doi.org/10.1145/42372.42381

        // Threshold for prefering the BINV algorithm. The paper suggests 10,
        // Ranlib uses 30, and GSL uses 14.
        const BINV_THRESHOLD: f64 = 10.;

        if (self.n as f64) * p < BINV_THRESHOLD &&
           self.n <= (::std::i32::MAX as u64) {
            // Use the BINV algorithm.
            let s = p / q;
            let a = ((self.n + 1) as f64) * s;
            let mut r = q.powi(self.n as i32);
            let mut u: f64 = rng.gen();
            let mut x = 0;
            while u > r as f64 {
                u -= r;
                x += 1;
                r *= a / (x as f64) - s;
            }
            result = x;
        } else {
            // Use the BTPE algorithm.

            // Threshold for using the squeeze algorithm. This can be freely
            // chosen based on performance. Ranlib and GSL use 20.
            const SQUEEZE_THRESHOLD: i64 = 20;

            // Step 0: Calculate constants as functions of `n` and `p`.
            let n = self.n as f64;
            let f_m = n*p + p;
            let m = f_m as i64;
            // radius of triangle region, since height=1 also area of region
            let p1 = (2.195 * (n*p*q).sqrt() - 4.6 * q).floor() + 0.5;
            // tip of triangle
            let x_m = (m as f64) + 0.5;
            // left edge of triangle
            let x_l = x_m - p1;
            // right edge of triangle
            let x_r = x_m + p1;
            let c = 0.134 + 20.5 / (15.3 + (m as f64));
            // p1 + area of parallelogram region
            let p2 = p1 * (1. + 2. * c);

            fn lambda(a: f64) -> f64 {
                a * (1. + 0.5 * a)
            }

            let lambda_l = lambda((f_m - x_l) / (f_m - x_l * p));
            let lambda_r = lambda((x_r - f_m) / (x_r * q));
            // p1 + area of left tail
            let p3 = p2 + c / lambda_l;
            // p1 + area of right tail
            let p4 = p3 + c / lambda_r;

            // return value
            let mut y: i64;

            loop {
                // Step 1: Generate `u` for selecting the region. If region 1 is
                // selected, generate a triangularly distributed variate.
                let u = rng.gen_range(0., p4);
                let mut v = rng.gen_range(0., 1.);
                if !(u > p1) {
                    y = (x_m - p1 * v + u) as i64;
                    break;
                }

                if !(u > p2) {
                    // Step 2: Region 2, parallelograms. Check if region 2 is
                    // used. If so, generate `y`.
                    let x = x_l + (u - p1) / c;
                    v = v * c + 1.0 - (x - x_m).abs() / p1;
                    if v > 1. {
                        continue;
                    } else {
                        y = x as i64;
                    }
                } else if !(u > p3) {
                    // Step 3: Region 3, left exponential tail.
                    y = (x_l + v.ln() / lambda_l) as i64;
                    if y < 0 {
                        continue;
                    } else {
                        v *= (u - p2) * lambda_l;
                    }
                } else {
                    // Step 4: Region 4, right exponential tail.
                    y = (x_r - v.ln() / lambda_r) as i64;
                    if y > 0 && (y as u64) > self.n {
                        continue;
                    } else {
                        v *= (u - p3) * lambda_r;
                    }
                }

                // Step 5: Acceptance/rejection comparison.

                // Step 5.0: Test for appropriate method of evaluating f(y).
                let k = (y - m).abs();
                if !(k > SQUEEZE_THRESHOLD && (k as f64) < 0.5 * n*p*q - 1.) {
                    // Step 5.1: Evaluate f(y) via the recursive relationship. Start the
                    // search from the mode.
                    let s = p / q;
                    let a = s * (n + 1.);
                    let mut f = 1.0;
                    if m < y {
                        let mut i = m;
                        loop {
                            i += 1;
                            f *= a / (i as f64) - s;
                            if i == y {
                                break;
                            }
                        }
                    } else if m > y {
                        let mut i = y;
                        loop {
                            i += 1;
                            f /= a / (i as f64) - s;
                            if i == m {
                                break;
                            }
                        }
                    }
                    if v > f {
                        continue;
                    } else {
                        break;
                    }
                }

                // Step 5.2: Squeezing. Check the value of ln(v) againts upper and
                // lower bound of ln(f(y)).
                let k = k as f64;
                let rho = (k / (n*p*q)) * ((k * (k / 3. + 0.625) + 1./6.) / (n*p*q) + 0.5);
                let t = -0.5 * k*k / (n*p*q);
                let alpha = v.ln();
                if alpha < t - rho {
                    break;
                }
                if alpha > t + rho {
                    continue;
                }

                // Step 5.3: Final acceptance/rejection test.
                let x1 = (y + 1) as f64;
                let f1 = (m + 1) as f64;
                let z = ((n as i64) + 1 - m) as f64;
                let w = ((n as i64) - y + 1) as f64;

                fn stirling(a: f64) -> f64 {
                    let a2 = a * a;
                    (13860. - (462. - (132. - (99. - 140. / a2) / a2) / a2) / a2) / a / 166320.
                }

                if alpha > x_m * (f1 / x1).ln()
                    + (n - (m as f64) + 0.5) * (z / w).ln()
                    + ((y - m) as f64) * (w * p / (x1 * q)).ln()
                    + stirling(f1) + stirling(z) + stirling(x1) + stirling(w)
                {
                    continue;
                }

                break;
            }
            assert!(y >= 0);
            result = y as u64;
        }

        // Invert the result for p < 0.5.
        if p != self.p {
            self.n - result
        } else {
            result
        }
    }
}

#[cfg(test)]
mod test {
    use Rng;
    use distributions::Distribution;
    use super::Binomial;

    fn test_binomial_mean_and_variance<R: Rng>(n: u64, p: f64, rng: &mut R) {
        let binomial = Binomial::new(n, p);

        let expected_mean = n as f64 * p;
        let expected_variance = n as f64 * p * (1.0 - p);

        let mut results = [0.0; 1000];
        for i in results.iter_mut() { *i = binomial.sample(rng) as f64; }

        let mean = results.iter().sum::<f64>() / results.len() as f64;
        assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0,
                "mean: {}, expected_mean: {}", mean, expected_mean);

        let variance =
            results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>()
            / results.len() as f64;
        assert!((variance - expected_variance).abs() < expected_variance / 10.0,
                "variance: {}, expected_variance: {}", variance, expected_variance);
    }

    #[test]
    fn test_binomial() {
        let mut rng = ::test::rng(351);
        test_binomial_mean_and_variance(150, 0.1, &mut rng);
        test_binomial_mean_and_variance(70, 0.6, &mut rng);
        test_binomial_mean_and_variance(40, 0.5, &mut rng);
        test_binomial_mean_and_variance(20, 0.7, &mut rng);
        test_binomial_mean_and_variance(20, 0.5, &mut rng);
    }

    #[test]
    fn test_binomial_end_points() {
        let mut rng = ::test::rng(352);
        assert_eq!(rng.sample(Binomial::new(20, 0.0)), 0);
        assert_eq!(rng.sample(Binomial::new(20, 1.0)), 20);
    }

    #[test]
    #[should_panic]
    fn test_binomial_invalid_lambda_neg() {
        Binomial::new(20, -10.0);
    }
}
<commit_msg>Binomial: Eliminate some common subexpressions<commit_after>// Copyright 2018 Developers of the Rand project.
// Copyright 2016-2017 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The binomial distribution.

use Rng;
use distributions::Distribution;

/// The binomial distribution `Binomial(n, p)`.
///
/// This distribution has density function:
/// `f(k) = n!/(k! (n-k)!) p^k (1-p)^(n-k)` for `k >= 0`.
///
/// # Example
///
/// ```
/// use rand::distributions::{Binomial, Distribution};
///
/// let bin = Binomial::new(20, 0.3);
/// let v = bin.sample(&mut rand::thread_rng());
/// println!("{} is from a binomial distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Binomial {
    /// Number of trials.
    n: u64,
    /// Probability of success.
    p: f64,
}

impl Binomial {
    /// Construct a new `Binomial` with the given shape parameters `n` (number
    /// of trials) and `p` (probability of success).
    ///
    /// Panics if `p < 0` or `p > 1`.
    pub fn new(n: u64, p: f64) -> Binomial {
        assert!(p >= 0.0, "Binomial::new called with p < 0");
        assert!(p <= 1.0, "Binomial::new called with p > 1");
        Binomial { n, p }
    }
}

impl Distribution<u64> for Binomial {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
        // Handle these values directly.
        if self.p == 0.0 {
            return 0;
        } else if self.p == 1.0 {
            return self.n;
        }

        // The binomial distribution is symmetrical with respect to p -> 1-p,
        // k -> n-k switch p so that it is less than 0.5 - this allows for lower
        // expected values we will just invert the result at the end
        let p = if self.p <= 0.5 {
            self.p
        } else {
            1.0 - self.p
        };

        let result;
        let q = 1. - p;

        // For small n * min(p, 1 - p), the BINV algorithm based on the inverse
        // transformation of the binomial distribution is efficient. Otherwise,
        // the BTPE algorithm is used.
        //
        // Voratas Kachitvichyanukul and Bruce W. Schmeiser. 1988. Binomial
        // random variate generation. Commun. ACM 31, 2 (February 1988),
        // 216-222. http://dx.doi.org/10.1145/42372.42381

        // Threshold for prefering the BINV algorithm. The paper suggests 10,
        // Ranlib uses 30, and GSL uses 14.
        const BINV_THRESHOLD: f64 = 10.;

        if (self.n as f64) * p < BINV_THRESHOLD &&
           self.n <= (::std::i32::MAX as u64) {
            // Use the BINV algorithm.
            let s = p / q;
            let a = ((self.n + 1) as f64) * s;
            let mut r = q.powi(self.n as i32);
            let mut u: f64 = rng.gen();
            let mut x = 0;
            while u > r as f64 {
                u -= r;
                x += 1;
                r *= a / (x as f64) - s;
            }
            result = x;
        } else {
            // Use the BTPE algorithm.

            // Threshold for using the squeeze algorithm. This can be freely
            // chosen based on performance. Ranlib and GSL use 20.
            const SQUEEZE_THRESHOLD: i64 = 20;

            // Step 0: Calculate constants as functions of `n` and `p`.
            let n = self.n as f64;
            let np = n * p;
            let npq = np * q;
            let f_m = np + p;
            let m = f_m as i64;
            // radius of triangle region, since height=1 also area of region
            let p1 = (2.195 * npq.sqrt() - 4.6 * q).floor() + 0.5;
            // tip of triangle
            let x_m = (m as f64) + 0.5;
            // left edge of triangle
            let x_l = x_m - p1;
            // right edge of triangle
            let x_r = x_m + p1;
            let c = 0.134 + 20.5 / (15.3 + (m as f64));
            // p1 + area of parallelogram region
            let p2 = p1 * (1. + 2. * c);

            fn lambda(a: f64) -> f64 {
                a * (1. + 0.5 * a)
            }

            let lambda_l = lambda((f_m - x_l) / (f_m - x_l * p));
            let lambda_r = lambda((x_r - f_m) / (x_r * q));
            // p1 + area of left tail
            let p3 = p2 + c / lambda_l;
            // p1 + area of right tail
            let p4 = p3 + c / lambda_r;

            // return value
            let mut y: i64;

            loop {
                // Step 1: Generate `u` for selecting the region. If region 1 is
                // selected, generate a triangularly distributed variate.
                let u = rng.gen_range(0., p4);
                let mut v = rng.gen_range(0., 1.);
                if !(u > p1) {
                    y = (x_m - p1 * v + u) as i64;
                    break;
                }

                if !(u > p2) {
                    // Step 2: Region 2, parallelograms. Check if region 2 is
                    // used. If so, generate `y`.
                    let x = x_l + (u - p1) / c;
                    v = v * c + 1.0 - (x - x_m).abs() / p1;
                    if v > 1. {
                        continue;
                    } else {
                        y = x as i64;
                    }
                } else if !(u > p3) {
                    // Step 3: Region 3, left exponential tail.
                    y = (x_l + v.ln() / lambda_l) as i64;
                    if y < 0 {
                        continue;
                    } else {
                        v *= (u - p2) * lambda_l;
                    }
                } else {
                    // Step 4: Region 4, right exponential tail.
                    y = (x_r - v.ln() / lambda_r) as i64;
                    if y > 0 && (y as u64) > self.n {
                        continue;
                    } else {
                        v *= (u - p3) * lambda_r;
                    }
                }

                // Step 5: Acceptance/rejection comparison.

                // Step 5.0: Test for appropriate method of evaluating f(y).
                let k = (y - m).abs();
                if !(k > SQUEEZE_THRESHOLD && (k as f64) < 0.5 * npq - 1.) {
                    // Step 5.1: Evaluate f(y) via the recursive relationship. Start the
                    // search from the mode.
                    let s = p / q;
                    let a = s * (n + 1.);
                    let mut f = 1.0;
                    if m < y {
                        let mut i = m;
                        loop {
                            i += 1;
                            f *= a / (i as f64) - s;
                            if i == y {
                                break;
                            }
                        }
                    } else if m > y {
                        let mut i = y;
                        loop {
                            i += 1;
                            f /= a / (i as f64) - s;
                            if i == m {
                                break;
                            }
                        }
                    }
                    if v > f {
                        continue;
                    } else {
                        break;
                    }
                }

                // Step 5.2: Squeezing. Check the value of ln(v) againts upper and
                // lower bound of ln(f(y)).
                let k = k as f64;
                let rho = (k / npq) * ((k * (k / 3. + 0.625) + 1./6.) / npq + 0.5);
                let t = -0.5 * k*k / npq;
                let alpha = v.ln();
                if alpha < t - rho {
                    break;
                }
                if alpha > t + rho {
                    continue;
                }

                // Step 5.3: Final acceptance/rejection test.
                let x1 = (y + 1) as f64;
                let f1 = (m + 1) as f64;
                let z = ((n as i64) + 1 - m) as f64;
                let w = ((n as i64) - y + 1) as f64;

                fn stirling(a: f64) -> f64 {
                    let a2 = a * a;
                    (13860. - (462. - (132. - (99. - 140. / a2) / a2) / a2) / a2) / a / 166320.
                }

                if alpha > x_m * (f1 / x1).ln()
                    + (n - (m as f64) + 0.5) * (z / w).ln()
                    + ((y - m) as f64) * (w * p / (x1 * q)).ln()
                    + stirling(f1) + stirling(z) + stirling(x1) + stirling(w)
                {
                    continue;
                }

                break;
            }
            assert!(y >= 0);
            result = y as u64;
        }

        // Invert the result for p < 0.5.
        if p != self.p {
            self.n - result
        } else {
            result
        }
    }
}

#[cfg(test)]
mod test {
    use Rng;
    use distributions::Distribution;
    use super::Binomial;

    fn test_binomial_mean_and_variance<R: Rng>(n: u64, p: f64, rng: &mut R) {
        let binomial = Binomial::new(n, p);

        let expected_mean = n as f64 * p;
        let expected_variance = n as f64 * p * (1.0 - p);

        let mut results = [0.0; 1000];
        for i in results.iter_mut() { *i = binomial.sample(rng) as f64; }

        let mean = results.iter().sum::<f64>() / results.len() as f64;
        assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0,
                "mean: {}, expected_mean: {}", mean, expected_mean);

        let variance =
            results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>()
            / results.len() as f64;
        assert!((variance - expected_variance).abs() < expected_variance / 10.0,
                "variance: {}, expected_variance: {}", variance, expected_variance);
    }

    #[test]
    fn test_binomial() {
        let mut rng = ::test::rng(351);
        test_binomial_mean_and_variance(150, 0.1, &mut rng);
        test_binomial_mean_and_variance(70, 0.6, &mut rng);
        test_binomial_mean_and_variance(40, 0.5, &mut rng);
        test_binomial_mean_and_variance(20, 0.7, &mut rng);
        test_binomial_mean_and_variance(20, 0.5, &mut rng);
    }

    #[test]
    fn test_binomial_end_points() {
        let mut rng = ::test::rng(352);
        assert_eq!(rng.sample(Binomial::new(20, 0.0)), 0);
        assert_eq!(rng.sample(Binomial::new(20, 1.0)), 20);
    }

    #[test]
    #[should_panic]
    fn test_binomial_invalid_lambda_neg() {
        Binomial::new(20, -10.0);
    }
}
<|endoftext|>