<commit_before>use Prefab;
use calx::{self, hex_neighbors, Dijkstra};
use euclid::vec2;
use form::{self, Form};
use image::{self, GenericImage, SubImage};
use location::{Location, Portal, Sector};
use mapgen::{self, DigCavesGen, Size2D};
use rand::{self, Rand, Rng, SeedableRng};
use serde;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::io::Cursor;
use std::slice;
use terrain::Terrain;
use world::Loadout;

/// Static generated world.
pub struct Worldgen {
    seed: u32,
    terrain: HashMap<Location, Terrain>,
    portals: HashMap<Location, Portal>,
    spawns: Vec<(Location, Loadout)>,
    player_entry: Location,
}

impl Worldgen {
    pub fn new(seed: u32) -> Worldgen {
        let mut ret = Worldgen {
            seed: seed,
            terrain: HashMap::new(),
            portals: HashMap::new(),
            spawns: Vec::new(),
            player_entry: Location::new(25, 0, 0),
        };

        let mut rng: ::Rng = SeedableRng::from_seed([seed, seed, seed, seed]);

        let mut cave_entrance = Location::new(9, 0, 0);
        ret.cave_entrance(cave_entrance);
        ret.terrain
            .insert(cave_entrance + vec2(1, 1), Terrain::Gate);

        for depth in 1..11 {
            let up_stairs;
            let down_stairs;
            let mut spawns = Vec::new();

            {
                let mut digger = SectorDigger::new(&mut ret, Sector::new(0, 0, depth as i16));
                let gen = DigCavesGen::new(digger.domain());
                gen.dig(&mut rng, &mut digger);

                up_stairs = digger.up_portal.expect("Mapgen didn't create stairs up");
                down_stairs = digger
                    .down_portal
                    .expect("Mapgen didn't create stairs down");

                // Spawn
                digger.clear_spawns_near_entrance(12);

                let mut spawn_locs = rand::sample(&mut rng, digger.spawn_region.iter(), 20);
                let n_spawns = spawn_locs.len();

                let items = Form::filter(|f| f.is_item() && f.at_depth(depth));
                for &loc in spawn_locs.drain(0..n_spawns / 2) {
                    spawns.push((
                        loc,
                        form::rand(&mut rng, &items)
                            .expect("No item spawn")
                            .loadout
                            .clone(),
                    ))
                }

                let mobs = Form::filter(|f| f.is_mob() && f.at_depth(depth));
                for &loc in spawn_locs {
                    spawns.push((
                        loc,
                        form::rand(&mut rng, &mobs)
                            .expect("No mob spawn")
                            .loadout
                            .clone(),
                    ))
                }
            }

            ret.spawns.extend(spawns.into_iter());
            ret.portal(cave_entrance, up_stairs + vec2(1, 1));
            ret.portal(up_stairs, cave_entrance - vec2(1, 1));
            cave_entrance = down_stairs;

            // TODO: Generator needs an option to not generate stairs down on bottom level
        }

        ret.load_map_bitmap(
            Location::new(0, 0, 0),
            include_bytes!("../assets/overland.png"),
        );

        ret
    }

    fn load_prefab(&mut self, origin: Location, prefab: &Prefab) {
        for (&p, &(ref terrain, ref entities)) in prefab.iter() {
            let loc = origin + p;

            self.terrain.insert(loc, *terrain);

            for spawn in entities.iter() {
                if spawn == "player" {
                    self.player_entry = loc;
                } else {
                    let form = Form::named(spawn)
                        .expect(&format!("Bad prefab: Form '{}' not found!", spawn));
                    self.spawns.push((loc, form.loadout.clone()));
                }
            }
        }
    }

    fn load_map_bitmap(&mut self, origin: Location, data: &[u8]) {
        let mut image = image::load(Cursor::new(data), image::ImageFormat::PNG).unwrap();
        // Skip the bottom horizontal line, it's used to store metadata pixels.
        let (w, h) = (image.width(), image.height());
        let input_map = SubImage::new(&mut image, 0, 0, w, h - 1);

        let prefab = calx::Prefab::parse(input_map).expect("Invalid overworld image map");

        self.terrain.extend(
            prefab
                .into_iter()
                .filter_map(|(p, c)| Terrain::from_color(c).map(|t| (origin + p, t))),
        );
    }

    pub fn seed(&self) -> u32 { self.seed }

    pub fn get_terrain(&self, loc: Location) -> Terrain {
        if let Some(&t) = self.terrain.get(&loc) {
            t
        } else {
            self.default_terrain(loc)
        }
    }

    fn default_terrain(&self, loc: Location) -> Terrain {
        use Terrain::*;
        if loc.z == 0 {
            /*
            match loc.noise() {
                n if n > 0.8 => Tree,
                n if n > -0.8 => Grass,
                _ => Water,
            }
            */
            Ground
        } else {
            Rock
        }
    }

    pub fn get_portal(&self, loc: Location) -> Option<Location> {
        self.portals.get(&loc).map(|&p| loc + p)
    }

    pub fn spawns(&self) -> slice::Iter<(Location, Loadout)> { self.spawns.iter() }

    pub fn player_entry(&self) -> Location { self.player_entry }

    /// Make a cave entrance going down.
    fn cave_entrance(&mut self, loc: Location) {
        const DOWNBOUND_ENCLOSURE: [(i32, i32); 5] = [(1, 0), (0, 1), (2, 1), (1, 2), (2, 2)];

        for &v in &DOWNBOUND_ENCLOSURE {
            let loc = loc + vec2(v.0, v.1);
            self.terrain.insert(loc, Terrain::Rock);
        }

        self.terrain.insert(loc, Terrain::Ground);
        self.terrain.insert(loc + vec2(1, 1), Terrain::Ground);
    }

    /// Punch a (one-way) portal between two points.
    fn portal(&mut self, origin: Location, destination: Location) {
        self.portals
            .insert(origin, Portal::new(origin, destination));
    }
}

impl serde::Serialize for Worldgen {
    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
        self.seed.serialize(s)
    }
}

impl<'a> serde::Deserialize<'a> for Worldgen {
    fn deserialize<D: serde::Deserializer<'a>>(d: D) -> Result<Self, D::Error> {
        Ok(Worldgen::new(serde::Deserialize::deserialize(d)?))
    }
}

struct SectorDigger<'a> {
    worldgen: &'a mut Worldgen,
    sector: Sector,
    up_portal: Option<Location>,
    down_portal: Option<Location>,
    spawn_region: BTreeSet<Location>,
}

impl<'a> SectorDigger<'a> {
    fn new(worldgen: &'a mut Worldgen, sector: Sector) -> SectorDigger<'a> {
        SectorDigger {
            worldgen,
            sector,
            up_portal: None,
            down_portal: None,
            spawn_region: BTreeSet::new(),
        }
    }

    // TODO return impl
    fn domain(&self) -> Vec<mapgen::Point2D> {
        let mut ret = Vec::new();

        let sector_origin = self.sector.origin();
        let mapgen_origin = mapgen::Point2D::zero();

        for loc in self.sector.iter() {
            let pos = mapgen_origin
                + sector_origin
                    .v2_at(loc)
                    .expect("Sector points are not Euclidean");

            // Okay, this part is a bit hairy, hang on.
            // Sectors are arranged in a rectangular grid, so there are some cells where you can
            // step diagonally across two sector boundaries. However, we want to pretend that
            // sectors are only connected in the four cardinal directions, so we omit these cells
            // from the domain to prevent surprise holes between two diagonally adjacent shafts.
            if Self::is_next_to_diagonal_sector(loc) {
                continue;
            }

            ret.push(pos);
        }

        ret
    }

    fn loc(&self, pos: mapgen::Point2D) -> Location { self.sector.origin() + pos.to_vector() }

    fn set(&mut self, loc: Location, terrain: Terrain) {
        if terrain != Terrain::Gate {
            debug_assert_ne!(self.up_portal, Some(loc), "mapgen overwriting exit");
            debug_assert_ne!(self.down_portal, Some(loc), "mapgen overwriting exit");
        }

        // FIXME: Mapgen currently trips this with exit placement
        // debug_assert!(
        //     !Self::is_next_to_diagonal_sector(loc) &&
        //     self.sector.iter().find(|x| x == &loc).is_some(),
        //     "Setting a location outsides sector area");

        self.worldgen.terrain.insert(loc, terrain);
    }

    fn clear_spawns_near_entrance(&mut self, range: u32) {
        for loc in Dijkstra::new(
            vec![self.up_portal.expect("No entrance generated")],
            |loc| {
                self.worldgen
                    .terrain
                    .get(loc)
                    .cloned()
                    .unwrap_or(Terrain::Empty)
                    .is_open()
            },
            10_000,
        ).weights
            .iter()
            .filter_map(|(loc, &w)| if w <= range { Some(loc) } else { None })
        {
            self.spawn_region.remove(loc);
        }
    }

    fn is_next_to_diagonal_sector(loc: Location) -> bool {
        hex_neighbors(loc)
            .map(|x| x.sector().taxicab_distance(loc.sector()))
            .any(|d| d > 1)
    }
}

impl<'a> mapgen::Dungeon for SectorDigger<'a> {
    type Prefab = Room;

    /// Return a random prefab room.
    fn sample_prefab<R: Rng>(&mut self, rng: &mut R) -> Self::Prefab { rng.gen() }

    /// Add a large open continuous region to dungeon.
    fn dig_chamber<I: IntoIterator<Item = mapgen::Point2D>>(&mut self, area: I) {
        for pos in area {
            let loc = self.loc(pos);
            self.spawn_region.insert(loc);
            self.set(loc, Terrain::Ground);
        }
    }

    /// Add a narrow corridor to dungeon.
    fn dig_corridor<I: IntoIterator<Item = mapgen::Point2D>>(&mut self, path: I) {
        for pos in path {
            let loc = self.loc(pos);
            // Do not spawn things in corridors.
            self.spawn_region.remove(&loc);
            self.set(loc, Terrain::Ground);
        }
    }

    fn add_prefab(&mut self, prefab: &Self::Prefab, pos: mapgen::Point2D) {
        // TODO
        unimplemented!();
    }

    fn add_door(&mut self, pos: mapgen::Point2D) {
        let loc = self.loc(pos);
        self.spawn_region.remove(&loc);
        self.set(loc, Terrain::Door);
    }

    fn add_up_stairs(&mut self, pos: mapgen::Point2D) {
        let loc = self.loc(pos);
        self.up_portal = Some(loc);
        self.spawn_region.remove(&loc);
        // TODO: Add debug assertion that up-stairs are dug in the upper wall of a room
        self.set(loc, Terrain::Gate);
    }

    fn add_down_stairs(&mut self, pos: mapgen::Point2D) {
        let loc = self.loc(pos);
        self.down_portal = Some(loc);
        self.spawn_region.remove(&loc);
        self.set(loc, Terrain::Gate);

        // Visual hack to make the down-stairs show up better, carve out the rock blob that would
        // be drawn partially in front of it on screen.
        self.set(loc + vec2(1, 1), Terrain::Empty);
        // TODO: Add a debug assertion here to check that the exit location has proper enclosure to
        // begin with, both it and the rock cell in front of it need to be surrounded by undug
        // front tiles.
        self.spawn_region.remove(&(loc + vec2(1, 1))); // Just in case
    }
}

struct Room {
    size: Size2D,
}

impl mapgen::Prefab for Room {
    fn contains(&self, pos: mapgen::Point2D) -> bool {
        pos.x >= 0 && pos.y >= 0 && pos.x < self.size.width && pos.y < self.size.height
    }

    fn can_make_door(&self, pos: mapgen::Point2D) -> bool {
        let on_x_wall = pos.x == 0 || pos.x == self.size.width - 1;
        let on_y_wall = pos.y == 0 || pos.y == self.size.height - 1;

        // Must touch one wall, but touching both makes it a corner and we don't want rooms there.
        on_x_wall ^ on_y_wall
    }

