<commit_before>use super::{GeoClass, Scratchpad};
use crate::StrError;
use plotpy::{Canvas, Plot, PolyCode, Text};
use russell_lab::Vector;

impl Scratchpad {
    /// Draws the geometric shape
    ///
    /// # Input
    ///
    /// * `plot` -- plot struct
    /// * `pad` -- the scratchpad with coordinates set already
    /// * `edge_color` -- color to draw edges; an empty string "" means use default color
    /// * `with_ids` -- draw IDs of nodes
    /// * `set_range` -- sets axis range
    pub fn draw_shape(
        &self,
        plot: &mut Plot,
        edge_color: &str,
        with_ids: bool,
        set_range: bool,
    ) -> Result<(), StrError> {
        if !self.ok_xxt {
            return Err("all components of the coordinates matrix must be set first");
        }
        let space_ndim = self.xmin.len();
        let mut canvas = Canvas::new();
        if self.kind.ndim() == 1 {
            let mut lin_pad = self.clone();
            draw_edge(&mut canvas, &mut lin_pad, edge_color)?;
        } else {
            let mut edge_pad = Scratchpad::new(space_ndim, self.kind.edge_kind().unwrap())?;
            for e in 0..self.kind.nedge() {
                for i in 0..self.kind.edge_nnode() {
                    let m = self.kind.edge_node_id(e, i);
                    for j in 0..space_ndim {
                        edge_pad.set_xx(i, j, self.xxt[j][m]);
                    }
                }
                draw_edge(&mut canvas, &mut edge_pad, edge_color)?;
            }
        }
        plot.add(&canvas);
        if with_ids {
            let mut labels_corner = Text::new();
            let mut labels_middle = Text::new();
            labels_corner
                .set_align_horizontal("center")
                .set_align_vertical("center")
                .set_bbox(true)
                .set_bbox_style("circle,pad=0.3")
                .set_bbox_facecolor("white")
                .set_bbox_edgecolor("red");
            labels_middle
                .set_align_horizontal("center")
                .set_align_vertical("center")
                .set_fontsize(8.0)
                .set_bbox(true)
                .set_bbox_style("circle,pad=0.2")
                .set_bbox_facecolor("white")
                .set_bbox_edgecolor("gold");
            if space_ndim == 2 {
                let m_corner_max = match self.kind.class() {
                    GeoClass::Lin => 2,
                    GeoClass::Tri => 3,
                    GeoClass::Qua => 4,
                    GeoClass::Tet => 4,
                    GeoClass::Hex => 8,
                };
                for m in 0..self.kind.nnode() {
                    if m < m_corner_max {
                        labels_corner.draw(self.xxt[0][m], self.xxt[1][m], format!("{}", m).as_str());
                    } else {
                        labels_middle.draw(self.xxt[0][m], self.xxt[1][m], format!("{}", m).as_str());
                    }
                }
            } else {
                let m_corner_max = if self.kind.is_tri_or_tet() { 4 } else { 8 };
                for m in 0..self.kind.nnode() {
                    if m < m_corner_max {
                        labels_corner.draw_3d(
                            self.xxt[0][m],
                            self.xxt[1][m],
                            self.xxt[2][m],
                            format!("{}", m).as_str(),
                        );
                    } else {
                        labels_middle.draw_3d(
                            self.xxt[0][m],
                            self.xxt[1][m],
                            self.xxt[2][m],
                            format!("{}", m).as_str(),
                        );
                    }
                }
            }
            plot.add(&labels_corner).add(&labels_middle);
        }
        if set_range {
            let dx = self.xmax[0] - self.xmin[0];
            let dy = self.xmax[1] - self.xmin[1];
            const PCT: f64 = 2.0 / 100.0;
            if space_ndim == 2 {
                plot.set_range(
                    self.xmin[0] - dx * PCT,
                    self.xmax[0] + dx * PCT,
                    self.xmin[1] - dy * PCT,
                    self.xmax[1] + dy * PCT,
                );
            } else {
                let dz = self.xmax[2] - self.xmin[2];
                plot.set_range_3d(
                    self.xmin[0] - dx * PCT,
                    self.xmax[0] + dx * PCT,
                    self.xmin[1] - dy * PCT,
                    self.xmax[1] + dy * PCT,
                    self.xmin[2] - dz * PCT,
                    self.xmax[2] + dz * PCT,
                );
            }
        }
        Ok(())
    }

    /// Draws shape (simple version)
    pub fn draw_shape_simple(&self, filename: &str) -> Result<(), StrError> {
        let mut plot = Plot::new();
        self.draw_shape(&mut plot, "", true, true)?;
        plot.set_equal_axes(true)
            .set_figure_size_points(600.0, 600.0)
            .save(filename)?;
        Ok(())
    }
}

