<commit_before>//! The Actual RethinkDB Commands

use errors::*;
use conn::{Connection, ConnectOpts};
use serde_json;
use ql2::proto;
use std::io::Write;
use byteorder::{WriteBytesExt, LittleEndian};
use std::io::Read;
use std::str;
use protobuf::ProtobufEnum;
use byteorder::ReadBytesExt;
use super::session::Client;
use super::{Result, r};

pub struct RootCommand(Result<String>);
pub struct Command;
pub struct Query;

impl Client {
    pub fn connection(&self) -> ConnectOpts {
        Self::config().read().clone()
    }

    pub fn db_create(&self, name: &str) -> RootCommand {
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DB_CREATE,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     None)))
    }

    pub fn db_drop(&self, name: &str) -> RootCommand {
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DB_DROP,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     None)))
    }

    pub fn db(&self, name: &str) -> RootCommand {
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DB,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     None)))
    }

    pub fn table_create(&self, name: &str) -> RootCommand {
        let config = Client::config().read();
        r.db(config.db).table_create(name)
    }

    pub fn table(&self, name: &str) -> RootCommand {
        let config = Client::config().read();
        r.db(config.db).table(name)
    }

    pub fn object(&self) -> serde_json::builder::ObjectBuilder {
        serde_json::builder::ObjectBuilder::new()
    }

    pub fn array(&self) -> serde_json::builder::ArrayBuilder {
        serde_json::builder::ArrayBuilder::new()
    }
}

impl RootCommand {
    pub fn table_create(self, name: &str) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::TABLE_CREATE,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     Some(&commands))))
    }

    pub fn table(self, name: &str) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::TABLE,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     Some(&commands))))
    }

    pub fn insert(self, expr: serde_json::Value) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        let data = match serde_json::to_string(&expr) {
            Ok(f) => f,
            Err(e) => return RootCommand(Err(From::from(DriverError::Json(e)))),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::INSERT,
                                     Some(&data),
                                     None,
                                     Some(&commands))))
    }

    pub fn delete(self) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DELETE, None, None, Some(&commands))))
    }

    pub fn filter(self, expr: serde_json::Value) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        let filter = match serde_json::to_string(&expr) {
            Ok(f) => f,
            Err(e) => return RootCommand(Err(From::from(DriverError::Json(e)))),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::FILTER,
                                     Some(&filter),
                                     None,
                                     Some(&commands))))
    }

    pub fn run(self) -> Result<String> {
        let logger = Client::logger().read();
        trace!(logger, "Calling r.run()");
        let commands = try!(self.0);
        let cfg = Client::config().read();
        let query = Query::wrap(proto::Query_QueryType::START, Some(&commands), None);
        debug!(logger, "{}", query);
        let mut conn = try!(Client::conn());
        // Try sending the query
        {
            let mut i = 0;
            let mut write = true;
            let mut connect = false;
            while i < cfg.retries {
                debug!(logger, "Getting connection...");
                if connect {
                    drop(&mut conn);
                    conn = match Client::conn() {
                        Ok(c) => c,
                        Err(error) => {
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(From::from(error));
                            } else {
                                debug!(logger, "Failed getting a connection. Retrying...");
                                i += 1;
                                continue;
                            }
                        }
                    };
                }
                debug!(logger, "Connection aquired.");
                conn.token += 1;
                if write {
                    debug!(logger, "Writing query...");
                    if let Err(error) = Query::write(&query, &mut conn) {
                        connect = true;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to write query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                    debug!(logger, "Query written...");
                    connect = false;
                }
                debug!(logger, "Reading query...");
                match Query::read(&mut conn) {
                    Ok(resp) => {
                        let result = try!(str::from_utf8(&resp));
                        debug!(logger, "{}", result);
                        // If the write operation failed, retry it
                        // {"t":18,"e":4100000,"r":["Cannot perform write: primary replica for
                        // shard [\"\", +inf) not available"],"b":[]}
                        let msg = r#"{"t":18,"e":4100000,"r":["Cannot perform write: primary replica for shard"#;
                        if result.starts_with(msg) {
                            write = true;
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(
                                From::from(
                                    AvailabilityError::OpFailed(
                                        String::from("Not available")
                                        )));
                            } else {
                                debug!(logger, "Write operation failed. Retrying...");
                                i += 1;
                                continue;
                            }
                        } else {
                            debug!(logger, "Query successfully read.");
                            // This is a successful operation
                            break;
                        }
                    }
                    Err(error) => {
                        write = false;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to read query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                }
            }
        }
        Ok(String::new())
    }
}

impl Command {
    pub fn wrap(command: proto::Term_TermType,
                arguments: Option<&str>,
                options: Option<&str>,
                commands: Option<&str>)
                -> String {
        let mut cmds = format!("[{},", command.value());
        let mut args = String::new();
        if let Some(commands) = commands {
            args.push_str(commands);
            if arguments.is_some() {
                args.push_str(",");
            }
        }
        if let Some(arguments) = arguments {
            args.push_str(arguments);
        }
        cmds.push_str(&format!("[{}]", args));
        if let Some(options) = options {
            cmds.push_str(&format!(",{{{}}}", options));
        }
        cmds.push(']');
        cmds
    }
}

impl Query {
    pub fn wrap(query_type: proto::Query_QueryType,
                query: Option<&str>,
                options: Option<&str>)
                -> String {
        let mut qry = format!("[{}", query_type.value());
        if let Some(query) = query {
            qry.push_str(&format!(",{}", query));
        }
        if let Some(options) = options {
            qry.push_str(&format!(",{}", options));
        }
        qry.push_str("]");
        qry
    }

    pub fn write(query: &str, conn: &mut Connection) -> Result<()> {
        let query = query.as_bytes();
        let token = conn.token;
        if let Err(error) = conn.stream.write_u64::<LittleEndian>(token) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_u32::<LittleEndian>(query.len() as u32) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_all(query) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.flush() {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(())
    }

    pub fn read(conn: &mut Connection) -> Result<Vec<u8>> {
        // @TODO use response_token to implement parallel reads and writes?
        // let response_token = try!(conn.stream.read_u64::<LittleEndian>());
        let _ = match conn.stream.read_u64::<LittleEndian>() {
            Ok(token) => token,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let len = match conn.stream.read_u32::<LittleEndian>() {
            Ok(len) => len,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let mut resp = vec![0u8; len as usize];
        if let Err(error) = conn.stream.read_exact(&mut resp) {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(resp)
    }
}
<commit_msg>String impl<commit_after>//! The Actual RethinkDB Commands

use errors::*;
use conn::{Connection, ConnectOpts};
use serde_json;
use ql2::proto;
use std::io::Write;
use byteorder::{WriteBytesExt, LittleEndian};
use std::io::Read;
use std::str;
use protobuf::ProtobufEnum;
use byteorder::ReadBytesExt;
use super::session::Client;
use super::{Result, r};
use std::fmt::Debug;

pub struct RootCommand(Result<String>);
pub struct Command;
pub struct Query;

pub trait IntoCommandArg {
    fn to_arg(&self) -> Result<String>;
}

impl<'a> IntoCommandArg for &'a str {
    fn to_arg(&self) -> Result<String> {
        Ok(self.to_string())
    }
}

impl<'a> IntoCommandArg for &'a String {
    fn to_arg(&self) -> Result<String> {
        let ref arg = **self;
        Ok(*arg)
    }
}

impl IntoCommandArg for serde_json::Value {
    fn to_arg(&self) -> Result<String> {
        serde_json::to_string(self).map_err(|e| From::from(DriverError::Json(e)))
    }
}

impl Client {
    pub fn connection(&self) -> ConnectOpts {
        Self::config().read().clone()
    }

    pub fn db_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DB_CREATE,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     None)))
    }

    pub fn db_drop<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DB_DROP,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     None)))
    }

    pub fn db<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DB,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     None)))
    }

    pub fn table_create<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let config = Client::config().read();
        r.db(config.db).table_create(name)
    }

    pub fn table<T>(&self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let config = Client::config().read();
        r.db(config.db).table(name)
    }

    pub fn object(&self) -> serde_json::builder::ObjectBuilder {
        serde_json::builder::ObjectBuilder::new()
    }

    pub fn array(&self) -> serde_json::builder::ArrayBuilder {
        serde_json::builder::ArrayBuilder::new()
    }
}

impl RootCommand {
    pub fn table_create<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::TABLE_CREATE,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     Some(&commands))))
    }

    pub fn table<T>(self, name: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::TABLE,
                                     Some(&format!("{:?}", name)),
                                     None,
                                     Some(&commands))))
    }

    pub fn insert<T>(self, data: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::INSERT,
                                     Some(data),
                                     None,
                                     Some(&commands))))
    }

    pub fn delete(self) -> RootCommand {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::DELETE, None as Option<&str>, None as Option<&str>, Some(&commands))))
    }

    pub fn filter<T>(self, filter: T) -> RootCommand
        where T: IntoCommandArg + Debug
    {
        let commands = match self.0 {
            Ok(t) => t,
            Err(e) => return RootCommand(Err(e)),
        };
        RootCommand(Ok(Command::wrap(proto::Term_TermType::FILTER,
                                     Some(filter),
                                     None,
                                     Some(&commands))))
    }

    pub fn run(self) -> Result<String> {
        let logger = Client::logger().read();
        trace!(logger, "Calling r.run()");
        let commands = try!(self.0);
        let cfg = Client::config().read();
        let query = Query::wrap(proto::Query_QueryType::START, Some(&commands), None);
        debug!(logger, "{}", query);
        let mut conn = try!(Client::conn());
        // Try sending the query
        {
            let mut i = 0;
            let mut write = true;
            let mut connect = false;
            while i < cfg.retries {
                debug!(logger, "Getting connection...");
                if connect {
                    drop(&mut conn);
                    conn = match Client::conn() {
                        Ok(c) => c,
                        Err(error) => {
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(From::from(error));
                            } else {
                                debug!(logger, "Failed getting a connection. Retrying...");
                                i += 1;
                                continue;
                            }
                        }
                    };
                }
                debug!(logger, "Connection aquired.");
                conn.token += 1;
                if write {
                    debug!(logger, "Writing query...");
                    if let Err(error) = Query::write(&query, &mut conn) {
                        connect = true;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to write query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                    debug!(logger, "Query written...");
                    connect = false;
                }
                debug!(logger, "Reading query...");
                match Query::read(&mut conn) {
                    Ok(resp) => {
                        let result = try!(str::from_utf8(&resp));
                        debug!(logger, "{}", result);
                        // If the write operation failed, retry it
                        // {"t":18,"e":4100000,"r":["Cannot perform write: primary replica for
                        // shard [\"\", +inf) not available"],"b":[]}
                        let msg = r#"{"t":18,"e":4100000,"r":["Cannot perform write: primary replica for shard"#;
                        if result.starts_with(msg) {
                            write = true;
                            if i == cfg.retries - 1 {
                                // The last error
                                return Err(
                                From::from(
                                    AvailabilityError::OpFailed(
                                        String::from("Not available")
                                        )));
                            } else {
                                debug!(logger, "Write operation failed. Retrying...");
                                i += 1;
                                continue;
                            }
                        } else {
                            debug!(logger, "Query successfully read.");
                            // This is a successful operation
                            break;
                        }
                    }
                    Err(error) => {
                        write = false;
                        if i == cfg.retries - 1 {
                            // The last error
                            return Err(From::from(error));
                        } else {
                            debug!(logger, "Failed to read query. Retrying...");
                            i += 1;
                            continue;
                        }
                    }
                }
            }
        }
        Ok(String::new())
    }
}

