<commit_before>extern crate resp;
extern crate rusqlite;

use self::resp::{Value};
use self::rusqlite::{Connection};
use std::sync::{Arc, Mutex, MutexGuard};
use std::cmp;

type CommandResult = Result<Value, String>;

struct Command<'a> {
    name: String,
    arguments: Vec<&'a String>,
    connection_mutex: &'a Arc<Mutex<Connection>>
}

struct CommandSettings {
    name: &'static str,
    argument_count: i32,
    //handler: fn(&'a Command<'a>) -> CommandResult
}

enum Direction {
    Left,
    Right
}

const COMMAND_SETTINGS: [CommandSettings; 8] = [
    CommandSettings { name: "LLEN", argument_count: 1 }, //, handler: Command::llen }
    CommandSettings { name: "LPOP", argument_count: 1 },
    CommandSettings { name: "RPOP", argument_count: 1 },
    CommandSettings { name: "LPUSH", argument_count: -2 },
    CommandSettings { name: "RPUSH", argument_count: -2 },
    CommandSettings { name: "LRANGE", argument_count: 3 },
    CommandSettings { name: "LTRIM", argument_count: 3 },
    CommandSettings { name: "RPOPLPUSH", argument_count: 2 },
];

pub fn handle_input(ref value: Value, connection_mutex: &Arc<Mutex<Connection>>) -> (Value, bool) {
    match parse_command(value, connection_mutex) {
        Ok(mut command) => { command.handle_command() }
        Err(error)  => { (Value::Error(format!("ERR {}", error)), false) }
    }
}

fn parse_command<'a>(value: &'a Value, connection_mutex: &'a Arc<Mutex<Connection>>) -> Result<Command<'a>, &'static str> {
    if let &Value::Array(ref array) = value {
        parse_command_array(array, connection_mutex)
    }
    else {
        Err("expected array")
    }
}

fn parse_command_array<'a>(array: &'a Vec<Value>, connection_mutex: &'a Arc<Mutex<Connection>>) -> Result<Command<'a>, &'static str> {
    let iter = array.iter().map(|value|
        match *value {
            Value::String(ref string) | Value::Bulk(ref string) => Ok(string),
            _ => Err("all arguments should be strings")
        }
    );

    let strings = iter.collect::<Result<Vec<&String>, &'static str>>()?;
    let (head, tail) = strings.split_at(1);

    Ok(Command {
        name: head[0].to_uppercase(),
        arguments: tail.to_vec(),
        connection_mutex: connection_mutex
    })
}

impl<'a> Command<'a> {
    fn handle_command(&mut self) -> (Value, bool) {
        print!("command {}\n", self.name);

        match self.name.as_ref() {
            "QUIT" => (Value::String("OK".to_string()), true),
            _      => (self.handle_nonterminal_command(), false)
        }
    }

    fn valid_argument_count(&self, command: &CommandSettings) -> bool {
        (command.argument_count < 0 && self.arguments.len() as i32 >= -command.argument_count) ||
            (command.argument_count >= 0 && self.arguments.len() as i32 == command.argument_count)
    }

    fn handle_nonterminal_command(&mut self) -> Value {
        let all = COMMAND_SETTINGS;
        let settings = all.iter().find(|settings| settings.name == self.name);

        match settings {
            Some(settings) => {
                if self.valid_argument_count(settings) {
                    let result = match settings.name {
                        "LLEN" => self.llen(),
                        "LPUSH" => self.lpush(),
                        "RPUSH" => self.rpush(),
                        "LPOP" => self.lpop(),
                        "RPOP" => self.rpop(),
                        "LRANGE" => self.lrange(),
                        "LTRIM" => self.ltrim(),
                        "RPOPLPUSH" => self.rpoplpush(),
                        _ => unimplemented!(),
                    };

                    match result {
                        Ok(value)  => value,
                        Err(error) => Value::Error(format!("ERR {}", error))
                    }
                }
                else {
                    Value::Error("ERR wrong number of arguments".to_string())
                }
            }

            None => Value::Error("ERR unsupported".to_string())
        }
    }

    fn llen(&self) -> CommandResult {
        let key = self.arguments[0];

        let connection = self.lock_connection();
        self.count_list_items_value(&*connection, key)
    }

    fn lpop(&self) -> CommandResult {
        let connection = self.lock_connection();

        match self.pop(&*connection, self.arguments[0], Direction::Left) {
            Some(data) => Ok(Value::Bulk(data)),
            None       => Ok(Value::NullArray)
        }
    }

    fn rpop(&self) -> CommandResult {
        let connection = self.lock_connection();

        match self.pop(&*connection, self.arguments[0], Direction::Right) {
            Some(data) => Ok(Value::Bulk(data)),
            None       => Ok(Value::NullArray)
        }
    }

    fn lpush(&self) -> CommandResult {
        let key = self.arguments[0];
        let mut connection = self.lock_connection();

        self.push(&mut *connection, key, Direction::Left, self.arguments.iter().skip(1));

        self.count_list_items_value(&*connection, key)
    }

    fn rpush(&self) -> CommandResult {
        let key = self.arguments[0];
        let mut connection = self.lock_connection();

        self.push(&mut *connection, key, Direction::Right, self.arguments.iter().skip(1));

        self.count_list_items_value(&*connection, key)
    }

    fn lrange(&self) -> CommandResult {
        let key = self.arguments[0];
        let mut start: i64 = self.arguments[1].parse().map_err(|_| "start must be an integer")?;
        let mut stop: i64 = self.arguments[2].parse().map_err(|_| "stop must be an integer")?;

        let connection = self.lock_connection();

        // LATER : maybe use find_position_boundaries and parse_positions instead of this
        // this needs only one select rather than two in the case where start and stop >= 0
        // but then it uses order & limit which are probably slower on a huge table
        // investigate performance at a later date

        if start < 0 || stop < -1 {
            let count = self.count_list_items(&connection, key);

            if start < 0  { start = cmp::max(0, count + start) }

            if stop < -1 {
                stop = count + stop;
                if stop < 0 { return Ok(Value::Array(vec![])); }
            }
        }

        if stop != -1 && start > stop { return Ok(Value::Array(vec![])); }

        let sql = match (start, stop) {
            (0, -1) => "".to_string(),
            (a, -1) => format!("LIMIT -1 OFFSET {}", a),
            (0, b)  => format!("LIMIT {}", b + 1),
            (a, b)  => format!("LIMIT {} OFFSET {}", b - a + 1, a),
        };

        let mut statement = connection.prepare(&format!("SELECT value FROM list_items WHERE key = ?1 ORDER BY position {}", sql)).unwrap();
        let rows = statement.query_map(&[key], |row| row.get(0)).unwrap();
        let result: Result<Vec<String>, _> = rows.collect();
        let values = result.unwrap().iter().map(|value| Value::String(value.clone())).collect();

        Ok(Value::Array(values))
    }

    fn ltrim(&self) -> CommandResult {
        let key = self.arguments[0];
        let start: i64 = self.arguments[1].parse().map_err(|_| "start must be an integer")?;
        let stop: i64 = self.arguments[2].parse().map_err(|_| "stop must be an integer")?;

        if start != 0 || stop != -1 {
            let connection = self.lock_connection();

            let boundaries = self.find_position_boundaries(&*connection, key);
            let (start_position, stop_position) = self.parse_positions(boundaries, (start, stop));

            connection.execute("DELETE FROM list_items WHERE key = ?1 AND (position < ?2 OR position > ?3)", &[key, &start_position, &stop_position]).unwrap();
        }

        Ok(Value::String("OK".to_string()))
    }

    fn rpoplpush(&self) -> CommandResult {
        let source = self.arguments[0];
        let destination = self.arguments[1];

        let mut connection = self.lock_connection();

        match self.pop(&*connection, source, Direction::Right) {
            Some(data) => {
                self.push(&mut *connection, destination, Direction::Left, [&data].iter());
                Ok(Value::Bulk(data))
            }

            None => Ok(Value::NullArray)
        }
    }

    // private

    fn pop(&self, connection: &Connection, key: &String, direction: Direction) -> Option<String> {
        let order = match direction { Direction::Left => "ASC", Direction::Right => "DESC" };
        let mut statement = connection.prepare(&format!("SELECT id, value FROM list_items WHERE key = ?1 ORDER BY position {} LIMIT 1", order)).unwrap();

        match statement.query_row(&[key], |row| (row.get(0), row.get(1))) {
            Ok(result) => {
                let (id, value): (i64, String) = result;

                connection.execute("DELETE FROM list_items WHERE id = ?1", &[&id]).unwrap();
                Some(value)
            }

            Err(rusqlite::Error::QueryReturnedNoRows) => None,

            Err(e) => Err(e).unwrap()
        }
    }

    fn push<'b, I>(&self, connection: &mut Connection, key: &String, direction: Direction, iterator: I) -> ()
        where I: Iterator<Item=&'b &'b String>
    {
        let tx = connection.transaction().unwrap();

        let next_position_sql = match direction {
            Direction::Left  => "coalesce(MIN(position), 0) - 1",
            Direction::Right => "coalesce(MAX(position), 0) + 1"
        };

        let sql = format!("INSERT INTO list_items (key, value, position) SELECT ?1, ?2, {} FROM list_items WHERE key = ?1", next_position_sql);

        iterator.map(|value| tx.execute(&sql, &[key, *value])).collect::<Result<Vec<_>, _>>().unwrap();

        tx.commit().unwrap();
    }

    fn lock_connection(&self) -> MutexGuard<Connection> {
        (*self.connection_mutex).lock().unwrap()
    }

    fn find_position_boundaries(&self, connection: &Connection, key: &String) -> (i64, i64) {
        let mut statement = connection.prepare("SELECT MIN(position), MAX(position) AS c FROM list_items WHERE key = ?1").unwrap();
        statement.query_row(&[key], |row| (row.get(0), row.get(1))).unwrap()
    }

    fn parse_positions(&self, (first_position, last_position) : (i64, i64), (start, stop): (i64, i64)) -> (i64, i64) {
        let start_position = match start {
            position if position < 0 => cmp::max(0, position + last_position + 1),
            position                 => position + first_position
        };

        let stop_position = match stop {
            position if position < 0 => cmp::max(0, position + last_position + 1),
            position                 => position + first_position
        };

        (start_position, stop_position)
    }

    fn count_list_items(&self, connection: &Connection, key: &String) -> i64 {
        let mut statement = connection.prepare("SELECT COUNT(*) AS c FROM list_items WHERE key = ?1").unwrap();
        statement.query_row(&[key], |row| row.get(0)).unwrap()
    }

    fn count_list_items_value(&self, connection: &Connection, key: &String) -> CommandResult {
        Ok(Value::Integer(self.count_list_items(connection, key)))
    }
}
<commit_msg>lindex<commit_after>extern crate resp;
extern crate rusqlite;

use self::resp::{Value};
use self::rusqlite::{Connection};
use std::sync::{Arc, Mutex, MutexGuard};
use std::cmp;

type CommandResult = Result<Value, String>;

