<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Code to handle initial setup steps for a pool.
// Initial setup steps are steps that do not alter the environment.

use std::{
    collections::{HashMap, HashSet},
    fs::OpenOptions,
    path::{Path, PathBuf},
};

use chrono::{DateTime, Utc};
use libudev;
use serde_json;

use devicemapper::{devnode_to_devno, Device, Sectors};

use crate::{
    engine::{
        strat_engine::{
            backstore::{
                blockdev::StratBlockDev,
                device::DevOwnership,
                metadata::{device_identifiers, BDA},
            },
            device::blkdev_size,
            serde_structs::{BackstoreSave, BaseBlockDevSave, PoolSave},
            udev::{block_enumerator, decide_ownership, is_multipath_member, stratis_enumerator},
        },
        types::{BlockDevTier, DevUuid, PoolUuid},
    },
    stratis::{ErrorEnum, StratisError, StratisResult},
};

/// Retrieve all Stratis block devices that should be made use of by the
/// Stratis engine. This excludes Stratis block devices that are multipath
/// members.
///
/// Includes a fallback path, which is used if no Stratis block devices are
/// found using the obvious udev property- and enumerator-based approach.
/// This fallback path is more expensive, because it must search all block
/// devices via udev, and may also attempt to directly read Stratis
/// metadata from the device.
fn get_stratis_block_devices() -> StratisResult<Vec<PathBuf>> {
    let devices = {
        let context = libudev::Context::new()?;
        let mut enumerator = stratis_enumerator(&context)?;

        enumerator
            .scan_devices()?
            .filter(|dev| dev.is_initialized())
            .filter(|dev| !is_multipath_member(dev).unwrap_or(true))
            .filter_map(|i| i.devnode().map(|d| d.into()))
            .collect::<Vec<PathBuf>>()
    };

    if devices.is_empty() {
        // No Stratis devices have been found, possible reasons are:
        // 1. There are none
        // 2. There are some but libblkid version is less than 2.32, so
        // Stratis devices are not recognized by udev.
        // 3. There are many incomplete udev entries, because this code is
        // being run before the udev database is populated.
        //
        // Try again to find Stratis block devices, but this time enumerate
        // all block devices, not just all the ones that can be identified
        // as Stratis blockdevs by udev, and then scrutinize each one
        // using various methods.
        let context = libudev::Context::new()?;
        let mut enumerator = block_enumerator(&context)?;
        Ok(enumerator
            .scan_devices()?
            .filter(|dev| dev.is_initialized())
            .filter_map(|dev| {
                dev.devnode().and_then(|devnode| {
                    decide_ownership(&dev)
                        .and_then(|decision| DevOwnership::from_udev_ownership(&decision, devnode))
                        .ok()
                        .and_then(|ownership| ownership.stratis_identifiers())
                        .map(|_| devnode.into())
                })
            })
            .collect())
    } else {
        Ok(devices)
    }
}

/// Find all Stratis devices.
///
/// Returns a map of pool uuids to a map of devices to devnodes for each pool.
pub fn find_all() -> StratisResult<HashMap<PoolUuid, HashMap<Device, PathBuf>>> {
    let mut pool_map = HashMap::new();

    for devnode in get_stratis_block_devices()? {
        match devnode_to_devno(&devnode)? {
            None => continue,
            Some(devno) => {
                if let Some((pool_uuid, _)) =
                    device_identifiers(&mut OpenOptions::new().read(true).open(&devnode)?)?
                {
                    pool_map
                        .entry(pool_uuid)
                        .or_insert_with(HashMap::new)
                        .insert(Device::from(devno), devnode);
                }
            }
        }
    }
    Ok(pool_map)
}

/// Get the most recent metadata from a set of Devices for a given pool UUID.
/// Returns None if no metadata found for this pool.
#[allow(clippy::implicit_hasher)]
pub fn get_metadata(
    pool_uuid: PoolUuid,
    devnodes: &HashMap<Device, PathBuf>,
) -> StratisResult<Option<(DateTime<Utc>, PoolSave)>> {
    // Get pairs of device nodes and matching BDAs
    // If no BDA, or BDA UUID does not match pool UUID, skip.
    // If there is an error reading the BDA, error. There could have been
    // vital information on that BDA, for example, it may have contained
    // the newest metadata.
    let mut bdas = Vec::new();
    for devnode in devnodes.values() {
        let bda = BDA::load(&mut OpenOptions::new().read(true).open(devnode)?)?;
        if let Some(bda) = bda {
            if bda.pool_uuid() == pool_uuid {
                bdas.push((devnode, bda));
            }
        }
    }

    // Most recent time should never be None if this was a properly
    // created pool; this allows for the method to be called in other
    // circumstances.
    let most_recent_time = {
        match bdas
            .iter()
            .filter_map(|&(_, ref bda)| bda.last_update_time())
            .max()
        {
            Some(time) => time,
            None => return Ok(None),
        }
    };

    // Try to read from all available devnodes that could contain most
    // recent metadata. In the event of errors, continue to try until all are
    // exhausted.
    bdas.iter()
        .filter_map(|&(devnode, ref bda)| {
            if bda.last_update_time() == Some(most_recent_time) {
                OpenOptions::new()
                    .read(true)
                    .open(devnode)
                    .ok()
                    .and_then(|mut f| bda.load_state(&mut f).unwrap_or(None))
                    .and_then(|data| serde_json::from_slice(&data).ok())
            } else {
                None
            }
        })
        .next()
        .ok_or_else(|| {
            StratisError::Engine(
                ErrorEnum::NotFound,
                "timestamp indicates data was written, but no data successfully read".into(),
            )
        })
        .map(|psave| Some((*most_recent_time, psave)))
}

/// Get all the blockdevs corresponding to this pool that can be obtained from
/// the given devices. Sort the blockdevs in the order in which they were
/// recorded in the metadata.
/// Returns an error if a BDA can not be read or can not be found on any
/// blockdev in devnodes.
/// Returns an error if the blockdevs obtained do not match the metadata.
/// Returns a tuple, of which the first are the data devs, and the second
/// are the devs that support the cache tier.
/// Precondition: Every device in devnodes has already been determined to
/// belong to the pool with the specified pool uuid.
#[allow(clippy::implicit_hasher)]
pub fn get_blockdevs(
    pool_uuid: PoolUuid,
    backstore_save: &BackstoreSave,
    devnodes: &HashMap<Device, PathBuf>,
) -> StratisResult<(Vec<StratBlockDev>, Vec<StratBlockDev>)> {
    let recorded_data_map: HashMap<DevUuid, (usize, &BaseBlockDevSave)> = backstore_save
        .data_tier
        .blockdev
        .devs
        .iter()
        .enumerate()
        .map(|(i, bds)| (bds.uuid, (i, bds)))
        .collect();

    let recorded_cache_map: HashMap<DevUuid, (usize, &BaseBlockDevSave)> =
        match backstore_save.cache_tier {
            Some(ref cache_tier) => cache_tier
                .blockdev
                .devs
                .iter()
                .enumerate()
                .map(|(i, bds)| (bds.uuid, (i, bds)))
                .collect(),
            None => HashMap::new(),
        };

    let mut segment_table: HashMap<DevUuid, Vec<(Sectors, Sectors)>> = HashMap::new();
    for seg in &backstore_save.data_tier.blockdev.allocs[0] {
        segment_table
            .entry(seg.parent)
            .or_insert_with(Vec::default)
            .push((seg.start, seg.length))
    }

    if let Some(ref cache_tier) = backstore_save.cache_tier {
        for seg in cache_tier.blockdev.allocs.iter().flat_map(|i| i.iter()) {
            segment_table
                .entry(seg.parent)
                .or_insert_with(Vec::default)
                .push((seg.start, seg.length))
        }
    }

    // Construct a single StratBlockDev. Return the tier to which the
    // blockdev has been found to belong. Returns an error if the block
    // device has shrunk, no metadata can be found for the block device,
    // or it is impossible to set up the device because the recorded
    // allocation information is impossible.
    fn get_blockdev(
        device: Device,
        devnode: &Path,
        bda: BDA,
        data_map: &HashMap<DevUuid, (usize, &BaseBlockDevSave)>,
        cache_map: &HashMap<DevUuid, (usize, &BaseBlockDevSave)>,
        segment_table: &HashMap<DevUuid, Vec<(Sectors, Sectors)>>,
    ) -> StratisResult<(BlockDevTier, StratBlockDev)> {
        // Return an error if apparent size of Stratis block device appears to
        // have decreased since metadata was recorded or if size of block
        // device could not be obtained.
        blkdev_size(&OpenOptions::new().read(true).open(devnode)?).and_then(|actual_size| {
            let actual_size_sectors = actual_size.sectors();
            let recorded_size = bda.dev_size().sectors();
            if actual_size_sectors < recorded_size {
                let err_msg = format!(
                    "Stratis device with device number {}, devnode {}, pool UUID {} and device UUID {} had recorded size ({}), but actual size is less at ({})",
                    device,
                    devnode.display(),
                    bda.pool_uuid(),
                    bda.dev_uuid(),
                    recorded_size,
                    actual_size
                );
                Err(StratisError::Engine(ErrorEnum::Error, err_msg))
            } else {
                Ok(())
            }
        })?;

        let dev_uuid = bda.dev_uuid();

        // Locate the device in the metadata using its uuid. Return the device
        // metadata and whether it was a cache or a datadev.
        let (tier, &(_, bd_save)) = data_map
            .get(&dev_uuid)
            .map(|bd_save| (BlockDevTier::Data, bd_save))
            .or_else(|| {
                cache_map
                    .get(&dev_uuid)
                    .map(|bd_save| (BlockDevTier::Cache, bd_save))
            })
            .ok_or_else(|| {
                let err_msg = format!(
                        "Stratis device with device number {}, devnode {}, pool UUID {} and device UUID {} had no record in pool metadata",
                        device,
                        devnode.display(),
                        bda.pool_uuid(),
                        bda.dev_uuid()
                    );
                StratisError::Engine(ErrorEnum::NotFound, err_msg)
            })?;

        // This should always succeed since the actual size is at
        // least the recorded size, so all segments should be
        // available to be allocated. If this fails, the most likely
        // conclusion is metadata corruption.
        let segments = segment_table.get(&dev_uuid);
        Ok((
            tier,
            StratBlockDev::new(
                device,
                devnode.to_owned(),
                bda,
                segments.unwrap_or(&vec![]),
                bd_save.user_info.clone(),
                bd_save.hardware_info.clone(),
            )?,
        ))
    }

    let (mut datadevs, mut cachedevs): (Vec<StratBlockDev>, Vec<StratBlockDev>) = (vec![], vec![]);
    for (device, devnode) in devnodes {
        let bda = BDA::load(&mut OpenOptions::new().read(true).open(devnode)?)?.ok_or_else(|| {
            StratisError::Engine(ErrorEnum::NotFound,
                                                 format!("Device {} with devnode {} was previously determined to belong to pool with uuid {} but no BDA was found",
                                                 device,
                                                 devnode.display(),
                                                 pool_uuid))
        })?;

        get_blockdev(
            *device,
            devnode,
            bda,
            &recorded_data_map,
            &recorded_cache_map,
            &segment_table,
        )
        .map(|(tier, blockdev)| {
            match tier {
                BlockDevTier::Data => &mut datadevs,
                BlockDevTier::Cache => &mut cachedevs,
            }
            .push(blockdev)
        })?;
    }

    // Verify that devices located are congruent with the metadata recorded
    // and generally consistent with expectations. If all seems correct,
    // sort the devices according to their order in the metadata.
    fn check_and_sort_devs(
        mut devs: Vec<StratBlockDev>,
        dev_map: &HashMap<DevUuid, (usize, &BaseBlockDevSave)>,
    ) -> StratisResult<Vec<StratBlockDev>> {
        let mut uuids = HashSet::new();
        let mut duplicate_uuids = Vec::new();
        for dev in &devs {
            let dev_uuid = dev.uuid();
            if !uuids.insert(dev_uuid) {
                duplicate_uuids.push(dev_uuid);
            }
        }

        if !duplicate_uuids.is_empty() {
            let err_msg = format!(
                "The following list of Stratis UUIDs were each claimed by more than one Stratis device: {}",
                duplicate_uuids.iter().map(|u| u.to_simple_ref().to_string()).collect::<Vec<_>>().join(", ")
            );
            return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
        }

        let recorded_uuids: HashSet<_> = dev_map.keys().cloned().collect();
        if uuids != recorded_uuids {
            let err_msg = format!(
                "UUIDs of devices found ({}) did not correspond with UUIDs specified in the metadata for this group of devices ({})",
                uuids.iter().map(|u| u.to_simple_ref().to_string()).collect::<Vec<_>>().join(", "),
                recorded_uuids.iter().map(|u| u.to_simple_ref().to_string()).collect::<Vec<_>>().join(", "),
            );
            return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
        }

        // Sort the devices according to their original location in the
        // metadata. Use a faster unstable sort, because the order of
        // devs before the sort is arbitrary and does not need to be
        // preserved.
        devs.sort_unstable_by_key(|dev| dev_map[&dev.uuid()].0);
        Ok(devs)
    }

    let datadevs = check_and_sort_devs(datadevs, &recorded_data_map)?;
    let cachedevs = check_and_sort_devs(cachedevs, &recorded_cache_map)?;

    Ok((datadevs, cachedevs))
}
<commit_msg>Add a log entry to show if forced to use fallback search<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Code to handle initial setup steps for a pool.
// Initial setup steps are steps that do not alter the environment.