/// Draws edge
fn draw_edge(canvas: &mut Canvas, edge_pad: &mut Scratchpad, edge_color: &str) -> Result<(), StrError> {
    const N: usize = 11; // number of points along edge (to handle nonlinear edges)
    let (ksi_min, ksi_del) = (-1.0, 2.0); // all Lin shapes go from -1 to +1
    let space_ndim = edge_pad.xmax.len();
    let mut x = Vector::new(space_ndim);
    canvas.set_face_color("None").set_line_width(3.0);
    if edge_color != "" {
        canvas.set_edge_color(edge_color);
    }
    if space_ndim == 2 {
        canvas.polycurve_begin();
        for i in 0..N {
            let kval = ksi_min + ksi_del * (i as f64) / ((N - 1) as f64);
            edge_pad.calc_coords(&mut x, &[kval])?;
            let code = if i == 0 { PolyCode::MoveTo } else { PolyCode::LineTo };
            canvas.polycurve_add(x[0], x[1], code);
        }
        canvas.polycurve_end(false);
    } else {
        canvas.polyline_3d_begin();
        for i in 0..N {
            let kval = ksi_min + ksi_del * (i as f64) / ((N - 1) as f64);
            edge_pad.calc_coords(&mut x, &[kval])?;
            canvas.polyline_3d_add(x[0], x[1], x[2]);
        }
        canvas.polyline_3d_end();
    }
    Ok(())
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#[cfg(test)]
mod tests {
    use crate::shapes::scratchpad_testing::aux;
    use crate::shapes::GeoKind;
    use crate::StrError;
    use plotpy::Plot;

    #[test]
    fn draw_shape_works() -> Result<(), StrError> {
        // select all kinds
        let kinds = vec![
            // Lin
            GeoKind::Lin2,
            GeoKind::Lin3,
            GeoKind::Lin4,
            GeoKind::Lin5,
            // Tri
            GeoKind::Tri3,
            GeoKind::Tri6,
            GeoKind::Tri10,
            GeoKind::Tri15,
            // Qua
            GeoKind::Qua4,
            GeoKind::Qua8,
            GeoKind::Qua9,
            GeoKind::Qua12,
            GeoKind::Qua16,
            GeoKind::Qua17,
            // Tet
            GeoKind::Tet4,
            GeoKind::Tet10,
            GeoKind::Tet20,
            // Hex
            GeoKind::Hex8,
            GeoKind::Hex20,
            GeoKind::Hex32,
        ];
        for kind in kinds {
            let space_ndim = usize::max(2, kind.ndim());
            let pad = aux::gen_scratchpad_with_coords(space_ndim, kind);
            let mut plot = Plot::new();
            if space_ndim == 2 {
                let canvas = aux::gen_canvas_mapping();
                plot.add(&canvas);
                plot.set_range(1.0, 10.0, 1.0, 10.0);
            }
            pad.draw_shape(&mut plot, "", true, true)?;
            if space_ndim == 2 {
                plot.set_range(1.0, 10.0, 1.0, 10.0);
            }
            if false {
                plot.set_figure_size_points(600.0, 600.0)
                    .set_equal_axes(true)
                    .save(format!("/tmp/gemlab/test_draw_shape_{}.svg", kind.to_string()).as_str())?;
            }
        }
        Ok(())
    }
}
<commit_msg>Fix set_range in draw_shape<commit_after>use super::{GeoClass, Scratchpad};
use crate::StrError;
use plotpy::{Canvas, Plot, PolyCode, Text};
use russell_lab::Vector;

impl Scratchpad {
    /// Draws the geometric shape
    ///
    /// # Input
    ///
    /// * `plot` -- plot struct
    /// * `pad` -- the scratchpad with coordinates set already
    /// * `edge_color` -- color to draw edges; an empty string "" means use default color
    /// * `with_ids` -- draw IDs of nodes
    /// * `set_range` -- sets axis range
    pub fn draw_shape(
        &self,
        plot: &mut Plot,
        edge_color: &str,
        with_ids: bool,
        set_range: bool,
    ) -> Result<(), StrError> {
        if !self.ok_xxt {
            return Err("all components of the coordinates matrix must be set first");
        }
        let space_ndim = self.xmin.len();
        let mut canvas = Canvas::new();
        if self.kind.ndim() == 1 {
            let mut lin_pad = self.clone();
            draw_edge(&mut canvas, &mut lin_pad, edge_color)?;
        } else {
            let mut edge_pad = Scratchpad::new(space_ndim, self.kind.edge_kind().unwrap())?;
            for e in 0..self.kind.nedge() {
                for i in 0..self.kind.edge_nnode() {
                    let m = self.kind.edge_node_id(e, i);
                    for j in 0..space_ndim {
                        edge_pad.set_xx(i, j, self.xxt[j][m]);
                    }
                }
                draw_edge(&mut canvas, &mut edge_pad, edge_color)?;
            }
        }
        plot.add(&canvas);
        if with_ids {
            let mut labels_corner = Text::new();
            let mut labels_middle = Text::new();
            labels_corner
                .set_align_horizontal("center")
                .set_align_vertical("center")
                .set_bbox(true)
                .set_bbox_style("circle,pad=0.3")
                .set_bbox_facecolor("white")
                .set_bbox_edgecolor("red");
            labels_middle
                .set_align_horizontal("center")
                .set_align_vertical("center")
                .set_fontsize(8.0)
                .set_bbox(true)
                .set_bbox_style("circle,pad=0.2")
                .set_bbox_facecolor("white")
                .set_bbox_edgecolor("gold");
            if space_ndim == 2 {
                let m_corner_max = match self.kind.class() {
                    GeoClass::Lin => 2,
                    GeoClass::Tri => 3,
                    GeoClass::Qua => 4,
                    GeoClass::Tet => 4,
                    GeoClass::Hex => 8,
                };
                for m in 0..self.kind.nnode() {
                    if m < m_corner_max {
                        labels_corner.draw(self.xxt[0][m], self.xxt[1][m], format!("{}", m).as_str());
                    } else {
                        labels_middle.draw(self.xxt[0][m], self.xxt[1][m], format!("{}", m).as_str());
                    }
                }
            } else {
                let m_corner_max = if self.kind.is_tri_or_tet() { 4 } else { 8 };
                for m in 0..self.kind.nnode() {
                    if m < m_corner_max {
                        labels_corner.draw_3d(
                            self.xxt[0][m],
                            self.xxt[1][m],
                            self.xxt[2][m],
                            format!("{}", m).as_str(),
                        );
                    } else {
                        labels_middle.draw_3d(
                            self.xxt[0][m],
                            self.xxt[1][m],
                            self.xxt[2][m],
                            format!("{}", m).as_str(),
                        );
                    }
                }
            }
            plot.add(&labels_corner).add(&labels_middle);
        }
        if set_range {
            let mut dx = self.xmax[0] - self.xmin[0];
            let mut dy = self.xmax[1] - self.xmin[1];
            if dx <= 0.0 {
                dx = 1.0
            }
            if dy <= 0.0 {
                dy = 1.0
            }
            const PCT: f64 = 2.0 / 100.0;
            if space_ndim == 2 {
                plot.set_range(
                    self.xmin[0] - dx * PCT,
                    self.xmax[0] + dx * PCT,
                    self.xmin[1] - dy * PCT,
                    self.xmax[1] + dy * PCT,
                );
            } else {
                let dz = self.xmax[2] - self.xmin[2];
                plot.set_range_3d(
                    self.xmin[0] - dx * PCT,
                    self.xmax[0] + dx * PCT,
                    self.xmin[1] - dy * PCT,
                    self.xmax[1] + dy * PCT,
                    self.xmin[2] - dz * PCT,
                    self.xmax[2] + dz * PCT,
                );
            }
        }
        Ok(())
    }

    /// Draws shape (simple version)
    pub fn draw_shape_simple(&self, filename: &str) -> Result<(), StrError> {
        let mut plot = Plot::new();
        self.draw_shape(&mut plot, "", true, true)?;
        plot.set_equal_axes(true)
            .set_figure_size_points(600.0, 600.0)
            .save(filename)?;
        Ok(())
    }
}

/// Draws edge
fn draw_edge(canvas: &mut Canvas, edge_pad: &mut Scratchpad, edge_color: &str) -> Result<(), StrError> {
    const N: usize = 11; // number of points along edge (to handle nonlinear edges)
    let (ksi_min, ksi_del) = (-1.0, 2.0); // all Lin shapes go from -1 to +1
    let space_ndim = edge_pad.xmax.len();
    let mut x = Vector::new(space_ndim);
    canvas.set_face_color("None").set_line_width(3.0);
    if edge_color != "" {
        canvas.set_edge_color(edge_color);
    }
    if space_ndim == 2 {
        canvas.polycurve_begin();
        for i in 0..N {
            let kval = ksi_min + ksi_del * (i as f64) / ((N - 1) as f64);
            edge_pad.calc_coords(&mut x, &[kval])?;
            let code = if i == 0 { PolyCode::MoveTo } else { PolyCode::LineTo };
            canvas.polycurve_add(x[0], x[1], code);
        }
        canvas.polycurve_end(false);
    } else {
        canvas.polyline_3d_begin();
        for i in 0..N {
            let kval = ksi_min + ksi_del * (i as f64) / ((N - 1) as f64);
            edge_pad.calc_coords(&mut x, &[kval])?;
            canvas.polyline_3d_add(x[0], x[1], x[2]);
        }
        canvas.polyline_3d_end();
    }
    Ok(())
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#[cfg(test)]
mod tests {
    use crate::shapes::scratchpad_testing::aux;
    use crate::shapes::GeoKind;
    use crate::StrError;
    use plotpy::Plot;

    #[test]
    fn draw_shape_works() -> Result<(), StrError> {
        // select all kinds
        let kinds = vec![
            // Lin
            GeoKind::Lin2,
            GeoKind::Lin3,
            GeoKind::Lin4,
            GeoKind::Lin5,
            // Tri
            GeoKind::Tri3,
            GeoKind::Tri6,
            GeoKind::Tri10,
            GeoKind::Tri15,
            // Qua
            GeoKind::Qua4,
            GeoKind::Qua8,
            GeoKind::Qua9,
            GeoKind::Qua12,
            GeoKind::Qua16,
            GeoKind::Qua17,
            // Tet
            GeoKind::Tet4,
            GeoKind::Tet10,
            GeoKind::Tet20,
            // Hex
            GeoKind::Hex8,
            GeoKind::Hex20,
            GeoKind::Hex32,
        ];
        for kind in kinds {
            let space_ndim = usize::max(2, kind.ndim());
            let pad = aux::gen_scratchpad_with_coords(space_ndim, kind);
            let mut plot = Plot::new();
            if space_ndim == 2 {
                let canvas = aux::gen_canvas_mapping();
                plot.add(&canvas);
                plot.set_range(1.0, 10.0, 1.0, 10.0);
            }
            pad.draw_shape(&mut plot, "", true, true)?;
            if space_ndim == 2 {
                plot.set_range(1.0, 10.0, 1.0, 10.0);
            }
            if false {
                plot.set_figure_size_points(600.0, 600.0)
                    .set_equal_axes(true)
                    .save(format!("/tmp/gemlab/test_draw_shape_{}.svg", kind.to_string()).as_str())?;
            }
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>use core::fmt;
use core::ops::{Index, IndexMut};

use super::PhysFrame;
use addr::PhysAddr;

use usize_conversions::usize_from;
use ux::*;

#[derive(Clone)]
#[repr(transparent)]
pub struct PageTableEntry(u64);

impl PageTableEntry {
    pub fn is_unused(&self) -> bool {
        self.0 == 0
    }

    pub fn set_unused(&mut self) {
        self.0 = 0;
    }

    pub fn flags(&self) -> PageTableFlags {
        PageTableFlags::from_bits_truncate(self.0)
    }

    pub fn frame(&self) -> Option<PhysFrame> {
        if self.flags().contains(PageTableFlags::PRESENT) {
            let addr = PhysAddr::new(self.0 & 0x000fffff_fffff000);
            Some(PhysFrame::containing_address(addr))
        } else {
            None
        }
    }

    pub fn set(&mut self, frame: PhysFrame, flags: PageTableFlags) {
        self.0 = (frame.start_address().as_u64()) | flags.bits();
    }
}

impl fmt::Debug for PageTableEntry {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut f = f.debug_struct("PageTableEntry");
        f.field("frame", &self.frame());
        f.field("flags", &self.flags());
        f.finish()
    }
}

bitflags! {
    pub struct PageTableFlags: u64 {
        const PRESENT =         1 << 0;
        const WRITABLE =        1 << 1;
        const USER_ACCESSIBLE = 1 << 2;
        const WRITE_THROUGH =   1 << 3;
        const NO_CACHE =        1 << 4;
        const ACCESSED =        1 << 5;
        const DIRTY =           1 << 6;
        const HUGE_PAGE =       1 << 7;
        const GLOBAL =          1 << 8;
        const BIT_9 =           1 << 9;
        const BIT_10 =          1 << 10;
        const BIT_11 =          1 << 11;
        const BIT_52 =          1 << 52;
        const BIT_53 =          1 << 53;
        const BIT_54 =          1 << 54;
        const BIT_55 =          1 << 55;
        const BIT_56 =          1 << 56;
        const BIT_57 =          1 << 57;
        const BIT_58 =          1 << 58;
        const BIT_59 =          1 << 59;
        const BIT_60 =          1 << 60;
        const BIT_61 =          1 << 61;
        const BIT_62 =          1 << 62;
        const NO_EXECUTE =      1 << 63;
    }
}

const ENTRY_COUNT: usize = 512;

#[repr(transparent)]
pub struct PageTable {
    entries: [PageTableEntry; ENTRY_COUNT],
}

impl PageTable {
    pub fn zero(&mut self) {
        for entry in self.entries.iter_mut() {
            entry.set_unused();
        }
    }
}

impl Index<usize> for PageTable {
    type Output = PageTableEntry;

    fn index(&self, index: usize) -> &PageTableEntry {
        &self.entries[index]
    }
}

impl IndexMut<usize> for PageTable {
    fn index_mut(&mut self, index: usize) -> &mut PageTableEntry {
        &mut self.entries[index]
    }
}

impl Index<u9> for PageTable {
    type Output = PageTableEntry;

    fn index(&self, index: u9) -> &PageTableEntry {
        &self.entries[usize_from(u16::from(index))]
    }
}

impl IndexMut<u9> for PageTable {
    fn index_mut(&mut self, index: u9) -> &mut PageTableEntry {
        &mut self.entries[usize_from(u16::from(index))]
    }
}

impl fmt::Debug for PageTable {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.entries[..].fmt(f)
    }
}
<commit_msg>Update page_table module with some renames and new methods<commit_after>use core::fmt;
use core::ops::{Index, IndexMut};

use super::{PhysFrame, PageSize, Size4KB};
use addr::PhysAddr;

use usize_conversions::usize_from;
use ux::*;

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum FrameError {
    FrameNotPresent,
    HugeFrame,
}

#[derive(Clone)]
#[repr(transparent)]
pub struct PageTableEntry {
    entry: u64,
}

impl PageTableEntry {
    pub fn is_unused(&self) -> bool {
        self.entry == 0
    }

    pub fn set_unused(&mut self) {
        self.entry = 0;
    }

    pub fn flags(&self) -> PageTableFlags {
        PageTableFlags::from_bits_truncate(self.entry)
    }

    pub fn addr(&self) -> PhysAddr {
        PhysAddr::new(self.entry & 0x000fffff_fffff000)
    }

    pub fn frame(&self) -> Result<PhysFrame, FrameError> {
        if !self.flags().contains(PageTableFlags::PRESENT) {
            Err(FrameError::FrameNotPresent)
        } else if self.flags().contains(PageTableFlags::HUGE_PAGE) {
            Err(FrameError::HugeFrame)
        } else {
            Ok(PhysFrame::containing_address(self.addr()))
        }
    }

    pub fn set_addr(&mut self, addr: PhysAddr, flags: PageTableFlags) {
        assert!(addr.is_aligned(Size4KB::SIZE));
        self.entry = (addr.as_u64()) | flags.bits();
    }

    pub fn set_frame(&mut self, frame: PhysFrame, flags: PageTableFlags) {
        assert!(!flags.contains(PageTableFlags::HUGE_PAGE));
        self.set_addr(frame.start_address(), flags)
    }
}

impl fmt::Debug for PageTableEntry {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut f = f.debug_struct("PageTableEntry");
        f.field("addr", &self.addr());
        f.field("flags", &self.flags());
        f.finish()
    }
}

bitflags! {
    pub struct PageTableFlags: u64 {
        const PRESENT =         1 << 0;
        const WRITABLE =        1 << 1;
        const USER_ACCESSIBLE = 1 << 2;
        const WRITE_THROUGH =   1 << 3;
        const NO_CACHE =        1 << 4;
        const ACCESSED =        1 << 5;
        const DIRTY =           1 << 6;
        const HUGE_PAGE =       1 << 7;
        const GLOBAL =          1 << 8;
        const BIT_9 =           1 << 9;
        const BIT_10 =          1 << 10;
        const BIT_11 =          1 << 11;
        const BIT_52 =          1 << 52;
        const BIT_53 =          1 << 53;
        const BIT_54 =          1 << 54;
        const BIT_55 =          1 << 55;
        const BIT_56 =          1 << 56;
        const BIT_57 =          1 << 57;
        const BIT_58 =          1 << 58;
        const BIT_59 =          1 << 59;
        const BIT_60 =          1 << 60;
        const BIT_61 =          1 << 61;
        const BIT_62 =          1 << 62;
        const NO_EXECUTE =      1 << 63;
    }
}

const ENTRY_COUNT: usize = 512;

#[repr(transparent)]
pub struct PageTable {
    entries: [PageTableEntry; ENTRY_COUNT],
}

impl PageTable {
    pub fn zero(&mut self) {
        for entry in self.entries.iter_mut() {
            entry.set_unused();
        }
    }
}

impl Index<usize> for PageTable {
    type Output = PageTableEntry;

    fn index(&self, index: usize) -> &Self::Output {
        &self.entries[index]
    }
}

impl IndexMut<usize> for PageTable {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        &mut self.entries[index]
    }
}

impl Index<u9> for PageTable {
    type Output = PageTableEntry;

    fn index(&self, index: u9) -> &Self::Output {
        &self.entries[usize_from(u16::from(index))]
    }
}

impl IndexMut<u9> for PageTable {
    fn index_mut(&mut self, index: u9) -> &mut Self::Output {
        &mut self.entries[usize_from(u16::from(index))]
    }
}

impl fmt::Debug for PageTable {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.entries[..].fmt(f)
    }
}
<|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.