struct Command<'a> {
    name: String,
    arguments: Vec<&'a String>,
    connection_mutex: &'a Arc<Mutex<Connection>>
}

struct CommandSettings {
    name: &'static str,
    argument_count: i32,
    //handler: fn(&'a Command<'a>) -> CommandResult
}

enum Direction {
    Left,
    Right
}

const COMMAND_SETTINGS: [CommandSettings; 9] = [
    CommandSettings { name: "LLEN", argument_count: 1 }, //, handler: Command::llen }
    CommandSettings { name: "LPOP", argument_count: 1 },
    CommandSettings { name: "RPOP", argument_count: 1 },
    CommandSettings { name: "LPUSH", argument_count: -2 },
    CommandSettings { name: "RPUSH", argument_count: -2 },
    CommandSettings { name: "LRANGE", argument_count: 3 },
    CommandSettings { name: "LTRIM", argument_count: 3 },
    CommandSettings { name: "RPOPLPUSH", argument_count: 2 },
    CommandSettings { name: "LINDEX", argument_count: 2 },
];

pub fn handle_input(ref value: Value, connection_mutex: &Arc<Mutex<Connection>>) -> (Value, bool) {
    match parse_command(value, connection_mutex) {
        Ok(mut command) => { command.handle_command() }
        Err(error)  => { (Value::Error(format!("ERR {}", error)), false) }
    }
}

fn parse_command<'a>(value: &'a Value, connection_mutex: &'a Arc<Mutex<Connection>>) -> Result<Command<'a>, &'static str> {
    if let &Value::Array(ref array) = value {
        parse_command_array(array, connection_mutex)
    }
    else {
        Err("expected array")
    }
}

fn parse_command_array<'a>(array: &'a Vec<Value>, connection_mutex: &'a Arc<Mutex<Connection>>) -> Result<Command<'a>, &'static str> {
    let iter = array.iter().map(|value|
        match *value {
            Value::String(ref string) | Value::Bulk(ref string) => Ok(string),
            _ => Err("all arguments should be strings")
        }
    );

    let strings = iter.collect::<Result<Vec<&String>, &'static str>>()?;
    let (head, tail) = strings.split_at(1);

    Ok(Command {
        name: head[0].to_uppercase(),
        arguments: tail.to_vec(),
        connection_mutex: connection_mutex
    })
}

impl<'a> Command<'a> {
    fn handle_command(&mut self) -> (Value, bool) {
        print!("command {}\n", self.name);

        match self.name.as_ref() {
            "QUIT" => (Value::String("OK".to_string()), true),
            _      => (self.handle_nonterminal_command(), false)
        }
    }

    fn valid_argument_count(&self, command: &CommandSettings) -> bool {
        (command.argument_count < 0 && self.arguments.len() as i32 >= -command.argument_count) ||
            (command.argument_count >= 0 && self.arguments.len() as i32 == command.argument_count)
    }

    fn handle_nonterminal_command(&mut self) -> Value {
        let all = COMMAND_SETTINGS;
        let settings = all.iter().find(|settings| settings.name == self.name);

        match settings {
            Some(settings) => {
                if self.valid_argument_count(settings) {
                    let result = match settings.name {
                        "LLEN" => self.llen(),
                        "LPUSH" => self.lpush(),
                        "RPUSH" => self.rpush(),
                        "LPOP" => self.lpop(),
                        "RPOP" => self.rpop(),
                        "LRANGE" => self.lrange(),
                        "LTRIM" => self.ltrim(),
                        "RPOPLPUSH" => self.rpoplpush(),
                        "LINDEX" => self.lindex(),
                        _ => unimplemented!(),
                    };

                    match result {
                        Ok(value)  => value,
                        Err(error) => Value::Error(format!("ERR {}", error))
                    }
                }
                else {
                    Value::Error("ERR wrong number of arguments".to_string())
                }
            }

            None => Value::Error("ERR unsupported".to_string())
        }
    }

    fn llen(&self) -> CommandResult {
        let key = self.arguments[0];

        let connection = self.lock_connection();
        self.count_list_items_value(&*connection, key)
    }

    fn lpop(&self) -> CommandResult {
        let connection = self.lock_connection();

        match self.pop(&*connection, self.arguments[0], Direction::Left) {
            Some(data) => Ok(Value::Bulk(data)),
            None       => Ok(Value::NullArray)
        }
    }

    fn rpop(&self) -> CommandResult {
        let connection = self.lock_connection();

        match self.pop(&*connection, self.arguments[0], Direction::Right) {
            Some(data) => Ok(Value::Bulk(data)),
            None       => Ok(Value::NullArray)
        }
    }

    fn lpush(&self) -> CommandResult {
        let key = self.arguments[0];
        let mut connection = self.lock_connection();

        self.push(&mut *connection, key, Direction::Left, self.arguments.iter().skip(1));

        self.count_list_items_value(&*connection, key)
    }

    fn rpush(&self) -> CommandResult {
        let key = self.arguments[0];
        let mut connection = self.lock_connection();

        self.push(&mut *connection, key, Direction::Right, self.arguments.iter().skip(1));

        self.count_list_items_value(&*connection, key)
    }

    fn lrange(&self) -> CommandResult {
        let key = self.arguments[0];
        let mut start: i64 = self.arguments[1].parse().map_err(|_| "start must be an integer")?;
        let mut stop: i64 = self.arguments[2].parse().map_err(|_| "stop must be an integer")?;

        let connection = self.lock_connection();

        // LATER : maybe use find_position_boundaries and parse_indexes instead of this
        // this needs only one select rather than two in the case where start and stop >= 0
        // but then it uses order & limit which are probably slower on a huge table
        // investigate performance at a later date

        if start < 0 || stop < -1 {
            let count = self.count_list_items(&connection, key);

            if start < 0  { start = cmp::max(0, count + start) }

            if stop < -1 {
                stop = count + stop;
                if stop < 0 { return Ok(Value::Array(vec![])); }
            }
        }

        if stop != -1 && start > stop { return Ok(Value::Array(vec![])); }

        let sql = match (start, stop) {
            (0, -1) => "".to_string(),
            (a, -1) => format!("LIMIT -1 OFFSET {}", a),
            (0, b)  => format!("LIMIT {}", b + 1),
            (a, b)  => format!("LIMIT {} OFFSET {}", b - a + 1, a),
        };

        let mut statement = connection.prepare(&format!("SELECT value FROM list_items WHERE key = ?1 ORDER BY position {}", sql)).unwrap();
        let rows = statement.query_map(&[key], |row| row.get(0)).unwrap();
        let result: Result<Vec<String>, _> = rows.collect();
        let values = result.unwrap().iter().map(|value| Value::String(value.clone())).collect();

        Ok(Value::Array(values))
    }

    fn ltrim(&self) -> CommandResult {
        let key = self.arguments[0];
        let start: i64 = self.arguments[1].parse().map_err(|_| "start must be an integer")?;
        let stop: i64 = self.arguments[2].parse().map_err(|_| "stop must be an integer")?;

        if start != 0 || stop != -1 {
            let connection = self.lock_connection();

            let boundaries = self.find_position_boundaries(&*connection, key);
            let (start_position, stop_position) = self.parse_indexes(boundaries, (start, stop));

            connection.execute("DELETE FROM list_items WHERE key = ?1 AND (position < ?2 OR position > ?3)", &[key, &start_position, &stop_position]).unwrap();
        }

        Ok(Value::String("OK".to_string()))
    }

    fn rpoplpush(&self) -> CommandResult {
        let source = self.arguments[0];
        let destination = self.arguments[1];

        let mut connection = self.lock_connection();

        match self.pop(&*connection, source, Direction::Right) {
            Some(data) => {
                self.push(&mut *connection, destination, Direction::Left, [&data].iter());
                Ok(Value::Bulk(data))
            }

            None => Ok(Value::NullArray)
        }
    }

    fn lindex(&self) -> CommandResult {
        let key = self.arguments[0];
        let index: i64 = self.arguments[1].parse().map_err(|_| "index must be an integer")?;

        let connection = self.lock_connection();

        let boundaries = self.find_position_boundaries(&*connection, key);
        let position = self.parse_index(boundaries, index);

        let mut statement = connection.prepare("SELECT value FROM list_items WHERE key = ?1 AND position = ?2 LIMIT 1").unwrap();

        match statement.query_row(&[key, &position], |row| row.get(0)) {
            Ok(data)                                  => Ok(Value::Bulk(data)),
            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(Value::NullArray),
            Err(e)                                    => panic!(e)
        }
    }

    // private

    fn pop(&self, connection: &Connection, key: &String, direction: Direction) -> Option<String> {
        let order = match direction { Direction::Left => "ASC", Direction::Right => "DESC" };
        let mut statement = connection.prepare(&format!("SELECT id, value FROM list_items WHERE key = ?1 ORDER BY position {} LIMIT 1", order)).unwrap();

        match statement.query_row(&[key], |row| (row.get(0), row.get(1))) {
            Ok(result) => {
                let (id, value): (i64, String) = result;

                connection.execute("DELETE FROM list_items WHERE id = ?1", &[&id]).unwrap();
                Some(value)
            }

            Err(rusqlite::Error::QueryReturnedNoRows) => None,

            Err(e) => Err(e).unwrap()
        }
    }

    fn push<'b, I>(&self, connection: &mut Connection, key: &String, direction: Direction, iterator: I) -> ()
        where I: Iterator<Item=&'b &'b String>
    {
        let tx = connection.transaction().unwrap();

        let next_position_sql = match direction {
            Direction::Left  => "coalesce(MIN(position), 0) - 1",
            Direction::Right => "coalesce(MAX(position), 0) + 1"
        };

        let sql = format!("INSERT INTO list_items (key, value, position) SELECT ?1, ?2, {} FROM list_items WHERE key = ?1", next_position_sql);

        iterator.map(|value| tx.execute(&sql, &[key, *value])).collect::<Result<Vec<_>, _>>().unwrap();

        tx.commit().unwrap();
    }

    fn lock_connection(&self) -> MutexGuard<Connection> {
        (*self.connection_mutex).lock().unwrap()
    }

    fn find_position_boundaries(&self, connection: &Connection, key: &String) -> (i64, i64) {
        let mut statement = connection.prepare("SELECT MIN(position), MAX(position) AS c FROM list_items WHERE key = ?1").unwrap();
        statement.query_row(&[key], |row| (row.get(0), row.get(1))).unwrap()
    }

    fn parse_indexes(&self, boundaries: (i64, i64), (start, stop): (i64, i64)) -> (i64, i64) {
        (self.parse_index(boundaries, start), self.parse_index(boundaries, stop))
    }

    fn parse_index(&self, (first_position, last_position): (i64, i64), index: i64) -> i64 {
        if index < 0 { index + last_position + 1 } else { index + first_position }
    }

    fn count_list_items(&self, connection: &Connection, key: &String) -> i64 {
        let mut statement = connection.prepare("SELECT COUNT(*) AS c FROM list_items WHERE key = ?1").unwrap();
        statement.query_row(&[key], |row| row.get(0)).unwrap()
    }