use std::{
    collections::{HashMap, HashSet},
    fs::OpenOptions,
    path::{Path, PathBuf},
};

use chrono::{DateTime, Utc};
use libudev;
use serde_json;

use devicemapper::{devnode_to_devno, Device, Sectors};

use crate::{
    engine::{
        strat_engine::{
            backstore::{
                blockdev::StratBlockDev,
                device::DevOwnership,
                metadata::{device_identifiers, BDA},
            },
            device::blkdev_size,
            serde_structs::{BackstoreSave, BaseBlockDevSave, PoolSave},
            udev::{block_enumerator, decide_ownership, is_multipath_member, stratis_enumerator},
        },
        types::{BlockDevTier, DevUuid, PoolUuid},
    },
    stratis::{ErrorEnum, StratisError, StratisResult},
};

/// Retrieve all Stratis block devices that should be made use of by the
/// Stratis engine. This excludes Stratis block devices that are multipath
/// members.
///
/// Includes a fallback path, which is used if no Stratis block devices are
/// found using the obvious udev property- and enumerator-based approach.
/// This fallback path is more expensive, because it must search all block
/// devices via udev, and may also attempt to directly read Stratis
/// metadata from the device.
fn get_stratis_block_devices() -> StratisResult<Vec<PathBuf>> {
    let devices = {
        let context = libudev::Context::new()?;
        let mut enumerator = stratis_enumerator(&context)?;

        enumerator
            .scan_devices()?
            .filter(|dev| dev.is_initialized())
            .filter(|dev| !is_multipath_member(dev).unwrap_or(true))
            .filter_map(|i| i.devnode().map(|d| d.into()))
            .collect::<Vec<PathBuf>>()
    };

    if devices.is_empty() {
        // No Stratis devices have been found, possible reasons are:
        // 1. There are none
        // 2. There are some but libblkid version is less than 2.32, so
        // Stratis devices are not recognized by udev.
        // 3. There are many incomplete udev entries, because this code is
        // being run before the udev database is populated.
        //
        // Try again to find Stratis block devices, but this time enumerate
        // all block devices, not just all the ones that can be identified
        // as Stratis blockdevs by udev, and then scrutinize each one
        // using various methods.

        info!("Could not identify any Stratis devices by a udev search for devices with ID_FS_TYPE=\"stratis\"; using fallback search mechanism");

        let context = libudev::Context::new()?;
        let mut enumerator = block_enumerator(&context)?;
        Ok(enumerator
            .scan_devices()?
            .filter(|dev| dev.is_initialized())
            .filter_map(|dev| {
                dev.devnode().and_then(|devnode| {
                    decide_ownership(&dev)
                        .and_then(|decision| DevOwnership::from_udev_ownership(&decision, devnode))
                        .ok()
                        .and_then(|ownership| ownership.stratis_identifiers())
                        .map(|_| devnode.into())
                })
            })
            .collect())
    } else {
        Ok(devices)
    }
}

/// Find all Stratis devices.
///
/// Returns a map of pool uuids to a map of devices to devnodes for each pool.
pub fn find_all() -> StratisResult<HashMap<PoolUuid, HashMap<Device, PathBuf>>> {
    let mut pool_map = HashMap::new();

    for devnode in get_stratis_block_devices()? {
        match devnode_to_devno(&devnode)? {
            None => continue,
            Some(devno) => {
                if let Some((pool_uuid, _)) =
                    device_identifiers(&mut OpenOptions::new().read(true).open(&devnode)?)?
                {
                    pool_map
                        .entry(pool_uuid)
                        .or_insert_with(HashMap::new)
                        .insert(Device::from(devno), devnode);
                }
            }
        }
    }
    Ok(pool_map)
}

/// Get the most recent metadata from a set of Devices for a given pool UUID.
/// Returns None if no metadata found for this pool.
#[allow(clippy::implicit_hasher)]
pub fn get_metadata(
    pool_uuid: PoolUuid,
    devnodes: &HashMap<Device, PathBuf>,
) -> StratisResult<Option<(DateTime<Utc>, PoolSave)>> {
    // Get pairs of device nodes and matching BDAs
    // If no BDA, or BDA UUID does not match pool UUID, skip.
    // If there is an error reading the BDA, error. There could have been
    // vital information on that BDA, for example, it may have contained
    // the newest metadata.
    let mut bdas = Vec::new();
    for devnode in devnodes.values() {
        let bda = BDA::load(&mut OpenOptions::new().read(true).open(devnode)?)?;
        if let Some(bda) = bda {
            if bda.pool_uuid() == pool_uuid {
                bdas.push((devnode, bda));
            }
        }
    }

    // Most recent time should never be None if this was a properly
    // created pool; this allows for the method to be called in other
    // circumstances.
    let most_recent_time = {
        match bdas
            .iter()
            .filter_map(|&(_, ref bda)| bda.last_update_time())
            .max()
        {
            Some(time) => time,
            None => return Ok(None),
        }
    };

    // Try to read from all available devnodes that could contain most
    // recent metadata. In the event of errors, continue to try until all are
    // exhausted.
    bdas.iter()
        .filter_map(|&(devnode, ref bda)| {
            if bda.last_update_time() == Some(most_recent_time) {
                OpenOptions::new()
                    .read(true)
                    .open(devnode)
                    .ok()
                    .and_then(|mut f| bda.load_state(&mut f).unwrap_or(None))
                    .and_then(|data| serde_json::from_slice(&data).ok())
            } else {
                None
            }
        })
        .next()
        .ok_or_else(|| {
            StratisError::Engine(
                ErrorEnum::NotFound,
                "timestamp indicates data was written, but no data successfully read".into(),
            )
        })
        .map(|psave| Some((*most_recent_time, psave)))
}