impl Command {
    pub fn wrap<T>(command: proto::Term_TermType,
                arguments: Option<T>,
                options: Option<T>,
                commands: Option<&str>)
                -> String
        where T: IntoCommandArg
    {
        let mut cmds = format!("[{},", command.value());
        let mut args = String::new();
        if let Some(commands) = commands {
            args.push_str(commands);
            if arguments.is_some() {
                args.push_str(",");
            }
        }
        if let Some(arguments) = arguments {
            args.push_str(&arguments.to_arg().unwrap());
        }
        cmds.push_str(&format!("[{}]", args));
        if let Some(options) = options {
            cmds.push_str(&format!(",{{{}}}", options.to_arg().unwrap()));
        }
        cmds.push(']');
        cmds
    }
}

impl Query {
    pub fn wrap(query_type: proto::Query_QueryType,
                query: Option<&str>,
                options: Option<&str>)
                -> String {
        let mut qry = format!("[{}", query_type.value());
        if let Some(query) = query {
            qry.push_str(&format!(",{}", query));
        }
        if let Some(options) = options {
            qry.push_str(&format!(",{}", options));
        }
        qry.push_str("]");
        qry
    }

    pub fn write(query: &str, conn: &mut Connection) -> Result<()> {
        let query = query.as_bytes();
        let token = conn.token;
        if let Err(error) = conn.stream.write_u64::<LittleEndian>(token) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_u32::<LittleEndian>(query.len() as u32) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.write_all(query) {
            conn.broken = true;
            return Err(From::from(error));
        }
        if let Err(error) = conn.stream.flush() {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(())
    }

    pub fn read(conn: &mut Connection) -> Result<Vec<u8>> {
        // @TODO use response_token to implement parallel reads and writes?
        // let response_token = try!(conn.stream.read_u64::<LittleEndian>());
        let _ = match conn.stream.read_u64::<LittleEndian>() {
            Ok(token) => token,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let len = match conn.stream.read_u32::<LittleEndian>() {
            Ok(len) => len,
            Err(error) => {
                conn.broken = true;
                return Err(From::from(error));
            }
        };
        let mut resp = vec![0u8; len as usize];
        if let Err(error) = conn.stream.read_exact(&mut resp) {
            conn.broken = true;
            return Err(From::from(error));
        }
        Ok(resp)
    }
}
<|endoftext|>


<commit_before>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

// ignore-fast

#[feature(asm)];

use std::io::Process;
use std::os;
use std::str;

// lifted from the test module
pub fn black_box<T>(dummy: T) { unsafe { asm!("" : : "r"(&dummy)) } }

fn silent_recurse() {
    let buf = [0, ..1000];
    black_box(buf);
    silent_recurse();
}

fn loud_recurse() {
    println!("hello!");
    loud_recurse();
}

fn main() {
    let args = os::args();
    if args.len() > 1 && args[1].as_slice() == "silent" {
        silent_recurse();
    } else if args.len() > 1 && args[1].as_slice() == "loud" {
        loud_recurse();
    } else {
        let silent = Process::output(args[0], [~"silent"]).unwrap();
        assert!(!silent.status.success());
        let error = str::from_utf8_lossy(silent.error);
        assert!(error.as_slice().contains("has overflowed its stack"));

        let loud = Process::output(args[0], [~"loud"]).unwrap();
        assert!(!loud.status.success());
        let error = str::from_utf8_lossy(silent.error);
        assert!(error.as_slice().contains("has overflowed its stack"));
    }
}
<commit_msg>test/run-pass/out-of-stack: prevent tco<commit_after>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

// ignore-fast

#[feature(asm)];

use std::io::Process;
use std::os;
use std::str;

// lifted from the test module
pub fn black_box<T>(dummy: T) { unsafe { asm!("" : : "r"(&dummy)) } }

fn silent_recurse() {
    let buf = [0, ..1000];
    black_box(buf);
    silent_recurse();
}

fn loud_recurse() {
    println!("hello!");
    loud_recurse();
    black_box(()); // don't optimize this into a tail call. please.
}

fn main() {
    let args = os::args();
    if args.len() > 1 && args[1].as_slice() == "silent" {
        silent_recurse();
    } else if args.len() > 1 && args[1].as_slice() == "loud" {
        loud_recurse();
    } else {
        let silent = Process::output(args[0], [~"silent"]).unwrap();
        assert!(!silent.status.success());
        let error = str::from_utf8_lossy(silent.error);
        assert!(error.as_slice().contains("has overflowed its stack"));

        let loud = Process::output(args[0], [~"loud"]).unwrap();
        assert!(!loud.status.success());
        let error = str::from_utf8_lossy(silent.error);
        assert!(error.as_slice().contains("has overflowed its stack"));
    }
}
<|endoftext|>


<commit_before>//! Types for user flags which can appear in rooms.
use std::fmt;

use serde::de::{Deserializer, Error, Unexpected, Visitor};

/// Single flag.
#[derive(Clone, Debug, Hash)]
pub struct Flag {
    /// The name of the flag, unique per user.
    pub name: String,
    /// The primary color of the flag.
    pub primary_color: FlagColor,
    /// The secondary color of the flag.
    pub secondary_color: FlagColor,
    /// The X position of the flag.
    pub x: i8,
    /// The Y position of the flag.
    pub y: i8,
}

/// All possible colors a flag can have.test
#[repr(u8)]
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub enum FlagColor {
    /// COLOR_RED = 1
    Red = 1,
    /// COLOR_PURPLE = 2
    Purple = 2,
    /// COLOR_BLUE = 3
    Blue = 3,
    /// COLOR_CYAN = 4
    Cyan = 4,
    /// COLOR_GREEN = 5
    Green = 5,
    /// COLOR_YELLOW = 6
    Yellow = 6,
    /// COLOR_ORANGE = 7
    Orange = 7,
    /// COLOR_BROWN = 8
    Brown = 8,
    /// COLOR_GREY = 9
    Grey = 9,
    /// COLOR_WHITE = 10
    White = 10,
}

/// An error resulting from trying to convert a FlagColor that's out of bounds.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct FlagColorError;

impl FlagColor {
    /// Converts an integer color code into a flag color.
    #[inline]
    pub fn from(v: u8) -> Result<Self, FlagColorError> {
        match v {
            1 => Ok(FlagColor::Red),
            2 => Ok(FlagColor::Purple),
            3 => Ok(FlagColor::Blue),
            4 => Ok(FlagColor::Cyan),
            5 => Ok(FlagColor::Green),
            6 => Ok(FlagColor::Yellow),
            7 => Ok(FlagColor::Orange),
            8 => Ok(FlagColor::Brown),
            9 => Ok(FlagColor::Grey),
            10 => Ok(FlagColor::White),
            _ => Err(FlagColorError),
        }
    }

    #[inline]
    fn from_serde<E: Error>(v: u8) -> Result<Self, E> {
        match v {
            1 => Ok(FlagColor::Red),
            2 => Ok(FlagColor::Purple),
            3 => Ok(FlagColor::Blue),
            4 => Ok(FlagColor::Cyan),
            5 => Ok(FlagColor::Green),
            6 => Ok(FlagColor::Yellow),
            7 => Ok(FlagColor::Orange),
            8 => Ok(FlagColor::Brown),
            9 => Ok(FlagColor::Grey),
            10 => Ok(FlagColor::White),
            other => Err(E::invalid_value(
                Unexpected::Unsigned(other as u64),
                &"an integer between 1 and 10",
            )),
        }
    }
}

struct FlagStringVisitor;

impl<'de> Visitor<'de> for FlagStringVisitor {
    type Value = Vec<Flag>;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("a string with flag formatting (`Flag571~2~3~14~7|`...)")
    }

    #[inline]
    fn visit_unit<E>(self) -> Result<Self::Value, E>
    where
        E: Error,
    {
        Ok(Vec::new())
    }

    #[inline]
    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
    where
        E: Error,
    {
        if v.is_empty() {
            return Ok(Vec::new());
        }

        // TODO: nocopy version of this maybe?
        v.split("|")
            .map(|flag_str| {
                let mut iter = flag_str.split("~");

                macro_rules! next {
                    () => ({
                        iter.next().ok_or_else(|| {
                            E::invalid_type(Unexpected::Str(flag_str),
                                            &"a string in the format of name~4~2~4~2")
                        })?
                    })
                }

                macro_rules! next_u8 {
                    () => ({
                        let next = next!();
                        next.parse().map_err(|_| E::invalid_type(Unexpected::Str(next),
                            &"an unsigned integer"))?
                    })
                }

                macro_rules! next_color {
                    () => ({
                        FlagColor::from_serde(next_u8!())?
                    })
                }

                Ok(Flag {
                    name: next!().to_owned(),
                    primary_color: next_color!(),
                    secondary_color: next_color!(),
                    x: next_u8!(),
                    y: next_u8!(),
                })
            })
            .collect::<Result<Vec<_>, _>>()
    }
}

pub(super) fn deserialize_flags<'de, D>(deserializer: D) -> Result<Vec<Flag>, D::Error>
where
    D: Deserializer<'de>,
{
    deserializer.deserialize_str(FlagStringVisitor)
}
<commit_msg>Update flag deserializer to handle 'null' passed as a flag list value.<commit_after>//! Types for user flags which can appear in rooms.
use std::fmt;

use serde::de::{Deserializer, Error, Unexpected, Visitor};

/// Single flag.
#[derive(Clone, Debug, Hash)]
pub struct Flag {
    /// The name of the flag, unique per user.
    pub name: String,
    /// The primary color of the flag.
    pub primary_color: FlagColor,
    /// The secondary color of the flag.
    pub secondary_color: FlagColor,
    /// The X position of the flag.
    pub x: i8,
    /// The Y position of the flag.
    pub y: i8,
}

/// All possible colors a flag can have.test
#[repr(u8)]
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub enum FlagColor {
    /// COLOR_RED = 1
    Red = 1,
    /// COLOR_PURPLE = 2
    Purple = 2,
    /// COLOR_BLUE = 3
    Blue = 3,
    /// COLOR_CYAN = 4
    Cyan = 4,
    /// COLOR_GREEN = 5
    Green = 5,
    /// COLOR_YELLOW = 6
    Yellow = 6,
    /// COLOR_ORANGE = 7
    Orange = 7,
    /// COLOR_BROWN = 8
    Brown = 8,
    /// COLOR_GREY = 9
    Grey = 9,
    /// COLOR_WHITE = 10
    White = 10,
}

/// An error resulting from trying to convert a FlagColor that's out of bounds.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct FlagColorError;

impl FlagColor {
    /// Converts an integer color code into a flag color.
    #[inline]
    pub fn from(v: u8) -> Result<Self, FlagColorError> {
        match v {
            1 => Ok(FlagColor::Red),
            2 => Ok(FlagColor::Purple),
            3 => Ok(FlagColor::Blue),
            4 => Ok(FlagColor::Cyan),
            5 => Ok(FlagColor::Green),
            6 => Ok(FlagColor::Yellow),
            7 => Ok(FlagColor::Orange),
            8 => Ok(FlagColor::Brown),
            9 => Ok(FlagColor::Grey),
            10 => Ok(FlagColor::White),
            _ => Err(FlagColorError),
        }
    }