//! System configuration loading for windows

use std::str::FromStr;
use std::net::SocketAddr;

use ipconfig::get_adapters;
use ipconfig::computer::{get_search_list, get_domain, is_round_robin_enabled};

use trust_dns_proto::rr::Name;

use config::{ResolverConfig, ResolverOpts, NameServerConfig, Protocol};
use error::*;

macro_rules! map_ipconfig_error {
    ($result:expr) => (
            $result.map_err(|e| ResolveErrorKind::Msg(format!("failed to read from registry: {}", e)).into())
        );
}

/// Returns the name servers of the computer (of all adapters)
fn get_name_servers() -> ResolveResult<Vec<NameServerConfig>> {
    let adapters = map_ipconfig_error!(get_adapters())?;
    let mut name_servers = vec![];

    for dns_server in adapters.iter().flat_map(|adapter| adapter.dns_servers().iter()) {
        let socket_addr = SocketAddr::new(*dns_server, 53);
        name_servers.push(NameServerConfig {
            socket_addr,
            protocol: Protocol::Udp,
        });
        name_servers.push(NameServerConfig {
            socket_addr,
            protocol: Protocol::Tcp,
        });
    };
    Ok(name_servers)
}

pub(crate) fn read_system_conf() -> ResolveResult<(ResolverConfig, ResolverOpts)> {
    let name_servers = get_name_servers()?;

    let search_list: Vec<Name> = map_ipconfig_error!(get_search_list())?
        .iter()
        .map(|x| Name::from_str(x))
        .collect::<Result<Vec<_>, _>>()?;

    let domain = match map_ipconfig_error!(get_domain())? {
        Some(domain) => Name::from_str(&domain)?,
        None => Name::root(),
    };

    let config = ResolverConfig::from_parts(domain, search_list, name_servers);

    let rotate = map_ipconfig_error!(is_round_robin_enabled())?;

    let opts = ResolverOpts {
        rotate,
        ..Default::default()
    };
    Ok((config, opts))
}
<commit_msg>make error explicit for ipconfig<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.

//! System configuration loading for windows

use std::str::FromStr;
use std::net::SocketAddr;

use ipconfig::get_adapters;
use ipconfig::computer::{get_search_list, get_domain, is_round_robin_enabled};

use trust_dns_proto::rr::Name;

use config::{ResolverConfig, ResolverOpts, NameServerConfig, Protocol};
use error::*;

macro_rules! map_ipconfig_error {
    // TODO: this should use error_chains chain facility
    ($result:expr) => (
            $result.map_err(|e| ResolveError::from(ResolveErrorKind::Msg(format!("failed to read from registry: {}", e))))
        );
}

/// Returns the name servers of the computer (of all adapters)
fn get_name_servers() -> ResolveResult<Vec<NameServerConfig>> {
    let adapters = map_ipconfig_error!(get_adapters())?;
    let mut name_servers = vec![];

    for dns_server in adapters.iter().flat_map(|adapter| adapter.dns_servers().iter()) {
        let socket_addr = SocketAddr::new(*dns_server, 53);
        name_servers.push(NameServerConfig {
            socket_addr,
            protocol: Protocol::Udp,
        });
        name_servers.push(NameServerConfig {
            socket_addr,
            protocol: Protocol::Tcp,
        });
    };
    Ok(name_servers)
}

pub(crate) fn read_system_conf() -> ResolveResult<(ResolverConfig, ResolverOpts)> {
    let name_servers = get_name_servers()?;

    let search_list: Vec<Name> = map_ipconfig_error!(get_search_list())?
        .iter()
        .map(|x| Name::from_str(x))
        .collect::<Result<Vec<_>, _>>()?;

    let domain = match map_ipconfig_error!(get_domain())? {
        Some(domain) => Name::from_str(&domain)?,
        None => Name::root(),
    };

    let config = ResolverConfig::from_parts(domain, search_list, name_servers);

    let rotate = map_ipconfig_error!(is_round_robin_enabled())?;

    let opts = ResolverOpts {
        rotate,
        ..Default::default()
    };
    Ok((config, opts))
}
<|endoftext|>


<commit_before>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.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 png;
extern crate std;
extern crate test;

use std::io;
use std::io::{File, Reader, Process};
use std::io::process::ExitStatus;
use std::os;
use std::str;
use test::{DynTestName, DynTestFn, TestDesc, TestOpts, TestDescAndFn};
use test::run_tests_console;

fn main() {
    let args = os::args();
    let mut parts = args.tail().split(|e| "--" == e.as_slice());

    let files = parts.next().unwrap();  // .split() is never empty
    let servo_args = parts.next().unwrap_or(&[]);

    if files.len() == 0 {
        fail!("error: at least one reftest list must be given");
    }

    let tests = parse_lists(files, servo_args);
    let test_opts = TestOpts {
        filter: None,
        run_ignored: false,
        logfile: None,
        run_tests: true,
        run_benchmarks: false,
        ratchet_noise_percent: None,
        ratchet_metrics: None,
        save_metrics: None,
        test_shard: None,
    };

    match run_tests_console(&test_opts, tests) {
        Ok(false) => os::set_exit_status(1), // tests failed
        Err(_) => os::set_exit_status(2),    // I/O-related failure
        _ => (),
    }
}

#[deriving(Eq)]
enum ReftestKind {
    Same,
    Different,
}

struct Reftest {
    name: ~str,
    kind: ReftestKind,
    files: [~str, ..2],
    id: uint,
    servo_args: ~[~str],
}

fn parse_lists(filenames: &[~str], servo_args: &[~str]) -> Vec<TestDescAndFn> {
    let mut tests = Vec::new();
    let mut next_id = 0;
    for file in filenames.iter() {
        let file_path = Path::new(file.clone());
        let contents = match File::open_mode(&file_path, io::Open, io::Read)
            .and_then(|mut f| {
                f.read_to_str()
            }) {
                Ok(s) => s,
                _ => fail!("Could not read file"),
            };

        for line in contents.lines() {
            // ignore comments
            if line.starts_with("#") {
                continue;
            }

            let parts: ~[&str] = line.split(' ').filter(|p| !p.is_empty()).collect();

            if parts.len() != 3 {
                fail!("reftest line: '{:s}' doesn't match 'KIND LEFT RIGHT'", line);
            }

            let kind = match parts[0] {
                "==" => Same,
                "!=" => Different,
                _ => fail!("reftest line: '{:s}' has invalid kind '{:s}'",
                           line, parts[0])
            };
            let src_path = file_path.dir_path();
            let src_dir = src_path.display().to_str();
            let file_left =  src_dir + "/" + parts[1];
            let file_right = src_dir + "/" + parts[2];

            let reftest = Reftest {
                name: parts[1] + " / " + parts[2],
                kind: kind,
                files: [file_left, file_right],
                id: next_id,
                servo_args: servo_args.to_owned(),
            };

            next_id += 1;

            tests.push(make_test(reftest));
        }
    }
    tests
}

fn make_test(reftest: Reftest) -> TestDescAndFn {
    let name = reftest.name.clone();
    TestDescAndFn {
        desc: TestDesc {
            name: DynTestName(name),
            ignore: false,
            should_fail: false,
        },
        testfn: DynTestFn(proc() {
            check_reftest(reftest);
        }),
    }
}

fn capture(reftest: &Reftest, side: uint) -> png::Image {
    let filename = format!("/tmp/servo-reftest-{:06u}-{:u}.png", reftest.id, side);
    let mut args = reftest.servo_args.clone();
    args.push_all_move(~["-f".to_owned(), "-o".to_owned(), filename.clone(), reftest.files[side].clone()]);

    let retval = match Process::status("./servo", args) {
        Ok(status) => status,
        Err(e) => fail!("failed to execute process: {}", e),
    };
    assert!(retval == ExitStatus(0));

    png::load_png(&from_str::<Path>(filename).unwrap()).unwrap()
}

fn check_reftest(reftest: Reftest) {
    let left  = capture(&reftest, 0);
    let right = capture(&reftest, 1);

    let pixels: ~[u8] = left.pixels.iter().zip(right.pixels.iter()).map(|(&a, &b)| {
            if a as i8 - b as i8 == 0 {
                // White for correct
                0xFF
            } else {
                // "1100" in the RGBA channel with an error for an incorrect value
                // This results in some number of C0 and FFs, which is much more
                // readable (and distinguishable) than the previous difference-wise
                // scaling but does not require reconstructing the actual RGBA pixel.
                0xC0
            }
        }).collect();

    if pixels.iter().any(|&a| a < 255) {
        let output = from_str::<Path>(format!("/tmp/servo-reftest-{:06u}-diff.png", reftest.id)).unwrap();

        let img = png::Image {
            width: left.width,
            height: left.height,
            color_type: png::RGBA8,
            pixels: pixels,
        };
        let res = png::store_png(&img, &output);
        assert!(res.is_ok());

        assert!(reftest.kind == Different);
    } else {
        assert!(reftest.kind == Same);
    }
}
<commit_msg>Print the location of the image diff when a ref test fails.<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.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 png;
extern crate std;
extern crate test;

use std::io;
use std::io::{File, Reader, Process};
use std::io::process::ExitStatus;
use std::os;
use std::str;
use test::{DynTestName, DynTestFn, TestDesc, TestOpts, TestDescAndFn};
use test::run_tests_console;

fn main() {
    let args = os::args();
    let mut parts = args.tail().split(|e| "--" == e.as_slice());

    let files = parts.next().unwrap();  // .split() is never empty
    let servo_args = parts.next().unwrap_or(&[]);

    if files.len() == 0 {
        fail!("error: at least one reftest list must be given");
    }

    let tests = parse_lists(files, servo_args);
    let test_opts = TestOpts {
        filter: None,
        run_ignored: false,
        logfile: None,
        run_tests: true,
        run_benchmarks: false,
        ratchet_noise_percent: None,
        ratchet_metrics: None,
        save_metrics: None,
        test_shard: None,
    };

    match run_tests_console(&test_opts, tests) {
        Ok(false) => os::set_exit_status(1), // tests failed
        Err(_) => os::set_exit_status(2),    // I/O-related failure
        _ => (),
    }
}

#[deriving(Eq)]
enum ReftestKind {
    Same,
    Different,
}

struct Reftest {
    name: ~str,
    kind: ReftestKind,
    files: [~str, ..2],
    id: uint,
    servo_args: ~[~str],
}