/// Get all the blockdevs corresponding to this pool that can be obtained from
/// the given devices. Sort the blockdevs in the order in which they were
/// recorded in the metadata.
/// Returns an error if a BDA can not be read or can not be found on any
/// blockdev in devnodes.
/// Returns an error if the blockdevs obtained do not match the metadata.
/// Returns a tuple, of which the first are the data devs, and the second
/// are the devs that support the cache tier.
/// Precondition: Every device in devnodes has already been determined to
/// belong to the pool with the specified pool uuid.
#[allow(clippy::implicit_hasher)]
pub fn get_blockdevs(
    pool_uuid: PoolUuid,
    backstore_save: &BackstoreSave,
    devnodes: &HashMap<Device, PathBuf>,
) -> StratisResult<(Vec<StratBlockDev>, Vec<StratBlockDev>)> {
    let recorded_data_map: HashMap<DevUuid, (usize, &BaseBlockDevSave)> = backstore_save
        .data_tier
        .blockdev
        .devs
        .iter()
        .enumerate()
        .map(|(i, bds)| (bds.uuid, (i, bds)))
        .collect();

    let recorded_cache_map: HashMap<DevUuid, (usize, &BaseBlockDevSave)> =
        match backstore_save.cache_tier {
            Some(ref cache_tier) => cache_tier
                .blockdev
                .devs
                .iter()
                .enumerate()
                .map(|(i, bds)| (bds.uuid, (i, bds)))
                .collect(),
            None => HashMap::new(),
        };

    let mut segment_table: HashMap<DevUuid, Vec<(Sectors, Sectors)>> = HashMap::new();
    for seg in &backstore_save.data_tier.blockdev.allocs[0] {
        segment_table
            .entry(seg.parent)
            .or_insert_with(Vec::default)
            .push((seg.start, seg.length))
    }

    if let Some(ref cache_tier) = backstore_save.cache_tier {
        for seg in cache_tier.blockdev.allocs.iter().flat_map(|i| i.iter()) {
            segment_table
                .entry(seg.parent)
                .or_insert_with(Vec::default)
                .push((seg.start, seg.length))
        }
    }

    // Construct a single StratBlockDev. Return the tier to which the
    // blockdev has been found to belong. Returns an error if the block
    // device has shrunk, no metadata can be found for the block device,
    // or it is impossible to set up the device because the recorded
    // allocation information is impossible.
    fn get_blockdev(
        device: Device,
        devnode: &Path,
        bda: BDA,
        data_map: &HashMap<DevUuid, (usize, &BaseBlockDevSave)>,
        cache_map: &HashMap<DevUuid, (usize, &BaseBlockDevSave)>,
        segment_table: &HashMap<DevUuid, Vec<(Sectors, Sectors)>>,
    ) -> StratisResult<(BlockDevTier, StratBlockDev)> {
        // Return an error if apparent size of Stratis block device appears to
        // have decreased since metadata was recorded or if size of block
        // device could not be obtained.
        blkdev_size(&OpenOptions::new().read(true).open(devnode)?).and_then(|actual_size| {
            let actual_size_sectors = actual_size.sectors();
            let recorded_size = bda.dev_size().sectors();
            if actual_size_sectors < recorded_size {
                let err_msg = format!(
                    "Stratis device with device number {}, devnode {}, pool UUID {} and device UUID {} had recorded size ({}), but actual size is less at ({})",
                    device,
                    devnode.display(),
                    bda.pool_uuid(),
                    bda.dev_uuid(),
                    recorded_size,
                    actual_size
                );
                Err(StratisError::Engine(ErrorEnum::Error, err_msg))
            } else {
                Ok(())
            }
        })?;

        let dev_uuid = bda.dev_uuid();

        // Locate the device in the metadata using its uuid. Return the device
        // metadata and whether it was a cache or a datadev.
        let (tier, &(_, bd_save)) = data_map
            .get(&dev_uuid)
            .map(|bd_save| (BlockDevTier::Data, bd_save))
            .or_else(|| {
                cache_map
                    .get(&dev_uuid)
                    .map(|bd_save| (BlockDevTier::Cache, bd_save))
            })
            .ok_or_else(|| {
                let err_msg = format!(
                        "Stratis device with device number {}, devnode {}, pool UUID {} and device UUID {} had no record in pool metadata",
                        device,
                        devnode.display(),
                        bda.pool_uuid(),
                        bda.dev_uuid()
                    );
                StratisError::Engine(ErrorEnum::NotFound, err_msg)
            })?;

        // This should always succeed since the actual size is at
        // least the recorded size, so all segments should be
        // available to be allocated. If this fails, the most likely
        // conclusion is metadata corruption.
        let segments = segment_table.get(&dev_uuid);
        Ok((
            tier,
            StratBlockDev::new(
                device,
                devnode.to_owned(),
                bda,
                segments.unwrap_or(&vec![]),
                bd_save.user_info.clone(),
                bd_save.hardware_info.clone(),
            )?,
        ))
    }

    let (mut datadevs, mut cachedevs): (Vec<StratBlockDev>, Vec<StratBlockDev>) = (vec![], vec![]);
    for (device, devnode) in devnodes {
        let bda = BDA::load(&mut OpenOptions::new().read(true).open(devnode)?)?.ok_or_else(|| {
            StratisError::Engine(ErrorEnum::NotFound,
                                                 format!("Device {} with devnode {} was previously determined to belong to pool with uuid {} but no BDA was found",
                                                 device,
                                                 devnode.display(),
                                                 pool_uuid))
        })?;

        get_blockdev(
            *device,
            devnode,
            bda,
            &recorded_data_map,
            &recorded_cache_map,
            &segment_table,
        )
        .map(|(tier, blockdev)| {
            match tier {
                BlockDevTier::Data => &mut datadevs,
                BlockDevTier::Cache => &mut cachedevs,
            }
            .push(blockdev)
        })?;
    }

    // Verify that devices located are congruent with the metadata recorded
    // and generally consistent with expectations. If all seems correct,
    // sort the devices according to their order in the metadata.
    fn check_and_sort_devs(
        mut devs: Vec<StratBlockDev>,
        dev_map: &HashMap<DevUuid, (usize, &BaseBlockDevSave)>,
    ) -> StratisResult<Vec<StratBlockDev>> {
        let mut uuids = HashSet::new();
        let mut duplicate_uuids = Vec::new();
        for dev in &devs {
            let dev_uuid = dev.uuid();
            if !uuids.insert(dev_uuid) {
                duplicate_uuids.push(dev_uuid);
            }
        }

        if !duplicate_uuids.is_empty() {
            let err_msg = format!(
                "The following list of Stratis UUIDs were each claimed by more than one Stratis device: {}",
                duplicate_uuids.iter().map(|u| u.to_simple_ref().to_string()).collect::<Vec<_>>().join(", ")
            );
            return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
        }

        let recorded_uuids: HashSet<_> = dev_map.keys().cloned().collect();
        if uuids != recorded_uuids {
            let err_msg = format!(
                "UUIDs of devices found ({}) did not correspond with UUIDs specified in the metadata for this group of devices ({})",
                uuids.iter().map(|u| u.to_simple_ref().to_string()).collect::<Vec<_>>().join(", "),
                recorded_uuids.iter().map(|u| u.to_simple_ref().to_string()).collect::<Vec<_>>().join(", "),
            );
            return Err(StratisError::Engine(ErrorEnum::Invalid, err_msg));
        }

        // Sort the devices according to their original location in the
        // metadata. Use a faster unstable sort, because the order of
        // devs before the sort is arbitrary and does not need to be
        // preserved.
        devs.sort_unstable_by_key(|dev| dev_map[&dev.uuid()].0);
        Ok(devs)
    }

    let datadevs = check_and_sort_devs(datadevs, &recorded_data_map)?;
    let cachedevs = check_and_sort_devs(cachedevs, &recorded_cache_map)?;

    Ok((datadevs, cachedevs))
}
<|endoftext|>


<commit_before>extern crate libc;
extern crate sys_util;
extern crate kvm_sys;
extern crate kvm;
extern crate kernel_loader;
extern crate x86_64;
extern crate clap;
extern crate devices;

pub mod machine;

use std::ffi::{CStr, CString};
use std::fs::File;
use std::io::{self, stdout, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Barrier, Mutex};
use std::thread;
use std::thread::{JoinHandle};
use kvm::*;
use kvm_sys::kvm_regs;
use sys_util::{register_signal_handler, EventFd, GuestAddress, GuestMemory,
               Killable, Pollable, Poller};
use machine::MachineCfg;

const KERNEL_START_OFFSET: usize = 0x200000;
const CMDLINE_OFFSET: usize = 0x20000;

pub enum Error {
    ConfigureSystem(x86_64::Error),
    EventFd(sys_util::Error),
    GuestMemory(sys_util::GuestMemoryError),
    Kernel(std::io::Error),
    KernelLoader(kernel_loader::Error),
    Kvm(sys_util::Error),
    Vcpu(sys_util::Error),
    VcpuSpawn(std::io::Error),
    Vm(sys_util::Error),
}

impl std::convert::From<kernel_loader::Error> for Error {
    fn from(e: kernel_loader::Error) -> Error {
        Error::KernelLoader(e)
    }
}

impl std::convert::From<x86_64::Error> for Error {
    fn from(e: x86_64::Error) -> Error {
        Error::ConfigureSystem(e)
    }
}

type Result<T> = std::result::Result<T, Error>;

pub fn boot_kernel(cfg: &MachineCfg) -> Result<()> {
    // FIXME branciog@ do not hardcode the vm mem size
    // Hardcoding the vm memory size to 128MB
    let mem_size = 128 << 20;
    let arch_mem_regions = x86_64::arch_memory_regions(mem_size);

    let mut kernel_file;
    match cfg.kernel_path {
        Some(ref kernel_path) => {
            kernel_file = File::open(kernel_path.as_path())
                    .map_err(Error::Kernel)?
        },
        None => {
            return Err(Error::Kernel(
                    io::Error::new(io::ErrorKind::NotFound,
                                   "missing kernel path")))
        }
    }

    let cmdline: CString = match cfg.kernel_cmdline {
        Some(ref v) => CString::new(v.as_bytes()).unwrap(),
        _ => return Err(Error::Kernel(
                io::Error::new(io::ErrorKind::NotFound,
                               "missing kernel cmdline")))
    };
    let cmdline: &CStr = &cmdline;
    let vcpu_count = 1;

    let kernel_start_addr = GuestAddress(KERNEL_START_OFFSET);
    let cmdline_addr = GuestAddress(CMDLINE_OFFSET);

    let guest_mem = GuestMemory::new(&arch_mem_regions)
            .map_err(Error::GuestMemory)?;

    let kvm = Kvm::new().map_err(Error::Kvm)?;
    let vm = Vm::new(&kvm, guest_mem).map_err(Error::Vm)?;

    let tss_addr = GuestAddress(0xfffbd000);
    vm.set_tss_addr(tss_addr).map_err(Error::Vm)?;
    vm.create_pit().map_err(Error::Vm)?;
    vm.create_irq_chip().map_err(Error::Vm)?;

    kernel_loader::load_kernel(vm.get_memory(), kernel_start_addr,
                               &mut kernel_file)?;
    kernel_loader::load_cmdline(vm.get_memory(), cmdline_addr, cmdline)?;

    x86_64::configure_system(vm.get_memory(),
                             kernel_start_addr,
                             cmdline_addr,
                             cmdline.to_bytes().len() + 1,
                             vcpu_count as u8)?;

    let mut io_bus = devices::Bus::new();
    let com_evt = EventFd::new().map_err(Error::EventFd)?;
    let stdio_serial =
        Arc::new(Mutex::new(
                    devices::Serial::new_out(com_evt, Box::new(stdout()))));
    io_bus.insert(stdio_serial, 0x3f8, 0x8).unwrap();

    let exit_evt = EventFd::new().map_err(Error::EventFd)?;
    let mut vcpu_handles = Vec::with_capacity(vcpu_count as usize);
    let vcpu_thread_barrier = Arc::new(Barrier::new((vcpu_count + 1) as usize));
    let kill_signaled = Arc::new(AtomicBool::new(false));

    for cpu_id in 0..vcpu_count {
        let io_bus = io_bus.clone();
        let kill_signaled = kill_signaled.clone();
        let vcpu_thread_barrier = vcpu_thread_barrier.clone();
        let vcpu_exit_evt = exit_evt.try_clone().map_err(Error::EventFd)?;

        let vcpu = Vcpu::new(cpu_id as libc::c_ulong, &kvm, &vm).map_err(Error::Vcpu)?;
        x86_64::configure_vcpu(vm.get_memory(), kernel_start_addr, &kvm,
                               &vcpu, cpu_id as u64, vcpu_count as u64)?;
        vcpu_handles.push(thread::Builder::new()
                .name(format!("fc_vcpu{}", cpu_id))
                .spawn(move || {
            unsafe {
                extern "C" fn handle_signal() {}
                // Our signal handler does nothing and is trivially async signal safe.
                register_signal_handler(0, handle_signal)
                        .expect("failed to register vcpu signal handler");
            }

            vcpu_thread_barrier.wait();

            loop {
                match vcpu.run() {
                    Ok(run) => match run {
                        VcpuExit::IoIn(addr, data) => {
                            io_bus.read(addr as u64, data);
                        },
                        VcpuExit::IoOut(addr, data) => {
                            io_bus.write(addr as u64, data);
                        },
                        VcpuExit::MmioRead(_, _) => {},
                        VcpuExit::MmioWrite(_, _) => {},
                        VcpuExit::Hlt => {
                            println!("KVM_EXIT_HLT");
                            break;
                        },
                        VcpuExit::Shutdown => {
                            println!("KVM_EXIT_SHUTDOWN");
                            break;
                        },
                        r => {
                            println!("unexpected exit reason: {:?}", r);
                            break;
                        }
                    },
                    Err(e) => {
                        match e.errno() {
                            libc::EAGAIN | libc::EINTR => {},
                            _ => {
                                println!("vcpu hit unknown error: {:?}", e);
                                break;
                            }
                        }
                    }
                }

                if kill_signaled.load(Ordering::SeqCst) {
                    break;
                }
            }

            vcpu_exit_evt
                .write(1)
                .expect("failed to signal vcpu exit eventfd");

        }).map_err(Error::VcpuSpawn)?);
    }

    vcpu_thread_barrier.wait();

    run_control(exit_evt, kill_signaled, vcpu_handles)
}