    fn size(&self) -> Size2D { self.size }
}

impl Rand for Room {
    fn rand<R: Rng>(rng: &mut R) -> Self {
        Room {
            size: Size2D::new(rng.gen_range(3, 10), rng.gen_range(3, 10)),
        }
    }
}

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

    #[test]
    fn test_determinism() {
        use rand::{self, Rand};

        let mut rng = rand::thread_rng();

        let seed = u32::rand(&mut rng);
        let mut gen = Worldgen::new(seed);

        // Build the value repeatedly using the same seed and see that they are all equal.
        for _ in 1..4 {
            let second = Worldgen::new(seed);

            assert_eq!(gen.seed, second.seed);
            // These can make huge printouts so don't use assert_eq that would try to print them to
            // stdout
            assert!(gen.terrain == second.terrain);
            assert!(gen.portals == second.portals);
            assert_eq!(gen.player_entry, second.player_entry);

            gen = second;
        }
    }
}
<commit_msg>Put dungeon entrance in overland map<commit_after>use Prefab;
use calx::{self, hex_neighbors, Dijkstra};
use euclid::vec2;
use form::{self, Form};
use image::{self, GenericImage, SubImage};
use location::{Location, Portal, Sector};
use mapgen::{self, DigCavesGen, Size2D};
use rand::{self, Rand, Rng, SeedableRng};
use serde;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::io::Cursor;
use std::slice;
use terrain::Terrain;
use world::Loadout;

/// Static generated world.
pub struct Worldgen {
    seed: u32,
    terrain: HashMap<Location, Terrain>,
    portals: HashMap<Location, Portal>,
    spawns: Vec<(Location, Loadout)>,
    player_entry: Location,
}

impl Worldgen {
    pub fn new(seed: u32) -> Worldgen {
        let mut ret = Worldgen {
            seed: seed,
            terrain: HashMap::new(),
            portals: HashMap::new(),
            spawns: Vec::new(),
            player_entry: Location::new(25, 0, 0),
        };

        ret.load_map_bitmap(
            Location::new(0, 0, 0),
            include_bytes!("../assets/overland.png"),
        );

        let mut rng: ::Rng = SeedableRng::from_seed([seed, seed, seed, seed]);

        let mut cave_entrance = Location::new(40, 12, 0);
        ret.cave_entrance(cave_entrance);
        ret.terrain
            .insert(cave_entrance, Terrain::Gate);
        ret.terrain
            .insert(cave_entrance + vec2(1, 1), Terrain::Empty);

        for depth in 1..11 {
            let up_stairs;
            let down_stairs;
            let mut spawns = Vec::new();

            {
                let mut digger = SectorDigger::new(&mut ret, Sector::new(0, 0, depth as i16));
                let gen = DigCavesGen::new(digger.domain());
                gen.dig(&mut rng, &mut digger);

                up_stairs = digger.up_portal.expect("Mapgen didn't create stairs up");
                down_stairs = digger
                    .down_portal
                    .expect("Mapgen didn't create stairs down");

                // Spawn
                digger.clear_spawns_near_entrance(12);

                let mut spawn_locs = rand::sample(&mut rng, digger.spawn_region.iter(), 20);
                let n_spawns = spawn_locs.len();

                let items = Form::filter(|f| f.is_item() && f.at_depth(depth));
                for &loc in spawn_locs.drain(0..n_spawns / 2) {
                    spawns.push((
                        loc,
                        form::rand(&mut rng, &items)
                            .expect("No item spawn")
                            .loadout
                            .clone(),
                    ))
                }

                let mobs = Form::filter(|f| f.is_mob() && f.at_depth(depth));
                for &loc in spawn_locs {
                    spawns.push((
                        loc,
                        form::rand(&mut rng, &mobs)
                            .expect("No mob spawn")
                            .loadout
                            .clone(),
                    ))
                }
            }

            ret.spawns.extend(spawns.into_iter());
            ret.portal(cave_entrance, up_stairs + vec2(1, 1));
            ret.portal(up_stairs, cave_entrance - vec2(1, 1));
            cave_entrance = down_stairs;

            // TODO: Generator needs an option to not generate stairs down on bottom level
        }

        ret
    }

    fn load_prefab(&mut self, origin: Location, prefab: &Prefab) {
        for (&p, &(ref terrain, ref entities)) in prefab.iter() {
            let loc = origin + p;

            self.terrain.insert(loc, *terrain);

            for spawn in entities.iter() {
                if spawn == "player" {
                    self.player_entry = loc;
                } else {
                    let form = Form::named(spawn)
                        .expect(&format!("Bad prefab: Form '{}' not found!", spawn));
                    self.spawns.push((loc, form.loadout.clone()));
                }
            }
        }
    }

    fn load_map_bitmap(&mut self, origin: Location, data: &[u8]) {
        let mut image = image::load(Cursor::new(data), image::ImageFormat::PNG).unwrap();
        // Skip the bottom horizontal line, it's used to store metadata pixels.
        let (w, h) = (image.width(), image.height());
        let input_map = SubImage::new(&mut image, 0, 0, w, h - 1);

        let prefab = calx::Prefab::parse(input_map).expect("Invalid overworld image map");

        self.terrain.extend(
            prefab
                .into_iter()
                .filter_map(|(p, c)| Terrain::from_color(c).map(|t| (origin + p, t))),
        );
    }

    pub fn seed(&self) -> u32 { self.seed }

    pub fn get_terrain(&self, loc: Location) -> Terrain {
        if let Some(&t) = self.terrain.get(&loc) {
            t
        } else {
            self.default_terrain(loc)
        }
    }

    fn default_terrain(&self, loc: Location) -> Terrain {
        use Terrain::*;
        if loc.z == 0 {
            /*
            match loc.noise() {
                n if n > 0.8 => Tree,
                n if n > -0.8 => Grass,
                _ => Water,
            }
            */
            Ground
        } else {
            Rock
        }
    }

    pub fn get_portal(&self, loc: Location) -> Option<Location> {
        self.portals.get(&loc).map(|&p| loc + p)
    }

    pub fn spawns(&self) -> slice::Iter<(Location, Loadout)> { self.spawns.iter() }

    pub fn player_entry(&self) -> Location { self.player_entry }

    /// Make a cave entrance going down.
    fn cave_entrance(&mut self, loc: Location) {
        const DOWNBOUND_ENCLOSURE: [(i32, i32); 5] = [(1, 0), (0, 1), (2, 1), (1, 2), (2, 2)];

        for &v in &DOWNBOUND_ENCLOSURE {
            let loc = loc + vec2(v.0, v.1);
            self.terrain.insert(loc, Terrain::Rock);
        }

        self.terrain.insert(loc, Terrain::Ground);
        self.terrain.insert(loc + vec2(1, 1), Terrain::Ground);
    }

    /// Punch a (one-way) portal between two points.
    fn portal(&mut self, origin: Location, destination: Location) {
        self.portals
            .insert(origin, Portal::new(origin, destination));
    }
}

impl serde::Serialize for Worldgen {
    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
        self.seed.serialize(s)
    }
}

impl<'a> serde::Deserialize<'a> for Worldgen {
    fn deserialize<D: serde::Deserializer<'a>>(d: D) -> Result<Self, D::Error> {
        Ok(Worldgen::new(serde::Deserialize::deserialize(d)?))
    }
}

struct SectorDigger<'a> {
    worldgen: &'a mut Worldgen,
    sector: Sector,
    up_portal: Option<Location>,
    down_portal: Option<Location>,
    spawn_region: BTreeSet<Location>,
}

impl<'a> SectorDigger<'a> {
    fn new(worldgen: &'a mut Worldgen, sector: Sector) -> SectorDigger<'a> {
        SectorDigger {
            worldgen,
            sector,
            up_portal: None,
            down_portal: None,
            spawn_region: BTreeSet::new(),
        }
    }

    // TODO return impl
    fn domain(&self) -> Vec<mapgen::Point2D> {
        let mut ret = Vec::new();

        let sector_origin = self.sector.origin();
        let mapgen_origin = mapgen::Point2D::zero();

        for loc in self.sector.iter() {
            let pos = mapgen_origin
                + sector_origin
                    .v2_at(loc)
                    .expect("Sector points are not Euclidean");

            // Okay, this part is a bit hairy, hang on.
            // Sectors are arranged in a rectangular grid, so there are some cells where you can
            // step diagonally across two sector boundaries. However, we want to pretend that
            // sectors are only connected in the four cardinal directions, so we omit these cells
            // from the domain to prevent surprise holes between two diagonally adjacent shafts.
            if Self::is_next_to_diagonal_sector(loc) {
                continue;
            }

            ret.push(pos);
        }

        ret
    }

    fn loc(&self, pos: mapgen::Point2D) -> Location { self.sector.origin() + pos.to_vector() }

    fn set(&mut self, loc: Location, terrain: Terrain) {
        if terrain != Terrain::Gate {
            debug_assert_ne!(self.up_portal, Some(loc), "mapgen overwriting exit");
            debug_assert_ne!(self.down_portal, Some(loc), "mapgen overwriting exit");
        }

        // FIXME: Mapgen currently trips this with exit placement
        // debug_assert!(
        //     !Self::is_next_to_diagonal_sector(loc) &&
        //     self.sector.iter().find(|x| x == &loc).is_some(),
        //     "Setting a location outsides sector area");

        self.worldgen.terrain.insert(loc, terrain);
    }

    fn clear_spawns_near_entrance(&mut self, range: u32) {
        for loc in Dijkstra::new(
            vec![self.up_portal.expect("No entrance generated")],
            |loc| {
                self.worldgen
                    .terrain
                    .get(loc)
                    .cloned()
                    .unwrap_or(Terrain::Empty)
                    .is_open()
            },
            10_000,
        ).weights
            .iter()
            .filter_map(|(loc, &w)| if w <= range { Some(loc) } else { None })
        {
            self.spawn_region.remove(loc);
        }
    }

    fn is_next_to_diagonal_sector(loc: Location) -> bool {
        hex_neighbors(loc)
            .map(|x| x.sector().taxicab_distance(loc.sector()))
            .any(|d| d > 1)
    }
}

impl<'a> mapgen::Dungeon for SectorDigger<'a> {
    type Prefab = Room;

    /// Return a random prefab room.
    fn sample_prefab<R: Rng>(&mut self, rng: &mut R) -> Self::Prefab { rng.gen() }

    /// Add a large open continuous region to dungeon.
    fn dig_chamber<I: IntoIterator<Item = mapgen::Point2D>>(&mut self, area: I) {
        for pos in area {
            let loc = self.loc(pos);
            self.spawn_region.insert(loc);
            self.set(loc, Terrain::Ground);
        }
    }

    /// Add a narrow corridor to dungeon.
    fn dig_corridor<I: IntoIterator<Item = mapgen::Point2D>>(&mut self, path: I) {
        for pos in path {
            let loc = self.loc(pos);
            // Do not spawn things in corridors.
            self.spawn_region.remove(&loc);
            self.set(loc, Terrain::Ground);
        }
    }

    fn add_prefab(&mut self, prefab: &Self::Prefab, pos: mapgen::Point2D) {
        // TODO
        unimplemented!();
    }

    fn add_door(&mut self, pos: mapgen::Point2D) {
        let loc = self.loc(pos);
        self.spawn_region.remove(&loc);
        self.set(loc, Terrain::Door);
    }

    fn add_up_stairs(&mut self, pos: mapgen::Point2D) {
        let loc = self.loc(pos);
        self.up_portal = Some(loc);
        self.spawn_region.remove(&loc);
        // TODO: Add debug assertion that up-stairs are dug in the upper wall of a room
        self.set(loc, Terrain::Gate);
    }