fn parse_lists(filenames: &[~str], servo_args: &[~str]) -> Vec<TestDescAndFn> {
    let mut tests = Vec::new();
    let mut next_id = 0;
    for file in filenames.iter() {
        let file_path = Path::new(file.clone());
        let contents = match File::open_mode(&file_path, io::Open, io::Read)
            .and_then(|mut f| {
                f.read_to_str()
            }) {
                Ok(s) => s,
                _ => fail!("Could not read file"),
            };

        for line in contents.lines() {
            // ignore comments
            if line.starts_with("#") {
                continue;
            }

            let parts: ~[&str] = line.split(' ').filter(|p| !p.is_empty()).collect();

            if parts.len() != 3 {
                fail!("reftest line: '{:s}' doesn't match 'KIND LEFT RIGHT'", line);
            }

            let kind = match parts[0] {
                "==" => Same,
                "!=" => Different,
                _ => fail!("reftest line: '{:s}' has invalid kind '{:s}'",
                           line, parts[0])
            };
            let src_path = file_path.dir_path();
            let src_dir = src_path.display().to_str();
            let file_left =  src_dir + "/" + parts[1];
            let file_right = src_dir + "/" + parts[2];

            let reftest = Reftest {
                name: parts[1] + " / " + parts[2],
                kind: kind,
                files: [file_left, file_right],
                id: next_id,
                servo_args: servo_args.to_owned(),
            };

            next_id += 1;

            tests.push(make_test(reftest));
        }
    }
    tests
}

fn make_test(reftest: Reftest) -> TestDescAndFn {
    let name = reftest.name.clone();
    TestDescAndFn {
        desc: TestDesc {
            name: DynTestName(name),
            ignore: false,
            should_fail: false,
        },
        testfn: DynTestFn(proc() {
            check_reftest(reftest);
        }),
    }
}

fn capture(reftest: &Reftest, side: uint) -> png::Image {
    let filename = format!("/tmp/servo-reftest-{:06u}-{:u}.png", reftest.id, side);
    let mut args = reftest.servo_args.clone();
    args.push_all_move(~["-f".to_owned(), "-o".to_owned(), filename.clone(), reftest.files[side].clone()]);

    let retval = match Process::status("./servo", args) {
        Ok(status) => status,
        Err(e) => fail!("failed to execute process: {}", e),
    };
    assert!(retval == ExitStatus(0));

    png::load_png(&from_str::<Path>(filename).unwrap()).unwrap()
}

fn check_reftest(reftest: Reftest) {
    let left  = capture(&reftest, 0);
    let right = capture(&reftest, 1);

    let pixels: ~[u8] = left.pixels.iter().zip(right.pixels.iter()).map(|(&a, &b)| {
            if a as i8 - b as i8 == 0 {
                // White for correct
                0xFF
            } else {
                // "1100" in the RGBA channel with an error for an incorrect value
                // This results in some number of C0 and FFs, which is much more
                // readable (and distinguishable) than the previous difference-wise
                // scaling but does not require reconstructing the actual RGBA pixel.
                0xC0
            }
        }).collect();

    if pixels.iter().any(|&a| a < 255) {
        let output_str = format!("/tmp/servo-reftest-{:06u}-diff.png", reftest.id);
        let output = from_str::<Path>(output_str).unwrap();

        let img = png::Image {
            width: left.width,
            height: left.height,
            color_type: png::RGBA8,
            pixels: pixels,
        };
        let res = png::store_png(&img, &output);
        assert!(res.is_ok());

        assert!(reftest.kind == Different, "rendering difference: {}", output_str);
    } else {
        assert!(reftest.kind == Same);
    }
}
<|endoftext|>


<commit_before>fn test_align_offset() {
    let d = Box::new([0u32; 4]);
    // Get u8 pointer to base
    let raw = d.as_ptr() as *const u8;

    assert_eq!(raw.align_offset(2), 0);
    assert_eq!(raw.align_offset(4), 0);
    assert_eq!(raw.align_offset(8), usize::max_value()); // requested alignment higher than allocation alignment

    assert_eq!(raw.wrapping_offset(1).align_offset(2), 1);
    assert_eq!(raw.wrapping_offset(1).align_offset(4), 3);
    assert_eq!(raw.wrapping_offset(1).align_offset(8), usize::max_value()); // requested alignment higher than allocation alignment

    assert_eq!(raw.wrapping_offset(2).align_offset(2), 0);
    assert_eq!(raw.wrapping_offset(2).align_offset(4), 2);
    assert_eq!(raw.wrapping_offset(2).align_offset(8), usize::max_value()); // requested alignment higher than allocation alignment
}

fn test_align_to() {
    const N: usize = 4;
    let d = Box::new([0u32; N]);
    // Get u8 slice covering the entire thing
    let s = unsafe { std::slice::from_raw_parts(d.as_ptr() as *const u8, 4 * N) };
    let raw = s.as_ptr();

    {
        let (l, m, r) = unsafe { s.align_to::<u32>() };
        assert_eq!(l.len(), 0);
        assert_eq!(r.len(), 0);
        assert_eq!(m.len(), N);
        assert_eq!(raw, m.as_ptr() as *const u8);
    }

    {
        let (l, m, r) = unsafe { s[1..].align_to::<u32>() };
        assert_eq!(l.len(), 3);
        assert_eq!(m.len(), N-1);
        assert_eq!(r.len(), 0);
        assert_eq!(raw.wrapping_offset(4), m.as_ptr() as *const u8);
    }

    {
        let (l, m, r) = unsafe { s[..4*N - 1].align_to::<u32>() };
        assert_eq!(l.len(), 0);
        assert_eq!(m.len(), N-1);
        assert_eq!(r.len(), 3);
        assert_eq!(raw, m.as_ptr() as *const u8);
    }

    {
        let (l, m, r) = unsafe { s[1..4*N - 1].align_to::<u32>() };
        assert_eq!(l.len(), 3);
        assert_eq!(m.len(), N-2);
        assert_eq!(r.len(), 3);
        assert_eq!(raw.wrapping_offset(4), m.as_ptr() as *const u8);
    }

    {
        #[repr(align(8))] struct Align8(u64);
        let (l, m, r) = unsafe { s.align_to::<Align8>() }; // requested alignment higher than allocation alignment
        assert_eq!(l.len(), 4*N);
        assert_eq!(r.len(), 0);
        assert_eq!(m.len(), 0);
    }
}

fn test_from_utf8() {
    const N: usize = 10;
    let vec = vec![0x4141414141414141u64; N];
    let content = unsafe { std::slice::from_raw_parts(vec.as_ptr() as *const u8, 8 * N) };
    println!("{:?}", std::str::from_utf8(content).unwrap());
}

fn main() {
    test_align_offset();
    test_align_to();
    test_from_utf8();
}
<commit_msg>test align_to example<commit_after>fn test_align_offset() {
    let d = Box::new([0u32; 4]);
    // Get u8 pointer to base
    let raw = d.as_ptr() as *const u8;

    assert_eq!(raw.align_offset(2), 0);
    assert_eq!(raw.align_offset(4), 0);
    assert_eq!(raw.align_offset(8), usize::max_value()); // requested alignment higher than allocation alignment

    assert_eq!(raw.wrapping_offset(1).align_offset(2), 1);
    assert_eq!(raw.wrapping_offset(1).align_offset(4), 3);
    assert_eq!(raw.wrapping_offset(1).align_offset(8), usize::max_value()); // requested alignment higher than allocation alignment

    assert_eq!(raw.wrapping_offset(2).align_offset(2), 0);
    assert_eq!(raw.wrapping_offset(2).align_offset(4), 2);
    assert_eq!(raw.wrapping_offset(2).align_offset(8), usize::max_value()); // requested alignment higher than allocation alignment
}

fn test_align_to() {
    const N: usize = 4;
    let d = Box::new([0u32; N]);
    // Get u8 slice covering the entire thing
    let s = unsafe { std::slice::from_raw_parts(d.as_ptr() as *const u8, 4 * N) };
    let raw = s.as_ptr();

    {
        let (l, m, r) = unsafe { s.align_to::<u32>() };
        assert_eq!(l.len(), 0);
        assert_eq!(r.len(), 0);
        assert_eq!(m.len(), N);
        assert_eq!(raw, m.as_ptr() as *const u8);
    }

    {
        let (l, m, r) = unsafe { s[1..].align_to::<u32>() };
        assert_eq!(l.len(), 3);
        assert_eq!(m.len(), N-1);
        assert_eq!(r.len(), 0);
        assert_eq!(raw.wrapping_offset(4), m.as_ptr() as *const u8);
    }

    {
        let (l, m, r) = unsafe { s[..4*N - 1].align_to::<u32>() };
        assert_eq!(l.len(), 0);
        assert_eq!(m.len(), N-1);
        assert_eq!(r.len(), 3);
        assert_eq!(raw, m.as_ptr() as *const u8);
    }

    {
        let (l, m, r) = unsafe { s[1..4*N - 1].align_to::<u32>() };
        assert_eq!(l.len(), 3);
        assert_eq!(m.len(), N-2);
        assert_eq!(r.len(), 3);
        assert_eq!(raw.wrapping_offset(4), m.as_ptr() as *const u8);
    }

    {
        #[repr(align(8))] struct Align8(u64);
        let (l, m, r) = unsafe { s.align_to::<Align8>() }; // requested alignment higher than allocation alignment
        assert_eq!(l.len(), 4*N);
        assert_eq!(r.len(), 0);
        assert_eq!(m.len(), 0);
    }
}

fn test_from_utf8() {
    const N: usize = 10;
    let vec = vec![0x4141414141414141u64; N];
    let content = unsafe { std::slice::from_raw_parts(vec.as_ptr() as *const u8, 8 * N) };
    println!("{:?}", std::str::from_utf8(content).unwrap());
}

fn test_u64_array() {
    #[derive(Default)]
    #[repr(align(8))]
    struct AlignToU64<T>(T);

    const BYTE_LEN: usize = std::mem::size_of::<[u64; 4]>();
    type Data = AlignToU64<[u8; BYTE_LEN]>;

    fn example(data: &Data) {
        let (head, u64_arrays, tail) = unsafe { data.0.align_to::<[u64; 4]>() };

        assert!(head.is_empty(), "buffer was not aligned for 64-bit numbers");
        assert_eq!(u64_arrays.len(), 1, "buffer was not long enough");
        assert!(tail.is_empty(), "buffer was too long");

        let u64_array = &u64_arrays[0];
        let _val = u64_array[0]; // make sure we can actually load this
    }

    example(&Data::default());
}

fn main() {
    test_align_offset();
    test_align_to();
    test_from_utf8();
    test_u64_array();
}
<|endoftext|>


<commit_before>// Copyright 2015 The Delix Project Authors. See the AUTHORS file at the top level directory.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

use std::collections::HashMap;
use std::result;
use std::sync::RwLock;

use node::{ID, request};
use transport::direct::{self, Link};

pub struct ServiceMap {
    balancer: Box<direct::Balancer>,
    entries: RwLock<HashMap<String, Entry>>,
}

struct Entry {
    local_handler: Option<Box<request::Handler>>,
    links: Vec<Link>,
    queue: Vec<Link>,
}

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

#[derive(Debug)]
pub enum Error {
    ServiceAlreadyExists,
    ServiceDoesNotExists,
    Connection(direct::ConnectionError),
    ConnectionMap(direct::ConnectionMapError),
}

impl ServiceMap {
    pub fn new(balancer: Box<direct::Balancer>) -> ServiceMap {
        ServiceMap {
            balancer: balancer,
            entries: RwLock::new(HashMap::new()),
        }
    }