    #[inline]
    fn from_serde<E: Error>(v: u8) -> Result<Self, E> {
        match v {
            1 => Ok(FlagColor::Red),
            2 => Ok(FlagColor::Purple),
            3 => Ok(FlagColor::Blue),
            4 => Ok(FlagColor::Cyan),
            5 => Ok(FlagColor::Green),
            6 => Ok(FlagColor::Yellow),
            7 => Ok(FlagColor::Orange),
            8 => Ok(FlagColor::Brown),
            9 => Ok(FlagColor::Grey),
            10 => Ok(FlagColor::White),
            other => Err(E::invalid_value(
                Unexpected::Unsigned(other as u64),
                &"an integer between 1 and 10",
            )),
        }
    }
}

struct FlagStringVisitor;

impl<'de> Visitor<'de> for FlagStringVisitor {
    type Value = Vec<Flag>;

    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str("a string with flag formatting (`Flag571~2~3~14~7|`...)")
    }

    #[inline]
    fn visit_unit<E>(self) -> Result<Self::Value, E>
    where
        E: Error,
    {
        Ok(Vec::new())
    }

    #[inline]
    fn visit_none<E>(self) -> Result<Self::Value, E>
    where
        E: Error,
    {
        Ok(Vec::new())
    }

    #[inline]
    fn visit_some<D>(self, d: D) -> Result<Self::Value, D::Error>
    where
        D: Deserializer<'de>,
    {
        d.deserialize_str(self)
    }

    #[inline]
    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
    where
        E: Error,
    {
        if v.is_empty() {
            return Ok(Vec::new());
        }

        // TODO: nocopy version of this maybe?
        v.split("|")
            .map(|flag_str| {
                let mut iter = flag_str.split("~");

                macro_rules! next {
                    () => ({
                        iter.next().ok_or_else(|| {
                            E::invalid_type(Unexpected::Str(flag_str),
                                            &"a string in the format of name~4~2~4~2")
                        })?
                    })
                }

                macro_rules! next_u8 {
                    () => ({
                        let next = next!();
                        next.parse().map_err(|_| E::invalid_type(Unexpected::Str(next),
                            &"an unsigned integer"))?
                    })
                }

                macro_rules! next_color {
                    () => ({
                        FlagColor::from_serde(next_u8!())?
                    })
                }

                Ok(Flag {
                    name: next!().to_owned(),
                    primary_color: next_color!(),
                    secondary_color: next_color!(),
                    x: next_u8!(),
                    y: next_u8!(),
                })
            })
            .collect::<Result<Vec<_>, _>>()
    }
}

pub(super) fn deserialize_flags<'de, D>(deserializer: D) -> Result<Vec<Flag>, D::Error>
where
    D: Deserializer<'de>,
{
    deserializer.deserialize_option(FlagStringVisitor)
}
<|endoftext|>


<commit_before>use regex::Regex;
use std::path::PathBuf;
use std::result::Result;
use subprocess::{Exec, Redirection};

const LOOP_BACKING_DEVICE_BASENAME: &'static str = "flaglock6-home";
const LOOP_DEVICE_PARTITION_LABEL: &'static str = "fl6";
const RSYNC_DOT_NET_REMOTE_DIR: &'static str = "usw-s007.rsync.net:./";

#[derive(Debug)]
pub enum AuricError {
    Invocation(String),
    Io(String),
    Subprocess(String),
    NotFound,
}

enum AuricOperationMode {
    Mount,
    Unmount,
}

impl From<subprocess::PopenError> for AuricError {
    fn from(err: subprocess::PopenError) -> AuricError {
        AuricError::Subprocess(err.to_string())
    }
}

struct SshfsManager {
    mountpoint: PathBuf,
}

struct LoopManager {
    target: Regex,
    backing: PathBuf,
}

struct LuksManager {
    device: PathBuf,
}

struct AuricImpl {
    sshfs_manager: SshfsManager,
    loop_manager: LoopManager,
    luks_manager: LuksManager,
}

fn exec_with_stdout(command: &str, args: &[&str], error_hint: &str) -> Result<String, AuricError> {
    let result = Exec::cmd(command)
        .args(args)
        .stdout(Redirection::Pipe)
        .stderr(Redirection::Pipe)
        .capture()?;
    if !result.success() {
        return Err(AuricError::Subprocess(format!(
            "{} failed: ``{}''",
            error_hint,
            result.stderr_str(),
        )));
    }
    Ok(result.stdout_str())
}

fn exec(command: &str, args: &[&str], error_hint: &str) -> Result<(), AuricError> {
    exec_with_stdout(command, args, error_hint).map(|_stdout_str| ())
}

impl SshfsManager {
    pub fn new() -> Result<SshfsManager, AuricError> {
        let xdg_runtime_dir = match std::env::var("XDG_RUNTIME_DIR") {
            Ok(val) => val,
            Err(e) => {
                return Err(AuricError::Invocation(format!(
                    "fatal error getting ${{XDG_RUNTIME_DIR}}: ``{}''",
                    e
                )))
            }
        };
        let mut mountpoint = PathBuf::new();
        mountpoint.push(xdg_runtime_dir);
        mountpoint.push("rsync.net-mountpoint");

        Ok(SshfsManager {
            mountpoint: mountpoint,
        })
    }

    fn ensure_mountpoint_is_dir(&self) -> Result<(), AuricError> {
        if self.mountpoint.is_dir() {
            return Ok(());
        }
        std::fs::create_dir(self.mountpoint.as_path())
            .map_err(|e| AuricError::Io(format!("bad mountpoint: ``{}''", e)))
    }

    // Does _not_ ensure that backing device exists.
    pub fn loop_backing_device_path(&self) -> PathBuf {
        let mut result: PathBuf = self.mountpoint.clone();
        result.push(LOOP_BACKING_DEVICE_BASENAME);
        result
    }

    fn loop_backing_device_present(&self) -> bool {
        self.loop_backing_device_path().is_file()
    }

    pub fn mount(&self) -> Result<(), AuricError> {
        self.ensure_mountpoint_is_dir()?;
        if self.loop_backing_device_present() {
            return Ok(());
        }

        // Invoke `sshfs`
        // *    with compression,
        // *    allowing other users to access the mounted filesystem
        //      (necessary to allow loop-mounts), and
        // *    attempting to reconnect if we lose the connection.
        exec(
            "sshfs",
            &[
                "-C",
                "-o",
                "allow_other,reconnect",
                RSYNC_DOT_NET_REMOTE_DIR,
                self.mountpoint.to_str().unwrap(),
            ],
            "sshfs mount",
        )
    }

    pub fn unmount(&self) -> Result<(), AuricError> {
        if !self.loop_backing_device_present() {
            return Ok(());
        }
        exec(
            "fusermount",
            &["-u", self.mountpoint.to_str().unwrap()],
            "fusermount -u",
        )
    }
}

impl LoopManager {
    pub fn new(backing: &PathBuf) -> LoopManager {
        let target = format!(
            r"^(/dev/loop\d)(\s+\d+){{4}}\s+({})\s+",
            backing.to_str().unwrap()
        );
        LoopManager {
            target: Regex::new(&target).unwrap(),
            backing: backing.clone(),
        }
    }

    pub fn find(&self) -> Result<PathBuf, AuricError> {
        let losetup_stdout = exec_with_stdout("losetup", &["-l"], "loop device enumeration")?;
        for line in losetup_stdout.lines() {
            if let Some(captures) = self.target.captures(line) {
                return Ok(PathBuf::from(captures.get(1).unwrap().as_str()));
            }
        }
        Err(AuricError::NotFound)
    }

    pub fn attach(&self) -> Result<(), AuricError> {
        if let Ok(_attached_loop) = self.find() {
            return Ok(());
        }

        exec(
            "udisksctl",
            &["loop-setup", "-f", self.backing.to_str().unwrap()],
            "loop device attach",
        )
    }

    // Note that the loop device will linger if the volume remains
    // unlocked. Prefer not to call this before locking.
    pub fn detach(&self) -> Result<(), AuricError> {
        let attached_loop = match self.find() {
            Ok(path) => path,
            Err(e) => {
                if let AuricError::NotFound = e {
                    return Ok(());
                }
                return Err(e);
            }
        };
        exec(
            "udisksctl",
            &["loop-delete", "-b", attached_loop.to_str().unwrap()],
            "loop device detach",
        )
    }
}

impl LuksManager {
    pub fn new() -> LuksManager {
        let mut device = PathBuf::from("/dev/disk/by-label");
        device.push(LOOP_DEVICE_PARTITION_LABEL);
        LuksManager { device: device }
    }