fn run_control(exit_evt: EventFd,
               kill_signaled: Arc<AtomicBool>,
               vcpu_handles: Vec<JoinHandle<()>>) -> Result<()> {
    const EXIT: u32 = 0;

    let mut pollables = Vec::new();
    pollables.push((EXIT, &exit_evt as &Pollable));

    let mut poller = Poller::new(pollables.len());

    'poll: loop {
        let tokens = match poller.poll(&pollables[..]) {
            Ok(v) => v,
            Err(e) => {
                println!("failed to poll: {:?}", e);
                break;
            }
        };

        for &token in tokens {
            match token {
                EXIT => {
                    println!("vcpu requested shutdown");
                    break 'poll;
                },
                _ => {}
            }
        }
    }

    kill_signaled.store(true, Ordering::SeqCst);
    for handle in vcpu_handles {
        match handle.kill(0) {
            Ok(_) => {
                if let Err(e) = handle.join() {
                    println!("failed to join vcpu thread: {:?}", e);
                }
            },
            Err(e) => println!("failed to kill vcpu thread: {:?}", e),
        }
    }

    Ok(())
}

pub fn run_x86_code() {
    // This example based on https://lwn.net/Articles/658511/
    let code = [
        0xba, 0xf8, 0x03, /* mov $0x3f8, %dx */
        0x00, 0xd8,       /* add %bl, %al */
        0x04, '0' as u8,  /* add $'0', %al */
        0xee,             /* out %al, (%dx) */
        0xb0, '\n' as u8, /* mov $'\n', %al */
        0xee,             /* out %al, (%dx) */
        0xf4,             /* hlt */
    ];

    let mem_size = 0x1000;
    let load_addr = GuestAddress(0x1000);
    let mem = GuestMemory::new(&vec![(load_addr, mem_size)]).unwrap();

    let kvm = Kvm::new().expect("new kvm failed");
    let vm = Vm::new(&kvm, mem).expect("new vm failed");
    let vcpu = Vcpu::new(0, &kvm, &vm).expect("new vcpu failed");

    vm.get_memory()
        .write_slice_at_addr(&code, load_addr)
        .expect("Writing code to memory failed.");

    let mut vcpu_sregs = vcpu.get_sregs().expect("get sregs failed");
    assert_ne!(vcpu_sregs.cs.base, 0);
    assert_ne!(vcpu_sregs.cs.selector, 0);
    vcpu_sregs.cs.base = 0;
    vcpu_sregs.cs.selector = 0;
    vcpu.set_sregs(&vcpu_sregs).expect("set sregs failed");

    let mut vcpu_regs: kvm_regs = unsafe { std::mem::zeroed() };
    vcpu_regs.rip = 0x1000;
    vcpu_regs.rax = 2;
    vcpu_regs.rbx = 2;
    vcpu_regs.rflags = 2;
    vcpu.set_regs(&vcpu_regs).expect("set regs failed");

    loop {
        match vcpu.run().expect("run failed") {
            VcpuExit::IoOut(0x3f8, data) => {
                assert_eq!(data.len(), 1);
                io::stdout().write(data).unwrap();
            },
            VcpuExit::Hlt => {
                io::stdout().write(b"KVM_EXIT_HLT\n").unwrap();
                break
            },
            r => panic!("unexpected exit reason: {:?}", r),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn run_code() {
        run_x86_code();
    }
}
<commit_msg>vmm: add serial devices and configure serial irqs<commit_after>extern crate libc;
extern crate sys_util;
extern crate kvm_sys;
extern crate kvm;
extern crate kernel_loader;
extern crate x86_64;
extern crate clap;
extern crate devices;

pub mod machine;

use std::ffi::{CStr, CString};
use std::fs::File;
use std::io::{self, stdout, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Barrier, Mutex};
use std::thread;
use std::thread::{JoinHandle};
use kvm::*;
use kvm_sys::kvm_regs;
use sys_util::{register_signal_handler, EventFd, GuestAddress, GuestMemory,
               Killable, Pollable, Poller};
use machine::MachineCfg;

const KERNEL_START_OFFSET: usize = 0x200000;
const CMDLINE_OFFSET: usize = 0x20000;

pub enum Error {
    ConfigureSystem(x86_64::Error),
    EventFd(sys_util::Error),
    GuestMemory(sys_util::GuestMemoryError),
    Irq(sys_util::Error),
    Kernel(std::io::Error),
    KernelLoader(kernel_loader::Error),
    Kvm(sys_util::Error),
    Vcpu(sys_util::Error),
    VcpuSpawn(std::io::Error),
    Vm(sys_util::Error),
}

impl std::convert::From<kernel_loader::Error> for Error {
    fn from(e: kernel_loader::Error) -> Error {
        Error::KernelLoader(e)
    }
}

impl std::convert::From<x86_64::Error> for Error {
    fn from(e: x86_64::Error) -> Error {
        Error::ConfigureSystem(e)
    }
}

type Result<T> = std::result::Result<T, Error>;

pub fn boot_kernel(cfg: &MachineCfg) -> Result<()> {
    // FIXME branciog@ do not hardcode the vm mem size
    // Hardcoding the vm memory size to 128MB
    let mem_size = 128 << 20;
    let arch_mem_regions = x86_64::arch_memory_regions(mem_size);

    let mut kernel_file;
    match cfg.kernel_path {
        Some(ref kernel_path) => {
            kernel_file = File::open(kernel_path.as_path())
                    .map_err(Error::Kernel)?
        },
        None => {
            return Err(Error::Kernel(
                    io::Error::new(io::ErrorKind::NotFound,
                                   "missing kernel path")))
        }
    }

    let cmdline: CString = match cfg.kernel_cmdline {
        Some(ref v) => CString::new(v.as_bytes()).unwrap(),
        _ => return Err(Error::Kernel(
                io::Error::new(io::ErrorKind::NotFound,
                               "missing kernel cmdline")))
    };
    let cmdline: &CStr = &cmdline;
    let vcpu_count = 1;

    let kernel_start_addr = GuestAddress(KERNEL_START_OFFSET);
    let cmdline_addr = GuestAddress(CMDLINE_OFFSET);

    let guest_mem = GuestMemory::new(&arch_mem_regions)
            .map_err(Error::GuestMemory)?;

    let kvm = Kvm::new().map_err(Error::Kvm)?;
    let vm = Vm::new(&kvm, guest_mem).map_err(Error::Vm)?;

    let tss_addr = GuestAddress(0xfffbd000);
    vm.set_tss_addr(tss_addr).map_err(Error::Vm)?;
    vm.create_pit().map_err(Error::Vm)?;
    vm.create_irq_chip().map_err(Error::Vm)?;

    kernel_loader::load_kernel(vm.get_memory(), kernel_start_addr,
                               &mut kernel_file)?;
    kernel_loader::load_cmdline(vm.get_memory(), cmdline_addr, cmdline)?;

    x86_64::configure_system(vm.get_memory(),
                             kernel_start_addr,
                             cmdline_addr,
                             cmdline.to_bytes().len() + 1,
                             vcpu_count as u8)?;

    let mut io_bus = devices::Bus::new();
    let com_evt_1_3 = EventFd::new().map_err(Error::EventFd)?;
    let com_evt_2_4 = EventFd::new().map_err(Error::EventFd)?;
    let stdio_serial = Arc::new(Mutex::new(devices::Serial::new_out(
            com_evt_1_3.try_clone().map_err(Error::EventFd)?,
            Box::new(stdout()))));

    io_bus.insert(stdio_serial, 0x3f8, 0x8).unwrap();
    io_bus.insert(Arc::new(Mutex::new(devices::Serial::new_sink(
            com_evt_2_4.try_clone().map_err(Error::EventFd)?))), 0x2f8, 0x8)
            .unwrap();
    io_bus.insert(Arc::new(Mutex::new(devices::Serial::new_sink(
            com_evt_1_3.try_clone().map_err(Error::EventFd)?))), 0x3e8, 0x8)
            .unwrap();
    io_bus.insert(Arc::new(Mutex::new(devices::Serial::new_sink(
            com_evt_2_4.try_clone().map_err(Error::EventFd)?))), 0x2e8, 0x8)
            .unwrap();

    vm.register_irqfd(&com_evt_1_3, 4)
            .map_err(Error::Irq)?;
    vm.register_irqfd(&com_evt_2_4, 3)
            .map_err(Error::Irq)?;

    let exit_evt = EventFd::new().map_err(Error::EventFd)?;
    let mut vcpu_handles = Vec::with_capacity(vcpu_count as usize);
    let vcpu_thread_barrier = Arc::new(Barrier::new((vcpu_count + 1) as usize));
    let kill_signaled = Arc::new(AtomicBool::new(false));

    for cpu_id in 0..vcpu_count {
        let io_bus = io_bus.clone();
        let kill_signaled = kill_signaled.clone();
        let vcpu_thread_barrier = vcpu_thread_barrier.clone();
        let vcpu_exit_evt = exit_evt.try_clone().map_err(Error::EventFd)?;

        let vcpu = Vcpu::new(cpu_id as libc::c_ulong, &kvm, &vm).map_err(Error::Vcpu)?;
        x86_64::configure_vcpu(vm.get_memory(), kernel_start_addr, &kvm,
                               &vcpu, cpu_id as u64, vcpu_count as u64)?;
        vcpu_handles.push(thread::Builder::new()
                .name(format!("fc_vcpu{}", cpu_id))
                .spawn(move || {
            unsafe {
                extern "C" fn handle_signal() {}
                // Our signal handler does nothing and is trivially async signal safe.
                register_signal_handler(0, handle_signal)
                        .expect("failed to register vcpu signal handler");
            }

            vcpu_thread_barrier.wait();

            loop {
                match vcpu.run() {
                    Ok(run) => match run {
                        VcpuExit::IoIn(addr, data) => {
                            io_bus.read(addr as u64, data);
                        },
                        VcpuExit::IoOut(addr, data) => {
                            io_bus.write(addr as u64, data);
                        },
                        VcpuExit::MmioRead(_, _) => {},
                        VcpuExit::MmioWrite(_, _) => {},
                        VcpuExit::Hlt => {
                            println!("KVM_EXIT_HLT");
                            break;
                        },
                        VcpuExit::Shutdown => {
                            println!("KVM_EXIT_SHUTDOWN");
                            break;
                        },
                        r => {
                            println!("unexpected exit reason: {:?}", r);
                            break;
                        }
                    },
                    Err(e) => {
                        match e.errno() {
                            libc::EAGAIN | libc::EINTR => {},
                            _ => {
                                println!("vcpu hit unknown error: {:?}", e);
                                break;
                            }
                        }
                    }
                }

                if kill_signaled.load(Ordering::SeqCst) {
                    break;
                }
            }

            vcpu_exit_evt
                .write(1)
                .expect("failed to signal vcpu exit eventfd");

        }).map_err(Error::VcpuSpawn)?);
    }

    vcpu_thread_barrier.wait();

    run_control(exit_evt, kill_signaled, vcpu_handles)
}

fn run_control(exit_evt: EventFd,
               kill_signaled: Arc<AtomicBool>,
               vcpu_handles: Vec<JoinHandle<()>>) -> Result<()> {
    const EXIT: u32 = 0;

    let mut pollables = Vec::new();
    pollables.push((EXIT, &exit_evt as &Pollable));

    let mut poller = Poller::new(pollables.len());

    'poll: loop {
        let tokens = match poller.poll(&pollables[..]) {
            Ok(v) => v,
            Err(e) => {
                println!("failed to poll: {:?}", e);
                break;
            }
        };

        for &token in tokens {
            match token {
                EXIT => {
                    println!("vcpu requested shutdown");
                    break 'poll;
                },
                _ => {}
            }
        }
    }

    kill_signaled.store(true, Ordering::SeqCst);
    for handle in vcpu_handles {
        match handle.kill(0) {
            Ok(_) => {
                if let Err(e) = handle.join() {
                    println!("failed to join vcpu thread: {:?}", e);
                }
            },
            Err(e) => println!("failed to kill vcpu thread: {:?}", e),
        }
    }

    Ok(())
}

pub fn run_x86_code() {
    // This example based on https://lwn.net/Articles/658511/
    let code = [
        0xba, 0xf8, 0x03, /* mov $0x3f8, %dx */
        0x00, 0xd8,       /* add %bl, %al */
        0x04, '0' as u8,  /* add $'0', %al */
        0xee,             /* out %al, (%dx) */
        0xb0, '\n' as u8, /* mov $'\n', %al */
        0xee,             /* out %al, (%dx) */
        0xf4,             /* hlt */
    ];

    let mem_size = 0x1000;
    let load_addr = GuestAddress(0x1000);
    let mem = GuestMemory::new(&vec![(load_addr, mem_size)]).unwrap();

    let kvm = Kvm::new().expect("new kvm failed");
    let vm = Vm::new(&kvm, mem).expect("new vm failed");
    let vcpu = Vcpu::new(0, &kvm, &vm).expect("new vcpu failed");

    vm.get_memory()
        .write_slice_at_addr(&code, load_addr)
        .expect("Writing code to memory failed.");

    let mut vcpu_sregs = vcpu.get_sregs().expect("get sregs failed");
    assert_ne!(vcpu_sregs.cs.base, 0);
    assert_ne!(vcpu_sregs.cs.selector, 0);
    vcpu_sregs.cs.base = 0;
    vcpu_sregs.cs.selector = 0;
    vcpu.set_sregs(&vcpu_sregs).expect("set sregs failed");

    let mut vcpu_regs: kvm_regs = unsafe { std::mem::zeroed() };
    vcpu_regs.rip = 0x1000;
    vcpu_regs.rax = 2;
    vcpu_regs.rbx = 2;
    vcpu_regs.rflags = 2;
    vcpu.set_regs(&vcpu_regs).expect("set regs failed");

    loop {
        match vcpu.run().expect("run failed") {
            VcpuExit::IoOut(0x3f8, data) => {
                assert_eq!(data.len(), 1);
                io::stdout().write(data).unwrap();
            },
            VcpuExit::Hlt => {
                io::stdout().write(b"KVM_EXIT_HLT\n").unwrap();
                break
            },
            r => panic!("unexpected exit reason: {:?}", r),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn run_code() {
        run_x86_code();
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

use std::fs::File;
use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
use std::{io, mem, result};

use libc::{c_void, dup, eventfd, read, write};

/// A safe wrapper around a Linux eventfd (man 2 eventfd).
///
/// An eventfd is useful because it is sendable across processes and can be used for signaling in
/// and out of the KVM API. They can also be polled like any other file descriptor.
pub struct EventFd {
    eventfd: File,
}

impl EventFd {
    /// Creates a new blocking EventFd with an initial value of 0.
    pub fn new() -> result::Result<EventFd, io::Error> {
        // This is safe because eventfd merely allocated an eventfd for our process and we handle
        // the error case.
        let ret = unsafe { eventfd(0, 0) };
        if ret < 0 {
            Err(io::Error::last_os_error())
        } else {
            // This is safe because we checked ret for success and know the kernel gave us an fd that we
            // own.
            Ok(EventFd {
                eventfd: unsafe { File::from_raw_fd(ret) },
            })
        }
    }

    /// Adds `v` to the eventfd's count, blocking until this won't overflow the count.
    pub fn write(&self, v: u64) -> result::Result<(), io::Error> {
        // This is safe because we made this fd and the pointer we pass can not overflow because we
        // give the syscall's size parameter properly.
        let ret = unsafe {
            write(
                self.as_raw_fd(),
                &v as *const u64 as *const c_void,
                mem::size_of::<u64>(),
            )
        };
        if ret <= 0 {
            Err(io::Error::last_os_error())
        } else {
            Ok(())
        }
    }

    /// Blocks until the the eventfd's count is non-zero, then resets the count to zero.
    pub fn read(&self) -> result::Result<u64, io::Error> {
        let mut buf: u64 = 0;
        let ret = unsafe {
            // This is safe because we made this fd and the pointer we pass can not overflow because
            // we give the syscall's size parameter properly.
            read(
                self.as_raw_fd(),
                &mut buf as *mut u64 as *mut c_void,
                mem::size_of::<u64>(),
            )
        };
        if ret <= 0 {
            Err(io::Error::last_os_error())
        } else {
            Ok(buf)
        }
    }

    /// Clones this EventFd, internally creating a new file descriptor. The new EventFd will share
    /// the same underlying count within the kernel.
    pub fn try_clone(&self) -> result::Result<EventFd, io::Error> {
        // This is safe because we made this fd and properly check that it returns without error.
        let ret = unsafe { dup(self.as_raw_fd()) };
        if ret < 0 {
            Err(io::Error::last_os_error())
        } else {
            // This is safe because we checked ret for success and know the kernel gave us an fd that we
            // own.
            Ok(EventFd {
                eventfd: unsafe { File::from_raw_fd(ret) },
            })
        }
    }
}

impl AsRawFd for EventFd {
    fn as_raw_fd(&self) -> RawFd {
        self.eventfd.as_raw_fd()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn new() {
        EventFd::new().unwrap();
    }

    #[test]
    fn read_write() {
        let evt = EventFd::new().unwrap();
        evt.write(55).unwrap();
        assert_eq!(evt.read().unwrap(), 55);
    }

    #[test]
    fn clone() {
        let evt = EventFd::new().unwrap();
        let evt_clone = evt.try_clone().unwrap();
        evt.write(923).unwrap();
        assert_eq!(evt_clone.read().unwrap(), 923);
    }
}
<commit_msg>- Makes eventfd read and write non-blocking - Add unit tests<commit_after>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

use std::fs::File;
use std::os::unix::io::{AsRawFd, FromRawFd, RawFd};
use std::{io, mem, result};

use libc::{c_void, dup, eventfd, read, write, EFD_NONBLOCK};

/// A safe wrapper around a Linux eventfd (man 2 eventfd).
///
/// An eventfd is useful because it is sendable across processes and can be used for signaling in
/// and out of the KVM API. They can also be polled like any other file descriptor.
pub struct EventFd {
    eventfd: File,
}

impl EventFd {
    /// Creates a new blocking EventFd with an initial value of 0.
    pub fn new() -> result::Result<EventFd, io::Error> {
        // This is safe because eventfd merely allocated an eventfd for our process and we handle
        // the error case.
        let ret = unsafe { eventfd(0, EFD_NONBLOCK) };
        if ret < 0 {
            Err(io::Error::last_os_error())
        } else {
            // This is safe because we checked ret for success and know the kernel gave us an fd that we
            // own.
            Ok(EventFd {
                eventfd: unsafe { File::from_raw_fd(ret) },
            })
        }
    }

    /// Adds `v` to the eventfd's count, does not block if the result will overflow the count
    pub fn write(&self, v: u64) -> result::Result<(), io::Error> {
        // This is safe because we made this fd and the pointer we pass can not overflow because we
        // give the syscall's size parameter properly.
        let ret = unsafe {
            write(
                self.as_raw_fd(),
                &v as *const u64 as *const c_void,
                mem::size_of::<u64>(),
            )
        };
        if ret <= 0 {
            Err(io::Error::last_os_error())
        } else {
            Ok(())
        }
    }

    /// Tries to read from the eventfd, does not block if the counter is zero
    pub fn read(&self) -> result::Result<u64, io::Error> {
        let mut buf: u64 = 0;
        let ret = unsafe {
            // This is safe because we made this fd and the pointer we pass can not overflow because
            // we give the syscall's size parameter properly.
            read(
                self.as_raw_fd(),
                &mut buf as *mut u64 as *mut c_void,
                mem::size_of::<u64>(),
            )
        };
        if ret < 0 {
            Err(io::Error::last_os_error())
        } else {
            Ok(buf)
        }
    }

    /// Clones this EventFd, internally creating a new file descriptor. The new EventFd will share
    /// the same underlying count within the kernel.
    pub fn try_clone(&self) -> result::Result<EventFd, io::Error> {
        // This is safe because we made this fd and properly check that it returns without error.
        let ret = unsafe { dup(self.as_raw_fd()) };
        if ret < 0 {
            Err(io::Error::last_os_error())
        } else {
            // This is safe because we checked ret for success and know the kernel gave us an fd that we
            // own.
            Ok(EventFd {
                eventfd: unsafe { File::from_raw_fd(ret) },
            })
        }
    }
}

impl AsRawFd for EventFd {
    fn as_raw_fd(&self) -> RawFd {
        self.eventfd.as_raw_fd()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn new() {
        EventFd::new().unwrap();
    }

    #[test]
    fn read_write() {
        let evt = EventFd::new().unwrap();
        evt.write(55).unwrap();
        assert_eq!(evt.read().unwrap(), 55);
    }

    #[test]
    fn write_overflow() {
        let evt = EventFd::new().unwrap();
        evt.write(std::u64::MAX - 1).unwrap();
        let r = evt.write(1);
        match r {
            Err(ref inner) if inner.kind() == io::ErrorKind::WouldBlock => (),
            _ => panic!("Unexpected"),
        }
    }

    #[test]
    fn read_nothing() {
        let evt = EventFd::new().unwrap();
        let r = evt.read();
        match r {
            Err(ref inner) if inner.kind() == io::ErrorKind::WouldBlock => (),
            _ => panic!("Unexpected"),
        }
    }

    #[test]
    fn clone() {
        let evt = EventFd::new().unwrap();
        let evt_clone = evt.try_clone().unwrap();
        evt.write(923).unwrap();
        assert_eq!(evt_clone.read().unwrap(), 923);
    }
}
<|endoftext|>


<commit_before>use specs::prelude::*;
use engine::components::{Color, Death, FactionId, Hp, HpChange, MovedFlag, Pos, Shape, UnitTypeTag};
use engine::{NeedsKeyInfo, Render};
use engine::resources::Skip;

use scaii_defs::protos::Entity as ScaiiEntity;

#[derive(SystemData)]
pub struct RenderSystemData<'a> {
    complete_rerender: Fetch<'a, NeedsKeyInfo>,
    color: ReadStorage<'a, Color>,
    pos: ReadStorage<'a, Pos>,
    shape: ReadStorage<'a, Shape>,
    hp: ReadStorage<'a, Hp>,
    hp_change: ReadStorage<'a, HpChange>,
    faction: ReadStorage<'a, FactionId>,
    moved: ReadStorage<'a, MovedFlag>,
    death: ReadStorage<'a, Death>,
    u_type: ReadStorage<'a, UnitTypeTag>,
    ids: Entities<'a>,
    skip: Fetch<'a, Skip>,

    out: FetchMut<'a, Render>,
}

#[derive(Default, Copy, Clone, Debug)]
pub struct RenderSystem {}

impl<'a> System<'a> for RenderSystem {
    type SystemData = RenderSystemData<'a>;

    fn run(&mut self, sys_data: Self::SystemData) {
        if sys_data.skip.0 {
            return;
        }
        if sys_data.complete_rerender.0 {
            self.render_all(sys_data);
        } else {
            self.render_delta(sys_data);
        }
    }
}

impl RenderSystem {
    fn render_delta(&mut self, mut sys_data: RenderSystemData) {
        let out = &mut sys_data.out.0;
        out.entities.clear();

        for (pos, id) in (&sys_data.pos, &*sys_data.ids).join() {
            if !(sys_data.moved.get(id).is_some() || sys_data.death.get(id).is_some()
                || sys_data.hp_change.get(id).is_some())
            {
                continue;
            }

            let mut entity = ScaiiEntity {
                id: id.id() as u64,
                pos: Some(pos.to_scaii_pos()),
                delete: sys_data.death.get(id).is_some(),
                shapes: vec![],

                ..ScaiiEntity::default()
            };

            if sys_data.hp_change.get(id).is_some() {
                let hit_points_string = format!("{}",sys_data.hp.get(id).unwrap().curr_hp);
                 println!(" sending hitpoints string in delta {}", hit_points_string);
                entity.float_string_metadata.insert("Hitpoints".to_string(), hit_points_string);
                //entity.float_metadata.insert("Hitpoints".to_string(), sys_data.hp.get(id).unwrap().curr_hp as f32);
            }

            out.entities.push(entity);
        }
    }

    fn render_all(&mut self, mut sys_data: RenderSystemData) {
        sys_data.out.0.entities.clear();

        for (color, pos, shape, id) in (
            &sys_data.color,
            &sys_data.pos,
            &sys_data.shape,
            &*sys_data.ids,
        ).join()
        {
            let entity = self.render_new(&sys_data, id, color, pos, shape);
            sys_data.out.0.entities.push(entity);
        }
    }

        fn render_new(&mut self, sys_data: &RenderSystemData, id: Entity, color: &Color, pos: &Pos, shape: &Shape) -> ScaiiEntity {
            println!("Rendering");
            use std::collections::HashMap;

            let mut scaii_shape = shape.to_scaii_shape(0);

            scaii_shape.color = Some(color.to_scaii_color());
            scaii_shape.relative_pos = Some(Pos::new(0.0, 0.0).to_scaii_pos());

            // TODO: maybe keep these around in a pool to reduce allocations?
            let float_metadata : HashMap<String, f32> = HashMap::with_capacity(1);
            //let mut bool_metadata : HashMap<String, bool> = HashMap::with_capacity(2);
            let mut string_metadata : HashMap<String, String> = HashMap::with_capacity(1);
            let mut bool_string_metadata : HashMap<String, String> = HashMap::with_capacity(2);
            let mut float_string_metadata : HashMap<String, String> = HashMap::with_capacity(1);
            let is_friend = sys_data.faction.get(id).unwrap().0 == 0;
            let hp = sys_data.hp.get(id).unwrap();
            let tag = sys_data.u_type.get(id).unwrap();

            let hit_points_string = format!("{}",hp.curr_hp);
            println!("entity : {:?} HitPoints string : {}", id, hit_points_string);
            float_string_metadata.insert("Hitpoints".to_string(), hit_points_string);
            //float_metadata.insert("Hitpoints".to_string(), hp.curr_hp as f32);
           // bool_metadata.insert("Enemy?".to_string(), not_is_friend);
           // bool_metadata.insert("Friend?".to_string(), is_friend);
           let mut enemy_bool_string = "true".to_string();
           let mut friend_bool_string = "false".to_string();
            if is_friend {
                enemy_bool_string = "false".to_string();
                friend_bool_string = "true".to_string();
            }

            bool_string_metadata.insert("Enemy?".to_string(),enemy_bool_string);
            bool_string_metadata.insert("Friend?".to_string(),friend_bool_string);
            string_metadata.insert("Unit Type".to_string(), tag.0.clone());
            ScaiiEntity {
                shapes: vec![scaii_shape],
                id: id.id() as u64,
                delete: sys_data.death.get(id).is_some(),
                pos: Some(pos.to_scaii_pos()),
                float_string_metadata,
                string_metadata,
                bool_string_metadata,
                ..ScaiiEntity::default()
            }
        }
}<commit_msg>code formatting change<commit_after>use specs::prelude::*;
use engine::components::{Color, Death, FactionId, Hp, HpChange, MovedFlag, Pos, Shape, UnitTypeTag};
use engine::{NeedsKeyInfo, Render};
use engine::resources::Skip;

use scaii_defs::protos::Entity as ScaiiEntity;

#[derive(SystemData)]
pub struct RenderSystemData<'a> {
    complete_rerender: Fetch<'a, NeedsKeyInfo>,
    color: ReadStorage<'a, Color>,
    pos: ReadStorage<'a, Pos>,
    shape: ReadStorage<'a, Shape>,
    hp: ReadStorage<'a, Hp>,
    hp_change: ReadStorage<'a, HpChange>,
    faction: ReadStorage<'a, FactionId>,
    moved: ReadStorage<'a, MovedFlag>,
    death: ReadStorage<'a, Death>,
    u_type: ReadStorage<'a, UnitTypeTag>,
    ids: Entities<'a>,
    skip: Fetch<'a, Skip>,

    out: FetchMut<'a, Render>,
}

#[derive(Default, Copy, Clone, Debug)]
pub struct RenderSystem {}

impl<'a> System<'a> for RenderSystem {
    type SystemData = RenderSystemData<'a>;

    fn run(&mut self, sys_data: Self::SystemData) {
        if sys_data.skip.0 {
            return;
        }
        if sys_data.complete_rerender.0 {
            self.render_all(sys_data);
        } else {
            self.render_delta(sys_data);
        }
    }
}

impl RenderSystem {
    fn render_delta(&mut self, mut sys_data: RenderSystemData) {
        let out = &mut sys_data.out.0;
        out.entities.clear();

        for (pos, id) in (&sys_data.pos, &*sys_data.ids).join() {
            if !(sys_data.moved.get(id).is_some() || sys_data.death.get(id).is_some()
                || sys_data.hp_change.get(id).is_some())
            {
                continue;
            }

            let mut entity = ScaiiEntity {
                id: id.id() as u64,
                pos: Some(pos.to_scaii_pos()),
                delete: sys_data.death.get(id).is_some(),
                shapes: vec![],

                ..ScaiiEntity::default()
            };

            if sys_data.hp_change.get(id).is_some() {
                let hit_points_string = format!("{}",sys_data.hp.get(id).unwrap().curr_hp);
                 println!(" sending hitpoints string in delta {}", hit_points_string);
                entity.float_string_metadata.insert("Hitpoints".to_string(), hit_points_string);
                //entity.float_metadata.insert("Hitpoints".to_string(), sys_data.hp.get(id).unwrap().curr_hp as f32);
            }

            out.entities.push(entity);
        }
    }

    fn render_all(&mut self, mut sys_data: RenderSystemData) {
        sys_data.out.0.entities.clear();

        for (color, pos, shape, id) in (
            &sys_data.color,
            &sys_data.pos,
            &sys_data.shape,
            &*sys_data.ids,
        ).join()
        {
            let entity = self.render_new(&sys_data, id, color, pos, shape);
            sys_data.out.0.entities.push(entity);
        }
    }

        fn render_new(&mut self, sys_data: &RenderSystemData, id: Entity, color: &Color, pos: &Pos, shape: &Shape) -> ScaiiEntity {
            println!("Rendering");
            use std::collections::HashMap;

            let mut scaii_shape = shape.to_scaii_shape(0);

            scaii_shape.color = Some(color.to_scaii_color());
            scaii_shape.relative_pos = Some(Pos::new(0.0, 0.0).to_scaii_pos());

            // TODO: maybe keep these around in a pool to reduce allocations?
            //let float_metadata : HashMap<String, f32> = HashMap::with_capacity(1);
            //let mut bool_metadata : HashMap<String, bool> = HashMap::with_capacity(2);
            let mut string_metadata : HashMap<String, String> = HashMap::with_capacity(1);
            let mut bool_string_metadata : HashMap<String, String> = HashMap::with_capacity(2);
            let mut float_string_metadata : HashMap<String, String> = HashMap::with_capacity(1);
            let is_friend = sys_data.faction.get(id).unwrap().0 == 0;
            let hp = sys_data.hp.get(id).unwrap();
            let tag = sys_data.u_type.get(id).unwrap();

            let hit_points_string = format!("{}",hp.curr_hp);
            println!("entity : {:?} HitPoints string : {}", id, hit_points_string);
            float_string_metadata.insert("Hitpoints".to_string(), hit_points_string);
            //float_metadata.insert("Hitpoints".to_string(), hp.curr_hp as f32);
            // bool_metadata.insert("Enemy?".to_string(), not_is_friend);
            // bool_metadata.insert("Friend?".to_string(), is_friend);
            let mut enemy_bool_string = "true".to_string();
            let mut friend_bool_string = "false".to_string();
            if is_friend {
                enemy_bool_string = "false".to_string();
                friend_bool_string = "true".to_string();
            }

            bool_string_metadata.insert("Enemy?".to_string(),enemy_bool_string);
            bool_string_metadata.insert("Friend?".to_string(),friend_bool_string);
            string_metadata.insert("Unit Type".to_string(), tag.0.clone());
            ScaiiEntity {
                shapes: vec![scaii_shape],
                id: id.id() as u64,
                delete: sys_data.death.get(id).is_some(),
                pos: Some(pos.to_scaii_pos()),
                float_string_metadata,
                string_metadata,
                bool_string_metadata,
                ..ScaiiEntity::default()
            }
        }
}<|endoftext|>


<commit_before>//! Optional handlers for the Iron framework. Requires the `iron-handlers` feature enabled.

use iron::prelude::*;
use iron::middleware::Handler;
use iron::mime::Mime;
use iron::status;
use iron::method;

use std::collections::{HashMap, BTreeMap};

use rustc_serialize::json::{ToJson, Json};

use ::{InputValue, GraphQLType, RootNode, execute};

/// Handler that executes GraphQL queries in the given schema
///
/// The handler responds to GET requests and POST requests only. In GET
/// requests, the query should be supplied in the `query` URL parameter, e.g.
/// `http://localhost:3000/graphql?query={hero{name}}`.
///
/// POST requests support both queries and variables. POST a JSON document to
/// this endpoint containing the field `"query"` and optionally `"variables"`.
/// The variables should be a JSON object containing the variable to value
/// mapping.
pub struct GraphQLHandler<'a, CtxFactory, Query, Mutation, CtxT>
    where CtxFactory: Fn(&mut Request) -> CtxT + Send + Sync + 'static,
          CtxT: 'static,
          Query: GraphQLType<Context=CtxT> + Send + Sync + 'static,
          Mutation: GraphQLType<Context=CtxT> + Send + Sync + 'static,
{
    context_factory: CtxFactory,
    root_node: RootNode<'a, Query, Mutation>,
}

/// Handler that renders GraphiQL - a graphical query editor interface
pub struct GraphiQLHandler {
    graphql_url: String,
}

impl<'a, CtxFactory, Query, Mutation, CtxT>
    GraphQLHandler<'a, CtxFactory, Query, Mutation, CtxT>
    where CtxFactory: Fn(&mut Request) -> CtxT + Send + Sync + 'static,
          CtxT: 'static,
          Query: GraphQLType<Context=CtxT> + Send + Sync + 'static,
          Mutation: GraphQLType<Context=CtxT> + Send + Sync + 'static,
{
    /// Build a new GraphQL handler
    ///
    /// The context factory will receive the Iron request object and is
    /// expected to construct a context object for the given schema. This can
    /// be used to construct e.g. database connections or similar data that
    /// the schema needs to execute the query.
    pub fn new(context_factory: CtxFactory, query: Query, mutation: Mutation) -> Self {
        GraphQLHandler {
            context_factory: context_factory,
            root_node: RootNode::new(query, mutation),
        }
    }

    fn handle_get(&self, req: &mut Request) -> IronResult<Response> {
        let url = req.url.clone().into_generic_url();

        let mut query = None;
        let variables = HashMap::new();

        for (k, v) in url.query_pairs() {
            if k == "query" {
                query = Some(v.into_owned());
            }
        }

        let query = iexpect!(query);

        self.execute(req, &query, &variables)
    }

    fn handle_post(&self, req: &mut Request) -> IronResult<Response> {
        let json_data = itry!(Json::from_reader(&mut req.body));

        let json_obj = match json_data {
            Json::Object(o) => o,
            _ => return Ok(Response::with((status::BadRequest, "No JSON object was decoded"))),
        };

        let mut query = None;
        let mut variables = HashMap::new();

        for (k, v) in json_obj.into_iter() {
            if k == "query" {
                query = v.as_string().map(|s| s.to_owned());
            }
            else if k == "variables" {
                variables = match InputValue::from_json(v).to_object_value() {
                    Some(o) => o.into_iter().map(|(k, v)| (k.to_owned(), v.clone())).collect(),
                    _ => HashMap::new(),
                };
            }
        }

        let query = iexpect!(query);

        self.execute(req, &query, &variables)
    }

    fn execute(&self, req: &mut Request, query: &str, variables: &HashMap<String, InputValue>) -> IronResult<Response> {
        let context = (self.context_factory)(req);
        let result = execute(query, None, &self.root_node, variables, &context);

        let content_type = "application/json".parse::<Mime>().unwrap();
        let mut map = BTreeMap::new();

        match result {
            Ok((result, errors)) => {
                map.insert("data".to_owned(), result.to_json());

                if !errors.is_empty() {
                    map.insert("errors".to_owned(), errors.to_json());
                }

                let data = Json::Object(map);
                let json = data.pretty();

                Ok(Response::with((content_type, status::Ok, json.to_string())))
            }

            Err(err) => {
                map.insert("errors".to_owned(), err.to_json());

                let data = Json::Object(map);
                let json = data.pretty();

                Ok(Response::with((content_type, status::BadRequest, json.to_string())))
            }
        }
    }
}

impl GraphiQLHandler {
    /// Build a new GraphiQL handler targeting the specified URL.
    ///
    /// The provided URL should point to the URL of the attached `GraphQLHandler`. It can be
    /// relative, so a common value could be `"/graphql"`.
    pub fn new(graphql_url: &str) -> GraphiQLHandler {
        GraphiQLHandler {
            graphql_url: graphql_url.to_owned(),
        }
    }
}

impl<'a, CtxFactory, Query, Mutation, CtxT>
    Handler
    for GraphQLHandler<'a, CtxFactory, Query, Mutation, CtxT>
    where CtxFactory: Fn(&mut Request) -> CtxT + Send + Sync + 'static,
          CtxT: 'static,
          Query: GraphQLType<Context=CtxT> + Send + Sync + 'static,
          Mutation: GraphQLType<Context=CtxT> + Send + Sync + 'static, 'a: 'static,
{
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        match req.method {
            method::Get => self.handle_get(req),
            method::Post => self.handle_post(req),
            _ => Ok(Response::with((status::MethodNotAllowed)))
        }
    }
}

impl Handler for GraphiQLHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        let content_type = "text/html".parse::<Mime>().unwrap();

        let stylesheet_source = r#"

        "#;

        let fetcher_source = r#"

        "#;

        let source = format!(r#"
<!DOCTYPE html>

        <title>GraphQL</title>
        {stylesheet_source}
        <link rel="stylesheet" type="text/css" href="//cdnjs.cloudflare.com/ajax/libs/graphiql/0.8.1/graphiql.css">

        {fetcher_source}

"#,
            graphql_url = self.graphql_url,
            stylesheet_source = stylesheet_source,
            fetcher_source = fetcher_source);

        Ok(Response::with((content_type, status::Ok, source)))
    }
}

#[cfg(test)]
mod tests {
    use rustc_serialize::json::Json;

    use iron::prelude::*;
    use iron::status;
    use iron::headers;
    use iron_test::{request, response};
    use iron::{Handler, Headers};

    use ::tests::model::Database;
    use types::scalars::EmptyMutation;

    use super::GraphQLHandler;

    fn context_factory(_: &mut Request) -> Database {
        Database::new()
    }

    fn make_handler() -> Box<Handler> {
        Box::new(GraphQLHandler::new(
            context_factory,
            Database::new(),
            EmptyMutation::<Database>::new(),
        ))
    }

    fn unwrap_json_response(resp: Response) -> Json {
        let result = response::extract_body_to_string(resp);

        Json::from_str(&result).expect("Could not parse JSON object")
    }

    #[test]
    fn test_simple_get() {
        let response = request::get(
            "http://localhost:3000/?query={hero{name}}",
            Headers::new(),
            &make_handler())
            .expect("Unexpected IronError");

        assert_eq!(response.status, Some(status::Ok));
        assert_eq!(response.headers.get::<headers::ContentType>(),
                   Some(&headers::ContentType::json()));

        let json = unwrap_json_response(response);

        assert_eq!(
            json,
            Json::from_str(r#"{"data": {"hero": {"name": "R2-D2"}}}"#)
                .expect("Invalid JSON constant in test"));
    }

    #[test]
    fn test_simple_post() {
        let response = request::post(
            "http://localhost:3000/",
            Headers::new(),
            r#"{"query": "{hero{name}}"}"#,
            &make_handler())
            .expect("Unexpected IronError");

        assert_eq!(response.status, Some(status::Ok));
        assert_eq!(response.headers.get::<headers::ContentType>(),
                   Some(&headers::ContentType::json()));

        let json = unwrap_json_response(response);

        assert_eq!(
            json,
            Json::from_str(r#"{"data": {"hero": {"name": "R2-D2"}}}"#)
                .expect("Invalid JSON constant in test"));
    }

    #[test]
    fn test_unsupported_method() {
        let response = request::options(
            "http://localhost:3000/?query={hero{name}}",
            Headers::new(),
            &make_handler())
            .expect("Unexpected IronError");

        assert_eq!(response.status, Some(status::MethodNotAllowed));
    }
}
<commit_msg>Fix graphQLFetcher body serialization<commit_after>//! Optional handlers for the Iron framework. Requires the `iron-handlers` feature enabled.

use iron::prelude::*;
use iron::middleware::Handler;
use iron::mime::Mime;
use iron::status;
use iron::method;

use std::collections::{HashMap, BTreeMap};

use rustc_serialize::json::{ToJson, Json};

use ::{InputValue, GraphQLType, RootNode, execute};

/// Handler that executes GraphQL queries in the given schema
///
/// The handler responds to GET requests and POST requests only. In GET
/// requests, the query should be supplied in the `query` URL parameter, e.g.
/// `http://localhost:3000/graphql?query={hero{name}}`.
///
/// POST requests support both queries and variables. POST a JSON document to
/// this endpoint containing the field `"query"` and optionally `"variables"`.
/// The variables should be a JSON object containing the variable to value
/// mapping.
pub struct GraphQLHandler<'a, CtxFactory, Query, Mutation, CtxT>
    where CtxFactory: Fn(&mut Request) -> CtxT + Send + Sync + 'static,
          CtxT: 'static,
          Query: GraphQLType<Context=CtxT> + Send + Sync + 'static,
          Mutation: GraphQLType<Context=CtxT> + Send + Sync + 'static,
{
    context_factory: CtxFactory,
    root_node: RootNode<'a, Query, Mutation>,
}

/// Handler that renders GraphiQL - a graphical query editor interface
pub struct GraphiQLHandler {
    graphql_url: String,
}

impl<'a, CtxFactory, Query, Mutation, CtxT>
    GraphQLHandler<'a, CtxFactory, Query, Mutation, CtxT>
    where CtxFactory: Fn(&mut Request) -> CtxT + Send + Sync + 'static,
          CtxT: 'static,
          Query: GraphQLType<Context=CtxT> + Send + Sync + 'static,
          Mutation: GraphQLType<Context=CtxT> + Send + Sync + 'static,
{
    /// Build a new GraphQL handler
    ///
    /// The context factory will receive the Iron request object and is
    /// expected to construct a context object for the given schema. This can
    /// be used to construct e.g. database connections or similar data that
    /// the schema needs to execute the query.
    pub fn new(context_factory: CtxFactory, query: Query, mutation: Mutation) -> Self {
        GraphQLHandler {
            context_factory: context_factory,
            root_node: RootNode::new(query, mutation),
        }
    }

    fn handle_get(&self, req: &mut Request) -> IronResult<Response> {
        let url = req.url.clone().into_generic_url();

        let mut query = None;
        let variables = HashMap::new();

        for (k, v) in url.query_pairs() {
            if k == "query" {
                query = Some(v.into_owned());
            }
        }

        let query = iexpect!(query);

        self.execute(req, &query, &variables)
    }

    fn handle_post(&self, req: &mut Request) -> IronResult<Response> {
        let json_data = itry!(Json::from_reader(&mut req.body));

        let json_obj = match json_data {
            Json::Object(o) => o,
            _ => return Ok(Response::with((status::BadRequest, "No JSON object was decoded"))),
        };

        let mut query = None;
        let mut variables = HashMap::new();

        for (k, v) in json_obj.into_iter() {
            if k == "query" {
                query = v.as_string().map(|s| s.to_owned());
            }
            else if k == "variables" {
                variables = match InputValue::from_json(v).to_object_value() {
                    Some(o) => o.into_iter().map(|(k, v)| (k.to_owned(), v.clone())).collect(),
                    _ => HashMap::new(),
                };
            }
        }

        let query = iexpect!(query);

        self.execute(req, &query, &variables)
    }

    fn execute(&self, req: &mut Request, query: &str, variables: &HashMap<String, InputValue>) -> IronResult<Response> {
        let context = (self.context_factory)(req);
        let result = execute(query, None, &self.root_node, variables, &context);

        let content_type = "application/json".parse::<Mime>().unwrap();
        let mut map = BTreeMap::new();

        match result {
            Ok((result, errors)) => {
                map.insert("data".to_owned(), result.to_json());

                if !errors.is_empty() {
                    map.insert("errors".to_owned(), errors.to_json());
                }

                let data = Json::Object(map);
                let json = data.pretty();

                Ok(Response::with((content_type, status::Ok, json.to_string())))
            }

            Err(err) => {
                map.insert("errors".to_owned(), err.to_json());

                let data = Json::Object(map);
                let json = data.pretty();

                Ok(Response::with((content_type, status::BadRequest, json.to_string())))
            }
        }
    }
}

impl GraphiQLHandler {
    /// Build a new GraphiQL handler targeting the specified URL.
    ///
    /// The provided URL should point to the URL of the attached `GraphQLHandler`. It can be
    /// relative, so a common value could be `"/graphql"`.
    pub fn new(graphql_url: &str) -> GraphiQLHandler {
        GraphiQLHandler {
            graphql_url: graphql_url.to_owned(),
        }
    }
}

impl<'a, CtxFactory, Query, Mutation, CtxT>
    Handler
    for GraphQLHandler<'a, CtxFactory, Query, Mutation, CtxT>
    where CtxFactory: Fn(&mut Request) -> CtxT + Send + Sync + 'static,
          CtxT: 'static,
          Query: GraphQLType<Context=CtxT> + Send + Sync + 'static,
          Mutation: GraphQLType<Context=CtxT> + Send + Sync + 'static, 'a: 'static,
{
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        match req.method {
            method::Get => self.handle_get(req),
            method::Post => self.handle_post(req),
            _ => Ok(Response::with((status::MethodNotAllowed)))
        }
    }
}

impl Handler for GraphiQLHandler {
    fn handle(&self, _: &mut Request) -> IronResult<Response> {
        let content_type = "text/html".parse::<Mime>().unwrap();

        let stylesheet_source = r#"

        "#;

        let fetcher_source = r#"

        "#;

        let source = format!(r#"
<!DOCTYPE html>

        <title>GraphQL</title>
        {stylesheet_source}
        <link rel="stylesheet" type="text/css" href="//cdnjs.cloudflare.com/ajax/libs/graphiql/0.8.1/graphiql.css">

        {fetcher_source}

"#,
            graphql_url = self.graphql_url,
            stylesheet_source = stylesheet_source,
            fetcher_source = fetcher_source);

        Ok(Response::with((content_type, status::Ok, source)))
    }
}

#[cfg(test)]
mod tests {
    use rustc_serialize::json::Json;

    use iron::prelude::*;
    use iron::status;
    use iron::headers;
    use iron_test::{request, response};
    use iron::{Handler, Headers};

    use ::tests::model::Database;
    use types::scalars::EmptyMutation;

    use super::GraphQLHandler;

    fn context_factory(_: &mut Request) -> Database {
        Database::new()
    }

    fn make_handler() -> Box<Handler> {
        Box::new(GraphQLHandler::new(
            context_factory,
            Database::new(),
            EmptyMutation::<Database>::new(),
        ))
    }

    fn unwrap_json_response(resp: Response) -> Json {
        let result = response::extract_body_to_string(resp);

        Json::from_str(&result).expect("Could not parse JSON object")
    }

    #[test]
    fn test_simple_get() {
        let response = request::get(
            "http://localhost:3000/?query={hero{name}}",
            Headers::new(),
            &make_handler())
            .expect("Unexpected IronError");

        assert_eq!(response.status, Some(status::Ok));
        assert_eq!(response.headers.get::<headers::ContentType>(),
                   Some(&headers::ContentType::json()));

        let json = unwrap_json_response(response);

        assert_eq!(
            json,
            Json::from_str(r#"{"data": {"hero": {"name": "R2-D2"}}}"#)
                .expect("Invalid JSON constant in test"));
    }

    #[test]
    fn test_simple_post() {
        let response = request::post(
            "http://localhost:3000/",
            Headers::new(),
            r#"{"query": "{hero{name}}"}"#,
            &make_handler())
            .expect("Unexpected IronError");

        assert_eq!(response.status, Some(status::Ok));
        assert_eq!(response.headers.get::<headers::ContentType>(),
                   Some(&headers::ContentType::json()));

        let json = unwrap_json_response(response);

        assert_eq!(
            json,
            Json::from_str(r#"{"data": {"hero": {"name": "R2-D2"}}}"#)
                .expect("Invalid JSON constant in test"));
    }

    #[test]
    fn test_unsupported_method() {
        let response = request::options(
            "http://localhost:3000/?query={hero{name}}",
            Headers::new(),
            &make_handler())
            .expect("Unexpected IronError");

        assert_eq!(response.status, Some(status::MethodNotAllowed));
    }
}
<|endoftext|>