    pub fn insert_local(&self, name: &str, f: Box<request::Handler>) -> Result<()> {
        let mut entries = self.entries.write().unwrap();

        if !entries.contains_key(name) {
            entries.insert(name.to_string(), Entry::new());
        }
        let mut entry = entries.get_mut(name).unwrap();

        if entry.local_handler.is_some() {
            return Err(Error::ServiceAlreadyExists);
        }

        entry.local_handler = Some(f);
        entry.links.push(Link::Local);

        Ok(())
    }

    pub fn insert_remote(&self, name: &str, peer_node_id: ID) -> Result<()> {
        let mut entries = self.entries.write().unwrap();

        if !entries.contains_key(name) {
            entries.insert(name.to_string(), Entry::new());
        }
        let mut entry = entries.get_mut(name).unwrap();

        if let Some(_) = entry.links.iter().find(|&link| Link::is_remote(link, &peer_node_id)) {
            return Err(Error::ServiceAlreadyExists);
        }

        entry.links.push(Link::Remote(peer_node_id));

        Ok(())
    }

    pub fn select<L, R>(&self, name: &str, local_handler: L, remote_handler: R) -> request::Response
        where L: Fn(&Box<request::Handler>) -> request::Response,
              R: Fn(ID) -> request::Response
    {
        let mut entries = self.entries.write().unwrap();

        loop {
            let response_pair = {
                let mut entry = match entries.get_mut(name) {
                    Some(entry) => entry,
                    None => return Err(request::Error::ServiceDoesNotExists),
                };

                if entry.queue.is_empty() {
                    entry.queue.append(&mut self.balancer.build_round(name, &entry.links));
                    entry.queue.reverse();
                }
                let link = entry.queue.pop().expect("balancer did not build any round");

                match link {
                    Link::Local => (local_handler(entry.local_handler.as_ref().unwrap()), None),
                    Link::Remote(ref peer_node_id) => {
                        (remote_handler(*peer_node_id), Some(*peer_node_id))
                    }
                }
            };

            if let (Err(request::Error::ServiceDoesNotExists),
                    Some(peer_node_id)) = response_pair {
                let mut remove_entry = false;
                if let Some(entry) = entries.get_mut(name) {
                    entry.links.retain(|link| !Link::is_remote(&link, &peer_node_id));
                    remove_entry = entry.links.is_empty();
                }
                if remove_entry {
                    entries.remove(name);
                }
            } else {
                return response_pair.0;
            }
        }
    }

    pub fn select_local<L>(&self, name: &str, local_handler: L) -> request::Response
        where L: Fn(&Box<request::Handler>) -> request::Response
    {
        let entries = self.entries.read().unwrap();

        let entry = match entries.get(name) {
            Some(entry) => entry,
            None => return Err(request::Error::ServiceDoesNotExists),
        };

        let link = match entry.links.iter().find(|link| Link::is_local(link)) {
            Some(link) => link,
            None => return Err(request::Error::ServiceDoesNotExists),
        };

        if let Link::Local = *link {
            local_handler(entry.local_handler.as_ref().unwrap())
        } else {
            unreachable!();
        }
    }

    pub fn local_service_names(&self) -> Vec<String> {
        self.entries
            .read()
            .unwrap()
            .iter()
            .filter_map(|(name, links)| {
                links.links.iter().find(|link| Link::is_local(link)).and(Some(name.to_string()))
            })
            .collect()
    }

    pub fn len(&self) -> usize {
        self.entries.read().unwrap().len()
    }

    pub fn remove_local(&self, name: &str) -> Result<()> {
        let mut entries = self.entries.write().unwrap();
        let remove = {
            let mut entry = match entries.get_mut(name) {
                Some(entry) => entry,
                None => return Err(Error::ServiceDoesNotExists),
            };
            entry.local_handler = None;
            entry.links.retain(|link| !Link::is_local(&&link));
            entry.links.len() == 0
        };
        if remove {
            entries.remove(name);
        }
        Ok(())
    }

    pub fn remove_remote(&self, peer_node_id: &ID) {
        let mut entries = self.entries.write().unwrap();
        let mut names = Vec::new();
        for (name, entry) in entries.iter_mut() {
            entry.links.retain(|link| !Link::is_remote(link, peer_node_id));
            if entry.links.len() == 0 {
                names.push(name.to_string());
            }
        }
        for name in names {
            entries.remove(&name);
        }
    }
}

impl Entry {
    fn new() -> Entry {
        Entry {
            local_handler: None,
            links: Vec::new(),
            queue: Vec::new(),
        }
    }
}

unsafe impl Send for ServiceMap {}

unsafe impl Sync for ServiceMap {}

impl From<direct::ConnectionError> for Error {
    fn from(error: direct::ConnectionError) -> Self {
        Error::Connection(error)
    }
}

impl From<direct::ConnectionMapError> for Error {
    fn from(error: direct::ConnectionMapError) -> Self {
        Error::ConnectionMap(error)
    }
}

#[cfg(test)]
mod tests {

    use node::ID;
    use super::ServiceMap;
    use super::super::balancer;

    #[test]
    fn insert_local() {
        let balancer = Box::new(balancer::DynamicRoundRobin::new());
        let service_map = ServiceMap::new(balancer);

        assert!(service_map.insert_local("test", Box::new(|request| Ok(request.to_vec()))).is_ok());
        assert!(service_map.insert_local("test", Box::new(|request| Ok(request.to_vec())))
                           .is_err());
        assert!(service_map.insert_remote("test", ID::new_random()).is_ok());

        assert_eq!(vec!["test"], service_map.local_service_names());
    }

    #[test]
    fn insert_remote() {
        let balancer = Box::new(balancer::DynamicRoundRobin::new());
        let service_map = ServiceMap::new(balancer);
        let node_id = ID::new_random();

        assert!(service_map.insert_remote("test", node_id).is_ok());
        assert!(service_map.insert_remote("test", node_id).is_err());
        assert!(service_map.insert_remote("test", ID::new_random()).is_ok());
        assert!(service_map.insert_local("test", Box::new(|request| Ok(request.to_vec()))).is_ok());

        assert_eq!(vec!["test"], service_map.local_service_names());
    }

    #[test]
    fn remove_local() {
        let balancer = Box::new(balancer::DynamicRoundRobin::new());
        let service_map = ServiceMap::new(balancer);
        service_map.insert_local("test", Box::new(|request| Ok(request.to_vec()))).unwrap();
        service_map.insert_remote("test", ID::new_random()).unwrap();

        assert!(service_map.remove_local("test").is_ok());

        assert_eq!(1, service_map.len());
    }

    #[test]
    fn remove_local_and_clean_up() {
        let balancer = Box::new(balancer::DynamicRoundRobin::new());
        let service_map = ServiceMap::new(balancer);
        service_map.insert_local("test", Box::new(|request| Ok(request.to_vec()))).unwrap();

        assert!(service_map.remove_local("test").is_ok());

        assert_eq!(0, service_map.len());
    }

    #[test]
    fn remove_remote() {
        let balancer = Box::new(balancer::DynamicRoundRobin::new());
        let service_map = ServiceMap::new(balancer);
        let node_id = ID::new_random();
        service_map.insert_remote("test", node_id).unwrap();
        service_map.insert_local("test", Box::new(|request| Ok(request.to_vec()))).unwrap();

        service_map.remove_remote(&node_id);

        assert_eq!(1, service_map.len());
    }

    #[test]
    fn remove_remote_and_clean_up() {
        let balancer = Box::new(balancer::DynamicRoundRobin::new());
        let service_map = ServiceMap::new(balancer);
        let node_id = ID::new_random();
        service_map.insert_remote("test", node_id).unwrap();

        service_map.remove_remote(&node_id);

        assert_eq!(0, service_map.len());
    }

}
<commit_msg>clean up load balancer queue on service removal<commit_after>// Copyright 2015 The Delix Project Authors. See the AUTHORS file at the top level directory.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

use std::collections::HashMap;
use std::result;
use std::sync::RwLock;

use node::{ID, request};
use transport::direct::{self, Link};

pub struct ServiceMap {
    balancer: Box<direct::Balancer>,
    entries: RwLock<HashMap<String, Entry>>,
}

struct Entry {
    local_handler: Option<Box<request::Handler>>,
    links: Vec<Link>,
    queue: Vec<Link>,
}

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

#[derive(Debug)]
pub enum Error {
    ServiceAlreadyExists,
    ServiceDoesNotExists,
    Connection(direct::ConnectionError),
    ConnectionMap(direct::ConnectionMapError),
}

impl ServiceMap {
    pub fn new(balancer: Box<direct::Balancer>) -> ServiceMap {
        ServiceMap {
            balancer: balancer,
            entries: RwLock::new(HashMap::new()),
        }
    }

    pub fn insert_local(&self, name: &str, f: Box<request::Handler>) -> Result<()> {
        let mut entries = self.entries.write().unwrap();

        if !entries.contains_key(name) {
            entries.insert(name.to_string(), Entry::new());
        }
        let mut entry = entries.get_mut(name).unwrap();

        if entry.local_handler.is_some() {
            return Err(Error::ServiceAlreadyExists);
        }

        entry.local_handler = Some(f);
        entry.links.push(Link::Local);

        Ok(())
    }

    pub fn insert_remote(&self, name: &str, peer_node_id: ID) -> Result<()> {
        let mut entries = self.entries.write().unwrap();

        if !entries.contains_key(name) {
            entries.insert(name.to_string(), Entry::new());
        }
        let mut entry = entries.get_mut(name).unwrap();

        if let Some(_) = entry.links.iter().find(|&link| Link::is_remote(link, &peer_node_id)) {
            return Err(Error::ServiceAlreadyExists);
        }

        entry.links.push(Link::Remote(peer_node_id));

        Ok(())
    }

    pub fn select<L, R>(&self, name: &str, local_handler: L, remote_handler: R) -> request::Response
        where L: Fn(&Box<request::Handler>) -> request::Response,
              R: Fn(ID) -> request::Response
    {
        let mut entries = self.entries.write().unwrap();

        loop {
            let response_pair = {
                let mut entry = match entries.get_mut(name) {
                    Some(entry) => entry,
                    None => return Err(request::Error::ServiceDoesNotExists),
                };

                if entry.queue.is_empty() {
                    entry.queue.append(&mut self.balancer.build_round(name, &entry.links));
                    entry.queue.reverse();
                }
                let link = entry.queue.pop().expect("balancer did not build any round");

                match link {
                    Link::Local => (local_handler(entry.local_handler.as_ref().unwrap()), None),
                    Link::Remote(ref peer_node_id) => {
                        (remote_handler(*peer_node_id), Some(*peer_node_id))
                    }
                }
            };

            if let (Err(request::Error::ServiceDoesNotExists),
                    Some(peer_node_id)) = response_pair {
                let mut remove_entry = false;
                if let Some(entry) = entries.get_mut(name) {
                    entry.links.retain(|link| !Link::is_remote(&link, &peer_node_id));
                    remove_entry = entry.links.is_empty();
                }
                if remove_entry {
                    entries.remove(name);
                }
            } else {
                return response_pair.0;
            }
        }
    }

    pub fn select_local<L>(&self, name: &str, local_handler: L) -> request::Response
        where L: Fn(&Box<request::Handler>) -> request::Response
    {
        let entries = self.entries.read().unwrap();

        let entry = match entries.get(name) {
            Some(entry) => entry,
            None => return Err(request::Error::ServiceDoesNotExists),
        };

        let link = match entry.links.iter().find(|link| Link::is_local(link)) {
            Some(link) => link,
            None => return Err(request::Error::ServiceDoesNotExists),
        };

        if let Link::Local = *link {
            local_handler(entry.local_handler.as_ref().unwrap())
        } else {
            unreachable!();
        }
    }