    fn count_list_items_value(&self, connection: &Connection, key: &String) -> CommandResult {
        Ok(Value::Integer(self.count_list_items(connection, key)))
    }
}
<|endoftext|>


<commit_before>use time;
use time::{Timespec, Tm};
use std::fmt::{Display, Error, Formatter};

/// Default timespec, used to signal a non initialized time.
#[allow(dead_code)]
pub const DEFAULT_TIMESPEC: Timespec = Timespec { sec: 0, nsec: 0 };

/// Utility struct that implements Display in a pretty style
/// for some Tm instance.
pub struct PrettyDisplay {
    tm: Tm,
}

/// The format to be used to display a time.
/// It could be a local or an UTC time.
#[allow(dead_code)]
pub enum Format {
    Local,
    Utc,
}

impl Display for PrettyDisplay {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        write!(f, "{}", self.tm.rfc822z())
    }
}

/// Returns an object implementing Display as a pretty printed UTC time.
#[allow(dead_code)]
pub fn to_pretty_utc(ts: Timespec) -> PrettyDisplay {
    PrettyDisplay { tm: time::at_utc(ts) }
}

/// Returns an object implementing Display as a pretty printed local time.
pub fn to_pretty_local(ts: Timespec) -> PrettyDisplay {
    PrettyDisplay { tm: time::at(ts) }
}

/// Returns an obejct implementing Display as a pretty printed time.
/// This could be a local or an UTC time, depending on the format parameter.
#[allow(dead_code)]
pub fn to_pretty(ts: Timespec, format: Format) -> PrettyDisplay {
    match format {
        Format::Local => to_pretty_local(ts),
        Format::Utc => to_pretty_utc(ts),
    }
}

/// Parse a string representing a duplicity timestamp and returns a Timespec
/// if all goes well.
pub fn parse_time_str(s: &str) -> Option<Timespec> {
    time::strptime(s, "%Y%m%dt%H%M%S%Z").ok().map(|tm| tm.to_timespec())
}

/// Change the current time zone
#[cfg(test)]
pub mod test_utils {
    use std::sync::{Mutex, MutexGuard};

    lazy_static! {
        static ref TZLOCK: Mutex<i32> = Mutex::new(0);
    }

    pub fn set_time_zone(tz: &str) -> MutexGuard<i32> {
        use std::env;
        use time;

        let lock = TZLOCK.lock();
        env::set_var("TZ", tz);
        time::tzset();
        lock.unwrap()
    }
}

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

    #[test]
    fn parse() {
        parse_time_str("19881211t172000z").unwrap();
    }

    #[test]
    fn parse_display_utc_london() {
        let time = parse_time_str("19881211t152000z").unwrap();
        let _lock = set_time_zone("Europe/London");
        assert_eq!(format!("{}", to_pretty_utc(time)),
                   "Sun, 11 Dec 1988 15:20:00 -0000");
    }

    #[test]
    fn parse_display_utc_rome() {
        let _lock = set_time_zone("Europe/Rome");
        let time = parse_time_str("19881211t152000z").unwrap();
        assert_eq!(format!("{}", to_pretty_utc(time)),
                   "Sun, 11 Dec 1988 15:20:00 -0000");
    }

    #[test]
    fn parse_display_local_rome() {
        let _lock = set_time_zone("Europe/Rome");
        let time = parse_time_str("19881211t152000z").unwrap();
        assert_eq!(format!("{}", to_pretty_local(time)),
                   "Sun, 11 Dec 1988 16:20:00 +0100");
    }

    #[test]
    fn parse_display_local_london() {
        let _lock = set_time_zone("Europe/London");
        let time = parse_time_str("19881211t152000z").unwrap();
        assert_eq!(format!("{}", to_pretty_local(time)),
                   "Sun, 11 Dec 1988 15:20:00 -0000");
    }

    #[test]
    fn time_crate() {
        use time::{Tm, at_utc, strftime, strptime};

        let _lock = set_time_zone("Europe/Rome");
        // parse
        let tm = strptime("20150617t182545z", "%Y%m%dt%H%M%S%Z").unwrap();
        // format
        assert_eq!(strftime("%a %d/%m/%Y %H:%M:%S", &tm).unwrap(),
                   "Sun 17/06/2015 18:25:45");
        assert_eq!(format!("{}", tm.rfc3339()), "2015-06-17T18:25:45Z");
        // store in Timespec and restore in Tm
        let ts = tm.to_timespec();
        let tm1 = at_utc(ts);
        // somehow they don't have the same identical structure :(
        // tm_wday, tm_yday are missing. See rust-lang-deprecated/time#92
        //assert_eq!(tm, tm1);
        // test equally formatted
        let format_fn = |tm: &Tm| format!("{}", tm.rfc3339());
        assert_eq!(format_fn(&tm), format_fn(&tm1));
    }
}
<commit_msg>Added comment<commit_after>use time;
use time::{Timespec, Tm};
use std::fmt::{Display, Error, Formatter};

/// Default timespec, used to signal a non initialized time.
#[allow(dead_code)]
pub const DEFAULT_TIMESPEC: Timespec = Timespec { sec: 0, nsec: 0 };

/// Utility struct that implements Display in a pretty style
/// for some Tm instance.
pub struct PrettyDisplay {
    tm: Tm,
}

/// The format to be used to display a time.
/// It could be a local or an UTC time.
#[allow(dead_code)]
pub enum Format {
    Local,
    Utc,
}

impl Display for PrettyDisplay {
    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
        write!(f, "{}", self.tm.rfc822z())
    }
}

/// Returns an object implementing Display as a pretty printed UTC time.
#[allow(dead_code)]
pub fn to_pretty_utc(ts: Timespec) -> PrettyDisplay {
    PrettyDisplay { tm: time::at_utc(ts) }
}

/// Returns an object implementing Display as a pretty printed local time.
pub fn to_pretty_local(ts: Timespec) -> PrettyDisplay {
    PrettyDisplay { tm: time::at(ts) }
}

/// Returns an obejct implementing Display as a pretty printed time.
/// This could be a local or an UTC time, depending on the format parameter.
#[allow(dead_code)]
pub fn to_pretty(ts: Timespec, format: Format) -> PrettyDisplay {
    match format {
        Format::Local => to_pretty_local(ts),
        Format::Utc => to_pretty_utc(ts),
    }
}

/// Parse a string representing a duplicity timestamp and returns a Timespec
/// if all goes well.
pub fn parse_time_str(s: &str) -> Option<Timespec> {
    time::strptime(s, "%Y%m%dt%H%M%S%Z").ok().map(|tm| tm.to_timespec())
}

/// Change the current time zone
#[cfg(test)]
pub mod test_utils {
    use std::sync::{Mutex, MutexGuard};

    // A global mutex is needed because tests are run in parallel
    // We need to avoid tests change time zone concurrently
    lazy_static! {
        static ref TZLOCK: Mutex<i32> = Mutex::new(0);
    }

    pub fn set_time_zone(tz: &str) -> MutexGuard<i32> {
        use std::env;
        use time;

        let lock = TZLOCK.lock();
        env::set_var("TZ", tz);
        time::tzset();
        lock.unwrap()
    }
}

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

    #[test]
    fn parse() {
        parse_time_str("19881211t172000z").unwrap();
    }

    #[test]
    fn parse_display_utc_london() {
        let time = parse_time_str("19881211t152000z").unwrap();
        let _lock = set_time_zone("Europe/London");
        assert_eq!(format!("{}", to_pretty_utc(time)),
                   "Sun, 11 Dec 1988 15:20:00 -0000");
    }

    #[test]
    fn parse_display_utc_rome() {
        let _lock = set_time_zone("Europe/Rome");
        let time = parse_time_str("19881211t152000z").unwrap();
        assert_eq!(format!("{}", to_pretty_utc(time)),
                   "Sun, 11 Dec 1988 15:20:00 -0000");
    }

    #[test]
    fn parse_display_local_rome() {
        let _lock = set_time_zone("Europe/Rome");
        let time = parse_time_str("19881211t152000z").unwrap();
        assert_eq!(format!("{}", to_pretty_local(time)),
                   "Sun, 11 Dec 1988 16:20:00 +0100");
    }

    #[test]
    fn parse_display_local_london() {
        let _lock = set_time_zone("Europe/London");
        let time = parse_time_str("19881211t152000z").unwrap();
        assert_eq!(format!("{}", to_pretty_local(time)),
                   "Sun, 11 Dec 1988 15:20:00 -0000");
    }

    #[test]
    fn time_crate() {
        use time::{Tm, at_utc, strftime, strptime};

        let _lock = set_time_zone("Europe/Rome");
        // parse
        let tm = strptime("20150617t182545z", "%Y%m%dt%H%M%S%Z").unwrap();
        // format
        assert_eq!(strftime("%a %d/%m/%Y %H:%M:%S", &tm).unwrap(),
                   "Sun 17/06/2015 18:25:45");
        assert_eq!(format!("{}", tm.rfc3339()), "2015-06-17T18:25:45Z");
        // store in Timespec and restore in Tm
        let ts = tm.to_timespec();
        let tm1 = at_utc(ts);
        // somehow they don't have the same identical structure :(
        // tm_wday, tm_yday are missing. See rust-lang-deprecated/time#92
        //assert_eq!(tm, tm1);
        // test equally formatted
        let format_fn = |tm: &Tm| format!("{}", tm.rfc3339());
        assert_eq!(format_fn(&tm), format_fn(&tm1));
    }
}
<|endoftext|>


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

use std::borrow::Cow;
use std::cell::RefCell;
use std::fmt::Display;
use std::path::PathBuf;
use std::rc::Rc;
use std::sync::Arc;

use dbus;
use dbus::Connection;
use dbus::Message;
use dbus::MessageItem;
use dbus::NameFlag;

use dbus::tree::Factory;
use dbus::tree::MethodErr;
use dbus::tree::MethodFn;
use dbus::tree::MethodResult;
use dbus::tree::Property;
use dbus::tree::Tree;

use dbus_consts::*;

use engine::Engine;
use types::{StratisResult, StratisError};

#[derive(Debug, Clone)]
pub struct DbusContext<'a> {
    name_prop: Arc<Property<MethodFn<'a>>>,
    pub remaining_prop: Arc<Property<MethodFn<'a>>>,
    pub total_prop: Arc<Property<MethodFn<'a>>>,
    pub status_prop: Arc<Property<MethodFn<'a>>>,
    pub running_status_prop: Arc<Property<MethodFn<'a>>>,
    pub block_devices_prop: Arc<Property<MethodFn<'a>>>,
}

impl<'a> DbusContext<'a> {
    pub fn update_one(prop: &Arc<Property<MethodFn<'a>>>, m: MessageItem) -> StratisResult<()> {
        match prop.set_value(m) {
            Ok(_) => Ok(()), // TODO: return signals
            Err(()) => Err(StratisError::Dbus(dbus::Error::new_custom(
                "UpdateError", "Could not update property with value"))),
        }
    }
}