    pub fn unlock(&self, locked_device: &PathBuf) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["unlock", "-b", locked_device.to_str().unwrap()],
            format!("unlocking {}", locked_device.to_str().unwrap()).as_str(),
        )
    }

    pub fn lock(&self, underlying_device: &PathBuf) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["lock", "-b", underlying_device.to_str().unwrap()],
            format!("locking {}", underlying_device.to_str().unwrap()).as_str(),
        )
    }

    pub fn mount(&self) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["mount", "-b", self.device.to_str().unwrap()],
            format!("mounting {}", self.device.to_str().unwrap()).as_str(),
        )
    }

    pub fn unmount(&self) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["unmount", "-b", self.device.to_str().unwrap()],
            format!("unmounting {}", self.device.to_str().unwrap()).as_str(),
        )
    }
}

fn get_action() -> Result<AuricOperationMode, AuricError> {
    let mut args_iter = std::env::args();
    args_iter.next();
    if let Some(action) = args_iter.next() {
        match action.as_str() {
            "mount" => return Ok(AuricOperationMode::Mount),
            "unmount" => return Ok(AuricOperationMode::Unmount),
            _ => {
                return Err(AuricError::Invocation(format!(
                    "unknown action: ``{}''",
                    action
                )))
            }
        }
    }
    Err(AuricError::Invocation("no action given".to_string()))
}

impl AuricImpl {
    pub fn new() -> Result<AuricImpl, AuricError> {
        let sshfs_manager = SshfsManager::new()?;
        let loop_backing_device_path = sshfs_manager.loop_backing_device_path();
        Ok(AuricImpl {
            sshfs_manager: sshfs_manager,
            loop_manager: LoopManager::new(&loop_backing_device_path),
            luks_manager: LuksManager::new(),
        })
    }

    fn mount(&self) -> Result<(), AuricError> {
        println!("{}", "Mounting sshfs...");
        self.sshfs_manager.mount()?;
        println!("{}", "Attaching loop device...");
        self.loop_manager.attach()?;
        println!("{}", "Finding newly attached loop device...");
        let loop_device = self.loop_manager.find()?;
        println!("{}", "Unlocking volume...");
        self.luks_manager.unlock(&loop_device)?;
        println!("{}", "Mounting...");
        self.luks_manager.mount()
    }

    fn unmount(&self) -> Result<(), AuricError> {
        println!("{}", "Unmounting...");
        self.luks_manager.unmount()?;
        println!("{}", "Finding backing loop device...");
        let loop_device = self.loop_manager.find()?;
        println!("{}", "Locking volume...");
        self.luks_manager.lock(&loop_device)?;
        println!("{}", "Detaching loop device...");
        self.loop_manager.detach()?;
        println!("{}", "Unmounting sshfs...");
        self.sshfs_manager.unmount()
    }

    pub fn act(&self, mode: AuricOperationMode) -> Result<(), AuricError> {
        match mode {
            AuricOperationMode::Mount => return self.mount(),
            AuricOperationMode::Unmount => return self.unmount(),
        }
    }
}

fn main_impl() -> Result<(), AuricError> {
    let action = get_action()?;
    let auric_impl = AuricImpl::new()?;
    auric_impl.act(action)
}

fn main() {
    std::process::exit(match main_impl() {
        Ok(_) => 0,
        Err(e) => {
            eprintln!("{:?}", e);
            1
        }
    });
}
<commit_msg>apricot-auric: Amend log message<commit_after>use regex::Regex;
use std::path::PathBuf;
use std::result::Result;
use subprocess::{Exec, Redirection};

const LOOP_BACKING_DEVICE_BASENAME: &'static str = "flaglock6-home";
const LOOP_DEVICE_PARTITION_LABEL: &'static str = "fl6";
const RSYNC_DOT_NET_REMOTE_DIR: &'static str = "usw-s007.rsync.net:./";

#[derive(Debug)]
pub enum AuricError {
    Invocation(String),
    Io(String),
    Subprocess(String),
    NotFound,
}

enum AuricOperationMode {
    Mount,
    Unmount,
}

impl From<subprocess::PopenError> for AuricError {
    fn from(err: subprocess::PopenError) -> AuricError {
        AuricError::Subprocess(err.to_string())
    }
}

struct SshfsManager {
    mountpoint: PathBuf,
}

struct LoopManager {
    target: Regex,
    backing: PathBuf,
}

struct LuksManager {
    device: PathBuf,
}

struct AuricImpl {
    sshfs_manager: SshfsManager,
    loop_manager: LoopManager,
    luks_manager: LuksManager,
}

fn exec_with_stdout(command: &str, args: &[&str], error_hint: &str) -> Result<String, AuricError> {
    let result = Exec::cmd(command)
        .args(args)
        .stdout(Redirection::Pipe)
        .stderr(Redirection::Pipe)
        .capture()?;
    if !result.success() {
        return Err(AuricError::Subprocess(format!(
            "{} failed: ``{}''",
            error_hint,
            result.stderr_str(),
        )));
    }
    Ok(result.stdout_str())
}

fn exec(command: &str, args: &[&str], error_hint: &str) -> Result<(), AuricError> {
    exec_with_stdout(command, args, error_hint).map(|_stdout_str| ())
}

impl SshfsManager {
    pub fn new() -> Result<SshfsManager, AuricError> {
        let xdg_runtime_dir = match std::env::var("XDG_RUNTIME_DIR") {
            Ok(val) => val,
            Err(e) => {
                return Err(AuricError::Invocation(format!(
                    "fatal error getting ${{XDG_RUNTIME_DIR}}: ``{}''",
                    e
                )))
            }
        };
        let mut mountpoint = PathBuf::new();
        mountpoint.push(xdg_runtime_dir);
        mountpoint.push("rsync.net-mountpoint");

        Ok(SshfsManager {
            mountpoint: mountpoint,
        })
    }

    fn ensure_mountpoint_is_dir(&self) -> Result<(), AuricError> {
        if self.mountpoint.is_dir() {
            return Ok(());
        }
        std::fs::create_dir(self.mountpoint.as_path())
            .map_err(|e| AuricError::Io(format!("bad mountpoint: ``{}''", e)))
    }

    // Does _not_ ensure that backing device exists.
    pub fn loop_backing_device_path(&self) -> PathBuf {
        let mut result: PathBuf = self.mountpoint.clone();
        result.push(LOOP_BACKING_DEVICE_BASENAME);
        result
    }

    fn loop_backing_device_present(&self) -> bool {
        self.loop_backing_device_path().is_file()
    }

    pub fn mount(&self) -> Result<(), AuricError> {
        self.ensure_mountpoint_is_dir()?;
        if self.loop_backing_device_present() {
            return Ok(());
        }

        // Invoke `sshfs`
        // *    with compression,
        // *    allowing other users to access the mounted filesystem
        //      (necessary to allow loop-mounts), and
        // *    attempting to reconnect if we lose the connection.
        exec(
            "sshfs",
            &[
                "-C",
                "-o",
                "allow_other,reconnect",
                RSYNC_DOT_NET_REMOTE_DIR,
                self.mountpoint.to_str().unwrap(),
            ],
            "sshfs mount",
        )
    }

    pub fn unmount(&self) -> Result<(), AuricError> {
        if !self.loop_backing_device_present() {
            return Ok(());
        }
        exec(
            "fusermount",
            &["-u", self.mountpoint.to_str().unwrap()],
            "fusermount -u",
        )
    }
}

impl LoopManager {
    pub fn new(backing: &PathBuf) -> LoopManager {
        let target = format!(
            r"^(/dev/loop\d)(\s+\d+){{4}}\s+({})\s+",
            backing.to_str().unwrap()
        );
        LoopManager {
            target: Regex::new(&target).unwrap(),
            backing: backing.clone(),
        }
    }

    pub fn find(&self) -> Result<PathBuf, AuricError> {
        let losetup_stdout = exec_with_stdout("losetup", &["-l"], "loop device enumeration")?;
        for line in losetup_stdout.lines() {
            if let Some(captures) = self.target.captures(line) {
                return Ok(PathBuf::from(captures.get(1).unwrap().as_str()));
            }
        }
        Err(AuricError::NotFound)
    }

    pub fn attach(&self) -> Result<(), AuricError> {
        if let Ok(_attached_loop) = self.find() {
            return Ok(());
        }

        exec(
            "udisksctl",
            &["loop-setup", "-f", self.backing.to_str().unwrap()],
            "loop device attach",
        )
    }

    // Note that the loop device will linger if the volume remains
    // unlocked. Prefer not to call this before locking.
    pub fn detach(&self) -> Result<(), AuricError> {
        let attached_loop = match self.find() {
            Ok(path) => path,
            Err(e) => {
                if let AuricError::NotFound = e {
                    return Ok(());
                }
                return Err(e);
            }
        };
        exec(
            "udisksctl",
            &["loop-delete", "-b", attached_loop.to_str().unwrap()],
            "loop device detach",
        )
    }
}

impl LuksManager {
    pub fn new() -> LuksManager {
        let mut device = PathBuf::from("/dev/disk/by-label");
        device.push(LOOP_DEVICE_PARTITION_LABEL);
        LuksManager { device: device }
    }

    pub fn unlock(&self, locked_device: &PathBuf) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["unlock", "-b", locked_device.to_str().unwrap()],
            format!("unlocking {}", locked_device.to_str().unwrap()).as_str(),
        )
    }

    pub fn lock(&self, underlying_device: &PathBuf) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["lock", "-b", underlying_device.to_str().unwrap()],
            format!("locking {}", underlying_device.to_str().unwrap()).as_str(),
        )
    }

    pub fn mount(&self) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["mount", "-b", self.device.to_str().unwrap()],
            format!("mounting {}", self.device.to_str().unwrap()).as_str(),
        )
    }

    pub fn unmount(&self) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["unmount", "-b", self.device.to_str().unwrap()],
            format!("unmounting {}", self.device.to_str().unwrap()).as_str(),
        )
    }
}

fn get_action() -> Result<AuricOperationMode, AuricError> {
    let mut args_iter = std::env::args();
    args_iter.next();
    if let Some(action) = args_iter.next() {
        match action.as_str() {
            "mount" => return Ok(AuricOperationMode::Mount),
            "unmount" => return Ok(AuricOperationMode::Unmount),
            _ => {
                return Err(AuricError::Invocation(format!(
                    "unknown action: ``{}''",
                    action
                )))
            }
        }
    }
    Err(AuricError::Invocation("no action given".to_string()))
}

impl AuricImpl {
    pub fn new() -> Result<AuricImpl, AuricError> {
        let sshfs_manager = SshfsManager::new()?;
        let loop_backing_device_path = sshfs_manager.loop_backing_device_path();
        Ok(AuricImpl {
            sshfs_manager: sshfs_manager,
            loop_manager: LoopManager::new(&loop_backing_device_path),
            luks_manager: LuksManager::new(),
        })
    }