    fn add_down_stairs(&mut self, pos: mapgen::Point2D) {
        let loc = self.loc(pos);
        self.down_portal = Some(loc);
        self.spawn_region.remove(&loc);
        self.set(loc, Terrain::Gate);

        // Visual hack to make the down-stairs show up better, carve out the rock blob that would
        // be drawn partially in front of it on screen.
        self.set(loc + vec2(1, 1), Terrain::Empty);
        // TODO: Add a debug assertion here to check that the exit location has proper enclosure to
        // begin with, both it and the rock cell in front of it need to be surrounded by undug
        // front tiles.
        self.spawn_region.remove(&(loc + vec2(1, 1))); // Just in case
    }
}

struct Room {
    size: Size2D,
}

impl mapgen::Prefab for Room {
    fn contains(&self, pos: mapgen::Point2D) -> bool {
        pos.x >= 0 && pos.y >= 0 && pos.x < self.size.width && pos.y < self.size.height
    }

    fn can_make_door(&self, pos: mapgen::Point2D) -> bool {
        let on_x_wall = pos.x == 0 || pos.x == self.size.width - 1;
        let on_y_wall = pos.y == 0 || pos.y == self.size.height - 1;

        // Must touch one wall, but touching both makes it a corner and we don't want rooms there.
        on_x_wall ^ on_y_wall
    }

    fn size(&self) -> Size2D { self.size }
}

impl Rand for Room {
    fn rand<R: Rng>(rng: &mut R) -> Self {
        Room {
            size: Size2D::new(rng.gen_range(3, 10), rng.gen_range(3, 10)),
        }
    }
}

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

    #[test]
    fn test_determinism() {
        use rand::{self, Rand};

        let mut rng = rand::thread_rng();

        let seed = u32::rand(&mut rng);
        let mut gen = Worldgen::new(seed);

        // Build the value repeatedly using the same seed and see that they are all equal.
        for _ in 1..4 {
            let second = Worldgen::new(seed);

            assert_eq!(gen.seed, second.seed);
            // These can make huge printouts so don't use assert_eq that would try to print them to
            // stdout
            assert!(gen.terrain == second.terrain);
            assert!(gen.portals == second.portals);
            assert_eq!(gen.player_entry, second.player_entry);

            gen = second;
        }
    }
}
<|endoftext|>


<commit_before>use byteorder::{ReadBytesExt, WriteBytesExt};
use std::error::Error;
use std::io::prelude::*;

use backend::{Backend, HasRawValue};
use deserialize::{self, FromSql};
use serialize::{self, IsNull, Output, ToSql};
use sql_types;

impl<DB> FromSql<sql_types::SmallInt, DB> for i16
where
    DB: Backend + for<'a> HasRawValue<'a, RawValue = &'a [u8]>,
{
    fn from_sql(bytes: Option<&[u8]>) -> deserialize::Result<Self> {
        let mut bytes = not_none!(bytes);
        debug_assert!(
            bytes.len() <= 2,
            "Received more than 2 bytes decoding i16. \
             Was an Integer expression accidentally identified as SmallInt?"
        );
        debug_assert!(
            bytes.len() >= 2,
            "Received fewer than 2 bytes decoding i16. \
             Was an expression of a different type accidentally identified \
             as SmallInt?"
        );
        bytes
            .read_i16::<DB::ByteOrder>()
            .map_err(|e| Box::new(e) as Box<Error + Send + Sync>)
    }
}

impl<DB: Backend> ToSql<sql_types::SmallInt, DB> for i16 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, DB>) -> serialize::Result {
        out.write_i16::<DB::ByteOrder>(*self)
            .map(|_| IsNull::No)
            .map_err(|e| Box::new(e) as Box<Error + Send + Sync>)
    }
}

impl<DB> FromSql<sql_types::Integer, DB> for i32
where
    DB: Backend + for<'a> HasRawValue<'a, RawValue = &'a [u8]>,
{
    fn from_sql(bytes: Option<&[u8]>) -> deserialize::Result<Self> {
        let mut bytes = not_none!(bytes);
        debug_assert!(
            bytes.len() <= 4,
            "Received more than 4 bytes decoding i32. \
             Was a BigInteger expression accidentally identified as Integer?"
        );
        debug_assert!(
            bytes.len() >= 4,
            "Received fewer than 4 bytes decoding i32. \
             Was a SmallInteger expression accidentally identified as Integer?"
        );
        bytes
            .read_i32::<DB::ByteOrder>()
            .map_err(|e| Box::new(e) as Box<Error + Send + Sync>)
    }
}

impl<DB: Backend> ToSql<sql_types::Integer, DB> for i32 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, DB>) -> serialize::Result {
        out.write_i32::<DB::ByteOrder>(*self)
            .map(|_| IsNull::No)
            .map_err(|e| Box::new(e) as Box<Error + Send + Sync>)
    }
}

impl<DB> FromSql<sql_types::BigInt, DB> for i64
where
    DB: Backend + for<'a> HasRawValue<'a, RawValue = &'a [u8]>,
{
    fn from_sql(bytes: Option<&[u8]>) -> deserialize::Result<Self> {
        let mut bytes = not_none!(bytes);
        debug_assert!(
            bytes.len() <= 8,
            "Received more than 8 bytes decoding i64. \
             Was an expression of a different type misidentified as BigInteger?"
        );
        debug_assert!(
            bytes.len() >= 8,
            "Received fewer than 8 bytes decoding i64. \
             Was an Integer expression misidentified as BigInteger?"
        );
        bytes
            .read_i64::<DB::ByteOrder>()
            .map_err(|e| Box::new(e) as Box<Error + Send + Sync>)
    }
}

impl<DB: Backend> ToSql<sql_types::BigInt, DB> for i64 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, DB>) -> serialize::Result {
        out.write_i64::<DB::ByteOrder>(*self)
            .map(|_| IsNull::No)
            .map_err(|e| Box::new(e) as Box<Error + Send + Sync>)
    }
}
<commit_msg>Correct suggested type to match diesel::sql_types<commit_after>use byteorder::{ReadBytesExt, WriteBytesExt};
use std::error::Error;
use std::io::prelude::*;

use backend::{Backend, HasRawValue};
use deserialize::{self, FromSql};
use serialize::{self, IsNull, Output, ToSql};
use sql_types;

impl<DB> FromSql<sql_types::SmallInt, DB> for i16
where
    DB: Backend + for<'a> HasRawValue<'a, RawValue = &'a [u8]>,
{
    fn from_sql(bytes: Option<&[u8]>) -> deserialize::Result<Self> {
        let mut bytes = not_none!(bytes);
        debug_assert!(
            bytes.len() <= 2,
            "Received more than 2 bytes decoding i16. \
             Was an Integer expression accidentally identified as SmallInt?"
        );
        debug_assert!(
            bytes.len() >= 2,
            "Received fewer than 2 bytes decoding i16. \
             Was an expression of a different type accidentally identified \
             as SmallInt?"
        );
        bytes
            .read_i16::<DB::ByteOrder>()
            .map_err(|e| Box::new(e) as Box<Error + Send + Sync>)
    }
}

impl<DB: Backend> ToSql<sql_types::SmallInt, DB> for i16 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, DB>) -> serialize::Result {
        out.write_i16::<DB::ByteOrder>(*self)
            .map(|_| IsNull::No)
            .map_err(|e| Box::new(e) as Box<Error + Send + Sync>)
    }
}

impl<DB> FromSql<sql_types::Integer, DB> for i32
where
    DB: Backend + for<'a> HasRawValue<'a, RawValue = &'a [u8]>,
{
    fn from_sql(bytes: Option<&[u8]>) -> deserialize::Result<Self> {
        let mut bytes = not_none!(bytes);
        debug_assert!(
            bytes.len() <= 4,
            "Received more than 4 bytes decoding i32. \
             Was a BigInt expression accidentally identified as Integer?"
        );
        debug_assert!(
            bytes.len() >= 4,
            "Received fewer than 4 bytes decoding i32. \
             Was a SmallInt expression accidentally identified as Integer?"
        );
        bytes
            .read_i32::<DB::ByteOrder>()
            .map_err(|e| Box::new(e) as Box<Error + Send + Sync>)
    }
}

impl<DB: Backend> ToSql<sql_types::Integer, DB> for i32 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, DB>) -> serialize::Result {
        out.write_i32::<DB::ByteOrder>(*self)
            .map(|_| IsNull::No)
            .map_err(|e| Box::new(e) as Box<Error + Send + Sync>)
    }
}

impl<DB> FromSql<sql_types::BigInt, DB> for i64
where
    DB: Backend + for<'a> HasRawValue<'a, RawValue = &'a [u8]>,
{
    fn from_sql(bytes: Option<&[u8]>) -> deserialize::Result<Self> {
        let mut bytes = not_none!(bytes);
        debug_assert!(
            bytes.len() <= 8,
            "Received more than 8 bytes decoding i64. \
             Was an expression of a different type misidentified as BigInt?"
        );
        debug_assert!(
            bytes.len() >= 8,
            "Received fewer than 8 bytes decoding i64. \
             Was an Integer expression misidentified as BigInt?"
        );
        bytes
            .read_i64::<DB::ByteOrder>()
            .map_err(|e| Box::new(e) as Box<Error + Send + Sync>)
    }
}

impl<DB: Backend> ToSql<sql_types::BigInt, DB> for i64 {
    fn to_sql<W: Write>(&self, out: &mut Output<W, DB>) -> serialize::Result {
        out.write_i64::<DB::ByteOrder>(*self)
            .map(|_| IsNull::No)
            .map_err(|e| Box::new(e) as Box<Error + Send + Sync>)
    }
}
<|endoftext|>


<commit_before>use crate::route_guide::Feature;
use crate::route_guide::Point;
use crate::route_guide::Rectangle;
use crate::route_guide::RouteNote;
use crate::route_guide::RouteSummary;
use crate::route_guide_grpc::RouteGuide;

use futures::future;

use futures::stream;
use futures::TryFutureExt;
use futures::TryStreamExt;

use futures::stream::StreamExt;

use grpc::Metadata;
use grpc::ServerHandlerContext;
use grpc::ServerRequest;
use grpc::ServerRequestSingle;
use grpc::ServerResponseSink;
use grpc::ServerResponseUnarySink;
use json::JsonValue;
use std::collections::HashMap;
use std::f64;
use std::fs;
use std::io::Read;
use std::path::Path;
use std::sync::Arc;
use std::sync::Mutex;
use std::task::Poll;
use std::time::Instant;

// https://github.com/grpc/grpc-go/blob/master/examples/route_guide/server/server.go
#[derive(Default)]
pub struct RouteGuideImpl {
    saved_features: Arc<Vec<Feature>>,
    route_notes: Arc<Mutex<HashMap<String, Vec<RouteNote>>>>,
}

impl RouteGuideImpl {
    pub fn new_and_load_db() -> RouteGuideImpl {
        RouteGuideImpl {
            saved_features: Arc::new(load_features(Path::new(ROUTE_GUIDE_DB_PATH))),
            route_notes: Default::default(),
        }
    }
}

impl RouteGuide for RouteGuideImpl {
    fn get_feature(
        &self,
        _o: ServerHandlerContext,
        req: ServerRequestSingle<Point>,
        resp: ServerResponseUnarySink<Feature>,
    ) -> grpc::Result<()> {
        for feature in &*self.saved_features {
            if feature.get_location() == &req.message {
                return resp.finish(feature.clone());
            }
        }

        resp.finish(Feature {
            location: Some(req.message).into(),
            ..Default::default()
        })
    }