fn list_pools(m: &Message, engine: &Rc<RefCell<Engine>>) -> MethodResult {

    let result = engine.borrow().list_pools();

    let mut msg_vec = Vec::new();

    // TODO: deal with failure
    let pool_tree = result.unwrap();

    for (name, pool) in pool_tree {
        let entry = vec![MessageItem::Str(format!("{}", name)),
                         MessageItem::UInt16(0),
                         MessageItem::Str(String::from("Ok"))];
        msg_vec.push(MessageItem::Struct(entry));
    }

    let item_array = MessageItem::Array(msg_vec, Cow::Borrowed("(sqs)"));

    Ok(vec![m.method_return().append1(item_array)])
}

fn create_pool(m: &Message, engine: &Rc<RefCell<Engine>>) -> MethodResult {

    let message: (Option<MessageItem>, Option<MessageItem>, Option<MessageItem>) = m.get3();

    let item0: MessageItem = try!(message.0.ok_or_else(MethodErr::no_arg));
    let name: &String = try!(item0.inner().map_err(|_| MethodErr::invalid_arg(&item0)));

    let item1: MessageItem = try!(message.1.ok_or_else(MethodErr::no_arg));
    let devs: &Vec<MessageItem> = try!(item1.inner().map_err(|_| MethodErr::invalid_arg(&item1)));

    let item2: MessageItem = try!(message.2.ok_or_else(MethodErr::no_arg));
    let raid_level: u16 = try!(item2.inner().map_err(|_| MethodErr::invalid_arg(&item2)));

    let blockdevs = devs.into_iter()
        .map(|x| PathBuf::from(x.inner::<&str>().unwrap()))
        .collect::<Vec<_>>();

    let result = engine.borrow_mut().create_pool(&name, &blockdevs, raid_level);

    Ok(vec![m.method_return().append3("/dbus/newpool/path", 0, "Ok")])
}

fn destroy_pool(m: &Message, engine: &Rc<RefCell<Engine>>) -> MethodResult {

    let item: MessageItem = try!(m.get1().ok_or_else(MethodErr::no_arg));

    let name: &String = try!(item.inner().map_err(|_| MethodErr::invalid_arg(&item)));

    let result = engine.borrow_mut().destroy_pool(&name);

    let message = m.method_return();

    let msg = match result {
        Ok(_) => {
            message.append2(MessageItem::UInt16(0),
                            MessageItem::Str(format!("{}", "Ok")))
        }
        Err(err) => {
            message.append2(MessageItem::UInt16(0),
                            MessageItem::Str(format!("{}", "Ok")))
        }
    };

    Ok(vec![msg])

}

fn get_pool_object_path(m: &Message) -> MethodResult {

    Ok(vec![m.method_return().append3("/dbus/pool/path", 0, "Ok")])
}

fn get_volume_object_path(m: &Message) -> MethodResult {
    Ok(vec![m.method_return().append3("/dbus/volume/path", 0, "Ok")])
}

fn get_dev_object_path(m: &Message) -> MethodResult {
    Ok(vec![m.method_return().append3("/dbus/dev/path", 0, "Ok")])
}

fn get_cache_object_path(m: &Message) -> MethodResult {
    Ok(vec![m.method_return().append3("/dbus/cache/path", 0, "Ok")])
}

fn get_list_items<T, I>(m: &Message, iter: I) -> MethodResult
    where T: HasCodes + Display,
          I: Iterator<Item = T>
{
    let msg_vec = iter.map(|item| {
            MessageItem::Struct(vec![MessageItem::Str(format!("{}", item)),
                                     MessageItem::UInt16(item.get_error_int()),
                                     MessageItem::Str(format!("{}", item.get_error_string()))])
        })
        .collect::<Vec<MessageItem>>();

    let item_array = MessageItem::Array(msg_vec, Cow::Borrowed("(sqs)"));
    Ok(vec![m.method_return().append1(item_array)])
}

fn get_error_codes(m: &Message) -> MethodResult {
    get_list_items(m, StratisErrorEnum::iter_variants())
}

fn get_raid_levels(m: &Message) -> MethodResult {
    get_list_items(m, StratisRaidType::iter_variants())
}