    fn mount(&self) -> Result<(), AuricError> {
        println!("{}", "Mounting sshfs...");
        self.sshfs_manager.mount()?;
        println!("{}", "Attaching loop device...");
        self.loop_manager.attach()?;
        println!("{}", "Finding newly attached loop device...");
        let loop_device = self.loop_manager.find()?;
        println!("{}", "Unlocking volume...");
        self.luks_manager.unlock(&loop_device)?;
        println!("{}", "Mounting volume...");
        self.luks_manager.mount()
    }

    fn unmount(&self) -> Result<(), AuricError> {
        println!("{}", "Unmounting volume...");
        self.luks_manager.unmount()?;
        println!("{}", "Finding backing loop device...");
        let loop_device = self.loop_manager.find()?;
        println!("{}", "Locking volume...");
        self.luks_manager.lock(&loop_device)?;
        println!("{}", "Detaching loop device...");
        self.loop_manager.detach()?;
        println!("{}", "Unmounting sshfs...");
        self.sshfs_manager.unmount()
    }

    pub fn act(&self, mode: AuricOperationMode) -> Result<(), AuricError> {
        match mode {
            AuricOperationMode::Mount => return self.mount(),
            AuricOperationMode::Unmount => return self.unmount(),
        }
    }
}

fn main_impl() -> Result<(), AuricError> {
    let action = get_action()?;
    let auric_impl = AuricImpl::new()?;
    auric_impl.act(action)
}

fn main() {
    std::process::exit(match main_impl() {
        Ok(_) => 0,
        Err(e) => {
            eprintln!("{:?}", e);
            1
        }
    });
}
<|endoftext|>


<commit_before>extern crate dns_parser;
extern crate resolv_conf;

use std::time::{Duration, Instant};
use std::io::{self, Read};
use std::fs::File;
use std::net::{IpAddr, SocketAddr, UdpSocket};
use std::collections::HashMap;

pub struct Resolver {
    server: SocketAddr,
    cache: HashMap<String, CacheEntry>,
    queries: HashMap<u16, Query>,
    responses: HashMap<String, Vec<usize>>,
    buf: Vec<u8>,
    qnum: u16,
}

struct Query {
    resps: u8,
    domain: String,
    deadline: Instant,
}

struct CacheEntry {
    ip: IpAddr,
    deadline: Instant,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Response {
    pub id: usize,
    pub result: Result<IpAddr, Error>,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Error {
    NotFound,
    Timeout,
}

impl Resolver {
    pub fn new(server: SocketAddr) -> Resolver {
        let mut buf = Vec::with_capacity(512);
        unsafe {
            buf.set_len(512);
        }
        Resolver {
            server,
            queries: HashMap::new(),
            responses: HashMap::new(),
            cache: HashMap::new(),
            buf,
            qnum: 0,
        }
    }

    pub fn purge(&mut self) {
        self.cache.clear();
    }

    pub fn from_resolv() -> io::Result<Resolver> {
        let mut buf = Vec::with_capacity(512);
        unsafe {
            buf.set_len(512);
        }

        let mut conf = Vec::with_capacity(4096);
        let mut f = File::open("/etc/resolv.conf")?;
        f.read_to_end(&mut conf)?;
        let cfg = resolv_conf::Config::parse(&conf).map_err(|e| {
            io::Error::new(io::ErrorKind::Other, format!("invalid resolv.conf: {}", e))
        })?;
        let server = match cfg.nameservers.first().cloned() {
            Some(resolv_conf::ScopedIp::V4(ip)) => SocketAddr::new(IpAddr::V4(ip), 53),
            Some(resolv_conf::ScopedIp::V6(ip, _)) => SocketAddr::new(IpAddr::V6(ip), 53),
            None => {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    "No nameservers found in resolv.conf!",
                ))
            }
        };

        Ok(Resolver {
            server,
            queries: HashMap::new(),
            responses: HashMap::new(),
            cache: HashMap::new(),
            buf,
            qnum: 0,
        })
    }

    pub fn query(
        &mut self,
        sock: &mut UdpSocket,
        id: usize,
        domain: &str,
    ) -> io::Result<Option<IpAddr>> {
        if let Some(entry) = self.cache.get(domain) {
            return Ok(Some(entry.ip));
        }

        let qn = self.qnum;
        self.qnum.wrapping_add(1);

        let mut query = dns_parser::Builder::new_query(qn, true);
        query.add_question(domain, dns_parser::QueryType::A, dns_parser::QueryClass::IN);
        let packet = query.build().unwrap_or_else(|d| d);
        sock.send_to(&packet, self.server)?;

        let mut query = dns_parser::Builder::new_query(qn, true);
        query.add_question(
            domain,
            dns_parser::QueryType::AAAA,
            dns_parser::QueryClass::IN,
        );
        let packet = query.build().unwrap_or_else(|d| d);
        sock.send_to(&packet, self.server)?;

        if self.responses.get(domain).is_none() {
            self.responses.insert(domain.to_string(), vec![]);
            self.queries.insert(
                qn,
                Query {
                    resps: 0,
                    domain: domain.to_string(),
                    deadline: Instant::now() + Duration::from_secs(3),
                },
            );
        }
        self.responses.get_mut(domain).unwrap().push(id);
        Ok(None)
    }

    pub fn read<F: FnMut(Response)>(&mut self, sock: &mut UdpSocket, mut f: F) -> io::Result<()> {
        'process: loop {
            match sock.recv_from(&mut self.buf) {
                Ok((amnt, _)) => {
                    match dns_parser::Packet::parse(&self.buf[..amnt]) {
                        Ok(packet) => {
                            let qn = packet.header.id;
                            let (domain, mut resps, deadline) = match self.queries.remove(&qn) {
                                Some(q) => (q.domain, q.resps, q.deadline),
                                // This could happen if timeout is exceeeded but we eventually get
                                // a response, ignore.
                                None => continue,
                            };
                            let now = Instant::now();
                            for answer in packet.answers {
                                match answer.data {
                                    dns_parser::RRData::A(addr) => {
                                        for id in self.responses.remove(&domain).unwrap() {
                                            f(Response {
                                                id,
                                                result: Ok(addr.into()),
                                            });
                                        }
                                        self.cache.insert(
                                            domain,
                                            CacheEntry {
                                                ip: addr.into(),
                                                deadline: now
                                                    + Duration::from_secs(answer.ttl.into()),
                                            },
                                        );
                                        continue 'process;
                                    }
                                    dns_parser::RRData::AAAA(addr) => {
                                        for id in self.responses.remove(&domain).unwrap() {
                                            f(Response {
                                                id,
                                                result: Ok(addr.into()),
                                            });
                                        }
                                        self.cache.insert(
                                            domain,
                                            CacheEntry {
                                                ip: addr.into(),
                                                deadline: now
                                                    + Duration::from_secs(answer.ttl.into()),
                                            },
                                        );
                                        continue 'process;
                                    }
                                    _ => continue,
                                }
                            }
                            if resps == 0 {
                                self.queries.insert(
                                    qn,
                                    Query {
                                        domain,
                                        resps: resps + 1,
                                        deadline,
                                    },
                                );
                            } else {
                                for id in self.responses.remove(&domain).unwrap() {
                                    f(Response {
                                        id,
                                        result: Err(Error::NotFound),
                                    });
                                }
                            }
                        }
                        Err(e) => {
                            return Err(io::Error::new(
                                io::ErrorKind::Other,
                                format!("malformed dns packet received: {}", e),
                            ));
                        }
                    }
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(()),
                Err(e) => return Err(e),
            }
        }
    }

    pub fn tick<F: FnMut(Response)>(&mut self, mut f: F) {
        let now = Instant::now();
        let responses = &mut self.responses;
        self.queries.retain(|_, query| {
            if now > query.deadline {
                for id in responses.remove(&query.domain).unwrap() {
                    f(Response {
                        id,
                        result: Err(Error::Timeout),
                    });
                }
                false
            } else {
                true
            }
        });
        self.cache.retain(|_, entry| now < entry.deadline);
    }
}

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

    #[test]
    fn test_google() {
        let mut resolver = Resolver::new("8.8.8.8:53".parse().unwrap());
        let mut sock = UdpSocket::bind("0.0.0.0:0").unwrap();
        sock.set_nonblocking(true).unwrap();

        assert_eq!(resolver.query(&mut sock, 0, "google.com").unwrap(), None);
        assert_eq!(resolver.query(&mut sock, 1, "google.com").unwrap(), None);
        assert_eq!(resolver.responses.get("google.com").unwrap().len(), 2);
        std::thread::sleep(Duration::from_millis(100));
        resolver.tick(|_| panic!("timeout should not have occured yet!"));
        let mut count = 0;
        resolver
            .read(&mut sock, |resp| {
                count += 1;
                assert!(resp.result.is_ok());
            })
            .unwrap();
        assert_eq!(count, 2);

        assert!(
            resolver
                .query(&mut sock, 0, "google.com")
                .unwrap()
                .is_some()
        );

        resolver
            .query(&mut sock, 0, "thiswebsiteshouldexit12589t69.com")
            .unwrap();
        std::thread::sleep(Duration::from_millis(200));
        let mut processed = false;
        resolver
            .read(&mut sock, |resp| {
                processed = true;
                assert_eq!(resp.result, Err(Error::NotFound))
            })
            .unwrap();
        #[cfg(not(target_os = "macos"))]
        assert!(processed);
    }
}
<commit_msg>Only initiate DNS queries when needed<commit_after>extern crate dns_parser;
extern crate resolv_conf;

use std::time::{Duration, Instant};
use std::io::{self, Read};
use std::fs::File;
use std::net::{IpAddr, SocketAddr, UdpSocket};
use std::collections::HashMap;

pub struct Resolver {
    server: SocketAddr,
    cache: HashMap<String, CacheEntry>,
    queries: HashMap<u16, Query>,
    responses: HashMap<String, Vec<usize>>,
    buf: Vec<u8>,
    qnum: u16,
}

struct Query {
    resps: u8,
    domain: String,
    deadline: Instant,
}