    fn list_features(
        &self,
        o: ServerHandlerContext,
        mut req: ServerRequestSingle<Rectangle>,
        resp: ServerResponseSink<Feature>,
    ) -> grpc::Result<()> {
        let req = req.take_message();
        let saved_features = self.saved_features.clone();
        let stream = stream::iter(0..saved_features.len())
            .filter_map(move |i| {
                let feature = saved_features[i].clone();
                future::ready({
                    if in_range(feature.get_location(), &req) {
                        Some(feature)
                    } else {
                        None
                    }
                })
            })
            .map(Ok);
        o.pump(stream, resp);
        Ok(())
    }

    fn record_route(
        &self,
        o: ServerHandlerContext,
        req: ServerRequest<Point>,
        resp: ServerResponseUnarySink<RouteSummary>,
    ) -> grpc::Result<()> {
        let start_time = Instant::now();

        struct State {
            point_count: u32,
            feature_count: u32,
            distance: u32,
            last_point: Option<Point>,
        }

        let state = State {
            point_count: 0,
            feature_count: 0,
            distance: 0,
            last_point: None,
        };

        let saved_features = self.saved_features.clone();

        let f = req
            .into_stream()
            .try_fold(state, move |mut state, point| {
                state.point_count += 1;
                for feature in &saved_features[..] {
                    if feature.get_location() == &point {
                        state.feature_count += 1;
                    }
                }
                if let Some(last_point) = &state.last_point {
                    state.distance += calc_distance(last_point, &point);
                }
                state.last_point = Some(point);
                future::ok::<_, grpc::Error>(state)
            })
            .map_ok(move |state| RouteSummary {
                point_count: state.point_count as i32,
                feature_count: state.feature_count as i32,
                distance: state.distance as i32,
                elapsed_time: start_time.elapsed().as_secs() as i32,
                ..Default::default()
            });
        o.pump_future(f, resp);
        Ok(())
    }

    fn route_chat(
        &self,
        o: ServerHandlerContext,
        req: ServerRequest<RouteNote>,
        mut resp: ServerResponseSink<RouteNote>,
    ) -> grpc::Result<()> {
        let mut req = req.into_stream();

        let route_notes_map = self.route_notes.clone();

        o.spawn_poll_fn(move |cx| {
            loop {
                // Wait until resp is writable
                if let Poll::Pending = resp.poll(cx)? {
                    return Poll::Pending;
                }

                match req.poll_next_unpin(cx)? {
                    Poll::Pending => return Poll::Pending,
                    Poll::Ready(Some(note)) => {
                        let key = serialize(note.get_location());

                        let mut route_notes_map = route_notes_map.lock().unwrap();

                        let route_notes = route_notes_map.entry(key).or_insert(Vec::new());
                        route_notes.push(note);

                        for note in route_notes {
                            resp.send_data(note.clone())?;
                        }
                    }
                    Poll::Ready(None) => {
                        resp.send_trailers(Metadata::new())?;
                        return Poll::Ready(Ok(()));
                    }
                }
            }
        });
        Ok(())
    }
}

fn in_range(point: &Point, rect: &Rectangle) -> bool {
    let left = f64::min(
        rect.get_lo().longitude as f64,
        rect.get_hi().longitude as f64,
    );
    let right = f64::max(
        rect.get_lo().longitude as f64,
        rect.get_hi().longitude as f64,
    );
    let top = f64::max(rect.get_lo().latitude as f64, rect.get_hi().latitude as f64);
    let bottom = f64::min(rect.get_lo().latitude as f64, rect.get_hi().latitude as f64);

    point.longitude as f64 >= left
        && point.longitude as f64 <= right
        && point.latitude as f64 >= bottom
        && point.latitude as f64 <= top
}

fn to_radians(num: f64) -> f64 {
    num * f64::consts::PI / 180.
}

fn calc_distance(p1: &Point, p2: &Point) -> u32 {
    let cord_factor: f64 = 1e7;
    let r = 6371000.; // earth radius in metres
    let lat1 = to_radians(p1.latitude as f64 / cord_factor);
    let lat2 = to_radians(p2.latitude as f64 / cord_factor);
    let lng1 = to_radians(p1.longitude as f64 / cord_factor);
    let lng2 = to_radians(p2.longitude as f64 / cord_factor);
    let dlat = lat2 - lat1;
    let dlng = lng2 - lng1;

    let a = f64::sin(dlat / 2.) * f64::sin(dlat / 2.)
        + f64::cos(lat1) * f64::cos(lat2) * f64::sin(dlng / 2.) * f64::sin(dlng / 2.);
    let c = 2. * f64::atan2(f64::sqrt(a), f64::sqrt(1. - a));

    let distance = r * c;
    distance as u32
}

fn serialize(point: &Point) -> String {
    format!("{} {}", point.latitude, point.longitude)
}

const ROUTE_GUIDE_DB_PATH: &str = "testdata/route_guide_db.json";

fn load_features(path: &Path) -> Vec<Feature> {
    let mut file = fs::File::open(path).expect("open");
    let mut s = String::new();
    file.read_to_string(&mut s).expect("read");

    // TODO: use protobuf mapper when new version is released

    let json_value = json::parse(&s).expect("parse json");
    let array = match json_value {
        JsonValue::Array(array) => array,
        _ => panic!(),
    };

    array
        .into_iter()
        .map(|item| {
            let object = match item {
                JsonValue::Object(object) => object,
                _ => panic!(),
            };

            let location = match object.get("location").expect("location") {
                JsonValue::Object(object) => object,
                _ => panic!(),
            };

            Feature {
                name: object
                    .get("name")
                    .expect("name")
                    .as_str()
                    .expect("unwrap")
                    .to_owned(),
                location: Some(Point {
                    latitude: location
                        .get("latitude")
                        .expect("latitude")
                        .as_i32()
                        .unwrap(),
                    longitude: location
                        .get("longitude")
                        .expect("longitude")
                        .as_i32()
                        .unwrap(),
                    ..Default::default()
                })
                .into(),
                ..Default::default()
            }
        })
        .collect()
}

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

    #[test]
    fn test_load_features() {
        let features = load_features(Path::new(ROUTE_GUIDE_DB_PATH));
        assert!(features.len() > 0);
    }
}
<commit_msg>More async/await in route guide example<commit_after>use crate::route_guide::Feature;
use crate::route_guide::Point;
use crate::route_guide::Rectangle;
use crate::route_guide::RouteNote;
use crate::route_guide::RouteSummary;
use crate::route_guide_grpc::RouteGuide;

use futures::future;

use futures::stream;
use futures::TryFutureExt;
use futures::TryStreamExt;

use futures::stream::StreamExt;

use grpc::Metadata;
use grpc::ServerHandlerContext;
use grpc::ServerRequest;
use grpc::ServerRequestSingle;
use grpc::ServerResponseSink;
use grpc::ServerResponseUnarySink;
use json::JsonValue;
use std::collections::HashMap;
use std::f64;
use std::fs;
use std::io::Read;
use std::path::Path;
use std::sync::Arc;
use std::sync::Mutex;
use std::task::Poll;
use std::time::Instant;

// https://github.com/grpc/grpc-go/blob/master/examples/route_guide/server/server.go
#[derive(Default)]
pub struct RouteGuideImpl {
    saved_features: Arc<Vec<Feature>>,
    route_notes: Arc<Mutex<HashMap<String, Vec<RouteNote>>>>,
}

impl RouteGuideImpl {
    pub fn new_and_load_db() -> RouteGuideImpl {
        RouteGuideImpl {
            saved_features: Arc::new(load_features(Path::new(ROUTE_GUIDE_DB_PATH))),
            route_notes: Default::default(),
        }
    }
}

impl RouteGuide for RouteGuideImpl {
    fn get_feature(
        &self,
        _o: ServerHandlerContext,
        req: ServerRequestSingle<Point>,
        resp: ServerResponseUnarySink<Feature>,
    ) -> grpc::Result<()> {
        for feature in &*self.saved_features {
            if feature.get_location() == &req.message {
                return resp.finish(feature.clone());
            }
        }

        resp.finish(Feature {
            location: Some(req.message).into(),
            ..Default::default()
        })
    }

    fn list_features(
        &self,
        o: ServerHandlerContext,
        mut req: ServerRequestSingle<Rectangle>,
        mut resp: ServerResponseSink<Feature>,
    ) -> grpc::Result<()> {
        let req = req.take_message();
        let saved_features = self.saved_features.clone();
        o.spawn(|| {
            for feature in &saved_features[..] {
                resp.ready().await?;
                resp.send_data(feature.clone())?;
            }
            resp.send_trailers(Metadata::new())
        });
        Ok(())
    }

    fn record_route(
        &self,
        o: ServerHandlerContext,
        req: ServerRequest<Point>,
        resp: ServerResponseUnarySink<RouteSummary>,
    ) -> grpc::Result<()> {
        let start_time = Instant::now();

        struct State {
            point_count: u32,
            feature_count: u32,
            distance: u32,
            last_point: Option<Point>,
        }

        let mut state = State {
            point_count: 0,
            feature_count: 0,
            distance: 0,
            last_point: None,
        };

        let saved_features = self.saved_features.clone();

        o.spawn(|| {
            let mut stream = req.into_stream();
            while let Some(point) = stream.next() {
                let point = point?;
                state.point_count += 1;
                for feature in &saved_features[..] {
                    if feature.get_location() == &point {
                        state.feature_count += 1;
                    }
                }
                if let Some(last_point) = &state.last_point {
                    state.distance += calc_distance(last_point, &point);
                }
                state.last_point = Some(point);
            }

            resp.finish(RouteSummary {
                point_count: state.point_count as i32,
                feature_count: state.feature_count as i32,
                distance: state.distance as i32,
                elapsed_time: start_time.elapsed().as_secs() as i32,
                ..Default::default()
            })
        });

        Ok(())
    }

    fn route_chat(
        &self,
        o: ServerHandlerContext,
        req: ServerRequest<RouteNote>,
        mut resp: ServerResponseSink<RouteNote>,
    ) -> grpc::Result<()> {
        let mut req = req.into_stream();

        let route_notes_map = self.route_notes.clone();

        o.spawn_poll_fn(move |cx| {
            loop {
                // Wait until resp is writable
                if let Poll::Pending = resp.poll(cx)? {
                    return Poll::Pending;
                }

                match req.poll_next_unpin(cx)? {
                    Poll::Pending => return Poll::Pending,
                    Poll::Ready(Some(note)) => {
                        let key = serialize(note.get_location());

                        let mut route_notes_map = route_notes_map.lock().unwrap();

                        let route_notes = route_notes_map.entry(key).or_insert(Vec::new());
                        route_notes.push(note);

                        for note in route_notes {
                            resp.send_data(note.clone())?;
                        }
                    }
                    Poll::Ready(None) => {
                        resp.send_trailers(Metadata::new())?;
                        return Poll::Ready(Ok(()));
                    }
                }
            }
        });
        Ok(())
    }
}

fn in_range(point: &Point, rect: &Rectangle) -> bool {
    let left = f64::min(
        rect.get_lo().longitude as f64,
        rect.get_hi().longitude as f64,
    );
    let right = f64::max(
        rect.get_lo().longitude as f64,
        rect.get_hi().longitude as f64,
    );
    let top = f64::max(rect.get_lo().latitude as f64, rect.get_hi().latitude as f64);
    let bottom = f64::min(rect.get_lo().latitude as f64, rect.get_hi().latitude as f64);

    point.longitude as f64 >= left
        && point.longitude as f64 <= right
        && point.latitude as f64 >= bottom
        && point.latitude as f64 <= top
}

fn to_radians(num: f64) -> f64 {
    num * f64::consts::PI / 180.
}