fn get_dev_types(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let _name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

pub fn get_base_tree<'a>(c: &'a Connection,
                         engine: Rc<RefCell<Engine>>)
                         -> StratisResult<Tree<MethodFn<'a>>> {
    c.register_name(STRATIS_BASE_SERVICE, NameFlag::ReplaceExisting as u32).unwrap();

    let f = Factory::new_fn();

    let base_tree = f.tree();

    let engine_clone = engine.clone();

    let createpool_method = f.method(CREATE_POOL, move |m, _, _| create_pool(m, &engine_clone))
        .in_arg(("pool_name", "s"))
        .in_arg(("dev_list", "as"))
        .in_arg(("raid_type", "q"))
        .out_arg(("object_path", "s"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let engine_clone = engine.clone();

    let destroypool_method = f.method(DESTROY_POOL, move |m, _, _| destroy_pool(m, &engine_clone))
        .in_arg(("pool_name", "s"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let engine_clone = engine.clone();

    let listpools_method = f.method(LIST_POOLS, move |m, _, _| list_pools(m, &engine_clone))
        .out_arg(("pool_names", "as"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let getpoolobjectpath_method =
        f.method(GET_POOL_OBJECT_PATH, move |m, _, _| get_pool_object_path(m))
            .in_arg(("pool_name", "s"))
            .out_arg(("object_path", "o"))
            .out_arg(("return_code", "q"))
            .out_arg(("return_string", "s"));

    let getvolumeobjectpath_method = f.method(GET_VOLUME_OBJECT_PATH,
                move |m, _, _| get_volume_object_path(m))
        .in_arg(("pool_name", "s"))
        .in_arg(("volume_name", "s"))
        .out_arg(("object_path", "o"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let getdevobjectpath_method =
        f.method(GET_DEV_OBJECT_PATH, move |m, _, _| get_dev_object_path(m))
            .in_arg(("dev_name", "s"))
            .out_arg(("object_path", "o"))
            .out_arg(("return_code", "q"))
            .out_arg(("return_string", "s"));

    let getcacheobjectpath_method = f.method(GET_CACHE_OBJECT_PATH,
                move |m, _, _| get_cache_object_path(m))
        .in_arg(("cache_dev_name", "s"))
        .out_arg(("object_path", "o"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let geterrorcodes_method = f.method(GET_ERROR_CODES, move |m, _, _| get_error_codes(m))
        .out_arg(("error_codes", "a(sqs)"));

    let getraidlevels_method = f.method(GET_RAID_LEVELS, move |m, _, _| get_raid_levels(m))
        .out_arg(("error_codes", "a(sqs)"));

    let getdevtypes_method = f.method(GET_DEV_TYPES, move |m, _, _| get_dev_types(m));

    let obj_path = f.object_path(STRATIS_BASE_PATH)
        .introspectable()
        .add(f.interface(STRATIS_MANAGER_INTERFACE)
            .add_m(listpools_method)
            .add_m(createpool_method)
            .add_m(destroypool_method)
            .add_m(getpoolobjectpath_method)
            .add_m(getvolumeobjectpath_method)
            .add_m(getdevobjectpath_method)
            .add_m(getcacheobjectpath_method)
            .add_m(geterrorcodes_method)
            .add_m(getraidlevels_method)
            .add_m(getdevtypes_method));

    let base_tree = base_tree.add(obj_path);
    try!(base_tree.set_registered(c, true));

    Ok(base_tree)
}
<commit_msg>Translate internal errors into DBus-API-specific error values<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::borrow::Cow;
use std::cell::RefCell;
use std::fmt::Display;
use std::path::PathBuf;
use std::rc::Rc;
use std::sync::Arc;
use std::io::ErrorKind;

use dbus;
use dbus::Connection;
use dbus::Message;
use dbus::MessageItem;
use dbus::NameFlag;

use dbus::tree::Factory;
use dbus::tree::MethodErr;
use dbus::tree::MethodFn;
use dbus::tree::MethodResult;
use dbus::tree::Property;
use dbus::tree::Tree;

use dbus_consts::*;

use engine::Engine;
use types::{StratisResult, StratisError};

#[derive(Debug, Clone)]
pub struct DbusContext<'a> {
    name_prop: Arc<Property<MethodFn<'a>>>,
    pub remaining_prop: Arc<Property<MethodFn<'a>>>,
    pub total_prop: Arc<Property<MethodFn<'a>>>,
    pub status_prop: Arc<Property<MethodFn<'a>>>,
    pub running_status_prop: Arc<Property<MethodFn<'a>>>,
    pub block_devices_prop: Arc<Property<MethodFn<'a>>>,
}

impl<'a> DbusContext<'a> {
    pub fn update_one(prop: &Arc<Property<MethodFn<'a>>>, m: MessageItem) -> StratisResult<()> {
        match prop.set_value(m) {
            Ok(_) => Ok(()), // TODO: return signals
            Err(()) => Err(StratisError::Dbus(dbus::Error::new_custom(
                "UpdateError", "Could not update property with value"))),
        }
    }
}

fn internal_to_dbus_err(err: &StratisError) -> StratisErrorEnum {
    match *err {
        StratisError::Stratis(_) => StratisErrorEnum::STRATIS_ERROR,
        StratisError::Io(ref err) => match err.kind() {
            ErrorKind::NotFound => StratisErrorEnum::STRATIS_NOTFOUND,
            ErrorKind::AlreadyExists => StratisErrorEnum::STRATIS_ALREADY_EXISTS,
            _ => StratisErrorEnum::STRATIS_ERROR,
        },
        _ => StratisErrorEnum::STRATIS_ERROR,
    }
}

fn list_pools(m: &Message, engine: &Rc<RefCell<Engine>>) -> MethodResult {

    let result = engine.borrow().list_pools();

    let mut msg_vec = Vec::new();

    // TODO: deal with failure
    let pool_tree = result.unwrap();

    for (name, pool) in pool_tree {
        let entry = vec![MessageItem::Str(format!("{}", name)),
                         MessageItem::UInt16(0),
                         MessageItem::Str(String::from("Ok"))];
        msg_vec.push(MessageItem::Struct(entry));
    }

    let item_array = MessageItem::Array(msg_vec, Cow::Borrowed("(sqs)"));

    Ok(vec![m.method_return().append1(item_array)])
}

fn create_pool(m: &Message, engine: &Rc<RefCell<Engine>>) -> MethodResult {

    let message: (Option<MessageItem>, Option<MessageItem>, Option<MessageItem>) = m.get3();

    let item0: MessageItem = try!(message.0.ok_or_else(MethodErr::no_arg));
    let name: &String = try!(item0.inner().map_err(|_| MethodErr::invalid_arg(&item0)));

    let item1: MessageItem = try!(message.1.ok_or_else(MethodErr::no_arg));
    let devs: &Vec<MessageItem> = try!(item1.inner().map_err(|_| MethodErr::invalid_arg(&item1)));

    let item2: MessageItem = try!(message.2.ok_or_else(MethodErr::no_arg));
    let raid_level: u16 = try!(item2.inner().map_err(|_| MethodErr::invalid_arg(&item2)));

    let blockdevs = devs.into_iter()
        .map(|x| PathBuf::from(x.inner::<&str>().unwrap()))
        .collect::<Vec<_>>();

    match engine.borrow_mut().create_pool(&name, &blockdevs, raid_level) {
        Ok(_) => Ok(vec![m.method_return().append3("/dbus/newpool/path", 0, "Ok")]),
        Err(x) => {
            let dbus_err = internal_to_dbus_err(&x);
            Ok(vec![m.method_return().append3("", dbus_err.get_error_int()
                                              , dbus_err.get_error_string())])
        }
    }
}

fn destroy_pool(m: &Message, engine: &Rc<RefCell<Engine>>) -> MethodResult {

    let item: MessageItem = try!(m.get1().ok_or_else(MethodErr::no_arg));

    let name: &String = try!(item.inner().map_err(|_| MethodErr::invalid_arg(&item)));

    let result = engine.borrow_mut().destroy_pool(&name);

    let message = m.method_return();

    let msg = match result {
        Ok(_) => {
            message.append2(MessageItem::UInt16(0),
                            MessageItem::Str(format!("{}", "Ok")))
        }
        Err(err) => {
            message.append2(MessageItem::UInt16(0),
                            MessageItem::Str(format!("{}", "Ok")))
        }
    };

    Ok(vec![msg])

}

fn get_pool_object_path(m: &Message) -> MethodResult {

    Ok(vec![m.method_return().append3("/dbus/pool/path", 0, "Ok")])
}

fn get_volume_object_path(m: &Message) -> MethodResult {
    Ok(vec![m.method_return().append3("/dbus/volume/path", 0, "Ok")])
}

fn get_dev_object_path(m: &Message) -> MethodResult {
    Ok(vec![m.method_return().append3("/dbus/dev/path", 0, "Ok")])
}

fn get_cache_object_path(m: &Message) -> MethodResult {
    Ok(vec![m.method_return().append3("/dbus/cache/path", 0, "Ok")])
}

fn get_list_items<T, I>(m: &Message, iter: I) -> MethodResult
    where T: HasCodes + Display,
          I: Iterator<Item = T>
{
    let msg_vec = iter.map(|item| {
            MessageItem::Struct(vec![MessageItem::Str(format!("{}", item)),
                                     MessageItem::UInt16(item.get_error_int()),
                                     MessageItem::Str(format!("{}", item.get_error_string()))])
        })
        .collect::<Vec<MessageItem>>();

    let item_array = MessageItem::Array(msg_vec, Cow::Borrowed("(sqs)"));
    Ok(vec![m.method_return().append1(item_array)])
}

fn get_error_codes(m: &Message) -> MethodResult {
    get_list_items(m, StratisErrorEnum::iter_variants())
}

fn get_raid_levels(m: &Message) -> MethodResult {
    get_list_items(m, StratisRaidType::iter_variants())
}

fn get_dev_types(m: &Message) -> MethodResult {
    let mut items = m.get_items();
    if items.len() < 1 {
        return Err(MethodErr::no_arg());
    }

    let _name = try!(items.pop()
        .ok_or_else(MethodErr::no_arg)
        .and_then(|i| {
            i.inner::<&str>()
                .map_err(|_| MethodErr::invalid_arg(&i))
                .map(|i| i.to_owned())
        }));

    println!("method called");

    Ok(vec![m.method_return()])
}

pub fn get_base_tree<'a>(c: &'a Connection,
                         engine: Rc<RefCell<Engine>>)
                         -> StratisResult<Tree<MethodFn<'a>>> {
    c.register_name(STRATIS_BASE_SERVICE, NameFlag::ReplaceExisting as u32).unwrap();

    let f = Factory::new_fn();

    let base_tree = f.tree();

    let engine_clone = engine.clone();

    let createpool_method = f.method(CREATE_POOL, move |m, _, _| create_pool(m, &engine_clone))
        .in_arg(("pool_name", "s"))
        .in_arg(("dev_list", "as"))
        .in_arg(("raid_type", "q"))
        .out_arg(("object_path", "s"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let engine_clone = engine.clone();

    let destroypool_method = f.method(DESTROY_POOL, move |m, _, _| destroy_pool(m, &engine_clone))
        .in_arg(("pool_name", "s"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let engine_clone = engine.clone();

    let listpools_method = f.method(LIST_POOLS, move |m, _, _| list_pools(m, &engine_clone))
        .out_arg(("pool_names", "as"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let getpoolobjectpath_method =
        f.method(GET_POOL_OBJECT_PATH, move |m, _, _| get_pool_object_path(m))
            .in_arg(("pool_name", "s"))
            .out_arg(("object_path", "o"))
            .out_arg(("return_code", "q"))
            .out_arg(("return_string", "s"));

    let getvolumeobjectpath_method = f.method(GET_VOLUME_OBJECT_PATH,
                move |m, _, _| get_volume_object_path(m))
        .in_arg(("pool_name", "s"))
        .in_arg(("volume_name", "s"))
        .out_arg(("object_path", "o"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let getdevobjectpath_method =
        f.method(GET_DEV_OBJECT_PATH, move |m, _, _| get_dev_object_path(m))
            .in_arg(("dev_name", "s"))
            .out_arg(("object_path", "o"))
            .out_arg(("return_code", "q"))
            .out_arg(("return_string", "s"));

    let getcacheobjectpath_method = f.method(GET_CACHE_OBJECT_PATH,
                move |m, _, _| get_cache_object_path(m))
        .in_arg(("cache_dev_name", "s"))
        .out_arg(("object_path", "o"))
        .out_arg(("return_code", "q"))
        .out_arg(("return_string", "s"));

    let geterrorcodes_method = f.method(GET_ERROR_CODES, move |m, _, _| get_error_codes(m))
        .out_arg(("error_codes", "a(sqs)"));

    let getraidlevels_method = f.method(GET_RAID_LEVELS, move |m, _, _| get_raid_levels(m))
        .out_arg(("error_codes", "a(sqs)"));

    let getdevtypes_method = f.method(GET_DEV_TYPES, move |m, _, _| get_dev_types(m));

    let obj_path = f.object_path(STRATIS_BASE_PATH)
        .introspectable()
        .add(f.interface(STRATIS_MANAGER_INTERFACE)
            .add_m(listpools_method)
            .add_m(createpool_method)
            .add_m(destroypool_method)
            .add_m(getpoolobjectpath_method)
            .add_m(getvolumeobjectpath_method)
            .add_m(getdevobjectpath_method)
            .add_m(getcacheobjectpath_method)
            .add_m(geterrorcodes_method)
            .add_m(getraidlevels_method)
            .add_m(getdevtypes_method));

    let base_tree = base_tree.add(obj_path);
    try!(base_tree.set_registered(c, true));

    Ok(base_tree)
}
<|endoftext|>


<commit_before>use syntax::{abi, ast, ast_util, codemap, parse};
use syntax::parse::token;
use syntax::ptr::P;
use syntax::ext::base;
use syntax::ext::build::AstBuilder;

static BEFORE_EACH: &'static str = "before_each";
static AFTER_EACH:  &'static str = "after_each";
static BEFORE:      &'static str = "before";
static AFTER:       &'static str = "after";
static IT:          &'static str = "it";
static DESCRIBE:    &'static str = "describe";

/// Defines the state of a `describe!` macro as it is parsing.
struct DescribeState {
    name: Option<ast::Ident>,
    before: Option<P<ast::Block>>,
    after: Option<P<ast::Block>>,
    before_each: Option<P<ast::Block>>,
    after_each: Option<P<ast::Block>>,
    subblocks: Vec<SubBlock>
}

enum SubBlock {
    TestBlock(Test),
    DescribeBlock(P<ast::Item>)
}

/// A test as a description and associated block.
struct Test {
    description: String,
    block: P<ast::Block>
}

/// Defines the overarching `describe!` syntax extension.
///
/// All other macros in stainless are actually "fake" in the sense
/// that they are detected and expanded inside of the implementation
/// of `describe!`.
pub fn describe<'a>(cx: &'a mut base::ExtCtxt, sp: codemap::Span,
                name: ast::Ident, tokens: Vec<ast::TokenTree>) -> Box<base::MacResult + 'a> {
    // Parse a full DescribeState from the input, emitting errors if used incorrectly.
    let state = parse_describe(&mut parse::tts_to_parser(cx.parse_sess(), tokens, cx.cfg()), sp, cx, Some(name));

    // Export the new module.
    base::MacItems::new(Some(create_describe_item(state, sp, cx)).into_iter())
}

fn create_describe_item(state: DescribeState, sp: codemap::Span, cx: &mut base::ExtCtxt) -> P<ast::Item> {
    // Get the name of this mod.
    let name = state.name.clone().unwrap();

    // Create subblocks from a full DescribeState
    let subblocks = create_subblocks(state, sp, cx);

    // Get a glob import of all items in scope to the module that `describe!` is called from.
    //
    // This glob is `pub use super::*` so that nested `describe!` blocks (which will also contain
    // this glob) will be able to see all the symbols.
    let super_glob = cx.view_use_glob(sp, ast::Public, vec![cx.ident_of("super")]);

    // Generate the new module.
    cx.item_mod(sp, sp, name, vec![], vec![super_glob], subblocks)
}

fn create_subblocks(state: DescribeState, sp: codemap::Span, cx: &mut base::ExtCtxt) -> Vec<P<ast::Item>> {
    // FIXME(reem): Implement before and after.
    let (_before, _after) = (state.before, state.after);

    let (before_each, after_each) = (state.before_each, state.after_each);
    let subblocks = state.subblocks;

    // Create the #[test] attribute.
    let test_attribute = cx.attribute(sp,
                                      cx.meta_word(sp, token::InternedString::new("test")));

    subblocks.into_iter().map(|block| {
        match block {
            TestBlock(Test { description, block }) => {
                // Create the full test body by splicing in the statements and view items of the before and
                // after blocks if they are present.
                let test_body = match (&before_each, &after_each) {
                    (&None, &None) => block,

                    (&Some(ref before), &None) => {
                        P(ast::Block {
                            view_items: before.view_items + block.view_items,
                            stmts: before.stmts + block.stmts,
                            ..block.deref().clone()
                        })
                    },

                    (&None, &Some(ref after)) => {
                        P(ast::Block {
                            view_items: block.view_items + after.view_items,
                            stmts: block.stmts + after.stmts,
                            ..block.deref().clone()
                        })
                    },

                    (&Some(ref before), &Some(ref after)) => {
                        P(ast::Block {
                            view_items: before.view_items + block.view_items + after.view_items,
                            stmts: before.stmts + block.stmts + after.stmts,
                            ..block.deref().clone()
                        })
                    }
                };

                // Create the final Item that represents the test.
                P(ast::Item {
                    // Name it with a snake_case version of the description.
                    ident: cx.ident_of(description.replace(" ", "_").as_slice()),

                    // Add #[test]
                    attrs: vec![test_attribute.clone()],
                    id: ast::DUMMY_NODE_ID,
                    node: ast::ItemFn(
                        // Takes no arguments and returns ()
                        cx.fn_decl(vec![], cx.ty_nil()),

                        // All the usual types.
                        ast::NormalFn,
                        abi::Rust,
                        ast_util::empty_generics(),

                        // Add the body of the function.
                        test_body
                    ),
                    // Inherited visibility (not pub)
                    vis: ast::Inherited,
                    span: sp
                })
            },
            DescribeBlock(item) => item
        }
    }).collect()
}

fn parse_describe(parser: &mut parse::parser::Parser, sp: codemap::Span,
                  cx: &mut base::ExtCtxt, name: Option<ast::Ident>) -> DescribeState {
    let mut state = DescribeState {
        name: None, before: None, after: None,
        before_each: None, after_each: None, subblocks: vec![]
    };

    state.name = match name {
        // Top-level describe block.
        Some(name) => Some(name),
        // Nested describe block.
        None => {
            // Get the name of this describe block
            let name = parser.parse_ident();
            // Move past the opening {
            if token::LBRACE != parser.bump_and_get() {
                parser.fatal("Expected { after the name of a describe! block.");
            }
            Some(name)
        }
    };

    // Now parse all tests and subsections:
    while parser.token != token::RBRACE && parser.token != token::EOF {
        // Get the name of this block, must be either:
        //     - before_each
        //     - after_each
        //     - before
        //     - after
        //     - it
        //     - describe!
        //
        // Any other top-level idents are not allowed.
        let block_name = parser.parse_ident();

        match block_name.as_str() {
            BEFORE_EACH => {
                if state.before_each.is_some() { parser.fatal("Only one `before_each` block is allowed per `describe!` block.") }
                state.before_each = Some(parser.parse_block());
            },

            AFTER_EACH => {
                if state.after_each.is_some() { parser.fatal("Only one `after_each` block is allowed per `describe!` block.") }
                state.after_each = Some(parser.parse_block());
            },

            BEFORE => {
                if state.before.is_some() { parser.fatal("Only one `before` block is allowed per `describe!` block.") }
                state.before = Some(parser.parse_block());
            },

            AFTER => {
                if state.after.is_some() { parser.fatal("Only one `after` block is allowed per `describe!` block.") }
                state.after = Some(parser.parse_block());
            },

            IT => {
                // Description of this `it` block.
                let (description, _) = parser.parse_str();

                state.subblocks.push(TestBlock(Test {
                    // Get as a String
                    description: description.get().to_string(),

                    // The associated block
                    block: parser.parse_block()
                }));
            },

            // Nested `describe!` block.
            DESCRIBE => {
                // Skip over the !
                match parser.bump_and_get() {
                    token::NOT => {},
                    other => parser.fatal(format!("Expected ! but found `{}`", other).as_slice())
                };

                // Parse this sublock, generate new item.
                state.subblocks.push(DescribeBlock(create_describe_item(parse_describe(parser, sp, cx, None), sp, cx)));

                // Move past closing bracket and paren.
                //
                // This has to go in here because it two is EOF on the highest-level invocation.
                match parser.bump_and_get() {
                    token::RBRACE => {},
                    other => parser.fatal(format!("Expected }} to close `describe!` but found: `{}`", other).as_slice())
                }
            }

            banned => {
                // Illegal block name.
                let span = parser.span;
                parser.span_fatal(span, format!("Expected one of: `{}`, but found: `{}`",
                    format!("{}, {}, {}, {}, {}", BEFORE_EACH, AFTER_EACH, BEFORE, AFTER, IT).as_slice(),
                    banned).as_slice());
            }
        }
    }

    state
}

<commit_msg>(feat) Added support for failing tests.<commit_after>use syntax::{abi, ast, ast_util, codemap, parse};
use syntax::parse::token;
use syntax::ptr::P;
use syntax::ext::base;
use syntax::ext::build::AstBuilder;

static BEFORE_EACH: &'static str = "before_each";
static AFTER_EACH:  &'static str = "after_each";
static BEFORE:      &'static str = "before";
static AFTER:       &'static str = "after";
static IT:          &'static str = "it";
static DESCRIBE:    &'static str = "describe";
static FAILING:     &'static str = "failing";

/// Defines the state of a `describe!` macro as it is parsing.
struct DescribeState {
    name: Option<ast::Ident>,
    before: Option<P<ast::Block>>,
    after: Option<P<ast::Block>>,
    before_each: Option<P<ast::Block>>,
    after_each: Option<P<ast::Block>>,
    subblocks: Vec<SubBlock>
}

enum SubBlock {
    TestBlock(Test),
    FailingTest(Test),
    DescribeBlock(P<ast::Item>)
}

/// A test as a description and associated block.
struct Test {
    description: String,
    block: P<ast::Block>
}

/// Defines the overarching `describe!` syntax extension.
///
/// All other macros in stainless are actually "fake" in the sense
/// that they are detected and expanded inside of the implementation
/// of `describe!`.
pub fn describe<'a>(cx: &'a mut base::ExtCtxt, sp: codemap::Span,
                name: ast::Ident, tokens: Vec<ast::TokenTree>) -> Box<base::MacResult + 'a> {
    // Parse a full DescribeState from the input, emitting errors if used incorrectly.
    let state = parse_describe(&mut parse::tts_to_parser(cx.parse_sess(), tokens, cx.cfg()), sp, cx, Some(name));

    // Export the new module.
    base::MacItems::new(Some(create_describe_item(state, sp, cx)).into_iter())
}

fn create_describe_item(state: DescribeState, sp: codemap::Span, cx: &mut base::ExtCtxt) -> P<ast::Item> {
    // Get the name of this mod.
    let name = state.name.clone().unwrap();

    // Create subblocks from a full DescribeState
    let subblocks = create_subblocks(state, sp, cx);

    // Get a glob import of all items in scope to the module that `describe!` is called from.
    //
    // This glob is `pub use super::*` so that nested `describe!` blocks (which will also contain
    // this glob) will be able to see all the symbols.
    let super_glob = cx.view_use_glob(sp, ast::Public, vec![cx.ident_of("super")]);

    // Generate the new module.
    cx.item_mod(sp, sp, name, vec![], vec![super_glob], subblocks)
}

impl SubBlock {
    fn to_item(self, sp: codemap::Span, blocks: (&Option<P<ast::Block>>, &Option<P<ast::Block>>),
               cx: &mut base::ExtCtxt) -> P<ast::Item> {
        match self {
            TestBlock(test) => test.to_item(false, sp, blocks, cx),
            FailingTest(test) => test.to_item(true, sp, blocks, cx),
            DescribeBlock(item) => item
        }
    }
}

impl Test {
    fn to_item(self, failing: bool, sp: codemap::Span,
               (before_each, after_each): (&Option<P<ast::Block>>, &Option<P<ast::Block>>),
               cx: &mut base::ExtCtxt) -> P<ast::Item> {
        let Test { description, block } = self;

        // Create the #[test] attribute.
        let test_attribute = cx.attribute(sp, cx.meta_word(sp, token::InternedString::new("test")));

        // Create the #[should_fail] attribute.
        let should_fail = cx.attribute(sp, cx.meta_word(sp, token::InternedString::new("should_fail")));

        // Create the full test body by splicing in the statements and view items of the before and
        // after blocks if they are present.
        let test_body = match (before_each, after_each) {
            (&None, &None) => block,

            (&Some(ref before), &None) => {
                P(ast::Block {
                    view_items: before.view_items + block.view_items,
                    stmts: before.stmts + block.stmts,
                    ..block.deref().clone()
                })
            },

            (&None, &Some(ref after)) => {
                P(ast::Block {
                    view_items: block.view_items + after.view_items,
                    stmts: block.stmts + after.stmts,
                    ..block.deref().clone()
                })
            },

            (&Some(ref before), &Some(ref after)) => {
                P(ast::Block {
                    view_items: before.view_items + block.view_items + after.view_items,
                    stmts: before.stmts + block.stmts + after.stmts,
                    ..block.deref().clone()
                })
            }
        };

        // Create the final Item that represents the test.
        P(ast::Item {
            // Name it with a snake_case version of the description.
            ident: cx.ident_of(description.replace(" ", "_").as_slice()),

            // Add #[test] and possibly #[should_fail]
            attrs: if failing { vec![test_attribute, should_fail] } else { vec![test_attribute] },
            id: ast::DUMMY_NODE_ID,
            node: ast::ItemFn(
                // Takes no arguments and returns ()
                cx.fn_decl(vec![], cx.ty_nil()),

                // All the usual types.
                ast::NormalFn,
                abi::Rust,
                ast_util::empty_generics(),

                // Add the body of the function.
                test_body
            ),
            // Inherited visibility (not pub)
            vis: ast::Inherited,
            span: sp
        })
    }
}

fn create_subblocks(state: DescribeState, sp: codemap::Span, cx: &mut base::ExtCtxt) -> Vec<P<ast::Item>> {
    // FIXME(reem): Implement before and after.
    let (_before, _after) = (state.before, state.after);

    let blocks = (&state.before_each, &state.after_each);
    let subblocks = state.subblocks;

    subblocks.into_iter().map(|block| { block.to_item(sp, blocks, cx) }).collect()
}

fn parse_describe(parser: &mut parse::parser::Parser, sp: codemap::Span,
                  cx: &mut base::ExtCtxt, name: Option<ast::Ident>) -> DescribeState {
    let mut state = DescribeState {
        name: None, before: None, after: None,
        before_each: None, after_each: None, subblocks: vec![]
    };

    state.name = match name {
        // Top-level describe block.
        Some(name) => Some(name),
        // Nested describe block.
        None => {
            // Get the name of this describe block
            let name = parser.parse_ident();
            // Move past the opening {
            if token::LBRACE != parser.bump_and_get() {
                parser.fatal("Expected { after the name of a describe! block.");
            }
            Some(name)
        }
    };

    // Now parse all tests and subsections:
    while parser.token != token::RBRACE && parser.token != token::EOF {
        // Get the name of this block, must be either:
        //     - before_each
        //     - after_each
        //     - before
        //     - after
        //     - it
        //     - failing
        //     - describe!
        //
        // Any other top-level idents are not allowed.
        let block_name = parser.parse_ident();

        match block_name.as_str() {
            BEFORE_EACH => {
                if state.before_each.is_some() { parser.fatal("Only one `before_each` block is allowed per `describe!` block.") }
                state.before_each = Some(parser.parse_block());
            },

            AFTER_EACH => {
                if state.after_each.is_some() { parser.fatal("Only one `after_each` block is allowed per `describe!` block.") }
                state.after_each = Some(parser.parse_block());
            },

            BEFORE => {
                if state.before.is_some() { parser.fatal("Only one `before` block is allowed per `describe!` block.") }
                state.before = Some(parser.parse_block());
            },

            AFTER => {
                if state.after.is_some() { parser.fatal("Only one `after` block is allowed per `describe!` block.") }
                state.after = Some(parser.parse_block());
            },

            // Regular `#[test]`.
            IT => { state.subblocks.push(TestBlock(parse_test(parser))) },

            // `#[should_fail]` test.
            FAILING => { state.subblocks.push(FailingTest(parse_test(parser))) },

            // Nested `describe!` block.
            DESCRIBE => {
                // Skip over the !
                match parser.bump_and_get() {
                    token::NOT => {},
                    other => parser.fatal(format!("Expected ! but found `{}`", other).as_slice())
                };

                // Parse this sublock, generate new item.
                state.subblocks.push(DescribeBlock(create_describe_item(parse_describe(parser, sp, cx, None), sp, cx)));

                // Move past closing bracket and paren.
                //
                // This has to go in here because it two is EOF on the highest-level invocation.
                match parser.bump_and_get() {
                    token::RBRACE => {},
                    other => parser.fatal(format!("Expected }} to close `describe!` but found: `{}`", other).as_slice())
                }
            }

            banned => {
                // Illegal block name.
                let span = parser.span;
                parser.span_fatal(span, format!("Expected one of: `{}`, but found: `{}`",
                    format!("{}, {}, {}, {}, {}", BEFORE_EACH, AFTER_EACH, BEFORE, AFTER, IT).as_slice(),
                    banned).as_slice());
            }
        }
    }

    state
}

fn parse_test(parser: &mut parse::parser::Parser) -> Test {
    // Description of this test.
    let (description, _) = parser.parse_str();

    Test {
        // Get as a String
        description: description.get().to_string(),

        // The associated block
        block: parser.parse_block()
    }
}

<|endoftext|>


<commit_before>use std::cmp;

fn max_length(str1: &str, str2: &str) -> usize {
    if str1.len() > str2.len() {
      str1.len()
    } else {
      str2.len()
    }
}

fn get_common_chars(str1: &str, str2: &str, ordering: bool) -> Vec<char> {
    let mut common_chars = Vec::new();

    if str1.len() == 0 || str2.len() == 0 {
        return common_chars;
    }

    let max_length = max_length(str1, str2);
    let match_buffer = ( max_length / 2 ) - 1;
    let mut f = str1;
    let mut s = str2;

    if ordering {
        f = str2;
        s = str1;
    }

    for (i, c) in f.chars().enumerate() {
        let s_index: usize = if i < match_buffer {
            0
        } else {
            i - match_buffer
        };

        let e_index: usize = if s.len() <= (i + match_buffer) {
            s.len()
        } else {
            i + match_buffer + 1
        };

        let word_slice = &s[s_index..e_index];
        if word_slice.contains(c) {
            common_chars.push(c);
        }
    }

    common_chars
}

pub fn jaro_winkler_distance(str1: &str, str2: &str) -> f32 {
    let matches1 = get_common_chars(str1, str2, false);
    let matches2 = get_common_chars(str1, str2, true);

    if matches1.is_empty() || matches2.is_empty()
      || matches1.len() != matches2.len()
    {
      return 0.0;
    }

    let mut t = 0.0;
    for (a, b) in matches1.iter().zip(matches2.iter()) {
      if a != b {
        t += 0.5;
      }
    }

    let f1 = matches1.len() as f32 / str1.len() as f32;
    let f2 = matches2.len() as f32 / str2.len() as f32;
    let f3 = (matches1.len() as f32 - t) / matches1.len() as f32;
    (f1 + f2 + f3) / 3.0
}

/* pub fn levenshtein_distance(a: &str, b: &str) -> usize {
    // TODO: unimplemented
}*/
<commit_msg>Remove unused import<commit_after>fn max_length(str1: &str, str2: &str) -> usize {
    if str1.len() > str2.len() {
      str1.len()
    } else {
      str2.len()
    }
}

fn get_common_chars(str1: &str, str2: &str, ordering: bool) -> Vec<char> {
    let mut common_chars = Vec::new();

    if str1.len() == 0 || str2.len() == 0 {
        return common_chars;
    }

    let max_length = max_length(str1, str2);
    let match_buffer = ( max_length / 2 ) - 1;
    let mut f = str1;
    let mut s = str2;

    if ordering {
        f = str2;
        s = str1;
    }

    for (i, c) in f.chars().enumerate() {
        let s_index: usize = if i < match_buffer {
            0
        } else {
            i - match_buffer
        };

        let e_index: usize = if s.len() <= (i + match_buffer) {
            s.len()
        } else {
            i + match_buffer + 1
        };

        let word_slice = &s[s_index..e_index];
        if word_slice.contains(c) {
            common_chars.push(c);
        }
    }

    common_chars
}

pub fn jaro_winkler_distance(str1: &str, str2: &str) -> f32 {
    let matches1 = get_common_chars(str1, str2, false);
    let matches2 = get_common_chars(str1, str2, true);

    if matches1.is_empty() || matches2.is_empty()
      || matches1.len() != matches2.len()
    {
      return 0.0;
    }

    let mut t = 0.0;
    for (a, b) in matches1.iter().zip(matches2.iter()) {
      if a != b {
        t += 0.5;
      }
    }

    let f1 = matches1.len() as f32 / str1.len() as f32;
    let f2 = matches2.len() as f32 / str2.len() as f32;
    let f3 = (matches1.len() as f32 - t) / matches1.len() as f32;
    (f1 + f2 + f3) / 3.0
}

/* pub fn levenshtein_distance(a: &str, b: &str) -> usize {
    // TODO: unimplemented
}*/
<|endoftext|>


<commit_before>// Defined note types for GNU systems.

#[cfg(feature = "log")]
use log::debug;
#[cfg(feature = "alloc")]
use scroll::{IOread, IOwrite, Pread, Pwrite, SizeWith};

/// ABI information.
///
/// The descriptor consists of words:
///  * word 0: OS descriptor
///  * word 1: major version of the ABI
///  * word 2: minor version of the ABI
///  * word 3: subminor version of the ABI
pub const NT_GNU_ABI_TAG: u32 = 1;

/// Old name
pub const ELF_NOTE_ABI: u32 = NT_GNU_ABI_TAG;
// Known OSes.  These values can appear in word 0 of an
// `NT_GNU_ABI_TAG` note section entry.
pub const ELF_NOTE_OS_LINUX: u32 = 0;
pub const ELF_NOTE_OS_GNU: u32 = 1;
pub const ELF_NOTE_OS_SOLARIS2: u32 = 2;
pub const ELF_NOTE_OS_FREEBSD: u32 = 3;

/// Synthetic `hwcap` information.
///
/// The descriptor begins with two words:
///  * word 0: number of entries
///  * word 1: bitmask of enabled entries
///
/// Then follow variable-length entries, one byte followed by a '\0'-terminated
/// `hwcap` name string.  The byte gives the bit number to test if enabled,
/// `(1U << bit) & bitmask`.
pub const NT_GNU_HWCAP: u32 = 2;

/// Build ID bits as generated by ld --build-id.
///
/// The descriptor consists of any nonzero number of bytes.
pub const NT_GNU_BUILD_ID: u32 = 3;

/// Version note generated by GNU gold containing a version string.
pub const NT_GNU_GOLD_VERSION: u32 = 4;

///Contains copy of prstatus struct.
pub const NT_PRSTATUS: u32 = 1;

///Contains copy of prpsinfo struct.
pub const NT_PRPSINFO: u32 = 3;

///Fields of siginfo_t.
pub const NT_SIGINFO: u32 = 0x5349_4749;

///Description of mapped files.
pub const NT_FILE: u32 = 0x4649_4c45;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "alloc", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header
/// of a fixed form.
pub struct Nhdr32 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u32,
    /// Length of the note's descriptor
    pub n_descsz: u32,
    /// Type of the note
    pub n_type: u32,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "alloc", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header
/// of a fixed form.
pub struct Nhdr64 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u64,
    /// Length of the note's descriptor.
    pub n_descsz: u64,
    /// Type of the note.
    pub n_type: u64,
}

if_alloc! {
    use crate::error;
    use crate::container;
    use scroll::ctx;
    use alloc::vec::Vec;

    /// An iterator over ELF binary notes in a note section or segment
    pub struct NoteDataIterator<'a> {
        pub data: &'a [u8],
        pub size: usize,
        pub offset: usize,
        pub ctx: (usize, container::Ctx), // (alignment, ctx)
    }

    impl<'a> Iterator for NoteDataIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            if self.offset >= self.size {
                None
            } else {
                debug!("NoteIterator - {:#x}", self.offset);
                match self.data.gread_with(&mut self.offset, self.ctx) {
                    Ok(res) => Some(Ok(res)),
                    Err(e) => Some(Err(e))
                }
            }
        }
    }

    /// An iterator over ELF binary notes
    pub struct NoteIterator<'a> {
        pub iters: Vec<NoteDataIterator<'a>>,
        pub index: usize,
    }

    impl<'a> Iterator for NoteIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            while self.index < self.iters.len() {
                if let Some(note_result) = self.iters[self.index].next() {
                    return Some(note_result);
                }

                self.index += 1;
            }

            None
        }
    }

    #[derive(Debug)]
    struct NoteHeader {
        n_namesz: usize,
        n_descsz: usize,
        n_type: u32,
    }

    impl From<Nhdr32> for NoteHeader {
        fn from(header: Nhdr32) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type,
            }
        }
    }

    impl From<Nhdr64> for NoteHeader {
        fn from(header: Nhdr64) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type as u32,
            }
        }
    }

    fn align(alignment: usize, offset: &mut usize) {
        let diff = *offset % alignment;
        if diff != 0 {
            *offset += alignment - diff;
        }
    }

    /// A 32/64 bit Note struct, with the name and desc pre-parsed
    #[derive(Debug)]
    pub struct Note<'a> {
        /// The type of this note
        pub n_type: u32,
        /// NUL terminated string, where `namesz` includes the terminator
        pub name: &'a str, // needs padding such that namesz + padding % {wordsize} == 0
        /// arbitrary data of length `descsz`
        pub desc: &'a [u8], // needs padding such that descsz + padding % {wordsize} == 0
    }

    impl<'a> Note<'a> {
        pub fn type_to_str(&self) -> &'static str {
            match self.n_type {
                NT_GNU_ABI_TAG => "NT_GNU_ABI_TAG",
                NT_GNU_HWCAP => "NT_GNU_HWCAP",
                NT_GNU_BUILD_ID => "NT_GNU_BUILD_ID",
                NT_GNU_GOLD_VERSION => "NT_GNU_GOLD_VERSION",
                _ => "NT_UNKNOWN"
            }
        }
    }

    impl<'a> ctx::TryFromCtx<'a, (usize, container::Ctx)> for Note<'a> {
        type Error = error::Error;
        fn try_from_ctx(bytes: &'a [u8], (alignment, ctx): (usize, container::Ctx)) -> Result<(Self, usize), Self::Error> {
            let offset = &mut 0;
            let mut alignment = alignment;
            if alignment < 4 {
                alignment = 4;
            }
            let header: NoteHeader = {
                match alignment {
                    4|8 => bytes.gread_with::<Nhdr32>(offset, ctx.le)?.into(),
                    _ => return Err(error::Error::Malformed(format!("Notes has unimplemented alignment requirement: {:#x}", alignment)))
                }
            };
            debug!("{:?} - {:#x}", header, *offset);
            // -1 because includes \0 terminator
            let name = bytes.gread_with::<&'a str>(offset, ctx::StrCtx::Length(header.n_namesz - 1))?;
            *offset += 1;
            align(alignment, offset);
            debug!("note name {} - {:#x}", name, *offset);
            let desc = bytes.gread_with::<&'a [u8]>(offset, header.n_descsz)?;
            align(alignment, offset);
            debug!("desc {:?} - {:#x}", desc, *offset);
            Ok((Note {
                name,
                desc,
                n_type: header.n_type,
            }, *offset))
        }
    }

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

        static NOTE_DATA: [u8; 68] = [0x04, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00,
                                     0x01, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
                                     0x06, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00,
                                     0x04, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00,
                                     0x03, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0xbc, 0xfc, 0x66, 0xcd, 0xc7, 0xd5, 0x14, 0x7b,
                                     0x53, 0xb1, 0x10, 0x11, 0x94, 0x86, 0x8e, 0xf9,
                                     0x4f, 0xe8, 0xdd, 0xdb];

        static CONTEXT: (usize, container::Ctx) = (4, container::Ctx {
            container: container::Container::Big,
            le: ::scroll::Endian::Little,
        });

        fn make_note_iter(start: usize, end: usize) -> NoteDataIterator<'static> {
            NoteDataIterator {
                data: &NOTE_DATA,
                size: end,
                offset: start,
                ctx: CONTEXT,
            }
        }

        #[test]
        fn iter_single_section() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn iter_multiple_sections() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 32), make_note_iter(32, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn skip_empty_sections() {
            let mut notes = NoteIterator {
                iters: vec![
                    make_note_iter(0, 32),
                    make_note_iter(0, 0),
                    make_note_iter(32, 68),
                ],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn ignore_no_sections() {
            let mut notes = NoteIterator { iters: vec![], index: 0 };
            assert!(notes.next().is_none());
        }
    }
}
<commit_msg>elf: avoid overflow when header.n_namesz is 0 (#256)<commit_after>// Defined note types for GNU systems.

#[cfg(feature = "log")]
use log::debug;
#[cfg(feature = "alloc")]
use scroll::{IOread, IOwrite, Pread, Pwrite, SizeWith};

/// ABI information.
///
/// The descriptor consists of words:
///  * word 0: OS descriptor
///  * word 1: major version of the ABI
///  * word 2: minor version of the ABI
///  * word 3: subminor version of the ABI
pub const NT_GNU_ABI_TAG: u32 = 1;

/// Old name
pub const ELF_NOTE_ABI: u32 = NT_GNU_ABI_TAG;
// Known OSes.  These values can appear in word 0 of an
// `NT_GNU_ABI_TAG` note section entry.
pub const ELF_NOTE_OS_LINUX: u32 = 0;
pub const ELF_NOTE_OS_GNU: u32 = 1;
pub const ELF_NOTE_OS_SOLARIS2: u32 = 2;
pub const ELF_NOTE_OS_FREEBSD: u32 = 3;

/// Synthetic `hwcap` information.
///
/// The descriptor begins with two words:
///  * word 0: number of entries
///  * word 1: bitmask of enabled entries
///
/// Then follow variable-length entries, one byte followed by a '\0'-terminated
/// `hwcap` name string.  The byte gives the bit number to test if enabled,
/// `(1U << bit) & bitmask`.
pub const NT_GNU_HWCAP: u32 = 2;

/// Build ID bits as generated by ld --build-id.
///
/// The descriptor consists of any nonzero number of bytes.
pub const NT_GNU_BUILD_ID: u32 = 3;

/// Version note generated by GNU gold containing a version string.
pub const NT_GNU_GOLD_VERSION: u32 = 4;

///Contains copy of prstatus struct.
pub const NT_PRSTATUS: u32 = 1;

///Contains copy of prpsinfo struct.
pub const NT_PRPSINFO: u32 = 3;

///Fields of siginfo_t.
pub const NT_SIGINFO: u32 = 0x5349_4749;

///Description of mapped files.
pub const NT_FILE: u32 = 0x4649_4c45;

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "alloc", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header
/// of a fixed form.
pub struct Nhdr32 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u32,
    /// Length of the note's descriptor
    pub n_descsz: u32,
    /// Type of the note
    pub n_type: u32,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "alloc", derive(Pread, Pwrite, IOread, IOwrite, SizeWith))]