struct CacheEntry {
    ip: IpAddr,
    deadline: Instant,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Response {
    pub id: usize,
    pub result: Result<IpAddr, Error>,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Error {
    NotFound,
    Timeout,
}

impl Resolver {
    pub fn new(server: SocketAddr) -> Resolver {
        let mut buf = Vec::with_capacity(512);
        unsafe {
            buf.set_len(512);
        }
        Resolver {
            server,
            queries: HashMap::new(),
            responses: HashMap::new(),
            cache: HashMap::new(),
            buf,
            qnum: 0,
        }
    }

    pub fn purge(&mut self) {
        self.cache.clear();
    }

    pub fn from_resolv() -> io::Result<Resolver> {
        let mut buf = Vec::with_capacity(512);
        unsafe {
            buf.set_len(512);
        }

        let mut conf = Vec::with_capacity(4096);
        let mut f = File::open("/etc/resolv.conf")?;
        f.read_to_end(&mut conf)?;
        let cfg = resolv_conf::Config::parse(&conf).map_err(|e| {
            io::Error::new(io::ErrorKind::Other, format!("invalid resolv.conf: {}", e))
        })?;
        let server = match cfg.nameservers.first().cloned() {
            Some(resolv_conf::ScopedIp::V4(ip)) => SocketAddr::new(IpAddr::V4(ip), 53),
            Some(resolv_conf::ScopedIp::V6(ip, _)) => SocketAddr::new(IpAddr::V6(ip), 53),
            None => {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    "No nameservers found in resolv.conf!",
                ))
            }
        };

        Ok(Resolver {
            server,
            queries: HashMap::new(),
            responses: HashMap::new(),
            cache: HashMap::new(),
            buf,
            qnum: 0,
        })
    }

    pub fn query(
        &mut self,
        sock: &mut UdpSocket,
        id: usize,
        domain: &str,
    ) -> io::Result<Option<IpAddr>> {
        if let Some(entry) = self.cache.get(domain) {
            return Ok(Some(entry.ip));
        }
        if self.responses.get(domain).is_none() {
            let qn = self.qnum;
            self.qnum.wrapping_add(1);
            let mut query = dns_parser::Builder::new_query(qn, true);
            query.add_question(domain, dns_parser::QueryType::A, dns_parser::QueryClass::IN);
            let packet = query.build().unwrap_or_else(|d| d);
            sock.send_to(&packet, self.server)?;

            let mut query = dns_parser::Builder::new_query(qn, true);
            query.add_question(
                domain,
                dns_parser::QueryType::AAAA,
                dns_parser::QueryClass::IN,
            );
            let packet = query.build().unwrap_or_else(|d| d);
            sock.send_to(&packet, self.server)?;

            self.responses.insert(domain.to_string(), vec![]);
            self.queries.insert(
                qn,
                Query {
                    resps: 0,
                    domain: domain.to_string(),
                    deadline: Instant::now() + Duration::from_secs(3),
                },
            );
        }
        self.responses.get_mut(domain).unwrap().push(id);
        Ok(None)
    }

    pub fn read<F: FnMut(Response)>(&mut self, sock: &mut UdpSocket, mut f: F) -> io::Result<()> {
        'process: loop {
            match sock.recv_from(&mut self.buf) {
                Ok((amnt, _)) => {
                    match dns_parser::Packet::parse(&self.buf[..amnt]) {
                        Ok(packet) => {
                            let qn = packet.header.id;
                            let (domain, mut resps, deadline) = match self.queries.remove(&qn) {
                                Some(q) => (q.domain, q.resps, q.deadline),
                                // This could happen if timeout is exceeeded but we eventually get
                                // a response, ignore.
                                None => continue,
                            };
                            let now = Instant::now();
                            for answer in packet.answers {
                                match answer.data {
                                    dns_parser::RRData::A(addr) => {
                                        for id in self.responses.remove(&domain).unwrap() {
                                            f(Response {
                                                id,
                                                result: Ok(addr.into()),
                                            });
                                        }
                                        self.cache.insert(
                                            domain,
                                            CacheEntry {
                                                ip: addr.into(),
                                                deadline: now
                                                    + Duration::from_secs(answer.ttl.into()),
                                            },
                                        );
                                        continue 'process;
                                    }
                                    dns_parser::RRData::AAAA(addr) => {
                                        for id in self.responses.remove(&domain).unwrap() {
                                            f(Response {
                                                id,
                                                result: Ok(addr.into()),
                                            });
                                        }
                                        self.cache.insert(
                                            domain,
                                            CacheEntry {
                                                ip: addr.into(),
                                                deadline: now
                                                    + Duration::from_secs(answer.ttl.into()),
                                            },
                                        );
                                        continue 'process;
                                    }
                                    _ => continue,
                                }
                            }
                            if resps == 0 {
                                self.queries.insert(
                                    qn,
                                    Query {
                                        domain,
                                        resps: resps + 1,
                                        deadline,
                                    },
                                );
                            } else {
                                for id in self.responses.remove(&domain).unwrap() {
                                    f(Response {
                                        id,
                                        result: Err(Error::NotFound),
                                    });
                                }
                            }
                        }
                        Err(e) => {
                            return Err(io::Error::new(
                                io::ErrorKind::Other,
                                format!("malformed dns packet received: {}", e),
                            ));
                        }
                    }
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(()),
                Err(e) => return Err(e),
            }
        }
    }