fn calc_distance(p1: &Point, p2: &Point) -> u32 {
    let cord_factor: f64 = 1e7;
    let r = 6371000.; // earth radius in metres
    let lat1 = to_radians(p1.latitude as f64 / cord_factor);
    let lat2 = to_radians(p2.latitude as f64 / cord_factor);
    let lng1 = to_radians(p1.longitude as f64 / cord_factor);
    let lng2 = to_radians(p2.longitude as f64 / cord_factor);
    let dlat = lat2 - lat1;
    let dlng = lng2 - lng1;

    let a = f64::sin(dlat / 2.) * f64::sin(dlat / 2.)
        + f64::cos(lat1) * f64::cos(lat2) * f64::sin(dlng / 2.) * f64::sin(dlng / 2.);
    let c = 2. * f64::atan2(f64::sqrt(a), f64::sqrt(1. - a));

    let distance = r * c;
    distance as u32
}

fn serialize(point: &Point) -> String {
    format!("{} {}", point.latitude, point.longitude)
}

const ROUTE_GUIDE_DB_PATH: &str = "testdata/route_guide_db.json";

fn load_features(path: &Path) -> Vec<Feature> {
    let mut file = fs::File::open(path).expect("open");
    let mut s = String::new();
    file.read_to_string(&mut s).expect("read");

    // TODO: use protobuf mapper when new version is released

    let json_value = json::parse(&s).expect("parse json");
    let array = match json_value {
        JsonValue::Array(array) => array,
        _ => panic!(),
    };

    array
        .into_iter()
        .map(|item| {
            let object = match item {
                JsonValue::Object(object) => object,
                _ => panic!(),
            };

            let location = match object.get("location").expect("location") {
                JsonValue::Object(object) => object,
                _ => panic!(),
            };

            Feature {
                name: object
                    .get("name")
                    .expect("name")
                    .as_str()
                    .expect("unwrap")
                    .to_owned(),
                location: Some(Point {
                    latitude: location
                        .get("latitude")
                        .expect("latitude")
                        .as_i32()
                        .unwrap(),
                    longitude: location
                        .get("longitude")
                        .expect("longitude")
                        .as_i32()
                        .unwrap(),
                    ..Default::default()
                })
                .into(),
                ..Default::default()
            }
        })
        .collect()
}

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

    #[test]
    fn test_load_features() {
        let features = load_features(Path::new(ROUTE_GUIDE_DB_PATH));
        assert!(features.len() > 0);
    }
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

extern crate chrono;
extern crate env_logger;
extern crate futures;
extern crate openssl;
extern crate trust_dns;
extern crate trust_dns_compatibility;

use std::env;
use std::fs::File;
use std::io;
use std::io::Read;
#[cfg(not(feature = "none"))]
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::str::FromStr;

#[cfg(not(feature = "none"))]
use chrono::Duration;
use futures::Stream;
use openssl::rsa::Rsa;

#[cfg(not(feature = "none"))]
use trust_dns::client::Client;
use trust_dns::client::{ClientConnection, SyncClient};
#[cfg(not(feature = "none"))]
use trust_dns::udp::UdpClientConnection;
#[cfg(not(feature = "none"))]
use trust_dns::op::ResponseCode;
use trust_dns::rr::Name;
#[cfg(not(feature = "none"))]
use trust_dns::rr::{DNSClass, RData, Record, RecordType};
use trust_dns::rr::dnssec::{Algorithm, KeyPair, Signer};
use trust_dns::rr::rdata::key::{KeyUsage, KEY};
#[cfg(not(feature = "none"))]
use trust_dns_compatibility::named_process;

#[cfg(not(feature = "none"))]
#[test]
#[allow(unused)]
fn test_get() {
    let (process, port) = named_process();
    let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), port);
    let conn = UdpClientConnection::new(socket).unwrap();
    let client = SyncClient::new(conn);

    let name = Name::parse("www.example.com.", None).unwrap();
    let result = client
        .query(&name, DNSClass::IN, RecordType::A)
        .expect("query failed");
    assert_eq!(result.response_code(), ResponseCode::NoError);
    assert_eq!(result.answers().len(), 1);
    assert_eq!(result.answers()[0].rr_type(), RecordType::A);

    let rdata = result.answers()[0].rdata();
    if let &RData::A(address) = rdata {
        assert_eq!(address, Ipv4Addr::new(127, 0, 0, 1));
    } else {
        assert!(false);
    }
}

#[allow(unused)]
fn create_sig0_ready_client<CC>(conn: CC) -> SyncClient<CC>
where
    CC: ClientConnection,
    <CC as ClientConnection>::MessageStream: Stream<Item = Vec<u8>, Error = io::Error> + 'static,
{
    let server_path = env::var("TDNS_SERVER_SRC_ROOT").unwrap_or(".".to_owned());
    let mut pem = File::open(format!(
        "{}/../compatibility-tests/tests/conf/Kupdate.example.com.+008+56935.pem",
        server_path
    )).expect("could not find pem file");

    let mut pem_buf = Vec::<u8>::new();
    pem.read_to_end(&mut pem_buf).expect("failed to read pem");
    let rsa = Rsa::private_key_from_pem(&pem_buf).expect("something wrong with key from pem");
    let key = KeyPair::from_rsa(rsa).unwrap();
    let sig0key = KEY::new(
        Default::default(),
        KeyUsage::Entity,
        Default::default(),
        Default::default(),
        Algorithm::RSASHA256,
        key.to_public_bytes().unwrap(),
    );

    let signer = Signer::sig0(
        sig0key,
        key,
        Name::from_str("update.example.com").unwrap(),
    );

    assert_eq!(signer.calculate_key_tag().unwrap(), 56935_u16);

    SyncClient::with_signer(conn, signer)
}

#[cfg(not(feature = "none"))]
#[test]
#[allow(unused)]
fn test_create() {
    let (process, port) = named_process();
    let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), port);
    let conn = UdpClientConnection::new(socket).unwrap();

    let client = create_sig0_ready_client(conn);
    let origin = Name::from_labels(vec!["example", "com"]);

    // create a record
    let mut record = Record::with(
        Name::from_labels(vec!["new", "example", "com"]),
        RecordType::A,
        Duration::minutes(5).num_seconds() as u32,
    );
    record.set_rdata(RData::A(Ipv4Addr::new(100, 10, 100, 10)));

    let result = client
        .create(record.clone(), origin.clone())
        .expect("create failed");
    assert_eq!(result.response_code(), ResponseCode::NoError);
    let result = client
        .query(record.name(), record.dns_class(), record.rr_type())
        .expect("query failed");
    assert_eq!(result.response_code(), ResponseCode::NoError);
    assert_eq!(result.answers().len(), 1);
    assert_eq!(result.answers()[0], record);

    // trying to create again should error
    // TODO: it would be cool to make this
    let result = client
        .create(record.clone(), origin.clone())
        .expect("create failed");
    assert_eq!(result.response_code(), ResponseCode::YXRRSet);

    // will fail if already set and not the same value.
    let mut record = record.clone();
    record.set_rdata(RData::A(Ipv4Addr::new(101, 11, 101, 11)));

    let result = client
        .create(record.clone(), origin.clone())
        .expect("create failed");
    assert_eq!(result.response_code(), ResponseCode::YXRRSet);
}

// #[test]
// fn test_update() {
//     // named_process();
// }

// #[test]
// fn test_delete() {
//     // named_process();
// }
<commit_msg>fix compatibility tests<commit_after>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

extern crate chrono;
extern crate env_logger;
extern crate futures;
extern crate openssl;
extern crate trust_dns;
extern crate trust_dns_compatibility;

use std::env;
use std::fs::File;
use std::io;
use std::io::Read;
#[cfg(not(feature = "none"))]
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::str::FromStr;

#[cfg(not(feature = "none"))]
use chrono::Duration;
use futures::Stream;
use openssl::rsa::Rsa;

#[cfg(not(feature = "none"))]
use trust_dns::client::Client;
use trust_dns::client::{ClientConnection, SyncClient};
#[cfg(not(feature = "none"))]
use trust_dns::udp::UdpClientConnection;
#[cfg(not(feature = "none"))]
use trust_dns::op::ResponseCode;
use trust_dns::rr::Name;
#[cfg(not(feature = "none"))]
use trust_dns::rr::{DNSClass, RData, Record, RecordType};
use trust_dns::rr::dnssec::{Algorithm, KeyPair, Signer};
use trust_dns::rr::rdata::key::{KeyUsage, KEY};
#[cfg(not(feature = "none"))]
use trust_dns_compatibility::named_process;

#[cfg(not(feature = "none"))]
#[test]
#[allow(unused)]
fn test_get() {
    let (process, port) = named_process();
    let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), port);
    let conn = UdpClientConnection::new(socket).unwrap();
    let client = SyncClient::new(conn);

    let name = Name::from_utf8("www.example.com.").unwrap();
    let result = client
        .query(&name, DNSClass::IN, RecordType::A)
        .expect("query failed");
    assert_eq!(result.response_code(), ResponseCode::NoError);
    assert_eq!(result.answers().len(), 1);
    assert_eq!(result.answers()[0].rr_type(), RecordType::A);

    let rdata = result.answers()[0].rdata();
    if let &RData::A(address) = rdata {
        assert_eq!(address, Ipv4Addr::new(127, 0, 0, 1));
    } else {
        assert!(false);
    }
}

#[allow(unused)]
fn create_sig0_ready_client<CC>(conn: CC) -> SyncClient<CC>
where
    CC: ClientConnection,
    <CC as ClientConnection>::MessageStream: Stream<Item = Vec<u8>, Error = io::Error> + 'static,
{
    let server_path = env::var("TDNS_SERVER_SRC_ROOT").unwrap_or(".".to_owned());
    let mut pem = File::open(format!(
        "{}/../compatibility-tests/tests/conf/Kupdate.example.com.+008+56935.pem",
        server_path
    )).expect("could not find pem file");

    let mut pem_buf = Vec::<u8>::new();
    pem.read_to_end(&mut pem_buf).expect("failed to read pem");
    let rsa = Rsa::private_key_from_pem(&pem_buf).expect("something wrong with key from pem");
    let key = KeyPair::from_rsa(rsa).unwrap();
    let sig0key = KEY::new(
        Default::default(),
        KeyUsage::Entity,
        Default::default(),
        Default::default(),
        Algorithm::RSASHA256,
        key.to_public_bytes().unwrap(),
    );

    let signer = Signer::sig0(
        sig0key,
        key,
        Name::from_str("update.example.com").unwrap(),
    );

    assert_eq!(signer.calculate_key_tag().unwrap(), 56935_u16);

    SyncClient::with_signer(conn, signer)
}

#[cfg(not(feature = "none"))]
#[test]
#[allow(unused)]
fn test_create() {
    let (process, port) = named_process();
    let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), port);
    let conn = UdpClientConnection::new(socket).unwrap();

    let client = create_sig0_ready_client(conn);
    let origin = Name::from_utf8("example.com.").unwrap();

    // create a record
    let mut record = Record::with(
        Name::from_utf8("new.example.com.").unwrap(),
        RecordType::A,
        Duration::minutes(5).num_seconds() as u32,
    );
    record.set_rdata(RData::A(Ipv4Addr::new(100, 10, 100, 10)));

    let result = client
        .create(record.clone(), origin.clone())
        .expect("create failed");
    assert_eq!(result.response_code(), ResponseCode::NoError);
    let result = client
        .query(record.name(), record.dns_class(), record.rr_type())
        .expect("query failed");
    assert_eq!(result.response_code(), ResponseCode::NoError);
    assert_eq!(result.answers().len(), 1);
    assert_eq!(result.answers()[0], record);

    // trying to create again should error
    // TODO: it would be cool to make this
    let result = client
        .create(record.clone(), origin.clone())
        .expect("create failed");
    assert_eq!(result.response_code(), ResponseCode::YXRRSet);

    // will fail if already set and not the same value.
    let mut record = record.clone();
    record.set_rdata(RData::A(Ipv4Addr::new(101, 11, 101, 11)));

    let result = client
        .create(record.clone(), origin.clone())
        .expect("create failed");
    assert_eq!(result.response_code(), ResponseCode::YXRRSet);
}