#[repr(C)]
/// Note section contents. Each entry in the note section begins with a header
/// of a fixed form.
pub struct Nhdr64 {
    /// Length of the note's name (includes the terminator)
    pub n_namesz: u64,
    /// Length of the note's descriptor.
    pub n_descsz: u64,
    /// Type of the note.
    pub n_type: u64,
}

if_alloc! {
    use crate::error;
    use crate::container;
    use scroll::ctx;
    use alloc::vec::Vec;

    /// An iterator over ELF binary notes in a note section or segment
    pub struct NoteDataIterator<'a> {
        pub data: &'a [u8],
        pub size: usize,
        pub offset: usize,
        pub ctx: (usize, container::Ctx), // (alignment, ctx)
    }

    impl<'a> Iterator for NoteDataIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            if self.offset >= self.size {
                None
            } else {
                debug!("NoteIterator - {:#x}", self.offset);
                match self.data.gread_with(&mut self.offset, self.ctx) {
                    Ok(res) => Some(Ok(res)),
                    Err(e) => Some(Err(e))
                }
            }
        }
    }

    /// An iterator over ELF binary notes
    pub struct NoteIterator<'a> {
        pub iters: Vec<NoteDataIterator<'a>>,
        pub index: usize,
    }

    impl<'a> Iterator for NoteIterator<'a> {
        type Item = error::Result<Note<'a>>;
        fn next(&mut self) -> Option<Self::Item> {
            while self.index < self.iters.len() {
                if let Some(note_result) = self.iters[self.index].next() {
                    return Some(note_result);
                }

                self.index += 1;
            }

            None
        }
    }

    #[derive(Debug)]
    struct NoteHeader {
        n_namesz: usize,
        n_descsz: usize,
        n_type: u32,
    }

    impl From<Nhdr32> for NoteHeader {
        fn from(header: Nhdr32) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type,
            }
        }
    }

    impl From<Nhdr64> for NoteHeader {
        fn from(header: Nhdr64) -> Self {
            NoteHeader {
                n_namesz: header.n_namesz as usize,
                n_descsz: header.n_descsz as usize,
                n_type: header.n_type as u32,
            }
        }
    }

    fn align(alignment: usize, offset: &mut usize) {
        let diff = *offset % alignment;
        if diff != 0 {
            *offset += alignment - diff;
        }
    }

    /// A 32/64 bit Note struct, with the name and desc pre-parsed
    #[derive(Debug)]
    pub struct Note<'a> {
        /// The type of this note
        pub n_type: u32,
        /// NUL terminated string, where `namesz` includes the terminator
        pub name: &'a str, // needs padding such that namesz + padding % {wordsize} == 0
        /// arbitrary data of length `descsz`
        pub desc: &'a [u8], // needs padding such that descsz + padding % {wordsize} == 0
    }

    impl<'a> Note<'a> {
        pub fn type_to_str(&self) -> &'static str {
            match self.n_type {
                NT_GNU_ABI_TAG => "NT_GNU_ABI_TAG",
                NT_GNU_HWCAP => "NT_GNU_HWCAP",
                NT_GNU_BUILD_ID => "NT_GNU_BUILD_ID",
                NT_GNU_GOLD_VERSION => "NT_GNU_GOLD_VERSION",
                _ => "NT_UNKNOWN"
            }
        }
    }

    impl<'a> ctx::TryFromCtx<'a, (usize, container::Ctx)> for Note<'a> {
        type Error = error::Error;
        fn try_from_ctx(bytes: &'a [u8], (alignment, ctx): (usize, container::Ctx)) -> Result<(Self, usize), Self::Error> {
            let offset = &mut 0;
            let mut alignment = alignment;
            if alignment < 4 {
                alignment = 4;
            }
            let header: NoteHeader = {
                match alignment {
                    4|8 => bytes.gread_with::<Nhdr32>(offset, ctx.le)?.into(),
                    _ => return Err(error::Error::Malformed(format!("Notes has unimplemented alignment requirement: {:#x}", alignment)))
                }
            };
            debug!("{:?} - {:#x}", header, *offset);
            // -1 because includes \0 terminator
            let name = bytes.gread_with::<&'a str>(offset, ctx::StrCtx::Length(header.n_namesz.saturating_sub(1)))?;
            if header.n_namesz > 0 {
                *offset += 1;
            }
            align(alignment, offset);
            debug!("note name {} - {:#x}", name, *offset);
            let desc = bytes.gread_with::<&'a [u8]>(offset, header.n_descsz)?;
            align(alignment, offset);
            debug!("desc {:?} - {:#x}", desc, *offset);
            Ok((Note {
                name,
                desc,
                n_type: header.n_type,
            }, *offset))
        }
    }

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

        static NOTE_DATA: [u8; 68] = [0x04, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00, 0x00,
                                     0x01, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00,
                                     0x06, 0x00, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00,
                                     0x04, 0x00, 0x00, 0x00, 0x14, 0x00, 0x00, 0x00,
                                     0x03, 0x00, 0x00, 0x00, 0x47, 0x4e, 0x55, 0x00,
                                     0xbc, 0xfc, 0x66, 0xcd, 0xc7, 0xd5, 0x14, 0x7b,
                                     0x53, 0xb1, 0x10, 0x11, 0x94, 0x86, 0x8e, 0xf9,
                                     0x4f, 0xe8, 0xdd, 0xdb];

        static CONTEXT: (usize, container::Ctx) = (4, container::Ctx {
            container: container::Container::Big,
            le: ::scroll::Endian::Little,
        });

        fn make_note_iter(start: usize, end: usize) -> NoteDataIterator<'static> {
            NoteDataIterator {
                data: &NOTE_DATA,
                size: end,
                offset: start,
                ctx: CONTEXT,
            }
        }

        #[test]
        fn iter_single_section() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn iter_multiple_sections() {
            let mut notes = NoteIterator {
                iters: vec![make_note_iter(0, 32), make_note_iter(32, 68)],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn skip_empty_sections() {
            let mut notes = NoteIterator {
                iters: vec![
                    make_note_iter(0, 32),
                    make_note_iter(0, 0),
                    make_note_iter(32, 68),
                ],
                index: 0,
            };

            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_ABI_TAG);
            assert_eq!(notes.next().unwrap().unwrap().n_type, NT_GNU_BUILD_ID);
            assert!(notes.next().is_none());
        }

        #[test]
        fn ignore_no_sections() {
            let mut notes = NoteIterator { iters: vec![], index: 0 };
            assert!(notes.next().is_none());
        }
    }
}
<|endoftext|>