    pub fn tick<F: FnMut(Response)>(&mut self, mut f: F) {
        let now = Instant::now();
        let responses = &mut self.responses;
        self.queries.retain(|_, query| {
            if now > query.deadline {
                for id in responses.remove(&query.domain).unwrap() {
                    f(Response {
                        id,
                        result: Err(Error::Timeout),
                    });
                }
                false
            } else {
                true
            }
        });
        self.cache.retain(|_, entry| now < entry.deadline);
    }
}

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

    #[test]
    fn test_google() {
        let mut resolver = Resolver::new("8.8.8.8:53".parse().unwrap());
        let mut sock = UdpSocket::bind("0.0.0.0:0").unwrap();
        sock.set_nonblocking(true).unwrap();

        assert_eq!(resolver.query(&mut sock, 0, "google.com").unwrap(), None);
        assert_eq!(resolver.query(&mut sock, 1, "google.com").unwrap(), None);
        assert_eq!(resolver.responses.get("google.com").unwrap().len(), 2);
        std::thread::sleep(Duration::from_millis(100));
        resolver.tick(|_| panic!("timeout should not have occured yet!"));
        let mut count = 0;
        resolver
            .read(&mut sock, |resp| {
                count += 1;
                assert!(resp.result.is_ok());
            })
            .unwrap();
        assert_eq!(count, 2);

        assert!(
            resolver
                .query(&mut sock, 0, "google.com")
                .unwrap()
                .is_some()
        );

        resolver
            .query(&mut sock, 0, "thiswebsiteshouldexit12589t69.com")
            .unwrap();
        std::thread::sleep(Duration::from_millis(200));
        let mut processed = false;
        resolver
            .read(&mut sock, |resp| {
                processed = true;
                assert_eq!(resp.result, Err(Error::NotFound))
            })
            .unwrap();
        #[cfg(not(target_os = "macos"))]
        assert!(processed);
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use xref::{Xref};
use super::{Object, ObjectId, Dictionary};
use object_stream::ObjectStream;
use byref::ByRef;

/// PDF document.
#[derive(Debug)]
pub struct Document {
	/// The version of the PDF specification to which the file conforms.
	pub version: String,

	/// The trailer gives the location of the cross-reference table and of certain special objects.
	pub trailer: Dictionary,

	/// The cross-reference table contains locations of the indirect objects.
	pub reference_table: Xref,

	/// The objects that make up the document contained in the file.
	pub objects: BTreeMap<ObjectId, Object>,

	/// The object streams which contains compressed objects.
	pub streams: BTreeMap<u32, ObjectStream>,

	/// Current maximum object id within the document.
	pub max_id: u32,
}

impl Document {

	/// Create new PDF document.
	pub fn new() -> Document {
		Document {
			version: "1.4".to_string(),
			trailer: Dictionary::new(),
			reference_table: Xref::new(0),
			objects: BTreeMap::new(),
			streams: BTreeMap::new(),
			max_id: 0,
		}
	}

	/// Create new PDF document.
	pub fn with_version<S: Into<String>>(version: S) -> Document {
		let mut document = Self::new();
		document.version = version.into();
		document
	}

	/// Get object by object id, will recursively dereference a referenced object.
	pub fn get_object(&self, id: ObjectId) -> Option<&Object> {

		let mut ret_entry = None;

		if let Some(entry) = self.objects.get(&id) {
			match *entry {
				Object::Reference(obj_ref) => { ret_entry = self.get_object(obj_ref); }
				_ => { ret_entry = Some(entry); }
			}
		}

		return ret_entry;
	}

	/// Traverse objects from trailer recursively, return all referenced object IDs.
	pub fn traverse_objects<A: Fn(&mut Object) -> ()>(&mut self, action: A) -> Vec<ObjectId> {
		fn traverse_array<A: Fn(&mut Object) -> ()>(array: &mut Vec<Object>, action: &A, refs: &mut Vec<ObjectId>) {
			for item in array.iter_mut() {
				traverse_object(item, action, refs);
			}
		}
		fn traverse_dictionary<A: Fn(&mut Object) -> ()>(dict: &mut Dictionary, action: &A, refs: &mut Vec<ObjectId>) {
			for (_, v) in dict.iter_mut() {
				traverse_object(v, action, refs);
			}
		}
		fn traverse_object<A: Fn(&mut Object) -> ()>(object: &mut Object, action: &A, refs: &mut Vec<ObjectId>) {
			action(object);
			match *object {
				Object::Array(ref mut array) => traverse_array(array, action, refs),
				Object::Dictionary(ref mut dict) => traverse_dictionary(dict, action, refs),
				Object::Stream(ref mut stream) => traverse_dictionary(&mut stream.dict, action, refs),
				Object::Reference(id) => {
					if !refs.contains(&id) {
						refs.push(id);
					}
				},
				_ => {}
			}
		}
		let mut refs = vec![];
		traverse_dictionary(&mut self.trailer, &action, &mut refs);
		let mut index = 0;
		while index < refs.len() {
			if let Some(object) = self.objects.get_mut(&refs[index]) {
				traverse_object(object, &action, &mut refs);
			}
			index += 1;
		}
		refs
	}

	/// Get catalog dictionary.
	pub fn catalog(&self) -> Option<&Dictionary> {
		self.trailer.get("Root").get_dict_by_ref(self)
	}

	/// Get page numbers and corresponding object ids.
	pub fn get_pages(&self) -> BTreeMap<u32, ObjectId> {
		fn collect_pages(doc: &Document, page_tree_id: ObjectId, page_number: &mut u32, pages: &mut BTreeMap<u32, ObjectId>) {
			if let Some(kids) = doc.get_object(page_tree_id).and_then(|obj|obj.as_dict()).and_then(|page_tree|page_tree.get("Kids")).and_then(|obj|obj.as_array()) {
				for kid in kids {
					if let Some(kid_id) = kid.as_reference() {
						if let Some(type_name) = doc.get_object(kid_id).and_then(|obj|obj.as_dict()).and_then(|dict|dict.type_name()) {
							match type_name {
								"Page" => {
									pages.insert(*page_number, kid_id);
									*page_number += 1;
								}
								"Pages" => {
									collect_pages(doc, kid_id, page_number, pages);
								}
								_ => {}
							}
						}
					}
				}
			}
		}

		let mut pages = BTreeMap::new();
		let mut page_number = 1;
		if let Some(page_tree_id) = self.catalog().and_then(|cat|cat.get("Pages")).and_then(|pages|pages.as_reference()) {
			collect_pages(self, page_tree_id, &mut page_number, &mut pages);
		}
		pages
	}
}
<commit_msg>Revert "Bug: get_object() was searching in Xref table, not in the objects."<commit_after>use std::collections::BTreeMap;
use xref::{Xref, XrefEntry};
use super::{Object, ObjectId, Dictionary};
use object_stream::ObjectStream;
use byref::ByRef;

/// PDF document.
#[derive(Debug)]
pub struct Document {
	/// The version of the PDF specification to which the file conforms.
	pub version: String,

	/// The trailer gives the location of the cross-reference table and of certain special objects.
	pub trailer: Dictionary,

	/// The cross-reference table contains locations of the indirect objects.
	pub reference_table: Xref,

	/// The objects that make up the document contained in the file.
	pub objects: BTreeMap<ObjectId, Object>,

	/// The object streams which contains compressed objects.
	pub streams: BTreeMap<u32, ObjectStream>,

	/// Current maximum object id within the document.
	pub max_id: u32,
}

impl Document {

	/// Create new PDF document.
	pub fn new() -> Document {
		Document {
			version: "1.4".to_string(),
			trailer: Dictionary::new(),
			reference_table: Xref::new(0),
			objects: BTreeMap::new(),
			streams: BTreeMap::new(),
			max_id: 0,
		}
	}

	/// Create new PDF document.
	pub fn with_version<S: Into<String>>(version: S) -> Document {
		let mut document = Self::new();
		document.version = version.into();
		document
	}

	/// Get object by object id, will recursively dereference a referenced object.
	pub fn get_object(&self, id: ObjectId) -> Option<&Object> {
		if let Some(entry) = self.reference_table.get(id.0) {
			match *entry {
				XrefEntry::Normal { .. } => {
					if let Some(object) = self.objects.get(&id) {
						if let Some(id) = object.as_reference() {
							return self.get_object(id);
						} else {
							return Some(object);
						}
					}
				}
				XrefEntry::Compressed { container, index } => {
					if let Some(stream) = self.streams.get(&container) {
						if let Some(&(_id, ref object)) = stream.get_object(index as usize) {
							return Some(object);
						}
					}
				}
				_ => {},
			}
		}
		return None;
	}

	/// Traverse objects from trailer recursively, return all referenced object IDs.
	pub fn traverse_objects<A: Fn(&mut Object) -> ()>(&mut self, action: A) -> Vec<ObjectId> {
		fn traverse_array<A: Fn(&mut Object) -> ()>(array: &mut Vec<Object>, action: &A, refs: &mut Vec<ObjectId>) {
			for item in array.iter_mut() {
				traverse_object(item, action, refs);
			}
		}
		fn traverse_dictionary<A: Fn(&mut Object) -> ()>(dict: &mut Dictionary, action: &A, refs: &mut Vec<ObjectId>) {
			for (_, v) in dict.iter_mut() {
				traverse_object(v, action, refs);
			}
		}
		fn traverse_object<A: Fn(&mut Object) -> ()>(object: &mut Object, action: &A, refs: &mut Vec<ObjectId>) {
			action(object);
			match *object {
				Object::Array(ref mut array) => traverse_array(array, action, refs),
				Object::Dictionary(ref mut dict) => traverse_dictionary(dict, action, refs),
				Object::Stream(ref mut stream) => traverse_dictionary(&mut stream.dict, action, refs),
				Object::Reference(id) => {
					if !refs.contains(&id) {
						refs.push(id);
					}
				},
				_ => {}
			}
		}
		let mut refs = vec![];
		traverse_dictionary(&mut self.trailer, &action, &mut refs);
		let mut index = 0;
		while index < refs.len() {
			if let Some(object) = self.objects.get_mut(&refs[index]) {
				traverse_object(object, &action, &mut refs);
			}
			index += 1;
		}
		refs
	}

	/// Get catalog dictionary.
	pub fn catalog(&self) -> Option<&Dictionary> {
		self.trailer.get("Root").get_dict_by_ref(self)
	}

	/// Get page numbers and corresponding object ids.
	pub fn get_pages(&self) -> BTreeMap<u32, ObjectId> {
		fn collect_pages(doc: &Document, page_tree_id: ObjectId, page_number: &mut u32, pages: &mut BTreeMap<u32, ObjectId>) {
			if let Some(kids) = doc.get_object(page_tree_id).and_then(|obj|obj.as_dict()).and_then(|page_tree|page_tree.get("Kids")).and_then(|obj|obj.as_array()) {
				for kid in kids {
					if let Some(kid_id) = kid.as_reference() {
						if let Some(type_name) = doc.get_object(kid_id).and_then(|obj|obj.as_dict()).and_then(|dict|dict.type_name()) {
							match type_name {
								"Page" => {
									pages.insert(*page_number, kid_id);
									*page_number += 1;
								}
								"Pages" => {
									collect_pages(doc, kid_id, page_number, pages);
								}
								_ => {}
							}
						}
					}
				}
			}
		}

		let mut pages = BTreeMap::new();
		let mut page_number = 1;
		if let Some(page_tree_id) = self.catalog().and_then(|cat|cat.get("Pages")).and_then(|pages|pages.as_reference()) {
			collect_pages(self, page_tree_id, &mut page_number, &mut pages);
		}
		pages
	}
}
<|endoftext|>


<commit_before>use std::path::Path;
use std::fs::{create_dir, remove_dir_all, remove_file, rename};

use fs::Filesystem;
use result::RustyPlatterResult;

/// Implementation of a local filesystem
pub struct LocalFileSystem<'a> {
    base_path: &'a Path,
}

impl<'a> LocalFileSystem<'a> {
    pub fn new(base_path: &'a str) -> Self {
        LocalFileSystem { base_path: Path::new(base_path) }
    }
}

impl<'a> Filesystem for LocalFileSystem<'a> {
    fn mkdir(&self, path: &str) -> RustyPlatterResult<()> {
        Ok(create_dir(self.base_path.join(path))?)
    }
    fn mv(&self, from: &str, to: &str) -> RustyPlatterResult<()> {
        Ok(rename(self.base_path.join(from), self.base_path.join(to))?)
    }
    fn rm(&self, path: &str) -> RustyPlatterResult<()> {
        let path = self.base_path.join(path);
        if path.is_dir() {
            Ok(remove_dir_all(path)?)
        } else {
            Ok(remove_file(path)?)
        }
    }
    fn exists(&self, path: &str) -> bool {
        self.base_path.join(path).exists()
    }
}

#[cfg(test)]
mod tests {
    extern crate tempdir;

    use std::fs as std_fs;

    use self::tempdir::TempDir;

    use ::fs::Filesystem;
    use ::fs::local::LocalFileSystem;

    #[test]
    fn test_mkdir() {
        let temp = TempDir::new("test_mkdir").unwrap();
        let path = temp.path();
        let fs = LocalFileSystem::new(path.to_str().unwrap());
        fs.mkdir("test")?;
        assert!(path.join("test").exists());
    }

    #[test]
    fn test_exists() {
        let temp = TempDir::new("test_mkdir").unwrap();
        let path = temp.path();
        let fs = LocalFileSystem::new(path.to_str().unwrap());
        assert!(fs.exists("."));
        assert!(!fs.exists("abc"));
    }

    #[test]
    fn test_rm() {
        let temp = TempDir::new("test_mkdir").unwrap();
        let path = temp.path();
        let fs = LocalFileSystem::new(path.to_str().unwrap());

        std_fs::create_dir(path.join("dir")).unwrap();
        fs.rm("dir").unwrap();
        std_fs::File::create(path.join("file")).unwrap();
        fs.rm("file").unwrap();
    }

    #[test]
    fn test_mv() {
        let temp = TempDir::new("test_mkdir").unwrap();
        let path = temp.path();
        let fs = LocalFileSystem::new(path.to_str().unwrap());

        std_fs::create_dir(path.join("dir")).unwrap();
        fs.mv("dir", "dir2").unwrap();
        assert!(path.join("dir2").exists());

        std_fs::File::create(path.join("file")).unwrap();
        fs.mv("file", "file2").unwrap();
        assert!(path.join("file2").exists());
    }
}
<commit_msg>test: fix mkdir test<commit_after>use std::path::Path;
use std::fs::{create_dir, remove_dir_all, remove_file, rename};

use fs::Filesystem;
use result::RustyPlatterResult;

/// Implementation of a local filesystem
pub struct LocalFileSystem<'a> {
    base_path: &'a Path,
}

impl<'a> LocalFileSystem<'a> {
    pub fn new(base_path: &'a str) -> Self {
        LocalFileSystem { base_path: Path::new(base_path) }
    }
}

impl<'a> Filesystem for LocalFileSystem<'a> {
    fn mkdir(&self, path: &str) -> RustyPlatterResult<()> {
        Ok(create_dir(self.base_path.join(path))?)
    }
    fn mv(&self, from: &str, to: &str) -> RustyPlatterResult<()> {
        Ok(rename(self.base_path.join(from), self.base_path.join(to))?)
    }
    fn rm(&self, path: &str) -> RustyPlatterResult<()> {
        let path = self.base_path.join(path);
        if path.is_dir() {
            Ok(remove_dir_all(path)?)
        } else {
            Ok(remove_file(path)?)
        }
    }
    fn exists(&self, path: &str) -> bool {
        self.base_path.join(path).exists()
    }
}

#[cfg(test)]
mod tests {
    extern crate tempdir;

    use std::fs as std_fs;

    use self::tempdir::TempDir;

    use ::fs::Filesystem;
    use ::fs::local::LocalFileSystem;