// #[test]
// fn test_update() {
//     // named_process();
// }

// #[test]
// fn test_delete() {
//     // named_process();
// }
<|endoftext|>


<commit_before>#[cfg(windows)]
use std::env;
use std::{path::Path,
          result};

#[cfg(windows)]
use crate::common::cli::{DEFAULT_BINLINK_DIR,
                         FS_ROOT};
use crate::{common::ui::{UIReader,
                         UIWriter,
                         UI},
            hcore::{crypto::SigKeyPair,
                    env as henv,
                    package::ident,
                    Error::InvalidOrigin}};
#[cfg(windows)]
use std::ptr;
use url::Url;
#[cfg(windows)]
use widestring::WideCString;
#[cfg(windows)]
use winapi::shared::minwindef::LPARAM;
#[cfg(windows)]
use winapi::um::winuser::{self,
                          HWND_BROADCAST,
                          SMTO_ABORTIFHUNG,
                          WM_SETTINGCHANGE};
#[cfg(windows)]
use winreg::enums::{HKEY_LOCAL_MACHINE,
                    KEY_ALL_ACCESS,
                    KEY_READ};
#[cfg(windows)]
use winreg::RegKey;

use crate::{command,
            config,
            error::Result,
            AUTH_TOKEN_ENVVAR,
            BLDR_URL_ENVVAR,
            CTL_SECRET_ENVVAR,
            ORIGIN_ENVVAR};

pub fn start(ui: &mut UI, cache_path: &Path) -> Result<()> {
    ui.br()?;
    ui.title("Habitat CLI Setup")?;
    ui.para("Welcome to hab setup. Let's get started.")?;

    ui.heading("Habitat Builder Instance")?;
    ui.para(
        "Habitat packages can be stored in either the public builder instance \
         https://bldr.habitat.sh or in an on-premises builder depot instance. If \
         you do not set a builder URL now, the `hab` CLI will default to using \
         the public builder instance. This can be overridden at any time after setup.",
    )?;
    if ask_default_builder_instance(ui)? {
        ui.br()?;
        ui.para("Enter the url of your builder instance. The default is \
                 https://bldr.habitat.sh. The configured endpoint can be overridden any time \
                 with a `HAB_BLDR_URL` envvar or a --url flag on the cli.")?;
        let mut url = prompt_url(ui)?;

        while valid_url(&url).is_err() {
            ui.br()?;
            ui.fatal(&format!("{}: is invalid, please provide a valid url", url))?;
            ui.br()?;

            url = prompt_url(ui)?;
        }

        write_cli_config_bldr_url(&url)?;
    } else {
        ui.br()?;
        ui.para("No worries, should you need to use a different bldr instance you can set a \
                 `HAB_BLDR_URL` envvar or pass the `--url` flag to the cli!")?;
    }

    ui.heading("Set up a default origin")?;
    ui.para("Every package in Habitat belongs to an origin, which indicates the person or \
             organization responsible for maintaining that package. Each origin also has a key \
             used to cryptographically sign packages in that origin.")?;
    ui.para("Selecting a default origin tells package building operations such as 'hab pkg \
             build' what key should be used to sign the packages produced. If you do not set a \
             default origin now, you will have to tell package building commands each time what \
             origin to use.")?;
    ui.para("For more information on origins and how they are used in building packages, please \
             consult the docs at https://www.habitat.sh/docs/create-packages-build/")?;
    if ask_default_origin(ui)? {
        ui.br()?;
        ui.para("Enter the name of your origin. If you plan to publish your packages publicly, \
                 we recommend that you select one that is not already in use on the Habitat \
                 build service found at https://bldr.habitat.sh/.")?;
        ui.para("Origins must begin with a lowercase letter or number. Allowed characters \
                 include lowercase letters, numbers, _, -. No more than 255 characters.")?;
        let mut origin = prompt_origin(ui)?;

        while !ident::is_valid_origin_name(&origin) {
            ui.br()?;
            ui.fatal(&format!("{}", InvalidOrigin(origin)))?;
            ui.br()?;

            origin = prompt_origin(ui)?;
        }
        write_cli_config_origin(&origin)?;
        ui.br()?;
        if is_origin_in_cache(&origin, cache_path) {
            ui.para(&format!("You already have an origin key for {} created and installed. \
                              Great work!",
                             &origin))?;
        } else {
            ui.heading("Create origin key pair")?;
            ui.para(&format!("It doesn't look like you have a signing key for the origin `{}'. \
                              Without it, you won't be able to build new packages successfully.",
                             &origin))?;
            ui.para("You can either create a new signing key now, or, if you are building \
                     packages for an origin that already exists, ask the owner to give you the \
                     signing key.")?;
            ui.para("For more information on the use of origin keys, please consult the \
                     documentation at https://www.habitat.sh/docs/concepts-keys/#origin-keys")?;
            if ask_create_origin(ui, &origin)? {
                create_origin(ui, &origin, cache_path)?;
            } else {
                ui.para(&format!("You might want to create an origin key later with: `hab \
                                  origin key generate {}'",
                                 &origin))?;
            }
        }
    } else {
        ui.para("Okay, maybe another time.")?;
    }
    ui.heading("Habitat Personal Access Token")?;
    ui.para("While you can perform tasks like building and running Habitat packages without \
             needing to authenticate with Builder, some operations like uploading your packages \
             to Builder, or checking the status of your build jobs from the Habitat client will \
             require you to use an access token.")?;
    ui.para("The Habitat Personal Access Token can be generated via the Builder  Profile page \
             (https://bldr.habitat.sh/#/profile). Once you have generated your token, you can \
             enter it here.")?;
    ui.para("If you would like to save your token for use by the Habitat client, please enter \
             your access token. Otherwise, just enter No.")?;
    ui.para(
        "For more information on using Builder, please read the \
         documentation at https://www.habitat.sh/docs/using-builder/",
    )?;
    if ask_default_auth_token(ui)? {
        ui.br()?;
        ui.para("Enter your Habitat Personal Access Token.")?;
        let auth_token = prompt_auth_token(ui)?;
        write_cli_config_auth_token(&auth_token)?;
    } else {
        ui.para("Okay, maybe another time.")?;
    }
    if ask_default_ctl_secret(ui)? {
        ui.br()?;
        ui.para("Enter your Habitat Supervisor CtlGateway secret.")?;
        let ctl_secret = prompt_ctl_secret(ui)?;
        write_cli_config_ctl_secret(&ctl_secret)?;
    } else {
        ui.para("Okay, maybe another time.")?;
    }
    #[cfg(windows)]
    {
        let binlink_path =
            Path::new(&*FS_ROOT).join(Path::new(DEFAULT_BINLINK_DIR).strip_prefix("/").unwrap());
        ui.heading("Habitat Binlink Path")?;
        ui.para("The `hab` command-line tool can create binlinks for package binaries in the \
                 'PATH' when using the 'pkg binlink' or 'pkg install --binlink' commands. By \
                 default, Habitat will create these binlinks in the '/hab/bin' directory. This \
                 directory will always be included in your 'PATH' when inside a Studio \
                 environment. However, you will want this directory on your machine's \
                 persistent 'PATH' in order to access binlinks outside of a Studio.")?;
        if ui.prompt_yes_no("Add binlink directory to PATH?", Some(true))? {
            set_binlink_path(&binlink_path)?;
            ui.para(&format!("{} has been added to your path. You will need to open a new \
                              console window for this added entry to take effect.",
                             binlink_path.display()))?;
        } else {
            ui.para("Okay, maybe another time.")?;
        }
    }
    ui.heading("CLI Setup Complete")?;
    ui.para("That's all for now. Thanks for using Habitat!")?;
    Ok(())
}

fn ask_default_origin(ui: &mut UI) -> Result<bool> {
    Ok(ui.prompt_yes_no("Set up a default origin?", Some(true))?)
}

fn ask_default_builder_instance(ui: &mut UI) -> Result<bool> {
    Ok(ui.prompt_yes_no("Connect to an on-premises bldr instance?", Some(true))?)
}

fn ask_create_origin(ui: &mut UI, origin: &str) -> Result<bool> {
    Ok(ui.prompt_yes_no(&format!("Create an origin key for `{}'?", origin),
                        Some(true))?)
}

fn write_cli_config_origin(origin: &str) -> Result<()> {
    let mut config = config::load()?;
    config.origin = Some(origin.to_string());
    config::save(&config)
}

fn write_cli_config_bldr_url(url: &str) -> Result<()> {
    let mut config = config::load()?;
    config.bldr_url = Some(url.to_string());
    config::save(&config)
}

fn write_cli_config_auth_token(auth_token: &str) -> Result<()> {
    let mut config = config::load()?;
    config.auth_token = Some(auth_token.to_string());
    config::save(&config)
}

fn write_cli_config_ctl_secret(value: &str) -> Result<()> {
    let mut config = config::load()?;
    config.ctl_secret = Some(value.to_string());
    config::save(&config)
}

fn is_origin_in_cache(origin: &str, cache_path: &Path) -> bool {
    match SigKeyPair::get_latest_pair_for(origin, cache_path, None) {
        Ok(pair) => {
            match pair.secret() {
                Ok(_) => true,
                _ => false,
            }
        }
        _ => false,
    }
}

fn create_origin(ui: &mut UI, origin: &str, cache_path: &Path) -> Result<()> {
    let result = command::origin::key::generate::start(ui, &origin, cache_path);
    ui.br()?;
    result
}

fn prompt_origin(ui: &mut UI) -> Result<String> {
    let config = config::load()?;
    let default = match config.origin {
        Some(o) => {
            ui.para(&format!("You already have a default origin set up as `{}', but feel free \
                              to change it if you wish.",
                             &o))?;
            Some(o)
        }
        None => henv::var(ORIGIN_ENVVAR).or_else(|_| henv::var("USER")).ok(),
    };
    Ok(ui.prompt_ask("Default origin name", default.as_ref().map(|x| &**x))?)
}

fn ask_default_auth_token(ui: &mut UI) -> Result<bool> {
    Ok(ui.prompt_yes_no("Set up a default Habitat personal access token?",
                        Some(true))?)
}

fn ask_default_ctl_secret(ui: &mut UI) -> Result<bool> {
    Ok(ui.prompt_yes_no("Set up a default Habitat Supervisor CtlGateway secret?",
                        Some(true))?)
}

fn prompt_url(ui: &mut UI) -> Result<String> {
    let config = config::load()?;
    let default = match config.bldr_url {
        Some(u) => {
            ui.para("You already have a default builder url set up, but feel free to change it \
                     if you wish.")?;
            Some(u)
        }
        None => henv::var(BLDR_URL_ENVVAR).ok(),
    };
    Ok(ui.prompt_ask("Private builder url", default.as_ref().map(|x| &**x))?)
}

fn prompt_auth_token(ui: &mut UI) -> Result<String> {
    let config = config::load()?;
    let default = match config.auth_token {
        Some(o) => {
            ui.para("You already have a default auth token set up, but feel free to change it \
                     if you wish.")?;
            Some(o)
        }
        None => henv::var(AUTH_TOKEN_ENVVAR).ok(),
    };
    Ok(ui.prompt_ask("Habitat personal access token",
                     default.as_ref().map(|x| &**x))?)
}

fn prompt_ctl_secret(ui: &mut UI) -> Result<String> {
    let config = config::load()?;
    let default = match config.ctl_secret {
        Some(o) => {
            ui.para(
                    "You already have a default CtlGateway secret set up, but feel free to \
                     change it
                if you wish.",
            )?;
            Some(o)
        }
        None => henv::var(CTL_SECRET_ENVVAR).ok(),
    };
    Ok(ui.prompt_ask("Habitat Supervisor CtlGateway secret",
                     default.as_ref().map(|x| &**x))?)
}