    pub fn local_service_names(&self) -> Vec<String> {
        self.entries
            .read()
            .unwrap()
            .iter()
            .filter_map(|(name, links)| {
                links.links.iter().find(|link| Link::is_local(link)).and(Some(name.to_string()))
            })
            .collect()
    }

    pub fn len(&self) -> usize {
        self.entries.read().unwrap().len()
    }

    pub fn remove_local(&self, name: &str) -> Result<()> {
        let mut entries = self.entries.write().unwrap();
        let remove = {
            let mut entry = match entries.get_mut(name) {
                Some(entry) => entry,
                None => return Err(Error::ServiceDoesNotExists),
            };
            entry.local_handler = None;
            entry.links.retain(|link| !Link::is_local(&&link));
            entry.queue.retain(|link| !Link::is_local(&&link));
            entry.links.len() == 0
        };
        if remove {
            entries.remove(name);
        }
        Ok(())
    }

    pub fn remove_remote(&self, peer_node_id: &ID) {
        let mut entries = self.entries.write().unwrap();
        let mut names = Vec::new();
        for (name, entry) in entries.iter_mut() {
            entry.links.retain(|link| !Link::is_remote(link, peer_node_id));
            entry.queue.retain(|link| !Link::is_remote(link, peer_node_id));
            if entry.links.len() == 0 {
                names.push(name.to_string());
            }
        }
        for name in names {
            entries.remove(&name);
        }
    }
}

impl Entry {
    fn new() -> Entry {
        Entry {
            local_handler: None,
            links: Vec::new(),
            queue: Vec::new(),
        }
    }
}

unsafe impl Send for ServiceMap {}

unsafe impl Sync for ServiceMap {}

impl From<direct::ConnectionError> for Error {
    fn from(error: direct::ConnectionError) -> Self {
        Error::Connection(error)
    }
}

impl From<direct::ConnectionMapError> for Error {
    fn from(error: direct::ConnectionMapError) -> Self {
        Error::ConnectionMap(error)
    }
}

#[cfg(test)]
mod tests {

    use node::ID;
    use super::ServiceMap;
    use super::super::balancer;

    #[test]
    fn insert_local() {
        let balancer = Box::new(balancer::DynamicRoundRobin::new());
        let service_map = ServiceMap::new(balancer);

        assert!(service_map.insert_local("test", Box::new(|request| Ok(request.to_vec()))).is_ok());
        assert!(service_map.insert_local("test", Box::new(|request| Ok(request.to_vec())))
                           .is_err());
        assert!(service_map.insert_remote("test", ID::new_random()).is_ok());

        assert_eq!(vec!["test"], service_map.local_service_names());
    }

    #[test]
    fn insert_remote() {
        let balancer = Box::new(balancer::DynamicRoundRobin::new());
        let service_map = ServiceMap::new(balancer);
        let node_id = ID::new_random();

        assert!(service_map.insert_remote("test", node_id).is_ok());
        assert!(service_map.insert_remote("test", node_id).is_err());
        assert!(service_map.insert_remote("test", ID::new_random()).is_ok());
        assert!(service_map.insert_local("test", Box::new(|request| Ok(request.to_vec()))).is_ok());

        assert_eq!(vec!["test"], service_map.local_service_names());
    }

    #[test]
    fn remove_local() {
        let balancer = Box::new(balancer::DynamicRoundRobin::new());
        let service_map = ServiceMap::new(balancer);
        service_map.insert_local("test", Box::new(|request| Ok(request.to_vec()))).unwrap();
        service_map.insert_remote("test", ID::new_random()).unwrap();

        assert!(service_map.remove_local("test").is_ok());

        assert_eq!(1, service_map.len());
    }

    #[test]
    fn remove_local_and_clean_up() {
        let balancer = Box::new(balancer::DynamicRoundRobin::new());
        let service_map = ServiceMap::new(balancer);
        service_map.insert_local("test", Box::new(|request| Ok(request.to_vec()))).unwrap();

        assert!(service_map.remove_local("test").is_ok());

        assert_eq!(0, service_map.len());
    }

    #[test]
    fn remove_remote() {
        let balancer = Box::new(balancer::DynamicRoundRobin::new());
        let service_map = ServiceMap::new(balancer);
        let node_id = ID::new_random();
        service_map.insert_remote("test", node_id).unwrap();
        service_map.insert_local("test", Box::new(|request| Ok(request.to_vec()))).unwrap();

        service_map.remove_remote(&node_id);

        assert_eq!(1, service_map.len());
    }

    #[test]
    fn remove_remote_and_clean_up() {
        let balancer = Box::new(balancer::DynamicRoundRobin::new());
        let service_map = ServiceMap::new(balancer);
        let node_id = ID::new_random();
        service_map.insert_remote("test", node_id).unwrap();

        service_map.remove_remote(&node_id);

        assert_eq!(0, service_map.len());
    }

}
<|endoftext|>


<commit_before>extern crate getopts;

pub struct HelpText<'a> {
    pub name: &'a str,
    pub version: &'a str,
    pub syntax: &'a str,
    pub summary: &'a str,
    pub long_help: &'a str,
    pub display_usage: bool,
}

pub struct CoreOptions<'a> {
    options: getopts::Options,
    help_text: HelpText<'a>,
}

impl<'a> CoreOptions<'a> {
    pub fn new(help_text: HelpText<'a>) -> Self {
        let mut ret = CoreOptions {
            options: getopts::Options::new(),
            help_text,
        };
        ret.options
            .optflag("", "help", "print usage information")
            .optflag("", "version", "print name and version number");
        ret
    }
    pub fn optflagopt(
        &mut self,
        short_name: &str,
        long_name: &str,
        desc: &str,
        hint: &str,
    ) -> &mut CoreOptions<'a> {
        self.options.optflagopt(short_name, long_name, desc, hint);
        self
    }
    pub fn optflag(
        &mut self,
        short_name: &str,
        long_name: &str,
        desc: &str,
    ) -> &mut CoreOptions<'a> {
        self.options.optflag(short_name, long_name, desc);
        self
    }
    pub fn optflagmulti(
        &mut self,
        short_name: &str,
        long_name: &str,
        desc: &str,
    ) -> &mut CoreOptions<'a> {
        self.options.optflagmulti(short_name, long_name, desc);
        self
    }
    pub fn optopt(
        &mut self,
        short_name: &str,
        long_name: &str,
        desc: &str,
        hint: &str,
    ) -> &mut CoreOptions<'a> {
        self.options.optopt(short_name, long_name, desc, hint);
        self
    }
    pub fn optmulti(
        &mut self,
        short_name: &str,
        long_name: &str,
        desc: &str,
        hint: &str,
    ) -> &mut CoreOptions<'a> {
        self.options.optmulti(short_name, long_name, desc, hint);
        self
    }
    pub fn usage(&self, summary: &str) -> String {
        self.options.usage(summary)
    }
    pub fn parse(&mut self, args: Vec<String>) -> getopts::Matches {
        let matches = match self.options.parse(&args[1..]) {
            Ok(m) => Some(m),
            Err(f) => {
                eprint!("{}: error: ", self.help_text.name);
                eprintln!("{}", f);
                ::std::process::exit(1);
            }
        }
        .unwrap();
        if matches.opt_present("help") {
            let usage_str = if self.help_text.display_usage {
                format!(
                    "\n {}\n\n Reference\n",
                    self.options.usage(self.help_text.summary)
                )
                .replace("Options:", " Options:")
            } else {
                String::new()
            };
            print!(
                "
 {0} {1}

 {0} {2}
{3}{4}
",
                self.help_text.name,
                self.help_text.version,
                self.help_text.syntax,
                usage_str,
                self.help_text.long_help
            );
            crate::exit!(0);
        } else if matches.opt_present("version") {
            println!("{} {}", self.help_text.name, self.help_text.version);
            crate::exit!(0);
        }
        matches
    }
}

#[macro_export]
macro_rules! new_coreopts {
    ($syntax: expr, $summary: expr, $long_help: expr) => {
        uucore::coreopts::CoreOptions::new(uucore::coreopts::HelpText {
            name: executable!(),
            version: env!("CARGO_PKG_VERSION"),
            syntax: $syntax,
            summary: $summary,
            long_help: $long_help,
            display_usage: true,
        })
    };
    ($syntax: expr, $summary: expr, $long_help: expr, $display_usage: expr) => {
        uucore::coreopts::CoreOptions::new(uucore::coreopts::HelpText {
            name: executable!(),
            version: env!("CARGO_PKG_VERSION"),
            syntax: $syntax,
            summary: $summary,
            long_help: $long_help,
            display_usage: $display_usage,
        })
    };
}
<commit_msg>change/API! ~ `new_coreopts!()` => `app!()`<commit_after>extern crate getopts;

pub struct HelpText<'a> {
    pub name: &'a str,
    pub version: &'a str,
    pub syntax: &'a str,
    pub summary: &'a str,
    pub long_help: &'a str,
    pub display_usage: bool,
}

pub struct CoreOptions<'a> {
    options: getopts::Options,
    help_text: HelpText<'a>,
}

impl<'a> CoreOptions<'a> {
    pub fn new(help_text: HelpText<'a>) -> Self {
        let mut ret = CoreOptions {
            options: getopts::Options::new(),
            help_text,
        };
        ret.options
            .optflag("", "help", "print usage information")
            .optflag("", "version", "print name and version number");
        ret
    }
    pub fn optflagopt(
        &mut self,
        short_name: &str,
        long_name: &str,
        desc: &str,
        hint: &str,
    ) -> &mut CoreOptions<'a> {
        self.options.optflagopt(short_name, long_name, desc, hint);
        self
    }
    pub fn optflag(
        &mut self,
        short_name: &str,
        long_name: &str,
        desc: &str,
    ) -> &mut CoreOptions<'a> {
        self.options.optflag(short_name, long_name, desc);
        self
    }
    pub fn optflagmulti(
        &mut self,
        short_name: &str,
        long_name: &str,
        desc: &str,
    ) -> &mut CoreOptions<'a> {
        self.options.optflagmulti(short_name, long_name, desc);
        self
    }
    pub fn optopt(
        &mut self,
        short_name: &str,
        long_name: &str,
        desc: &str,
        hint: &str,
    ) -> &mut CoreOptions<'a> {
        self.options.optopt(short_name, long_name, desc, hint);
        self
    }
    pub fn optmulti(
        &mut self,
        short_name: &str,
        long_name: &str,
        desc: &str,
        hint: &str,
    ) -> &mut CoreOptions<'a> {
        self.options.optmulti(short_name, long_name, desc, hint);
        self
    }
    pub fn usage(&self, summary: &str) -> String {
        self.options.usage(summary)
    }
    pub fn parse(&mut self, args: Vec<String>) -> getopts::Matches {
        let matches = match self.options.parse(&args[1..]) {
            Ok(m) => Some(m),
            Err(f) => {
                eprint!("{}: error: ", self.help_text.name);
                eprintln!("{}", f);
                ::std::process::exit(1);
            }
        }
        .unwrap();
        if matches.opt_present("help") {
            let usage_str = if self.help_text.display_usage {
                format!(
                    "\n {}\n\n Reference\n",
                    self.options.usage(self.help_text.summary)
                )
                .replace("Options:", " Options:")
            } else {
                String::new()
            };
            print!(
                "
 {0} {1}

 {0} {2}
{3}{4}
",
                self.help_text.name,
                self.help_text.version,
                self.help_text.syntax,
                usage_str,
                self.help_text.long_help
            );
            crate::exit!(0);
        } else if matches.opt_present("version") {
            println!("{} {}", self.help_text.name, self.help_text.version);
            crate::exit!(0);
        }
        matches
    }
}