    #[test]
    fn test_mkdir() {
        let temp = TempDir::new("test_mkdir").unwrap();
        let path = temp.path();
        let fs = LocalFileSystem::new(path.to_str().unwrap());
        fs.mkdir("test").unwrap();
        assert!(path.join("test").exists());
    }

    #[test]
    fn test_exists() {
        let temp = TempDir::new("test_mkdir").unwrap();
        let path = temp.path();
        let fs = LocalFileSystem::new(path.to_str().unwrap());
        assert!(fs.exists("."));
        assert!(!fs.exists("abc"));
    }

    #[test]
    fn test_rm() {
        let temp = TempDir::new("test_mkdir").unwrap();
        let path = temp.path();
        let fs = LocalFileSystem::new(path.to_str().unwrap());

        std_fs::create_dir(path.join("dir")).unwrap();
        fs.rm("dir").unwrap();
        std_fs::File::create(path.join("file")).unwrap();
        fs.rm("file").unwrap();
    }

    #[test]
    fn test_mv() {
        let temp = TempDir::new("test_mkdir").unwrap();
        let path = temp.path();
        let fs = LocalFileSystem::new(path.to_str().unwrap());

        std_fs::create_dir(path.join("dir")).unwrap();
        fs.mv("dir", "dir2").unwrap();
        assert!(path.join("dir2").exists());

        std_fs::File::create(path.join("file")).unwrap();
        fs.mv("file", "file2").unwrap();
        assert!(path.join("file2").exists());
    }
}
<|endoftext|>


<commit_before>#![feature(test)]
extern crate test;

#[macro_use]
extern crate lazy_static;

extern crate brainfuck;

use test::Bencher;

use brainfuck::{precompile, Instruction, OptimizationLevel};

lazy_static! {
    // This program is trivial to run in both size and speed
    static ref TRIVIAL_SOURCE: Vec<u8> = include_bytes!("../examples/hello-world.bf").to_vec();
    // This is a very large program in size
    static ref LARGE_SOURCE: Vec<u8> = include_bytes!("../examples/hanoi.bf").to_vec();
    // This program is absurdly huge and is mostly a corner case for the time being
    static ref HUGE_SOURCE: Vec<u8> = include_bytes!("../examples/LostKingdom.bf").to_vec();
    // This program is reasonably small, non-trivial to run, and simple overall
    static ref SIMPLE_SOURCE: Vec<u8> = include_bytes!("../examples/99bottles.bf").to_vec();
    // This program, while only relatively large, takes a long time to run
    static ref SLOW_SOURCE: Vec<u8> = include_bytes!("../examples/mandel.bf").to_vec();
}

struct NullWrite;

impl std::io::Write for NullWrite {
    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
        Ok(bytes.len())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

fn interpret(program: Vec<Instruction>) {
    let mut inp: &[u8] = &[];
    brainfuck::interpret(&mut inp, NullWrite, program, |_| {});
}

#[bench]
fn b01_compile_trivial(b: &mut Bencher) {
    b.iter(|| precompile(TRIVIAL_SOURCE.iter(), OptimizationLevel::Off));
}

#[bench]
fn b01_compile_trivial_opt(b: &mut Bencher) {
    b.iter(|| precompile(TRIVIAL_SOURCE.iter(), OptimizationLevel::Speed));
}

#[bench]
fn b02_compile_large(b: &mut Bencher) {
    b.iter(|| precompile(LARGE_SOURCE.iter(), OptimizationLevel::Off));
}

#[bench]
fn b02_compile_large_opt(b: &mut Bencher) {
    b.iter(|| precompile(LARGE_SOURCE.iter(), OptimizationLevel::Speed));
}

#[bench]
fn b03_compile_huge(b: &mut Bencher) {
    b.iter(|| precompile(HUGE_SOURCE.iter(), OptimizationLevel::Off));
}

#[bench]
fn b03_compile_huge_opt(b: &mut Bencher) {
    b.iter(|| precompile(HUGE_SOURCE.iter(), OptimizationLevel::Speed));
}

#[bench]
fn b04_compile_simple(b: &mut Bencher) {
    b.iter(|| precompile(SIMPLE_SOURCE.iter(), OptimizationLevel::Off));
}

#[bench]
fn b04_compile_simple_opt(b: &mut Bencher) {
    b.iter(|| precompile(SIMPLE_SOURCE.iter(), OptimizationLevel::Speed));
}

#[bench]
fn b05_compile_slow(b: &mut Bencher) {
    b.iter(|| precompile(SLOW_SOURCE.iter(), OptimizationLevel::Off));
}

#[bench]
fn b05_compile_slow_opt(b: &mut Bencher) {
    b.iter(|| precompile(SLOW_SOURCE.iter(), OptimizationLevel::Speed));
}

#[bench]
fn b06_interpret_trivial(b: &mut Bencher) {
    let program = precompile(TRIVIAL_SOURCE.iter(), OptimizationLevel::Off);
    b.iter(|| interpret(program.clone()));
}

#[bench]
fn b06_interpret_trivial_opt(b: &mut Bencher) {
    let program = precompile(TRIVIAL_SOURCE.iter(), OptimizationLevel::Speed);
    b.iter(|| interpret(program.clone()));
}

#[bench]
fn b07_interpret_simple(b: &mut Bencher) {
    let program = precompile(SIMPLE_SOURCE.iter(), OptimizationLevel::Off);
    b.iter(|| interpret(program.clone()));
}

#[bench]
fn b07_interpret_simple_opt(b: &mut Bencher) {
    let program = precompile(SIMPLE_SOURCE.iter(), OptimizationLevel::Speed);
    b.iter(|| interpret(program.clone()));
}

#[bench]
#[ignore]
fn b08_interpret_slow(b: &mut Bencher) {
    let program = precompile(SLOW_SOURCE.iter(), OptimizationLevel::Off);
    b.iter(|| interpret(program.clone()));
}

#[bench]
#[ignore]
fn b08_interpret_slow_opt(b: &mut Bencher) {
    let program = precompile(SLOW_SOURCE.iter(), OptimizationLevel::Speed);
    b.iter(|| interpret(program.clone()));
}
<commit_msg>Replaced NullWrite with io::Sink (#21)<commit_after>#![feature(test)]
extern crate test;

#[macro_use]
extern crate lazy_static;

extern crate brainfuck;

use std::io;

use test::Bencher;

use brainfuck::{precompile, Instruction, OptimizationLevel};

lazy_static! {
    // This program is trivial to run in both size and speed
    static ref TRIVIAL_SOURCE: Vec<u8> = include_bytes!("../examples/hello-world.bf").to_vec();
    // This is a very large program in size
    static ref LARGE_SOURCE: Vec<u8> = include_bytes!("../examples/hanoi.bf").to_vec();
    // This program is absurdly huge and is mostly a corner case for the time being
    static ref HUGE_SOURCE: Vec<u8> = include_bytes!("../examples/LostKingdom.bf").to_vec();
    // This program is reasonably small, non-trivial to run, and simple overall
    static ref SIMPLE_SOURCE: Vec<u8> = include_bytes!("../examples/99bottles.bf").to_vec();
    // This program, while only relatively large, takes a long time to run
    static ref SLOW_SOURCE: Vec<u8> = include_bytes!("../examples/mandel.bf").to_vec();
}

fn interpret(program: Vec<Instruction>) {
    let mut inp: &[u8] = &[];
    brainfuck::interpret(&mut inp, io::sink(), program, |_| {});
}

#[bench]
fn b01_compile_trivial(b: &mut Bencher) {
    b.iter(|| precompile(TRIVIAL_SOURCE.iter(), OptimizationLevel::Off));
}

#[bench]
fn b01_compile_trivial_opt(b: &mut Bencher) {
    b.iter(|| precompile(TRIVIAL_SOURCE.iter(), OptimizationLevel::Speed));
}

#[bench]
fn b02_compile_large(b: &mut Bencher) {
    b.iter(|| precompile(LARGE_SOURCE.iter(), OptimizationLevel::Off));
}

#[bench]
fn b02_compile_large_opt(b: &mut Bencher) {
    b.iter(|| precompile(LARGE_SOURCE.iter(), OptimizationLevel::Speed));
}

#[bench]
fn b03_compile_huge(b: &mut Bencher) {
    b.iter(|| precompile(HUGE_SOURCE.iter(), OptimizationLevel::Off));
}

#[bench]
fn b03_compile_huge_opt(b: &mut Bencher) {
    b.iter(|| precompile(HUGE_SOURCE.iter(), OptimizationLevel::Speed));
}

#[bench]
fn b04_compile_simple(b: &mut Bencher) {
    b.iter(|| precompile(SIMPLE_SOURCE.iter(), OptimizationLevel::Off));
}

#[bench]
fn b04_compile_simple_opt(b: &mut Bencher) {
    b.iter(|| precompile(SIMPLE_SOURCE.iter(), OptimizationLevel::Speed));
}

#[bench]
fn b05_compile_slow(b: &mut Bencher) {
    b.iter(|| precompile(SLOW_SOURCE.iter(), OptimizationLevel::Off));
}

#[bench]
fn b05_compile_slow_opt(b: &mut Bencher) {
    b.iter(|| precompile(SLOW_SOURCE.iter(), OptimizationLevel::Speed));
}

#[bench]
fn b06_interpret_trivial(b: &mut Bencher) {
    let program = precompile(TRIVIAL_SOURCE.iter(), OptimizationLevel::Off);
    b.iter(|| interpret(program.clone()));
}

#[bench]
fn b06_interpret_trivial_opt(b: &mut Bencher) {
    let program = precompile(TRIVIAL_SOURCE.iter(), OptimizationLevel::Speed);
    b.iter(|| interpret(program.clone()));
}

#[bench]
fn b07_interpret_simple(b: &mut Bencher) {
    let program = precompile(SIMPLE_SOURCE.iter(), OptimizationLevel::Off);
    b.iter(|| interpret(program.clone()));
}

#[bench]
fn b07_interpret_simple_opt(b: &mut Bencher) {
    let program = precompile(SIMPLE_SOURCE.iter(), OptimizationLevel::Speed);
    b.iter(|| interpret(program.clone()));
}

#[bench]
#[ignore]
fn b08_interpret_slow(b: &mut Bencher) {
    let program = precompile(SLOW_SOURCE.iter(), OptimizationLevel::Off);
    b.iter(|| interpret(program.clone()));
}

#[bench]
#[ignore]
fn b08_interpret_slow_opt(b: &mut Bencher) {
    let program = precompile(SLOW_SOURCE.iter(), OptimizationLevel::Speed);
    b.iter(|| interpret(program.clone()));
}
<|endoftext|>