fn valid_url(val: &str) -> result::Result<(), String> {
    match Url::parse(&val) {
        Ok(_) => Ok(()),
        Err(_) => Err(format!("URL: '{}' is not valid", &val)),
    }
}

#[cfg(windows)]
fn binlink_is_on_path(binlink_path: &Path) -> bool {
    match RegKey::predef(HKEY_LOCAL_MACHINE).open_subkey_with_flags(
        r"System\CurrentControlSet\Control\Session Manager\Environment",
        KEY_READ,
    ) {
        Ok(env) => {
            let path: String = env
                .get_value("path")
                .expect("could not find a machine PATH");
            env::split_paths(&path).any(|p| p == binlink_path)
        }
        _ => false,
    }
}

/// this sets the permanent machine PATH and not
/// the path of this process. These are maintained
/// in the Windows registry
#[cfg(windows)]
fn set_binlink_path(binlink_path: &Path) -> Result<()> {
    if !binlink_is_on_path(binlink_path) {
        let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
        let env = hklm.open_subkey_with_flags(
            r"System\CurrentControlSet\Control\Session Manager\Environment",
            KEY_ALL_ACCESS,
        )?;
        let mut paths = vec![binlink_path.to_path_buf()];
        let path: String = env.get_value("path")?;
        paths.append(&mut env::split_paths(&path).collect());
        env.set_value("path", &env::join_paths(paths)?.to_str().unwrap())?;

        // After altering the machine environment, we must broadcast
        // a WM_SETTINGCHANGE message to all windows so the user
        // will not need to sign out/in for the new path to take effect
        unsafe {
            winuser::SendMessageTimeoutW(HWND_BROADCAST,
                                         WM_SETTINGCHANGE,
                                         0,
                                         WideCString::from_str("Environment").unwrap().as_ptr()
                                         as LPARAM,
                                         SMTO_ABORTIFHUNG,
                                         5000,
                                         ptr::null_mut());
        }
    }
    Ok(())
}
<commit_msg>do not attempt to set binlink path if not in an elevated console<commit_after>#[cfg(windows)]
use std::env;
use std::{path::Path,
          result};

#[cfg(windows)]
use crate::{common::cli::{DEFAULT_BINLINK_DIR,
                          FS_ROOT},
            hcore::fs};
use crate::{common::ui::{UIReader,
                         UIWriter,
                         UI},
            hcore::{crypto::SigKeyPair,
                    env as henv,
                    package::ident,
                    Error::InvalidOrigin}};
#[cfg(windows)]
use std::ptr;
use url::Url;
#[cfg(windows)]
use widestring::WideCString;
#[cfg(windows)]
use winapi::shared::minwindef::LPARAM;
#[cfg(windows)]
use winapi::um::winuser::{self,
                          HWND_BROADCAST,
                          SMTO_ABORTIFHUNG,
                          WM_SETTINGCHANGE};
#[cfg(windows)]
use winreg::enums::{HKEY_LOCAL_MACHINE,
                    KEY_ALL_ACCESS,
                    KEY_READ};
#[cfg(windows)]
use winreg::RegKey;

use crate::{command,
            config,
            error::Result,
            AUTH_TOKEN_ENVVAR,
            BLDR_URL_ENVVAR,
            CTL_SECRET_ENVVAR,
            ORIGIN_ENVVAR};

pub fn start(ui: &mut UI, cache_path: &Path) -> Result<()> {
    ui.br()?;
    ui.title("Habitat CLI Setup")?;
    ui.para("Welcome to hab setup. Let's get started.")?;

    ui.heading("Habitat Builder Instance")?;
    ui.para(
        "Habitat packages can be stored in either the public builder instance \
         https://bldr.habitat.sh or in an on-premises builder depot instance. If \
         you do not set a builder URL now, the `hab` CLI will default to using \
         the public builder instance. This can be overridden at any time after setup.",
    )?;
    if ask_default_builder_instance(ui)? {
        ui.br()?;
        ui.para("Enter the url of your builder instance. The default is \
                 https://bldr.habitat.sh. The configured endpoint can be overridden any time \
                 with a `HAB_BLDR_URL` envvar or a --url flag on the cli.")?;
        let mut url = prompt_url(ui)?;

        while valid_url(&url).is_err() {
            ui.br()?;
            ui.fatal(&format!("{}: is invalid, please provide a valid url", url))?;
            ui.br()?;

            url = prompt_url(ui)?;
        }

        write_cli_config_bldr_url(&url)?;
    } else {
        ui.br()?;
        ui.para("No worries, should you need to use a different bldr instance you can set a \
                 `HAB_BLDR_URL` envvar or pass the `--url` flag to the cli!")?;
    }

    ui.heading("Set up a default origin")?;
    ui.para("Every package in Habitat belongs to an origin, which indicates the person or \
             organization responsible for maintaining that package. Each origin also has a key \
             used to cryptographically sign packages in that origin.")?;
    ui.para("Selecting a default origin tells package building operations such as 'hab pkg \
             build' what key should be used to sign the packages produced. If you do not set a \
             default origin now, you will have to tell package building commands each time what \
             origin to use.")?;
    ui.para("For more information on origins and how they are used in building packages, please \
             consult the docs at https://www.habitat.sh/docs/create-packages-build/")?;
    if ask_default_origin(ui)? {
        ui.br()?;
        ui.para("Enter the name of your origin. If you plan to publish your packages publicly, \
                 we recommend that you select one that is not already in use on the Habitat \
                 build service found at https://bldr.habitat.sh/.")?;
        ui.para("Origins must begin with a lowercase letter or number. Allowed characters \
                 include lowercase letters, numbers, _, -. No more than 255 characters.")?;
        let mut origin = prompt_origin(ui)?;

        while !ident::is_valid_origin_name(&origin) {
            ui.br()?;
            ui.fatal(&format!("{}", InvalidOrigin(origin)))?;
            ui.br()?;

            origin = prompt_origin(ui)?;
        }
        write_cli_config_origin(&origin)?;
        ui.br()?;
        if is_origin_in_cache(&origin, cache_path) {
            ui.para(&format!("You already have an origin key for {} created and installed. \
                              Great work!",
                             &origin))?;
        } else {
            ui.heading("Create origin key pair")?;
            ui.para(&format!("It doesn't look like you have a signing key for the origin `{}'. \
                              Without it, you won't be able to build new packages successfully.",
                             &origin))?;
            ui.para("You can either create a new signing key now, or, if you are building \
                     packages for an origin that already exists, ask the owner to give you the \
                     signing key.")?;
            ui.para("For more information on the use of origin keys, please consult the \
                     documentation at https://www.habitat.sh/docs/concepts-keys/#origin-keys")?;
            if ask_create_origin(ui, &origin)? {
                create_origin(ui, &origin, cache_path)?;
            } else {
                ui.para(&format!("You might want to create an origin key later with: `hab \
                                  origin key generate {}'",
                                 &origin))?;
            }
        }
    } else {
        ui.para("Okay, maybe another time.")?;
    }
    ui.heading("Habitat Personal Access Token")?;
    ui.para("While you can perform tasks like building and running Habitat packages without \
             needing to authenticate with Builder, some operations like uploading your packages \
             to Builder, or checking the status of your build jobs from the Habitat client will \
             require you to use an access token.")?;
    ui.para("The Habitat Personal Access Token can be generated via the Builder  Profile page \
             (https://bldr.habitat.sh/#/profile). Once you have generated your token, you can \
             enter it here.")?;
    ui.para("If you would like to save your token for use by the Habitat client, please enter \
             your access token. Otherwise, just enter No.")?;
    ui.para(
        "For more information on using Builder, please read the \
         documentation at https://www.habitat.sh/docs/using-builder/",
    )?;
    if ask_default_auth_token(ui)? {
        ui.br()?;
        ui.para("Enter your Habitat Personal Access Token.")?;
        let auth_token = prompt_auth_token(ui)?;
        write_cli_config_auth_token(&auth_token)?;
    } else {
        ui.para("Okay, maybe another time.")?;
    }
    if ask_default_ctl_secret(ui)? {
        ui.br()?;
        ui.para("Enter your Habitat Supervisor CtlGateway secret.")?;
        let ctl_secret = prompt_ctl_secret(ui)?;
        write_cli_config_ctl_secret(&ctl_secret)?;
    } else {
        ui.para("Okay, maybe another time.")?;
    }
    #[cfg(windows)]
    {
        let binlink_path =
            Path::new(&*FS_ROOT).join(Path::new(DEFAULT_BINLINK_DIR).strip_prefix("/").unwrap());
        ui.heading("Habitat Binlink Path")?;
        ui.para("The `hab` command-line tool can create binlinks for package binaries in the \
                 'PATH' when using the 'pkg binlink' or 'pkg install --binlink' commands. By \
                 default, Habitat will create these binlinks in the '/hab/bin' directory. This \
                 directory will always be included in your 'PATH' when inside a Studio \
                 environment. However, you will want this directory on your machine's \
                 persistent 'PATH' in order to access binlinks outside of a Studio.")?;
        if ui.prompt_yes_no("Add binlink directory to PATH?", Some(true))? {
            if !fs::am_i_root() {
                ui.fatal("You must be in an elevated console running as an administrator in \
                          order to set the machine's PATH. Please open a new console running as \
                          administrator and run setup again to add the Habitat binlink \
                          directory to your path.")?;
            } else {
                set_binlink_path(&binlink_path)?;
                ui.para(&format!("{} has been added to your path. You will need to open a new \
                                  console window for this added entry to take effect.",
                                 binlink_path.display()))?;
            }
        } else {
            ui.para("Okay, maybe another time.")?;
        }
    }
    ui.heading("CLI Setup Complete")?;
    ui.para("That's all for now. Thanks for using Habitat!")?;
    Ok(())
}

fn ask_default_origin(ui: &mut UI) -> Result<bool> {
    Ok(ui.prompt_yes_no("Set up a default origin?", Some(true))?)
}

fn ask_default_builder_instance(ui: &mut UI) -> Result<bool> {
    Ok(ui.prompt_yes_no("Connect to an on-premises bldr instance?", Some(true))?)
}

fn ask_create_origin(ui: &mut UI, origin: &str) -> Result<bool> {
    Ok(ui.prompt_yes_no(&format!("Create an origin key for `{}'?", origin),
                        Some(true))?)
}

fn write_cli_config_origin(origin: &str) -> Result<()> {
    let mut config = config::load()?;
    config.origin = Some(origin.to_string());
    config::save(&config)
}

fn write_cli_config_bldr_url(url: &str) -> Result<()> {
    let mut config = config::load()?;
    config.bldr_url = Some(url.to_string());
    config::save(&config)
}

fn write_cli_config_auth_token(auth_token: &str) -> Result<()> {
    let mut config = config::load()?;
    config.auth_token = Some(auth_token.to_string());
    config::save(&config)
}

fn write_cli_config_ctl_secret(value: &str) -> Result<()> {
    let mut config = config::load()?;
    config.ctl_secret = Some(value.to_string());
    config::save(&config)
}

fn is_origin_in_cache(origin: &str, cache_path: &Path) -> bool {
    match SigKeyPair::get_latest_pair_for(origin, cache_path, None) {
        Ok(pair) => {
            match pair.secret() {
                Ok(_) => true,
                _ => false,
            }
        }
        _ => false,
    }
}

fn create_origin(ui: &mut UI, origin: &str, cache_path: &Path) -> Result<()> {
    let result = command::origin::key::generate::start(ui, &origin, cache_path);
    ui.br()?;
    result
}