#[macro_export]
macro_rules! app {
    ($syntax: expr, $summary: expr, $long_help: expr) => {
        uucore::coreopts::CoreOptions::new(uucore::coreopts::HelpText {
            name: executable!(),
            version: env!("CARGO_PKG_VERSION"),
            syntax: $syntax,
            summary: $summary,
            long_help: $long_help,
            display_usage: true,
        })
    };
    ($syntax: expr, $summary: expr, $long_help: expr, $display_usage: expr) => {
        uucore::coreopts::CoreOptions::new(uucore::coreopts::HelpText {
            name: executable!(),
            version: env!("CARGO_PKG_VERSION"),
            syntax: $syntax,
            summary: $summary,
            long_help: $long_help,
            display_usage: $display_usage,
        })
    };
}
<|endoftext|>


<commit_before>// Copyright 2017 Mozilla Foundation
// Copyright 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::cell::RefCell;
use std::fmt;
use std::fs::{self, File};
use std::io::{self, Read};
use std::rc::Rc;
use std::time;

use cache::{
    Cache,
    CacheRead,
    CacheWrite,
    Storage,
};
use chrono;
use futures::future::Shared;
use futures::{future, Async, Future, Stream};
use hyper;
use hyper::header::{Authorization, Bearer, ContentType, ContentLength};
use hyper::Method;
use hyper::client::{Client, HttpConnector, Request};
use hyper_tls::HttpsConnector;
use jwt;
use openssl;
use serde_json;
use tokio_core::reactor::Handle;
use url::form_urlencoded;
use url::percent_encoding::{percent_encode, PATH_SEGMENT_ENCODE_SET, QUERY_ENCODE_SET};

use errors::*;

type HyperClient = Client<HttpsConnector<HttpConnector>>;

/// A GCS bucket
struct Bucket {
    name: String,
    base_url: String,
    client: HyperClient,
}

impl fmt::Display for Bucket {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Bucket(name={}, base_url={})", self.name, self.base_url)
    }
}

impl Bucket {
    pub fn new(name: String, base_url: String, handle: &Handle) -> Result<Bucket> {
        let client = Client::configure()
                        .connector(HttpsConnector::new(1, handle)?)
                        .build(handle);

        Ok(Bucket { name, base_url, client })
    }

    fn get(&self, key: &str, cred_provider: &GCSCredentialProvider) -> SFuture<Vec<u8>> {
        let url = format!("{}/download/storage/v1/b/{}/o/{}?alt=media",
                    self.base_url,
                    percent_encode(self.name.as_bytes(), PATH_SEGMENT_ENCODE_SET),
                    percent_encode(key.as_bytes(), PATH_SEGMENT_ENCODE_SET));

        let client = self.client.clone();

        Box::new(cred_provider.credentials(&self.client).and_then(move |creds| {
            let mut request = Request::new(Method::Get, url.parse().unwrap());
            request.headers_mut()
                .set(Authorization(Bearer { token: creds.token }));
            client.request(request).chain_err(move || {
                format!("failed GET: {}", url)
            }).and_then(|res| {
                if res.status().is_success() {
                    Ok(res.body())
                } else {
                    Err(ErrorKind::BadHTTPStatus(res.status().clone()).into())
                }
            }).and_then(|body| {
                body.fold(Vec::new(), |mut body, chunk| {
                    body.extend_from_slice(&chunk);
                    Ok::<_, hyper::Error>(body)
                }).chain_err(|| {
                    "failed to read HTTP body"
                })
            })
        }))
    }

    fn put(&self, key: &str, content: Vec<u8>, cred_provider: &GCSCredentialProvider) -> SFuture<()> {
        let url = format!("{}/upload/storage/v1/b/{}/o?name={}&uploadType=media",
                    self.base_url,
                    percent_encode(self.name.as_bytes(), PATH_SEGMENT_ENCODE_SET),
                    percent_encode(key.as_bytes(), QUERY_ENCODE_SET));

        let client = self.client.clone();

        Box::new(cred_provider.credentials(&client).and_then(move |creds| {
            let mut request = Request::new(Method::Post, url.parse().unwrap());
            {
                let mut headers = request.headers_mut();
                headers.set(Authorization(Bearer { token: creds.token }));
                headers.set(ContentType("application/octet-stream".parse().unwrap()));
                headers.set(ContentLength(content.len() as u64));
            }
            request.set_body(content);

            client.request(request).then(|result| {
                match result {
                    Ok(res) => {
                        if res.status().is_success() {
                            trace!("PUT succeeded");
                            Ok(())
                        } else {
                            trace!("PUT failed with HTTP status: {}", res.status());
                            Err(ErrorKind::BadHTTPStatus(res.status().clone()).into())
                        }
                    }
                    Err(e) => {
                        trace!("PUT failed with error: {:?}", e);
                        Err(e.into())
                    }
                }
            })
        }))
    }
}

pub struct GCSCredentialProvider {
    read_only: bool,
    credentials_path: String,
    cached_credentials: RefCell<Option<Shared<SFuture<GCSCredential>>>>,
}

#[derive(Debug, Deserialize)]
struct ServiceAccountKey {
    #[serde(rename = "type")]
    _type: String,
    project_id: String,
    private_key_id: String,
    private_key: String,
    client_email: String,
    client_id: String,
    auth_uri: String,
    token_uri: String,
    auth_provider_x509_cert_url: String,
}

#[derive(Serialize)]
struct JwtClaims {
    #[serde(rename = "iss")]
    issuer: String,
    scope: String,
    #[serde(rename = "aud")]
    audience: String,
    #[serde(rename = "exp")]
    expiration: i64,
    #[serde(rename = "iat")]
    issued_at: i64,
}

#[derive(Deserialize)]
struct TokenMsg {
    access_token: String,
    token_type: String,
    expires_in: usize,
}

#[derive(Clone)]
pub struct GCSCredential {
    token: String,
    expiration_time: chrono::DateTime<chrono::UTC>,
}

impl GCSCredentialProvider {
    pub fn new(read_only: bool, credentials_path: String) -> Self {
        GCSCredentialProvider {
            read_only,
            credentials_path,
            cached_credentials: RefCell::new(None),
        }
    }

    fn auth_request_jwt(&self, expire_at: &chrono::DateTime<chrono::UTC>) -> Result<String> {
        let metadata = fs::metadata(&self.credentials_path).chain_err(|| {
            "Couldn't stat GCS credentials file"
        })?;
        if !metadata.is_file() {
            bail!("Couldn't open GCS credentials file.");
        }
        let mut file = File::open(&self.credentials_path)?;
        let mut service_account_json = String::new();
        file.read_to_string(&mut service_account_json)?;
        let sa_key: ServiceAccountKey = serde_json::from_str(&service_account_json)?;

        let scope = (if self.read_only {
            "https://www.googleapis.com/auth/devstorage.readonly"
        } else {
            "https://www.googleapis.com/auth/devstorage.read_write"
        }).to_owned();

        let jwt_claims = JwtClaims {
            issuer: sa_key.client_email,
            scope: scope,
            audience: "https://www.googleapis.com/oauth2/v4/token".to_owned(),
            expiration: expire_at.timestamp(),
            issued_at: chrono::UTC::now().timestamp(),
        };

        let binary_key = openssl::rsa::Rsa::private_key_from_pem(
            sa_key.private_key.as_bytes()
        )?.private_key_to_der()?;

        let auth_request_jwt = jwt::encode(
            &jwt::Header::new(jwt::Algorithm::RS256),
            &jwt_claims,
            &binary_key,
        )?;

        Ok(auth_request_jwt)
    }

    pub fn credentials(&self, client: &HyperClient) -> SFuture<GCSCredential> {
        let mut future_opt = self.cached_credentials.borrow_mut();

        let needs_refresh = match Option::as_mut(&mut future_opt).map(|mut f| f.poll()) {
            None => true,
            Some(Ok(Async::Ready(ref creds))) => creds.expiration_time < chrono::UTC::now(),
            _ => false
        };

        if needs_refresh {
            let client = client.clone();
            let expires_at = chrono::UTC::now() + chrono::Duration::minutes(59);
            let auth_jwt = self.auth_request_jwt(&expires_at);
            let credentials: SFuture<_> = Box::new(future::result(auth_jwt).and_then(move |auth_jwt| {
                let url = "https://www.googleapis.com/oauth2/v4/token";
                let params = form_urlencoded::Serializer::new(String::new())
                    .append_pair("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer")
                    .append_pair("assertion", &auth_jwt)
                    .finish();

                let mut request = Request::new(Method::Post, url.parse().unwrap());
                {
                    let mut headers = request.headers_mut();
                    headers.set(ContentType("application/x-www-form-urlencoded".parse().unwrap()));
                    headers.set(ContentLength(params.len() as u64));
                }
                request.set_body(params);

                client.request(request).map_err(Into::into)
            }).and_then(move |res| {
                if res.status().is_success() {
                    Ok(res.body())
                } else {
                    Err(ErrorKind::BadHTTPStatus(res.status().clone()).into())
                }
            }).and_then(move |body| {
                body.fold(Vec::new(), |mut body, chunk| {
                    body.extend_from_slice(&chunk);
                    Ok::<_, hyper::Error>(body)
                }).chain_err(|| {
                    "failed to read HTTP body"
                })
            }).and_then(move |body| {
                let body_str = String::from_utf8(body)?;
                let token_msg: TokenMsg = serde_json::from_str(&body_str)?;
                Ok(GCSCredential {
                    token: token_msg.access_token,
                    expiration_time: expires_at,
                })
            }));

            *future_opt = Some(credentials.shared());
        };

        Box::new(Option::as_mut(&mut future_opt).unwrap().clone().then(|result| {
            match result {
                Ok(e) => Ok((*e).clone()),
                Err(e) => Err(e.to_string().into()),
            }
        }))
    }
}

/// A cache that stores entries in Google Cloud Storage
pub struct GCSCache {
    /// The GCS bucket
    bucket: Rc<Bucket>,
    /// Credential provider for GCS
    credential_provider: GCSCredentialProvider,
}

impl GCSCache {
    /// Create a new `GCSCache` storing data in `bucket`
    pub fn new(bucket: String,
               endpoint: String,
               credential_provider: GCSCredentialProvider,
               handle: &Handle) -> Result<GCSCache>
    {
        Ok(GCSCache {
            bucket: Rc::new(Bucket::new(bucket, endpoint, handle)?),
            credential_provider: credential_provider,
        })
    }
}

impl Storage for GCSCache {
    fn get(&self, key: &str) -> SFuture<Cache> {
        Box::new(self.bucket.get(&key, &self.credential_provider).then(|result| {
            match result {
                Ok(data) => {
                    let hit = CacheRead::from(io::Cursor::new(data))?;
                    Ok(Cache::Hit(hit))
                }
                Err(e) => {
                    warn!("Got GCS error: {:?}", e);
                    Ok(Cache::Miss)
                }
            }
        }))
    }