fn prompt_origin(ui: &mut UI) -> Result<String> {
    let config = config::load()?;
    let default = match config.origin {
        Some(o) => {
            ui.para(&format!("You already have a default origin set up as `{}', but feel free \
                              to change it if you wish.",
                             &o))?;
            Some(o)
        }
        None => henv::var(ORIGIN_ENVVAR).or_else(|_| henv::var("USER")).ok(),
    };
    Ok(ui.prompt_ask("Default origin name", default.as_ref().map(|x| &**x))?)
}

fn ask_default_auth_token(ui: &mut UI) -> Result<bool> {
    Ok(ui.prompt_yes_no("Set up a default Habitat personal access token?",
                        Some(true))?)
}

fn ask_default_ctl_secret(ui: &mut UI) -> Result<bool> {
    Ok(ui.prompt_yes_no("Set up a default Habitat Supervisor CtlGateway secret?",
                        Some(true))?)
}

fn prompt_url(ui: &mut UI) -> Result<String> {
    let config = config::load()?;
    let default = match config.bldr_url {
        Some(u) => {
            ui.para("You already have a default builder url set up, but feel free to change it \
                     if you wish.")?;
            Some(u)
        }
        None => henv::var(BLDR_URL_ENVVAR).ok(),
    };
    Ok(ui.prompt_ask("Private builder url", default.as_ref().map(|x| &**x))?)
}

fn prompt_auth_token(ui: &mut UI) -> Result<String> {
    let config = config::load()?;
    let default = match config.auth_token {
        Some(o) => {
            ui.para("You already have a default auth token set up, but feel free to change it \
                     if you wish.")?;
            Some(o)
        }
        None => henv::var(AUTH_TOKEN_ENVVAR).ok(),
    };
    Ok(ui.prompt_ask("Habitat personal access token",
                     default.as_ref().map(|x| &**x))?)
}

fn prompt_ctl_secret(ui: &mut UI) -> Result<String> {
    let config = config::load()?;
    let default = match config.ctl_secret {
        Some(o) => {
            ui.para(
                    "You already have a default CtlGateway secret set up, but feel free to \
                     change it
                if you wish.",
            )?;
            Some(o)
        }
        None => henv::var(CTL_SECRET_ENVVAR).ok(),
    };
    Ok(ui.prompt_ask("Habitat Supervisor CtlGateway secret",
                     default.as_ref().map(|x| &**x))?)
}

fn valid_url(val: &str) -> result::Result<(), String> {
    match Url::parse(&val) {
        Ok(_) => Ok(()),
        Err(_) => Err(format!("URL: '{}' is not valid", &val)),
    }
}

#[cfg(windows)]
fn binlink_is_on_path(binlink_path: &Path) -> bool {
    match RegKey::predef(HKEY_LOCAL_MACHINE).open_subkey_with_flags(
        r"System\CurrentControlSet\Control\Session Manager\Environment",
        KEY_READ,
    ) {
        Ok(env) => {
            let path: String = env
                .get_value("path")
                .expect("could not find a machine PATH");
            env::split_paths(&path).any(|p| p == binlink_path)
        }
        _ => false,
    }
}

/// this sets the permanent machine PATH and not
/// the path of this process. These are maintained
/// in the Windows registry
#[cfg(windows)]
fn set_binlink_path(binlink_path: &Path) -> Result<()> {
    if !binlink_is_on_path(binlink_path) {
        let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
        let env = hklm.open_subkey_with_flags(
            r"System\CurrentControlSet\Control\Session Manager\Environment",
            KEY_ALL_ACCESS,
        )?;
        let mut paths = vec![binlink_path.to_path_buf()];
        let path: String = env.get_value("path")?;
        paths.append(&mut env::split_paths(&path).collect());
        env.set_value("path", &env::join_paths(paths)?.to_str().unwrap())?;

        // After altering the machine environment, we must broadcast
        // a WM_SETTINGCHANGE message to all windows so the user
        // will not need to sign out/in for the new path to take effect
        unsafe {
            winuser::SendMessageTimeoutW(HWND_BROADCAST,
                                         WM_SETTINGCHANGE,
                                         0,
                                         WideCString::from_str("Environment").unwrap().as_ptr()
                                         as LPARAM,
                                         SMTO_ABORTIFHUNG,
                                         5000,
                                         ptr::null_mut());
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>#![allow(unstable)]

use syntax::ast;
use syntax::ext::base;
use syntax::ext::build::AstBuilder;
use syntax::ext::deriving::generic;
use syntax::codemap;
use syntax::parse::token;
use syntax::ptr::P;

use syntax::ext::deriving::generic::ty::{Ty,PtrTy};
use syntax::ast::Mutability::MutImmutable;

pub fn ply_data(ecx: &mut base::ExtCtxt, span: codemap::Span,
                meta_item: &ast::MetaItem, item: &ast::Item,
                mut push: Box<FnMut(P<ast::Item>)>)
{
    generic::TraitDef {
        span: span,
        attributes: Vec::new(),
        path: generic::ty::Path {
            path: vec!["ply", "PlyModel"],
            lifetime: None,
            params: Vec::new(),
            global: true,
        },
        additional_bounds: Vec::new(),
        generics: generic::ty::LifetimeBounds::empty(),
        methods: vec![
            generic::MethodDef {
                name: "new",
                generics: generic::ty::LifetimeBounds::empty(),
                explicit_self: None,
                args: vec![generic::ty::Ptr(Box::new(
                        generic::ty::Literal(
                            generic::ty::Path::new(vec!["ply","parser","PLY"])
                        )
                    ), PtrTy::Borrowed(None, MutImmutable)
                )],
                ret_ty: generic::ty::Literal( // Result<Self, &'static str>
                    generic::ty::Path {
                        path: vec!["std", "result", "Result"],
                        lifetime: None,
                        params: vec![
                            Box::new(Ty::Self),
                            Box::new(Ty::Ptr(Box::new(
                                    generic::ty::Literal(
                                        generic::ty::Path::new(vec!["str"])
                                    )
                                ), PtrTy::Borrowed(Some("'static"), MutImmutable)
                            ))],
                        global: true
                    }
                ),
                attributes: vec![/*
                    ecx.attribute(span.clone(), ecx.meta_list(span.clone(),
                        token::InternedString::new("allow"),
                        vec![ecx.meta_word(span.clone(),
                                token::InternedString::new("unused_assignments"))]
                    ))*/
                ],
                combine_substructure: generic::combine_substructure(Box::new(body)),
            },
        ],
    }.expand(ecx, meta_item, item, |i| push(i));
}

// To print the expanded form, use `puts --pretty expanded -Z unstable-options`.
fn body(ecx: &mut base::ExtCtxt, span: codemap::Span,
        substr: &generic::Substructure) -> P<ast::Expr>
{
    let ecx: &base::ExtCtxt = ecx;
    let self_ty = &substr.type_ident;

    match substr.fields {
        &generic::StaticStruct(ref definition, generic::Named(ref fields)) => {
            let field_count = definition.fields.len();

            let struct_expr = ecx.expr_struct_ident(span, *self_ty,
                fields.iter().zip(definition.fields.iter())
                             .map(|(&(field_ident, _), field_def)| {
                    let ref field_type = field_def.node.ty;
                    let ident_str = token::get_ident(field_ident);
                    let ident_str = ident_str.get();
                    ecx.field_imm(span, field_ident, quote_expr!(ecx, { // For each element

                        if let Some(e) = ply.elements.iter()
                                         .filter(|&e| e.name == $ident_str).next() {
                            try!(ply::ElementVec::check(None::<$field_type>, e));
                            let mut accum = vec![];
                            for line in e.data.iter() {
                                let res = try!(ply::Element::parse(line));
                                accum.push(res);
                            }
                            accum
                        } else {
                            return Err("Did not find a corresponding element name.");
                        }

                    }))
                }).collect()
            );
            quote_expr!(ecx, { // The function body

                let ply = __arg_0;
                // check
        		if ply.elements.len() != $field_count {
        			return Err("The number of elements in PlyModel is not correct.");
        		}
                let ret = $struct_expr;
                Ok(ret)

            })
        },
        _ => {
            ecx.span_err(span, "Unable to implement `PlyModel` on a non-structure");
            ecx.expr_int(span, 0)
        }
    }
}
<commit_msg>A minor cleanup<commit_after>#![allow(unstable)]

use syntax::ast;
use syntax::ast::Mutability::MutImmutable;
use syntax::codemap::Span;
use syntax::ext::base::ExtCtxt;
use syntax::ext::deriving::generic;
use syntax::ext::deriving::generic::ty;
use syntax::parse::token;
use syntax::ptr::P;

pub fn ply_data(ecx: &mut ExtCtxt, span: Span,
                meta_item: &ast::MetaItem, item: &ast::Item,
                mut push: Box<FnMut(P<ast::Item>)>)
{
    generic::TraitDef {
        span: span,
        attributes: Vec::new(),
        path: ty::Path {
            path: vec!["ply", "PlyModel"],
            lifetime: None,
            params: Vec::new(),
            global: true,
        },
        additional_bounds: Vec::new(),
        generics: ty::LifetimeBounds::empty(),
        methods: vec![
            generic::MethodDef {
                name: "new",
                generics: ty::LifetimeBounds::empty(),
                explicit_self: None,
                args: vec![ty::Ptr(Box::new(
                        ty::Literal(
                            ty::Path::new(vec!["ply","parser","PLY"])
                        )
                    ), ty::PtrTy::Borrowed(None, MutImmutable)
                )],
                ret_ty: ty::Literal( // Result<Self, &'static str>
                    ty::Path {
                        path: vec!["std", "result", "Result"],
                        lifetime: None,
                        params: vec![
                            Box::new(ty::Ty::Self),
                            Box::new(ty::Ty::Ptr(Box::new(
                                    ty::Literal(
                                        ty::Path::new(vec!["str"])
                                    )
                                ), ty::PtrTy::Borrowed(Some("'static"), MutImmutable)
                            ))],
                        global: true
                    }
                ),
                // If ever superfluous fields are allowed, check out how glium supresses warnings.
                attributes: vec![],
                combine_substructure: generic::combine_substructure(Box::new(body)),
            },
        ],
    }.expand(ecx, meta_item, item, |i| push(i));
}

// To print the expanded form, use `puts --pretty expanded -Z unstable-options`.
// Function body expansion
fn body(ecx: &mut ExtCtxt, span: Span,
        substr: &generic::Substructure) -> P<ast::Expr>
{
    let ecx: &ExtCtxt = ecx;
    let self_ty = &substr.type_ident;

    match substr.fields {
        &generic::StaticStruct(ref definition, generic::Named(ref fields)) => {
            let field_count = definition.fields.len();

            let struct_expr = ecx.expr_struct_ident(span, *self_ty,
                fields.iter().zip(definition.fields.iter())
                             .map(|(&(field_ident, _), field_def)| {
                    let ref field_type = field_def.node.ty;
                    let ident_str = token::get_ident(field_ident);
                    let ident_str = ident_str.get();
                    ecx.field_imm(span, field_ident, quote_expr!(ecx, { // For each element

                        if let Some(e) = ply.elements.iter()
                                         .filter(|&e| e.name == $ident_str).next() {
                            try!(ply::ElementVec::check(None::<$field_type>, e));
                            let mut accum = vec![];
                            for line in e.data.iter() {
                                let res = try!(ply::Element::parse(line));
                                accum.push(res);
                            }
                            accum
                        } else {
                            return Err("Did not find a corresponding element name.");
                        }

                    }))
                }).collect()
            );
            quote_expr!(ecx, { // The function body

                let ply = __arg_0;
                // check
        		if ply.elements.len() != $field_count {
        			return Err("The number of elements in PlyModel is not correct.");
        		}
                let ret = $struct_expr;
                Ok(ret)

            })
        },
        _ => {
            ecx.span_err(span, "Unable to implement `PlyModel` on a non-structure");
            ecx.expr_int(span, 0)
        }
    }
}
<|endoftext|>