    fn put(&self, key: &str, entry: CacheWrite) -> SFuture<time::Duration> {
        let start = time::Instant::now();
        let data = match entry.finish() {
            Ok(data) => data,
            Err(e) => return future::err(e.into()).boxed(),
        };
        let bucket = self.bucket.clone();
        let response = bucket.put(&key, data, &self.credential_provider).chain_err(|| {
            "failed to put cache entry in GCS"
        });

        Box::new(response.map(move |_| start.elapsed()))
    }

    fn location(&self) -> String {
        format!("GCS, bucket: {}", self.bucket)
    }

    fn current_size(&self) -> Option<usize> { None }
    fn max_size(&self) -> Option<usize> { None }
}
<commit_msg>Cleanup extra GCS fields<commit_after>// Copyright 2017 Mozilla Foundation
// Copyright 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::cell::RefCell;
use std::fmt;
use std::fs::{self, File};
use std::io::{self, Read};
use std::rc::Rc;
use std::time;

use cache::{
    Cache,
    CacheRead,
    CacheWrite,
    Storage,
};
use chrono;
use futures::future::Shared;
use futures::{future, Async, Future, Stream};
use hyper;
use hyper::header::{Authorization, Bearer, ContentType, ContentLength};
use hyper::Method;
use hyper::client::{Client, HttpConnector, Request};
use hyper_tls::HttpsConnector;
use jwt;
use openssl;
use serde_json;
use tokio_core::reactor::Handle;
use url::form_urlencoded;
use url::percent_encoding::{percent_encode, PATH_SEGMENT_ENCODE_SET, QUERY_ENCODE_SET};

use errors::*;

type HyperClient = Client<HttpsConnector<HttpConnector>>;

/// A GCS bucket
struct Bucket {
    name: String,
    base_url: String,
    client: HyperClient,
}

impl fmt::Display for Bucket {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Bucket(name={}, base_url={})", self.name, self.base_url)
    }
}

impl Bucket {
    pub fn new(name: String, base_url: String, handle: &Handle) -> Result<Bucket> {
        let client = Client::configure()
                        .connector(HttpsConnector::new(1, handle)?)
                        .build(handle);

        Ok(Bucket { name, base_url, client })
    }

    fn get(&self, key: &str, cred_provider: &GCSCredentialProvider) -> SFuture<Vec<u8>> {
        let url = format!("{}/download/storage/v1/b/{}/o/{}?alt=media",
                    self.base_url,
                    percent_encode(self.name.as_bytes(), PATH_SEGMENT_ENCODE_SET),
                    percent_encode(key.as_bytes(), PATH_SEGMENT_ENCODE_SET));

        let client = self.client.clone();

        Box::new(cred_provider.credentials(&self.client).and_then(move |creds| {
            let mut request = Request::new(Method::Get, url.parse().unwrap());
            request.headers_mut()
                .set(Authorization(Bearer { token: creds.token }));
            client.request(request).chain_err(move || {
                format!("failed GET: {}", url)
            }).and_then(|res| {
                if res.status().is_success() {
                    Ok(res.body())
                } else {
                    Err(ErrorKind::BadHTTPStatus(res.status().clone()).into())
                }
            }).and_then(|body| {
                body.fold(Vec::new(), |mut body, chunk| {
                    body.extend_from_slice(&chunk);
                    Ok::<_, hyper::Error>(body)
                }).chain_err(|| {
                    "failed to read HTTP body"
                })
            })
        }))
    }

    fn put(&self, key: &str, content: Vec<u8>, cred_provider: &GCSCredentialProvider) -> SFuture<()> {
        let url = format!("{}/upload/storage/v1/b/{}/o?name={}&uploadType=media",
                    self.base_url,
                    percent_encode(self.name.as_bytes(), PATH_SEGMENT_ENCODE_SET),
                    percent_encode(key.as_bytes(), QUERY_ENCODE_SET));

        let client = self.client.clone();

        Box::new(cred_provider.credentials(&client).and_then(move |creds| {
            let mut request = Request::new(Method::Post, url.parse().unwrap());
            {
                let mut headers = request.headers_mut();
                headers.set(Authorization(Bearer { token: creds.token }));
                headers.set(ContentType("application/octet-stream".parse().unwrap()));
                headers.set(ContentLength(content.len() as u64));
            }
            request.set_body(content);

            client.request(request).then(|result| {
                match result {
                    Ok(res) => {
                        if res.status().is_success() {
                            trace!("PUT succeeded");
                            Ok(())
                        } else {
                            trace!("PUT failed with HTTP status: {}", res.status());
                            Err(ErrorKind::BadHTTPStatus(res.status().clone()).into())
                        }
                    }
                    Err(e) => {
                        trace!("PUT failed with error: {:?}", e);
                        Err(e.into())
                    }
                }
            })
        }))
    }
}

pub struct GCSCredentialProvider {
    read_only: bool,
    credentials_path: String,
    cached_credentials: RefCell<Option<Shared<SFuture<GCSCredential>>>>,
}

#[derive(Debug, Deserialize)]
struct ServiceAccountKey {
    private_key: String,
    client_email: String,
}

#[derive(Serialize)]
struct JwtClaims {
    #[serde(rename = "iss")]
    issuer: String,
    scope: String,
    #[serde(rename = "aud")]
    audience: String,
    #[serde(rename = "exp")]
    expiration: i64,
    #[serde(rename = "iat")]
    issued_at: i64,
}

#[derive(Deserialize)]
struct TokenMsg {
    access_token: String,
}

#[derive(Clone)]
pub struct GCSCredential {
    token: String,
    expiration_time: chrono::DateTime<chrono::UTC>,
}

impl GCSCredentialProvider {
    pub fn new(read_only: bool, credentials_path: String) -> Self {
        GCSCredentialProvider {
            read_only,
            credentials_path,
            cached_credentials: RefCell::new(None),
        }
    }

    fn auth_request_jwt(&self, expire_at: &chrono::DateTime<chrono::UTC>) -> Result<String> {
        let metadata = fs::metadata(&self.credentials_path).chain_err(|| {
            "Couldn't stat GCS credentials file"
        })?;
        if !metadata.is_file() {
            bail!("Couldn't open GCS credentials file.");
        }
        let mut file = File::open(&self.credentials_path)?;
        let mut service_account_json = String::new();
        file.read_to_string(&mut service_account_json)?;
        let sa_key: ServiceAccountKey = serde_json::from_str(&service_account_json)?;

        let scope = (if self.read_only {
            "https://www.googleapis.com/auth/devstorage.readonly"
        } else {
            "https://www.googleapis.com/auth/devstorage.read_write"
        }).to_owned();

        let jwt_claims = JwtClaims {
            issuer: sa_key.client_email,
            scope: scope,
            audience: "https://www.googleapis.com/oauth2/v4/token".to_owned(),
            expiration: expire_at.timestamp(),
            issued_at: chrono::UTC::now().timestamp(),
        };

        let binary_key = openssl::rsa::Rsa::private_key_from_pem(
            sa_key.private_key.as_bytes()
        )?.private_key_to_der()?;

        let auth_request_jwt = jwt::encode(
            &jwt::Header::new(jwt::Algorithm::RS256),
            &jwt_claims,
            &binary_key,
        )?;

        Ok(auth_request_jwt)
    }

    pub fn credentials(&self, client: &HyperClient) -> SFuture<GCSCredential> {
        let mut future_opt = self.cached_credentials.borrow_mut();

        let needs_refresh = match Option::as_mut(&mut future_opt).map(|mut f| f.poll()) {
            None => true,
            Some(Ok(Async::Ready(ref creds))) => creds.expiration_time < chrono::UTC::now(),
            _ => false
        };

        if needs_refresh {
            let client = client.clone();
            let expires_at = chrono::UTC::now() + chrono::Duration::minutes(59);
            let auth_jwt = self.auth_request_jwt(&expires_at);
            let credentials: SFuture<_> = Box::new(future::result(auth_jwt).and_then(move |auth_jwt| {
                let url = "https://www.googleapis.com/oauth2/v4/token";
                let params = form_urlencoded::Serializer::new(String::new())
                    .append_pair("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer")
                    .append_pair("assertion", &auth_jwt)
                    .finish();

                let mut request = Request::new(Method::Post, url.parse().unwrap());
                {
                    let mut headers = request.headers_mut();
                    headers.set(ContentType("application/x-www-form-urlencoded".parse().unwrap()));
                    headers.set(ContentLength(params.len() as u64));
                }
                request.set_body(params);

                client.request(request).map_err(Into::into)
            }).and_then(move |res| {
                if res.status().is_success() {
                    Ok(res.body())
                } else {
                    Err(ErrorKind::BadHTTPStatus(res.status().clone()).into())
                }
            }).and_then(move |body| {
                body.fold(Vec::new(), |mut body, chunk| {
                    body.extend_from_slice(&chunk);
                    Ok::<_, hyper::Error>(body)
                }).chain_err(|| {
                    "failed to read HTTP body"
                })
            }).and_then(move |body| {
                let body_str = String::from_utf8(body)?;
                let token_msg: TokenMsg = serde_json::from_str(&body_str)?;
                Ok(GCSCredential {
                    token: token_msg.access_token,
                    expiration_time: expires_at,
                })
            }));

            *future_opt = Some(credentials.shared());
        };

        Box::new(Option::as_mut(&mut future_opt).unwrap().clone().then(|result| {
            match result {
                Ok(e) => Ok((*e).clone()),
                Err(e) => Err(e.to_string().into()),
            }
        }))
    }
}

/// A cache that stores entries in Google Cloud Storage
pub struct GCSCache {
    /// The GCS bucket
    bucket: Rc<Bucket>,
    /// Credential provider for GCS
    credential_provider: GCSCredentialProvider,
}

impl GCSCache {
    /// Create a new `GCSCache` storing data in `bucket`
    pub fn new(bucket: String,
               endpoint: String,
               credential_provider: GCSCredentialProvider,
               handle: &Handle) -> Result<GCSCache>
    {
        Ok(GCSCache {
            bucket: Rc::new(Bucket::new(bucket, endpoint, handle)?),
            credential_provider: credential_provider,
        })
    }
}

impl Storage for GCSCache {
    fn get(&self, key: &str) -> SFuture<Cache> {
        Box::new(self.bucket.get(&key, &self.credential_provider).then(|result| {
            match result {
                Ok(data) => {
                    let hit = CacheRead::from(io::Cursor::new(data))?;
                    Ok(Cache::Hit(hit))
                }
                Err(e) => {
                    warn!("Got GCS error: {:?}", e);
                    Ok(Cache::Miss)
                }
            }
        }))
    }

    fn put(&self, key: &str, entry: CacheWrite) -> SFuture<time::Duration> {
        let start = time::Instant::now();
        let data = match entry.finish() {
            Ok(data) => data,
            Err(e) => return future::err(e.into()).boxed(),
        };
        let bucket = self.bucket.clone();
        let response = bucket.put(&key, data, &self.credential_provider).chain_err(|| {
            "failed to put cache entry in GCS"
        });

        Box::new(response.map(move |_| start.elapsed()))
    }

    fn location(&self) -> String {
        format!("GCS, bucket: {}", self.bucket)
    }

    fn current_size(&self) -> Option<usize> { None }
    fn max_size(&self) -> Option<usize> { None }
}
<|endoftext|>