<commit_before>use std::collections::HashMap;
use std::cell::{RefCell, RefMut};

use reqwest::{Response};

use errors::*;
use client::MatrixClient;
use matrix_types::Event;

#[derive(Clone)]
pub struct RoomEvent<'a> {
    pub room_id: &'a str,
    pub raw_event: &'a Event,
}

pub struct Bot <'a, 'b> {
    client: RefCell<&'b mut MatrixClient>,
    root_services: Vec<&'a str>,
    all_services: HashMap<&'a str, RefCell<Box<Node<'a>>>>,

    //rooms: RefCell<Vec<&'a str>>,
}

impl<'a, 'b> Bot<'a, 'b> {
    pub fn new(client: &'b mut MatrixClient) -> Self{
        Bot {
            client: RefCell::new(client),
            root_services: Vec::new(),
            all_services: HashMap::new(),
            //rooms: RefCell::new(Vec::new()),
        }
    }

    pub fn join(&self, room_id: &str) -> Result<Response, RustixError>{
        println!("Joining room with id: {}", &room_id);
        self.client.borrow().join(room_id)

        //self.rooms.borrow_mut().push(room_id);
    }

    pub fn say(&self, room_id: &str, message: &str) -> Result<Response, RustixError> {
        self.client.borrow_mut().send_msg(room_id, message)
    }

    pub fn reply(&self, event: &RoomEvent, message: &str) -> Result<Response, RustixError> {
        self.say(event.room_id, message)
    }

    /*
    pub fn action(&self, room_id: &str, action: &str) {
    }

    pub fn set_name(&self, name: &str) {
    }
    */

    pub fn register_service(&mut self,
                            name: &'a str,
                            parent: Option<&'a str>,
                            mut service: Box<Node<'a>>) -> Option<&'a str> {
        match parent {
            Some(p) => {
                self.all_services.get_mut(p).expect("Invalid parent node")
                    .borrow_mut().register_child(name)
            },
            None => self.root_services.push(name),
        };

        service.on_load();

        self.all_services.insert(name, RefCell::new(service));

        Some(name)
    }

    pub fn get_service(&self, name: &str) -> RefMut<Box<Node<'a>>> {
        self.all_services.get(name).unwrap().borrow_mut()
    }

    pub fn propagate_event(&self, event: &RoomEvent) {
        for service in &self.root_services {
            self.all_services.get(service).unwrap()
                .borrow_mut().handle(self, event.clone());
        }
    }

    pub fn run(&mut self) {
        let monitor_room = self.client.borrow().get_public_room_id("test").unwrap();

        self.join(&monitor_room);
        self.say(&monitor_room, "Hello world from rust!");

        if let Some(rid) = self.client.borrow().get_public_room_id("#geeks") {
            self.join(&rid);
        }

        let mut next_batch: String = self.client.borrow().sync(None).unwrap().next_batch;

        loop {
            let sync_data;
            if let Ok(res) = self.client.borrow().sync(Some(&next_batch)) {
                sync_data = res;
            } else {
                continue;
            }

            if let Some(room) = sync_data.rooms.join.get(&monitor_room) {
                for raw_event in &room.timeline.events {
                    self.propagate_event(
                        &RoomEvent{room_id: &monitor_room,
                                   raw_event}
                    );
                }
            }

            next_batch = sync_data.next_batch;
        }
    }
}

pub trait Node<'a> {
    fn children(&self) -> Option<&Vec<&'a str>> {
        None
    }

    fn register_child(&mut self, name: &'a str) {
    }

    fn propagate_event(&self, bot: &Bot, event: RoomEvent) {
        if let Some(children) = self.children() {
            for child in children {
                bot.get_service(child).handle(bot, event.clone());
            }
        }
    }

    fn handle(&mut self, bot: &Bot, event: RoomEvent) {
        self.propagate_event(bot, event);
    }

    fn on_load(&mut self) { }

    fn on_exit(&self) { }
}

/*
if let Ok(x) = serde_json::to_string_pretty(&res.rooms.join) {
    println!("{}", x);
} */
<commit_msg>Bot monitors multiple rooms based on join<commit_after>use std::collections::HashMap;
use std::cell::{RefCell, RefMut};

use reqwest::{Response};

use errors::*;
use client::MatrixClient;
use matrix_types::Event;

#[derive(Clone)]
pub struct RoomEvent<'a> {
    pub room_id: &'a str,
    pub raw_event: &'a Event,
}

pub struct Bot <'a, 'b> {
    client: RefCell<&'b mut MatrixClient>,
    root_services: Vec<&'a str>,
    all_services: HashMap<&'a str, RefCell<Box<Node<'a>>>>,

    rooms: RefCell<Vec<String>>,
}

impl<'a, 'b> Bot<'a, 'b> {
    pub fn new(client: &'b mut MatrixClient) -> Self{
        Bot {
            client: RefCell::new(client),
            root_services: Vec::new(),
            all_services: HashMap::new(),
            rooms: RefCell::new(Vec::new()),
        }
    }

    pub fn join(&self, room_id: &str) -> Result<Response, RustixError>{
        self.rooms.borrow_mut().push(room_id.clone().to_string());

        self.client.borrow().join(room_id)
    }

    pub fn say(&self, room_id: &str, message: &str) -> Result<Response, RustixError> {
        self.client.borrow_mut().send_msg(room_id, message)
    }

    pub fn reply(&self, event: &RoomEvent, message: &str) -> Result<Response, RustixError> {
        self.say(event.room_id, message)
    }

    /*
    pub fn action(&self, room_id: &str, action: &str) {
    }

    pub fn set_name(&self, name: &str) {
    }
    */

    pub fn register_service(&mut self,
                            name: &'a str,
                            parent: Option<&'a str>,
                            mut service: Box<Node<'a>>) -> Option<&'a str> {
        match parent {
            Some(p) => {
                self.all_services.get_mut(p).expect("Invalid parent node")
                    .borrow_mut().register_child(name)
            },
            None => self.root_services.push(name),
        };

        service.on_load();

        self.all_services.insert(name, RefCell::new(service));

        Some(name)
    }

    pub fn get_service(&self, name: &str) -> RefMut<Box<Node<'a>>> {
        self.all_services.get(name).unwrap().borrow_mut()
    }

    pub fn propagate_event(&self, event: &RoomEvent) {
        for service in &self.root_services {
            self.all_services.get(service).unwrap()
                .borrow_mut().handle(self, event.clone());
        }
    }

    pub fn run(&mut self) {
        let initial_rooms = vec!["test", "test2", "#geeks"];
        for room in initial_rooms {
            if let Some(rid) = self.client.borrow().get_public_room_id(room) {
                println!("Joining {} id: {}", &room, &rid);
                self.join(&rid);
            }
        }

        let mut next_batch: String = self.client.borrow().sync(None).unwrap().next_batch;

        loop {
            let sync_data = self.client.borrow().sync(Some(&next_batch)).unwrap();

            for room_id in self.rooms.borrow().iter() {
                if let Some(room) = sync_data.rooms.join.get(room_id) {
                    for raw_event in &room.timeline.events {
                        self.propagate_event(
                            &RoomEvent{room_id, raw_event}
                        );
                    }
                }
            }

            next_batch = sync_data.next_batch;
        }
    }
}

pub trait Node<'a> {
    fn children(&self) -> Option<&Vec<&'a str>> {
        None
    }

    fn register_child(&mut self, name: &'a str) {
    }

    fn propagate_event(&self, bot: &Bot, event: RoomEvent) {
        if let Some(children) = self.children() {
            for child in children {
                bot.get_service(child).handle(bot, event.clone());
            }
        }
    }

    fn handle(&mut self, bot: &Bot, event: RoomEvent) {
        self.propagate_event(bot, event);
    }

    fn on_load(&mut self) { }

    fn on_exit(&self) { }
}

/*
if let Ok(x) = serde_json::to_string_pretty(&res.rooms.join) {
    println!("{}", x);
} */
<|endoftext|>


<commit_before>// Copyright (c) 2017-2018 Rene van der Meer
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

//! Raspberry Pi system-related tools.
//!
//! Use [`DeviceInfo`] to identify what Raspberry Pi model and SoC the software is
//! running on.
//!
//! [`DeviceInfo`]: struct.DeviceInfo.html

use std::fmt;
use std::fs;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::result;

const PERIPHERAL_BASE_RPI: u32 = 0x2000_0000;
const PERIPHERAL_BASE_RPI2: u32 = 0x3f00_0000;
const GPIO_OFFSET: u32 = 0x20_0000;

quick_error! {
/// Errors that can occur when trying to identify the Raspberry Pi hardware.
    #[derive(Debug)]
    pub enum Error {
/// Unknown model.
///
/// `DeviceInfo` was unable to identify the Raspberry Pi model based on the
/// contents of `/proc/cpuinfo`, `/sys/firmware/devicetree/base/compatible`
/// and `/sys/firmware/devicetree/base/model`.
        UnknownModel { description("unknown Raspberry Pi model") }
    }
}

/// Result type returned from methods that can have `system::Error`s.
pub type Result<T> = result::Result<T, Error>;

/// Identifiable Raspberry Pi models.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Model {
    RaspberryPiA,
    RaspberryPiAPlus,
    RaspberryPiBRev1,
    RaspberryPiBRev2,
    RaspberryPiBPlus,
    RaspberryPi2B,
    RaspberryPi3B,
    RaspberryPi3BPlus,
    RaspberryPi3APlus,
    RaspberryPiComputeModule,
    RaspberryPiComputeModule3,
    RaspberryPiZero,
    RaspberryPiZeroW,
}

impl fmt::Display for Model {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Model::RaspberryPiA => write!(f, "Raspberry Pi A"),
            Model::RaspberryPiAPlus => write!(f, "Raspberry Pi A+"),
            Model::RaspberryPiBRev1 => write!(f, "Raspberry Pi B Rev 1"),
            Model::RaspberryPiBRev2 => write!(f, "Raspberry Pi B Rev 2"),
            Model::RaspberryPiBPlus => write!(f, "Raspberry Pi B+"),
            Model::RaspberryPi2B => write!(f, "Raspberry Pi 2 B"),
            Model::RaspberryPi3B => write!(f, "Raspberry Pi 3 B"),
            Model::RaspberryPi3BPlus => write!(f, "Raspberry Pi 3 B+"),
            Model::RaspberryPi3APlus => write!(f, "Raspberry Pi 3 A+"),
            Model::RaspberryPiComputeModule => write!(f, "Raspberry Pi Compute Module"),
            Model::RaspberryPiComputeModule3 => write!(f, "Raspberry Pi Compute Module 3"),
            Model::RaspberryPiZero => write!(f, "Raspberry Pi Zero"),
            Model::RaspberryPiZeroW => write!(f, "Raspberry Pi Zero W"),
        }
    }
}

/// Identifiable Raspberry Pi SoCs.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum SoC {
    Bcm2835,
    Bcm2836,
    Bcm2837A1,
    Bcm2837B0,
}

impl fmt::Display for SoC {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            SoC::Bcm2835 => write!(f, "BCM2835"),
            SoC::Bcm2836 => write!(f, "BCM2836"),
            SoC::Bcm2837A1 => write!(f, "BCM2837A1"),
            SoC::Bcm2837B0 => write!(f, "BCM2837B0"),
        }
    }
}

// Identify Pi model based on /proc/cpuinfo
fn parse_proc_cpuinfo() -> Result<Model> {
    let proc_cpuinfo = BufReader::new(match File::open("/proc/cpuinfo") {
        Ok(file) => file,
        Err(_) => return Err(Error::UnknownModel),
    });

    let mut hardware: String = String::new();
    let mut revision: String = String::new();
    for line_result in proc_cpuinfo.lines() {
        if let Ok(line) = line_result {
            if line.starts_with("Hardware\t: ") {
                hardware = String::from(&line[11..]);
            } else if line.starts_with("Revision\t: ") {
                revision = String::from(&line[11..]).to_lowercase();
            }
        }
    }

    // Return an error if we don't recognize the SoC. This check is
    // done to prevent accidentally identifying a non-Pi SBC as a Pi
    // solely based on the revision field.
    match &hardware[..] {
        "BCM2708" | "BCM2835" | "BCM2709" | "BCM2836" | "BCM2710" | "BCM2837" | "BCM2837A1"
        | "BCM2837B0" => {}
        _ => return Err(Error::UnknownModel),
    }

    let model = if (revision.len() == 4) || (revision.len() == 8) {
        // Older revisions are 4 characters long, or 8 if they've been over-volted
        match &revision[revision.len() - 4..] {
            "0007" | "0008" | "0009" | "0015" => Model::RaspberryPiA,
            "Beta" | "0002" | "0003" => Model::RaspberryPiBRev1,
            "0004" | "0005" | "0006" | "000d" | "000e" | "000f" => Model::RaspberryPiBRev2,
            "0012" => Model::RaspberryPiAPlus,
            "0010" | "0013" => Model::RaspberryPiBPlus,
            "0011" | "0014" => Model::RaspberryPiComputeModule,
            _ => return Err(Error::UnknownModel),
        }
    } else if revision.len() >= 6 {
        // Newer revisions consist of at least 6 characters
        match &revision[revision.len() - 3..revision.len() - 1] {
            "00" => Model::RaspberryPiA,
            "01" => Model::RaspberryPiBRev2,
            "02" => Model::RaspberryPiAPlus,
            "03" => Model::RaspberryPiBPlus,
            "04" => Model::RaspberryPi2B,
            "06" => Model::RaspberryPiComputeModule,
            "08" => Model::RaspberryPi3B,
            "09" => Model::RaspberryPiZero,
            "0a" => Model::RaspberryPiComputeModule3,
            "0c" => Model::RaspberryPiZeroW,
            "0d" => Model::RaspberryPi3BPlus,
            "0e" => Model::RaspberryPi3APlus,
            _ => return Err(Error::UnknownModel),
        }
    } else {
        return Err(Error::UnknownModel);
    };

    Ok(model)
}

// Identify Pi model based on /sys/firmware/devicetree/base/compatible
fn parse_base_compatible() -> Result<Model> {
    let base_compatible = match fs::read_to_string("/sys/firmware/devicetree/base/compatible") {
        Ok(buffer) => buffer,
        Err(_) => return Err(Error::UnknownModel),
    };

    // Based on /arch/arm/boot/dts/ and /Documentation/devicetree/bindings/arm/bcm/
    for comp_id in base_compatible.split('\0') {
        let model = match comp_id {
            "raspberrypi,model-b-i2c0" => Model::RaspberryPiBRev1,
            "raspberrypi,model-b" => Model::RaspberryPiBRev1,
            "raspberrypi,model-a" => Model::RaspberryPiA,
            "raspberrypi,model-b-rev2" => Model::RaspberryPiBRev2,
            "raspberrypi,model-a-plus" => Model::RaspberryPiAPlus,
            "raspberrypi,model-b-plus" => Model::RaspberryPiBPlus,
            "raspberrypi,2-model-b" => Model::RaspberryPi2B,
            "raspberrypi,compute-module" => Model::RaspberryPiComputeModule,
            "raspberrypi,3-model-b" => Model::RaspberryPi3B,
            "raspberrypi,model-zero" => Model::RaspberryPiZero,
            "raspberrypi,3-compute-module" => Model::RaspberryPiComputeModule3,
            "raspberrypi,model-zero-w" => Model::RaspberryPiZeroW,
            "raspberrypi,3-model-b-plus" => Model::RaspberryPi3BPlus,
            "raspberrypi,3-model-a-plus" => Model::RaspberryPi3APlus,
            _ => continue,
        };

        return Ok(model);
    }

    Err(Error::UnknownModel)
}

// Identify Pi model based on /sys/firmware/devicetree/base/model
fn parse_base_model() -> Result<Model> {
    let mut base_model = match fs::read_to_string("/sys/firmware/devicetree/base/model") {
        Ok(mut buffer) => if let Some(idx) = buffer.find('\0') {
            buffer.truncate(idx);
            buffer
        } else {
            buffer
        },
        Err(_) => return Err(Error::UnknownModel),
    };

    // Check if this is a Pi B rev 2 before we remove the revision part, assuming the
    // PCB Revision numbers on https://elinux.org/RPi_HardwareHistory are correct, and
    // the installed distro appends the revision to the model name.
    match &base_model[..] {
        "Raspberry Pi Model B Rev 2.0" => return Ok(Model::RaspberryPiBRev2),
        "Raspberry Pi Model B rev2 Rev 2.0" => return Ok(Model::RaspberryPiBRev2),
        _ => (),
    }

    if let Some(idx) = base_model.find(" Rev ") {
        base_model.truncate(idx);
    }

    // Based on /arch/arm/boot/dts/ and /Documentation/devicetree/bindings/arm/bcm/
    let model = match &base_model[..] {
        "Raspberry Pi Model B (no P5)" => Model::RaspberryPiBRev1,
        "Raspberry Pi Model B" => Model::RaspberryPiBRev1,
        "Raspberry Pi Model A" => Model::RaspberryPiA,
        "Raspberry Pi Model B rev2" => Model::RaspberryPiBRev2,
        "Raspberry Pi Model A+" => Model::RaspberryPiAPlus,
        "Raspberry Pi Model A Plus" => Model::RaspberryPiAPlus,
        "Raspberry Pi Model B+" => Model::RaspberryPiBPlus,
        "Raspberry Pi Model B Plus" => Model::RaspberryPiBPlus,
        "Raspberry Pi 2 Model B" => Model::RaspberryPi2B,
        "Raspberry Pi Compute Module" => Model::RaspberryPiComputeModule,
        "Raspberry Pi 3 Model B" => Model::RaspberryPi3B,
        "Raspberry Pi Zero" => Model::RaspberryPiZero,
        "Raspberry Pi Compute Module 3" => Model::RaspberryPiComputeModule3,
        "Raspberry Pi Zero W" => Model::RaspberryPiZeroW,
        "Raspberry Pi 3 Model B+" => Model::RaspberryPi3BPlus,
        "Raspberry Pi 3 Model B Plus" => Model::RaspberryPi3BPlus,
        "Raspberry Pi 3 Model A Plus" => Model::RaspberryPi3APlus,
        _ => return Err(Error::UnknownModel),
    };

    Ok(model)
}

/// Retrieves Raspberry Pi device information.
#[derive(Debug, PartialEq, Copy, Clone)]
pub struct DeviceInfo {
    model: Model,
    soc: SoC,
    peripheral_base: u32,
    gpio_offset: u32,
}

impl DeviceInfo {
    /// Constructs a new `DeviceInfo`.
    ///
    /// `new` automatically identifies the Raspberry Pi model and SoC based on
    /// the contents of `/proc/cpuinfo`, `/sys/firmware/devicetree/base/compatible`
    /// and `/sys/firmware/devicetree/base/model`.
    pub fn new() -> Result<DeviceInfo> {
        // Parse order from most-detailed to least-detailed info
        let model = parse_proc_cpuinfo()
            .or_else(|_| parse_base_compatible().or_else(|_| parse_base_model()))?;

        // Set SoC and memory offsets based on model
        match model {
            Model::RaspberryPiA
            | Model::RaspberryPiAPlus
            | Model::RaspberryPiBRev1
            | Model::RaspberryPiBRev2
            | Model::RaspberryPiBPlus
            | Model::RaspberryPiComputeModule
            | Model::RaspberryPiZero
            | Model::RaspberryPiZeroW => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2835,
                peripheral_base: PERIPHERAL_BASE_RPI,
                gpio_offset: GPIO_OFFSET,
            }),
            Model::RaspberryPi2B => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2836,
                peripheral_base: PERIPHERAL_BASE_RPI2,
                gpio_offset: GPIO_OFFSET,
            }),
            Model::RaspberryPi3B | Model::RaspberryPiComputeModule3 => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2837A1,
                peripheral_base: PERIPHERAL_BASE_RPI2,
                gpio_offset: GPIO_OFFSET,
            }),
            Model::RaspberryPi3BPlus | Model::RaspberryPi3APlus => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2837B0,
                peripheral_base: PERIPHERAL_BASE_RPI2,
                gpio_offset: GPIO_OFFSET,
            }),
        }
    }

    /// Returns the Raspberry Pi's model.
    pub fn model(&self) -> Model {
        self.model
    }

    /// Returns the Raspberry Pi's SoC.
    pub fn soc(&self) -> SoC {
        self.soc
    }

    /// Returns the base memory address for the BCM283x peripherals.
    pub fn peripheral_base(&self) -> u32 {
        self.peripheral_base
    }

    /// Returns the offset from the base memory address for the GPIO section.
    pub fn gpio_offset(&self) -> u32 {
        self.gpio_offset
    }
}
<commit_msg>Compare full rev code for future compatibility<commit_after>// Copyright (c) 2017-2018 Rene van der Meer
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

//! Raspberry Pi system-related tools.
//!
//! Use [`DeviceInfo`] to identify what Raspberry Pi model and SoC the software is
//! running on.
//!
//! [`DeviceInfo`]: struct.DeviceInfo.html

use std::fmt;
use std::fs;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::result;

const PERIPHERAL_BASE_RPI: u32 = 0x2000_0000;
const PERIPHERAL_BASE_RPI2: u32 = 0x3f00_0000;
const GPIO_OFFSET: u32 = 0x20_0000;

quick_error! {
/// Errors that can occur when trying to identify the Raspberry Pi hardware.
    #[derive(Debug)]
    pub enum Error {
/// Unknown model.
///
/// `DeviceInfo` was unable to identify the Raspberry Pi model based on the
/// contents of `/proc/cpuinfo`, `/sys/firmware/devicetree/base/compatible`
/// and `/sys/firmware/devicetree/base/model`.
        UnknownModel { description("unknown Raspberry Pi model") }
    }
}

/// Result type returned from methods that can have `system::Error`s.
pub type Result<T> = result::Result<T, Error>;

/// Identifiable Raspberry Pi models.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Model {
    RaspberryPiA,
    RaspberryPiAPlus,
    RaspberryPiBRev1,
    RaspberryPiBRev2,
    RaspberryPiBPlus,
    RaspberryPi2B,
    RaspberryPi3B,
    RaspberryPi3BPlus,
    RaspberryPi3APlus,
    RaspberryPiComputeModule,
    RaspberryPiComputeModule3,
    RaspberryPiZero,
    RaspberryPiZeroW,
}

impl fmt::Display for Model {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Model::RaspberryPiA => write!(f, "Raspberry Pi A"),
            Model::RaspberryPiAPlus => write!(f, "Raspberry Pi A+"),
            Model::RaspberryPiBRev1 => write!(f, "Raspberry Pi B Rev 1"),
            Model::RaspberryPiBRev2 => write!(f, "Raspberry Pi B Rev 2"),
            Model::RaspberryPiBPlus => write!(f, "Raspberry Pi B+"),
            Model::RaspberryPi2B => write!(f, "Raspberry Pi 2 B"),
            Model::RaspberryPi3B => write!(f, "Raspberry Pi 3 B"),
            Model::RaspberryPi3BPlus => write!(f, "Raspberry Pi 3 B+"),
            Model::RaspberryPi3APlus => write!(f, "Raspberry Pi 3 A+"),
            Model::RaspberryPiComputeModule => write!(f, "Raspberry Pi Compute Module"),
            Model::RaspberryPiComputeModule3 => write!(f, "Raspberry Pi Compute Module 3"),
            Model::RaspberryPiZero => write!(f, "Raspberry Pi Zero"),
            Model::RaspberryPiZeroW => write!(f, "Raspberry Pi Zero W"),
        }
    }
}

/// Identifiable Raspberry Pi SoCs.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum SoC {
    Bcm2835,
    Bcm2836,
    Bcm2837A1,
    Bcm2837B0,
}

impl fmt::Display for SoC {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            SoC::Bcm2835 => write!(f, "BCM2835"),
            SoC::Bcm2836 => write!(f, "BCM2836"),
            SoC::Bcm2837A1 => write!(f, "BCM2837A1"),
            SoC::Bcm2837B0 => write!(f, "BCM2837B0"),
        }
    }
}

// Identify Pi model based on /proc/cpuinfo
fn parse_proc_cpuinfo() -> Result<Model> {
    let proc_cpuinfo = BufReader::new(match File::open("/proc/cpuinfo") {
        Ok(file) => file,
        Err(_) => return Err(Error::UnknownModel),
    });

    let mut hardware: String = String::new();
    let mut revision: String = String::new();
    for line_result in proc_cpuinfo.lines() {
        if let Ok(line) = line_result {
            if line.starts_with("Hardware\t: ") {
                hardware = String::from(&line[11..]);
            } else if line.starts_with("Revision\t: ") {
                revision = String::from(&line[11..]).to_lowercase();
            }
        }
    }

    // Return an error if we don't recognize the SoC. This check is
    // done to prevent accidentally identifying a non-Pi SBC as a Pi
    // solely based on the revision field.
    match &hardware[..] {
        "BCM2708" | "BCM2835" | "BCM2709" | "BCM2836" | "BCM2710" | "BCM2837" | "BCM2837A1"
        | "BCM2837B0" => {}
        _ => return Err(Error::UnknownModel),
    }

    let model = if (revision.len() == 4) || (revision.len() == 8) {
        // Older revisions are 4 characters long, or 8 if they've been over-volted
        match &revision[revision.len() - 4..] {
            "0007" | "0008" | "0009" | "0015" => Model::RaspberryPiA,
            "Beta" | "0002" | "0003" => Model::RaspberryPiBRev1,
            "0004" | "0005" | "0006" | "000d" | "000e" | "000f" => Model::RaspberryPiBRev2,
            "0012" => Model::RaspberryPiAPlus,
            "0010" | "0013" => Model::RaspberryPiBPlus,
            "0011" | "0014" => Model::RaspberryPiComputeModule,
            _ => return Err(Error::UnknownModel),
        }
    } else if revision.len() >= 6 {
        // Newer revisions consist of at least 6 characters
        match &revision[..] {
            "900021" => Model::RaspberryPiAPlus,
            "900032" => Model::RaspberryPiBPlus,
            "a01040" | "a01041" | "a21041" | "a22042" => Model::RaspberryPi2B,
            "a02082" | "a22082" | "a32082" | "a52082" => Model::RaspberryPi3B,
            "900092" | "920092" | "900093" | "920093" => Model::RaspberryPiZero,
            "a020a0" => Model::RaspberryPiComputeModule3,
            "9000c1" => Model::RaspberryPiZeroW,
            "a020d3" => Model::RaspberryPi3BPlus,
            "9020e0" => Model::RaspberryPi3APlus,
            _ => return Err(Error::UnknownModel),
        }
    } else {
        return Err(Error::UnknownModel);
    };

    Ok(model)
}

// Identify Pi model based on /sys/firmware/devicetree/base/compatible
fn parse_base_compatible() -> Result<Model> {
    let base_compatible = match fs::read_to_string("/sys/firmware/devicetree/base/compatible") {
        Ok(buffer) => buffer,
        Err(_) => return Err(Error::UnknownModel),
    };

    // Based on /arch/arm/boot/dts/ and /Documentation/devicetree/bindings/arm/bcm/
    for comp_id in base_compatible.split('\0') {
        let model = match comp_id {
            "raspberrypi,model-b-i2c0" => Model::RaspberryPiBRev1,
            "raspberrypi,model-b" => Model::RaspberryPiBRev1,
            "raspberrypi,model-a" => Model::RaspberryPiA,
            "raspberrypi,model-b-rev2" => Model::RaspberryPiBRev2,
            "raspberrypi,model-a-plus" => Model::RaspberryPiAPlus,
            "raspberrypi,model-b-plus" => Model::RaspberryPiBPlus,
            "raspberrypi,2-model-b" => Model::RaspberryPi2B,
            "raspberrypi,compute-module" => Model::RaspberryPiComputeModule,
            "raspberrypi,3-model-b" => Model::RaspberryPi3B,
            "raspberrypi,model-zero" => Model::RaspberryPiZero,
            "raspberrypi,3-compute-module" => Model::RaspberryPiComputeModule3,
            "raspberrypi,model-zero-w" => Model::RaspberryPiZeroW,
            "raspberrypi,3-model-b-plus" => Model::RaspberryPi3BPlus,
            "raspberrypi,3-model-a-plus" => Model::RaspberryPi3APlus,
            _ => continue,
        };

        return Ok(model);
    }

    Err(Error::UnknownModel)
}

// Identify Pi model based on /sys/firmware/devicetree/base/model
fn parse_base_model() -> Result<Model> {
    let mut base_model = match fs::read_to_string("/sys/firmware/devicetree/base/model") {
        Ok(mut buffer) => if let Some(idx) = buffer.find('\0') {
            buffer.truncate(idx);
            buffer
        } else {
            buffer
        },
        Err(_) => return Err(Error::UnknownModel),
    };

    // Check if this is a Pi B rev 2 before we remove the revision part, assuming the
    // PCB Revision numbers on https://elinux.org/RPi_HardwareHistory are correct, and
    // the installed distro appends the revision to the model name.
    match &base_model[..] {
        "Raspberry Pi Model B Rev 2.0" => return Ok(Model::RaspberryPiBRev2),
        "Raspberry Pi Model B rev2 Rev 2.0" => return Ok(Model::RaspberryPiBRev2),
        _ => (),
    }

    if let Some(idx) = base_model.find(" Rev ") {
        base_model.truncate(idx);
    }

    // Based on /arch/arm/boot/dts/ and /Documentation/devicetree/bindings/arm/bcm/
    let model = match &base_model[..] {
        "Raspberry Pi Model B (no P5)" => Model::RaspberryPiBRev1,
        "Raspberry Pi Model B" => Model::RaspberryPiBRev1,
        "Raspberry Pi Model A" => Model::RaspberryPiA,
        "Raspberry Pi Model B rev2" => Model::RaspberryPiBRev2,
        "Raspberry Pi Model A+" => Model::RaspberryPiAPlus,
        "Raspberry Pi Model A Plus" => Model::RaspberryPiAPlus,
        "Raspberry Pi Model B+" => Model::RaspberryPiBPlus,
        "Raspberry Pi Model B Plus" => Model::RaspberryPiBPlus,
        "Raspberry Pi 2 Model B" => Model::RaspberryPi2B,
        "Raspberry Pi Compute Module" => Model::RaspberryPiComputeModule,
        "Raspberry Pi 3 Model B" => Model::RaspberryPi3B,
        "Raspberry Pi Zero" => Model::RaspberryPiZero,
        "Raspberry Pi Compute Module 3" => Model::RaspberryPiComputeModule3,
        "Raspberry Pi Zero W" => Model::RaspberryPiZeroW,
        "Raspberry Pi 3 Model B+" => Model::RaspberryPi3BPlus,
        "Raspberry Pi 3 Model B Plus" => Model::RaspberryPi3BPlus,
        "Raspberry Pi 3 Model A Plus" => Model::RaspberryPi3APlus,
        _ => return Err(Error::UnknownModel),
    };

    Ok(model)
}

/// Retrieves Raspberry Pi device information.
#[derive(Debug, PartialEq, Copy, Clone)]
pub struct DeviceInfo {
    model: Model,
    soc: SoC,
    peripheral_base: u32,
    gpio_offset: u32,
}

impl DeviceInfo {
    /// Constructs a new `DeviceInfo`.
    ///
    /// `new` automatically identifies the Raspberry Pi model and SoC based on
    /// the contents of `/proc/cpuinfo`, `/sys/firmware/devicetree/base/compatible`
    /// and `/sys/firmware/devicetree/base/model`.
    pub fn new() -> Result<DeviceInfo> {
        // Parse order from most-detailed to least-detailed info
        let model = parse_proc_cpuinfo()
            .or_else(|_| parse_base_compatible().or_else(|_| parse_base_model()))?;

        // Set SoC and memory offsets based on model
        match model {
            Model::RaspberryPiA
            | Model::RaspberryPiAPlus
            | Model::RaspberryPiBRev1
            | Model::RaspberryPiBRev2
            | Model::RaspberryPiBPlus
            | Model::RaspberryPiComputeModule
            | Model::RaspberryPiZero
            | Model::RaspberryPiZeroW => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2835,
                peripheral_base: PERIPHERAL_BASE_RPI,
                gpio_offset: GPIO_OFFSET,
            }),
            Model::RaspberryPi2B => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2836,
                peripheral_base: PERIPHERAL_BASE_RPI2,
                gpio_offset: GPIO_OFFSET,
            }),
            Model::RaspberryPi3B | Model::RaspberryPiComputeModule3 => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2837A1,
                peripheral_base: PERIPHERAL_BASE_RPI2,
                gpio_offset: GPIO_OFFSET,
            }),
            Model::RaspberryPi3BPlus | Model::RaspberryPi3APlus => Ok(DeviceInfo {
                model,
                soc: SoC::Bcm2837B0,
                peripheral_base: PERIPHERAL_BASE_RPI2,
                gpio_offset: GPIO_OFFSET,
            }),
        }
    }

    /// Returns the Raspberry Pi's model.
    pub fn model(&self) -> Model {
        self.model
    }

    /// Returns the Raspberry Pi's SoC.
    pub fn soc(&self) -> SoC {
        self.soc
    }

    /// Returns the base memory address for the BCM283x peripherals.
    pub fn peripheral_base(&self) -> u32 {
        self.peripheral_base
    }

    /// Returns the offset from the base memory address for the GPIO section.
    pub fn gpio_offset(&self) -> u32 {
        self.gpio_offset
    }
}
<|endoftext|>


<commit_before>use std::env;
use std::collections::HashMap;
use error::*;
use source::Source;
use value::{Value, ValueKind};

#[derive(Clone, Debug)]
pub struct Environment {
    /// Optional prefix that will limit access to the environment to only keys that
    /// begin with the defined prefix.
    ///
    /// A prefix, followed by `_` (the seperator),
    /// is tested to be present on each key before its considered
    /// to be part of the source environment.
    ///
    /// For example, the key `CONFIG_DEBUG` would become `DEBUG` with a prefix of `config`.
    prefix: Option<String>,

    /// The character sequence that separates each key segment in an environment key pattern.
    /// Consider a nested configuration such as `redis.password`, a separator of `_` would allow
    /// an environment key of `REDIS_PASSWORD` to match.
    ///
    /// The default separator is `_`.
    separator: String,
}

impl Environment {
    pub fn new() -> Self {
        Environment::default()
    }

    pub fn with_prefix(s: &str) -> Self {
        Environment {
            prefix: Some(s.into()),
            ..Environment::default()
        }
    }

    pub fn prefix(&mut self, s: String) -> &mut Self {
        self.prefix = s.into();
        self
    }

    pub fn separator(&mut self, s: String) -> &mut Self {
        self.separator = s;
        self
    }
}

impl Default for Environment {
    fn default() -> Environment {
        Environment {
            prefix: None,
            separator: "_".into(),
        }
    }
}

impl Source for Environment {
    fn clone_into_box(&self) -> Box<Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<HashMap<String, Value>> {
        let mut m = HashMap::new();
        let uri: String = "the environment".into();

        // Define a prefiux pattern to test and exclude from keys
        let prefix_pattern = match self.prefix {
            Some(ref prefix) => Some(prefix.clone() + &self.separator),
            _ => None,
        };

        for (key, value) in env::vars() {
            let mut key = key.to_string();

            // Check for prefix
            if let Some(ref prefix_pattern) = prefix_pattern {
                if key.to_lowercase().starts_with(prefix_pattern) {
                    // Remove this prefix from the key
                    key = key[prefix_pattern.len()..].to_string();
                } else {
                    // Skip this key
                    continue;
                }
            }

            m.insert(
                key.to_lowercase(),
                Value::new(Some(&uri), ValueKind::String(value)),
            );
        }

        Ok(m)
    }
}
<commit_msg>Revert "remove unnecessary env key manipulation"<commit_after>use std::env;
use std::collections::HashMap;
use error::*;
use source::Source;
use value::{Value, ValueKind};

#[derive(Clone, Debug)]
pub struct Environment {
    /// Optional prefix that will limit access to the environment to only keys that
    /// begin with the defined prefix.
    ///
    /// A prefix, followed by `_` (the seperator),
    /// is tested to be present on each key before its considered
    /// to be part of the source environment.
    ///
    /// For example, the key `CONFIG_DEBUG` would become `DEBUG` with a prefix of `config`.
    prefix: Option<String>,

    /// The character sequence that separates each key segment in an environment key pattern.
    /// Consider a nested configuration such as `redis.password`, a separator of `_` would allow
    /// an environment key of `REDIS_PASSWORD` to match.
    ///
    /// The default separator is `_`.
    separator: String,
}

impl Environment {
    pub fn new() -> Self {
        Environment::default()
    }

    pub fn with_prefix(s: &str) -> Self {
        Environment {
            prefix: Some(s.into()),
            ..Environment::default()
        }
    }

    pub fn prefix(&mut self, s: String) -> &mut Self {
        self.prefix = s.into();
        self
    }

    pub fn separator(&mut self, s: String) -> &mut Self {
        self.separator = s;
        self
    }
}

impl Default for Environment {
    fn default() -> Environment {
        Environment {
            prefix: None,
            separator: "_".into(),
        }
    }
}

impl Source for Environment {
    fn clone_into_box(&self) -> Box<Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<HashMap<String, Value>> {
        let mut m = HashMap::new();
        let uri: String = "the environment".into();

        // Define a prefiux pattern to test and exclude from keys
        let prefix_pattern = match self.prefix {
            Some(ref prefix) => Some(prefix.clone() + &self.separator),
            _ => None,
        };

        for (key, value) in env::vars() {
            let mut key = key.to_string();

            // Check for prefix
            if let Some(ref prefix_pattern) = prefix_pattern {
                if key.to_lowercase().starts_with(prefix_pattern) {
                    // Remove this prefix from the key
                    key = key[prefix_pattern.len()..].to_string();
                } else {
                    // Skip this key
                    continue;
                }
            }

            // Replace `separator` with `.`
            key = key.replace(&self.separator, ".");

            m.insert(
                key.to_lowercase(),
                Value::new(Some(&uri), ValueKind::String(value)),
            );
        }

        Ok(m)
    }
}
<|endoftext|>


<commit_before>//! Defines types to use with the geospatial commands.

use super::{ErrorKind, RedisResult};
use types::{FromRedisValue, ToRedisArgs, Value};

macro_rules! invalid_type_error {
    ($v:expr, $det:expr) => ({
        fail!((ErrorKind::TypeError,
               "Response was of incompatible type",
               format!("{:?} (response was {:?})", $det, $v)));
    })
}

/// Units used by [`geo_dist`][1] and [`geo_radius`][2].
///
/// [1]: ../trait.Commands.html#method.geo_dist
/// [2]: ../trait.Commands.html#method.geo_radius
pub enum Unit {
    Meters,
    Kilometers,
    Miles,
    Feet,
}

impl ToRedisArgs for Unit {
    fn to_redis_args(&self) -> Vec<Vec<u8>> {
        let args = match *self {
            Unit::Meters => "m",
            Unit::Kilometers => "km",
            Unit::Miles => "mi",
            Unit::Feet => "ft",
        };
        vec![args.as_bytes().to_vec()]
    }
}

/// A coordinate (longitude, latitude). Can be used with [`geo_pos`][1]
/// to parse response from Redis.
///
/// [1]: ../trait.Commands.html#method.geo_pos
///
/// `T` is the type of the every value.
///
/// * You may want to use either `f64` or `f32` if you want to perform mathematical operations.
/// * To keep the raw value from Redis, use `String`.
#[derive(Debug, PartialEq)]
pub struct Coord<T> {
    pub longitude: T,
    pub latitude: T,
}

impl<T> Coord<T> {
    /// Create a new Coord with the (longitude, latitude)
    pub fn lon_lat(longitude: T, latitude: T) -> Coord<T> {
        Coord { longitude, latitude }
    }
}

impl<T: FromRedisValue> FromRedisValue for Coord<T> {
    fn from_redis_value(v: &Value) -> RedisResult<Self> {
        let values: Vec<T> = FromRedisValue::from_redis_value(v)?;
        if values.len() != 2 {
            invalid_type_error!(v, "Expect a pair of numbers");
        }
        let mut values = values.into_iter();
        Ok(Coord { longitude: values.next().unwrap(), latitude: values.next().unwrap() })
    }
}

impl<T: ToRedisArgs> ToRedisArgs for Coord<T> {
    fn to_redis_args(&self) -> Vec<Vec<u8>> {
        let mut args = Vec::new();
        for field in &[&self.longitude, &self.latitude] {
            args.extend(ToRedisArgs::to_redis_args(*field));
        }
        args
    }
}

/// Options to sort results from GEORADIUS and GEORADIUSBYMEMBER commands
pub enum RadiusOrder {
    /// Don't sort the results
    Unsorted,

    /// Sort returned items from the nearest to the farthest, relative to the center.
    Asc,

    /// Sort returned items from the farthest to the nearest, relative to the center.
    Desc,
}

impl Default for RadiusOrder {
    fn default() -> RadiusOrder {
        RadiusOrder::Unsorted
    }
}

/// Options for the GEORADIUS and GEORADIUSBYMEMBER commands
#[derive(Default)]
pub struct RadiusOptions {
    with_coord: bool,
    with_dist: bool,
    count: Option<usize>,
    order: RadiusOrder,
    store: Option<Vec<Vec<u8>>>,
    store_dist: Option<Vec<Vec<u8>>>,
}

impl RadiusOptions {

    /// Limit the results to the first N matching items.
    pub fn limit(mut self, n: usize) -> Self {
        self.count = Some(n);
        self
    }

    /// Return the distance of the returned items from the specified center.
    /// The distance is returned in the same unit as the unit specified as the
    /// radius argument of the command.
    pub fn with_dist(mut self) -> Self {
        self.with_dist = true;
        self
    }

    /// Return the longitude,latitude coordinates of the matching items.
    pub fn with_coord(mut self) -> Self {
        self.with_coord = true;
        self
    }

    /// Sort the returned items
    pub fn order(mut self, o: RadiusOrder) -> Self {
        self.order = o;
        self
    }

    /// Store the results in a sorted set at key, instead of returning them.
    ///
    /// This feature can't be used with any `with_*` method.
    pub fn store<K: ToRedisArgs>(mut self, key: K) -> Self {
        self.store = Some(ToRedisArgs::to_redis_args(&key));
        self
    }

    /// Similar to `store`, but includes the distance of every member from the
    /// center. This feature can't be used with any `with_*` method.
    pub fn store_dist<K: ToRedisArgs>(mut self, key: K) -> Self {
        self.store_dist = Some(ToRedisArgs::to_redis_args(&key));
        self
    }

}

impl ToRedisArgs for RadiusOptions {
    fn to_redis_args(&self) -> Vec<Vec<u8>> {
        let mut args = Vec::new();

        if self.with_coord {
            args.push("WITHCOORD".as_bytes().to_vec());
        }

        if self.with_dist {
            args.push("WITHDIST".as_bytes().to_vec());
        }

        if let Some(n) = self.count {
            args.push("COUNT".as_bytes().to_vec());
            args.push(format!("{}", n).into_bytes());
        }

        match self.order {
            RadiusOrder::Asc => args.push("ASC".as_bytes().to_vec()),
            RadiusOrder::Desc => args.push("DESC".as_bytes().to_vec()),
            _ => (),
        };

        if let Some(ref store) = self.store {
            args.push("STORE".as_bytes().to_vec());
            args.extend_from_slice(&store[..]);
        }

        if let Some(ref store_dist) = self.store_dist {
            args.push("STOREDIST".as_bytes().to_vec());
            args.extend_from_slice(&store_dist[..]);
        }

        args
    }
}

/// Contain an item returned by [`geo_radius`][1] and [`geo_radius_by_member`][2].
///
/// [1]: ../trait.Commands.html#method.geo_radius
/// [2]: ../trait.Commands.html#method.geo_radius_by_member
pub struct RadiusSearchResult {
    pub name: String,
    pub coord: Option<Coord<f64>>,
    pub dist: Option<f64>,
}

impl FromRedisValue for RadiusSearchResult {
    fn from_redis_value(v: &Value) -> RedisResult<Self> {

        // If we receive only the member name, it will be a plain string
        if let Ok(name) = FromRedisValue::from_redis_value(v) {
            return Ok(RadiusSearchResult {
                name: name,
                coord: None,
                dist: None,
            });
        }

        // Try to parse the result from multitple values
        if let Value::Bulk(ref items) = *v {
            if let Some(result) = RadiusSearchResult::parse_multi_values(items) {
                return Ok(result);
            }
        }

        invalid_type_error!(v, "Response type not RadiusSearchResult compatible.");
    }
}

impl RadiusSearchResult {
    fn parse_multi_values(items: &[Value]) -> Option<Self> {
        let mut iter = items.iter();

        // First item is always the member name
        let name: String = match iter.next().map(FromRedisValue::from_redis_value) {
            Some(Ok(n)) => n,
            _ => return None,
        };

        let mut next = iter.next();

        // Next element, if present, will be the distance.
        let dist = match next.map(FromRedisValue::from_redis_value) {
            Some(Ok(c)) => { next = iter.next(); Some(c) },
            _ => None,
        };

        // Finally, if present, the last item will be the coordinates

        let coord = match next.map(FromRedisValue::from_redis_value) {
            Some(Ok(c)) => Some(c),
            _ => None,
        };

        Some(RadiusSearchResult { name, coord, dist })
    }
}

#[cfg(test)]
mod tests {
    use types::ToRedisArgs;
    use super::{Coord, RadiusOptions, RadiusOrder};
    use std::str;

    macro_rules! assert_args {
        ($value:expr, $($args:expr),+) => {
            let args = $value.to_redis_args();
            let strings: Vec<_> = args.iter()
                                      .map(|a| str::from_utf8(a.as_ref()).unwrap())
                                      .collect();
            assert_eq!(strings, vec![$($args),+]);
        }
    }

    #[test]
    fn test_coord_to_args() {
        let member = ("Palermo", Coord::lon_lat("13.361389", "38.115556"));
        assert_args!(&member, "Palermo", "13.361389", "38.115556");
    }

    #[test]
    fn test_radius_options() {
        // Without options, should not generate any argument
        let empty = RadiusOptions::default();
        assert_eq!(ToRedisArgs::to_redis_args(&empty).len(), 0);

        // Some combinations with WITH* options
        let opts = || RadiusOptions::default();

        assert_args!(
            opts().with_coord().with_dist(),
            "WITHCOORD", "WITHDIST");

        assert_args!(opts().limit(50), "COUNT", "50");

        assert_args!(
            opts().limit(50).store("x"),
            "COUNT", "50", "STORE", "x");

        assert_args!(
            opts().limit(100).store_dist("y"),
            "COUNT", "100", "STOREDIST", "y");

        assert_args!(
            opts().order(RadiusOrder::Asc).limit(10).with_dist(),
            "WITHDIST", "COUNT", "10", "ASC");

    }
}
<commit_msg>Better documentation for RadiusOrder and RadiusOptions<commit_after>//! Defines types to use with the geospatial commands.

use super::{ErrorKind, RedisResult};
use types::{FromRedisValue, ToRedisArgs, Value};

macro_rules! invalid_type_error {
    ($v:expr, $det:expr) => ({
        fail!((ErrorKind::TypeError,
               "Response was of incompatible type",
               format!("{:?} (response was {:?})", $det, $v)));
    })
}

/// Units used by [`geo_dist`][1] and [`geo_radius`][2].
///
/// [1]: ../trait.Commands.html#method.geo_dist
/// [2]: ../trait.Commands.html#method.geo_radius
pub enum Unit {
    Meters,
    Kilometers,
    Miles,
    Feet,
}

impl ToRedisArgs for Unit {
    fn to_redis_args(&self) -> Vec<Vec<u8>> {
        let args = match *self {
            Unit::Meters => "m",
            Unit::Kilometers => "km",
            Unit::Miles => "mi",
            Unit::Feet => "ft",
        };
        vec![args.as_bytes().to_vec()]
    }
}

/// A coordinate (longitude, latitude). Can be used with [`geo_pos`][1]
/// to parse response from Redis.
///
/// [1]: ../trait.Commands.html#method.geo_pos
///
/// `T` is the type of the every value.
///
/// * You may want to use either `f64` or `f32` if you want to perform mathematical operations.
/// * To keep the raw value from Redis, use `String`.
#[derive(Debug, PartialEq)]
pub struct Coord<T> {
    pub longitude: T,
    pub latitude: T,
}

impl<T> Coord<T> {
    /// Create a new Coord with the (longitude, latitude)
    pub fn lon_lat(longitude: T, latitude: T) -> Coord<T> {
        Coord { longitude, latitude }
    }
}

impl<T: FromRedisValue> FromRedisValue for Coord<T> {
    fn from_redis_value(v: &Value) -> RedisResult<Self> {
        let values: Vec<T> = FromRedisValue::from_redis_value(v)?;
        if values.len() != 2 {
            invalid_type_error!(v, "Expect a pair of numbers");
        }
        let mut values = values.into_iter();
        Ok(Coord { longitude: values.next().unwrap(), latitude: values.next().unwrap() })
    }
}

impl<T: ToRedisArgs> ToRedisArgs for Coord<T> {
    fn to_redis_args(&self) -> Vec<Vec<u8>> {
        let mut args = Vec::new();
        for field in &[&self.longitude, &self.latitude] {
            args.extend(ToRedisArgs::to_redis_args(*field));
        }
        args
    }
}

/// Options to sort results from [GEORADIUS][1] and [GEORADIUSBYMEMBER][2] commands
///
/// [1]: https://redis.io/commands/georadius
/// [2]: https://redis.io/commands/georadiusbymember
pub enum RadiusOrder {
    /// Don't sort the results
    Unsorted,

    /// Sort returned items from the nearest to the farthest, relative to the center.
    Asc,

    /// Sort returned items from the farthest to the nearest, relative to the center.
    Desc,
}

impl Default for RadiusOrder {
    fn default() -> RadiusOrder {
        RadiusOrder::Unsorted
    }
}

/// Options for the [GEORADIUS][1] and [GEORADIUSBYMEMBER][2] commands
///
/// [1]: https://redis.io/commands/georadius
/// [2]: https://redis.io/commands/georadiusbymember
///
/// # Example
///
/// ```rust,no_run
/// use redis::{Commands, RedisResult};
/// use redis::geo::{RadiusSearchResult, RadiusOptions, RadiusOrder, Unit};
/// fn nearest_in_radius(
///     con: &redis::Connection,
///     key: &str,
///     longitude: f64,
///     latitude: f64,
///     meters: f64,
///     limit: usize,
/// ) -> RedisResult<Vec<RadiusSearchResult>> {
///     let opts = RadiusOptions::default()
///         .order(RadiusOrder::Asc)
///         .limit(limit);
///     con.geo_radius(key, longitude, latitude, meters, Unit::Meters, opts)
/// }
/// ```

#[derive(Default)]
pub struct RadiusOptions {
    with_coord: bool,
    with_dist: bool,
    count: Option<usize>,
    order: RadiusOrder,
    store: Option<Vec<Vec<u8>>>,
    store_dist: Option<Vec<Vec<u8>>>,
}

impl RadiusOptions {

    /// Limit the results to the first N matching items.
    pub fn limit(mut self, n: usize) -> Self {
        self.count = Some(n);
        self
    }

    /// Return the distance of the returned items from the specified center.
    /// The distance is returned in the same unit as the unit specified as the
    /// radius argument of the command.
    pub fn with_dist(mut self) -> Self {
        self.with_dist = true;
        self
    }

    /// Return the longitude,latitude coordinates of the matching items.
    pub fn with_coord(mut self) -> Self {
        self.with_coord = true;
        self
    }

    /// Sort the returned items
    pub fn order(mut self, o: RadiusOrder) -> Self {
        self.order = o;
        self
    }

    /// Store the results in a sorted set at key, instead of returning them.
    ///
    /// This feature can't be used with any `with_*` method.
    pub fn store<K: ToRedisArgs>(mut self, key: K) -> Self {
        self.store = Some(ToRedisArgs::to_redis_args(&key));
        self
    }

    /// Similar to `store`, but includes the distance of every member from the
    /// center. This feature can't be used with any `with_*` method.
    pub fn store_dist<K: ToRedisArgs>(mut self, key: K) -> Self {
        self.store_dist = Some(ToRedisArgs::to_redis_args(&key));
        self
    }

}

impl ToRedisArgs for RadiusOptions {
    fn to_redis_args(&self) -> Vec<Vec<u8>> {
        let mut args = Vec::new();

        if self.with_coord {
            args.push("WITHCOORD".as_bytes().to_vec());
        }

        if self.with_dist {
            args.push("WITHDIST".as_bytes().to_vec());
        }

        if let Some(n) = self.count {
            args.push("COUNT".as_bytes().to_vec());
            args.push(format!("{}", n).into_bytes());
        }

        match self.order {
            RadiusOrder::Asc => args.push("ASC".as_bytes().to_vec()),
            RadiusOrder::Desc => args.push("DESC".as_bytes().to_vec()),
            _ => (),
        };

        if let Some(ref store) = self.store {
            args.push("STORE".as_bytes().to_vec());
            args.extend_from_slice(&store[..]);
        }

        if let Some(ref store_dist) = self.store_dist {
            args.push("STOREDIST".as_bytes().to_vec());
            args.extend_from_slice(&store_dist[..]);
        }

        args
    }
}

/// Contain an item returned by [`geo_radius`][1] and [`geo_radius_by_member`][2].
///
/// [1]: ../trait.Commands.html#method.geo_radius
/// [2]: ../trait.Commands.html#method.geo_radius_by_member
pub struct RadiusSearchResult {
    pub name: String,
    pub coord: Option<Coord<f64>>,
    pub dist: Option<f64>,
}

impl FromRedisValue for RadiusSearchResult {
    fn from_redis_value(v: &Value) -> RedisResult<Self> {

        // If we receive only the member name, it will be a plain string
        if let Ok(name) = FromRedisValue::from_redis_value(v) {
            return Ok(RadiusSearchResult {
                name: name,
                coord: None,
                dist: None,
            });
        }

        // Try to parse the result from multitple values
        if let Value::Bulk(ref items) = *v {
            if let Some(result) = RadiusSearchResult::parse_multi_values(items) {
                return Ok(result);
            }
        }

        invalid_type_error!(v, "Response type not RadiusSearchResult compatible.");
    }
}

impl RadiusSearchResult {
    fn parse_multi_values(items: &[Value]) -> Option<Self> {
        let mut iter = items.iter();

        // First item is always the member name
        let name: String = match iter.next().map(FromRedisValue::from_redis_value) {
            Some(Ok(n)) => n,
            _ => return None,
        };

        let mut next = iter.next();

        // Next element, if present, will be the distance.
        let dist = match next.map(FromRedisValue::from_redis_value) {
            Some(Ok(c)) => { next = iter.next(); Some(c) },
            _ => None,
        };

        // Finally, if present, the last item will be the coordinates

        let coord = match next.map(FromRedisValue::from_redis_value) {
            Some(Ok(c)) => Some(c),
            _ => None,
        };

        Some(RadiusSearchResult { name, coord, dist })
    }
}

#[cfg(test)]
mod tests {
    use types::ToRedisArgs;
    use super::{Coord, RadiusOptions, RadiusOrder};
    use std::str;

    macro_rules! assert_args {
        ($value:expr, $($args:expr),+) => {
            let args = $value.to_redis_args();
            let strings: Vec<_> = args.iter()
                                      .map(|a| str::from_utf8(a.as_ref()).unwrap())
                                      .collect();
            assert_eq!(strings, vec![$($args),+]);
        }
    }

    #[test]
    fn test_coord_to_args() {
        let member = ("Palermo", Coord::lon_lat("13.361389", "38.115556"));
        assert_args!(&member, "Palermo", "13.361389", "38.115556");
    }

    #[test]
    fn test_radius_options() {
        // Without options, should not generate any argument
        let empty = RadiusOptions::default();
        assert_eq!(ToRedisArgs::to_redis_args(&empty).len(), 0);

        // Some combinations with WITH* options
        let opts = || RadiusOptions::default();

        assert_args!(
            opts().with_coord().with_dist(),
            "WITHCOORD", "WITHDIST");

        assert_args!(opts().limit(50), "COUNT", "50");

        assert_args!(
            opts().limit(50).store("x"),
            "COUNT", "50", "STORE", "x");

        assert_args!(
            opts().limit(100).store_dist("y"),
            "COUNT", "100", "STOREDIST", "y");

        assert_args!(
            opts().order(RadiusOrder::Asc).limit(10).with_dist(),
            "WITHDIST", "COUNT", "10", "ASC");

    }
}
<|endoftext|>


<commit_before>//! A macro for defining `#[cfg]` if-else statements.
//!
//! The macro provided by this crate, `cfg_if`, is similar to the `if/elif` C
//! preprocessor macro by allowing definition of a cascade of `#[cfg]` cases,
//! emitting the implementation which matches first.
//!
//! This allows you to conveniently provide a long list `#[cfg]`'d blocks of code
//! without having to rewrite each clause multiple times.
//!
//! # Example
//!
//! ```
//! cfg_if::cfg_if! {
//!     if #[cfg(unix)] {
//!         fn foo() { /* unix specific functionality */ }
//!     } else if #[cfg(target_pointer_width = "32")] {
//!         fn foo() { /* non-unix, 32-bit functionality */ }
//!     } else {
//!         fn foo() { /* fallback implementation */ }
//!     }
//! }
//!
//! # fn main() {}
//! ```

#![no_std]
#![doc(html_root_url = "https://docs.rs/cfg-if")]
#![deny(missing_docs)]
#![cfg_attr(test, deny(warnings))]

/// The main macro provided by this crate. See crate documentation for more
/// information.
#[macro_export]
macro_rules! cfg_if {
    // match if/else chains with a final `else`
    ($(
        if #[cfg($($meta:meta),*)] { $($it:item)* }
    ) else * else {
        $($it2:item)*
    }) => {
        $crate::cfg_if! {
            @__items
            () ;
            $( ( ($($meta),*) ($($it)*) ), )*
            ( () ($($it2)*) ),
        }
    };

    // match if/else chains lacking a final `else`
    (
        if #[cfg($($i_met:meta),*)] { $($i_it:item)* }
        $(
            else if #[cfg($($e_met:meta),*)] { $($e_it:item)* }
        )*
    ) => {
        $crate::cfg_if! {
            @__items
            () ;
            ( ($($i_met),*) ($($i_it)*) ),
            $( ( ($($e_met),*) ($($e_it)*) ), )*
            ( () () ),
        }
    };

    // Internal and recursive macro to emit all the items
    //
    // Collects all the negated cfgs in a list at the beginning and after the
    // semicolon is all the remaining items
    (@__items ($($not:meta,)*) ; ) => {};
    (@__items ($($not:meta,)*) ; ( ($($m:meta),*) ($($it:item)*) ), $($rest:tt)*) => {
        // Emit all items within one block, applying an approprate #[cfg]. The
        // #[cfg] will require all `$m` matchers specified and must also negate
        // all previous matchers.
        $crate::cfg_if! { @__apply cfg(all($($m,)* not(any($($not),*)))), $($it)* }

        // Recurse to emit all other items in `$rest`, and when we do so add all
        // our `$m` matchers to the list of `$not` matchers as future emissions
        // will have to negate everything we just matched as well.
        $crate::cfg_if! { @__items ($($not,)* $($m,)*) ; $($rest)* }
    };

    // Internal macro to Apply a cfg attribute to a list of items
    (@__apply $m:meta, $($it:item)*) => {
        $(#[$m] $it)*
    };
}

#[cfg(test)]
mod tests {
    cfg_if! {
        if #[cfg(test)] {
            use core::option::Option as Option2;
            fn works1() -> Option2<u32> { Some(1) }
        } else {
            fn works1() -> Option<u32> { None }
        }
    }

    cfg_if! {
        if #[cfg(foo)] {
            fn works2() -> bool { false }
        } else if #[cfg(test)] {
            fn works2() -> bool { true }
        } else {
            fn works2() -> bool { false }
        }
    }

    cfg_if! {
        if #[cfg(foo)] {
            fn works3() -> bool { false }
        } else {
            fn works3() -> bool { true }
        }
    }

    cfg_if! {
        if #[cfg(test)] {
            use core::option::Option as Option3;
            fn works4() -> Option3<u32> { Some(1) }
        }
    }

    cfg_if! {
        if #[cfg(foo)] {
            fn works5() -> bool { false }
        } else if #[cfg(test)] {
            fn works5() -> bool { true }
        }
    }

    #[test]
    fn it_works() {
        assert!(works1().is_some());
        assert!(works2());
        assert!(works3());
        assert!(works4().is_some());
        assert!(works5());
    }
}
<commit_msg>convert to token trees<commit_after>//! A macro for defining `#[cfg]` if-else statements.
//!
//! The macro provided by this crate, `cfg_if`, is similar to the `if/elif` C
//! preprocessor macro by allowing definition of a cascade of `#[cfg]` cases,
//! emitting the implementation which matches first.
//!
//! This allows you to conveniently provide a long list `#[cfg]`'d blocks of code
//! without having to rewrite each clause multiple times.
//!
//! # Example
//!
//! ```
//! cfg_if::cfg_if! {
//!     if #[cfg(unix)] {
//!         fn foo() { /* unix specific functionality */ }
//!     } else if #[cfg(target_pointer_width = "32")] {
//!         fn foo() { /* non-unix, 32-bit functionality */ }
//!     } else {
//!         fn foo() { /* fallback implementation */ }
//!     }
//! }
//!
//! # fn main() {}
//! ```

#![no_std]
#![doc(html_root_url = "https://docs.rs/cfg-if")]
#![deny(missing_docs)]
#![cfg_attr(test, deny(warnings))]

/// The main macro provided by this crate. See crate documentation for more
/// information.
#[macro_export]
macro_rules! cfg_if {
    // match if/else chains with a final `else`
    ($(
        if #[cfg($($meta:meta),*)] { $($tokens:tt)* }
    ) else * else {
        $($tokens2:tt)*
    }) => {
        $crate::cfg_if! {
            @__items
            () ;
            $( ( ($($meta),*) ($($tokens)*) ), )*
            ( () ($($tokens2)*) ),
        }
    };

    // match if/else chains lacking a final `else`
    (
        if #[cfg($($i_met:meta),*)] { $($i_tokens:tt)* }
        $(
            else if #[cfg($($e_met:meta),*)] { $($e_tokens:tt)* }
        )*
    ) => {
        $crate::cfg_if! {
            @__items
            () ;
            ( ($($i_met),*) ($($i_tokens)*) ),
            $( ( ($($e_met),*) ($($e_tokens)*) ), )*
            ( () () ),
        }
    };

    // Internal and recursive macro to emit all the items
    //
    // Collects all the negated cfgs in a list at the beginning and after the
    // semicolon is all the remaining items
    (@__items ($($not:meta,)*) ; ) => {};
    (@__items ($($not:meta,)*) ; ( ($($m:meta),*) ($($tokens:tt)*) ), $($rest:tt)*) => {
        // Emit all items within one block, applying an appropriate #[cfg]. The
        // #[cfg] will require all `$m` matchers specified and must also negate
        // all previous matchers.
        $crate::cfg_if! { @__apply cfg(all($($m,)* not(any($($not),*)))), $($tokens)* }

        // Recurse to emit all other items in `$rest`, and when we do so add all
        // our `$m` matchers to the list of `$not` matchers as future emissions
        // will have to negate everything we just matched as well.
        $crate::cfg_if! { @__items ($($not,)* $($m,)*) ; $($rest)* }
    };

    // Internal macro to Apply a cfg attribute to a list of items
    (@__apply $m:meta, $($tokens:tt)*) => {
        #[$m] $crate::cfg_if! { @__identity $($tokens)* }
    };

    // Internal macro to Apply a cfg attribute to a list of items
    (@__identity $($tokens:tt)*) => {
        $($tokens)*
    };
}

#[cfg(test)]
mod tests {
    cfg_if! {
        if #[cfg(test)] {
            use core::option::Option as Option2;
            fn works1() -> Option2<u32> { Some(1) }
        } else {
            fn works1() -> Option<u32> { None }
        }
    }

    cfg_if! {
        if #[cfg(foo)] {
            fn works2() -> bool { false }
        } else if #[cfg(test)] {
            fn works2() -> bool { true }
        } else {
            fn works2() -> bool { false }
        }
    }

    cfg_if! {
        if #[cfg(foo)] {
            fn works3() -> bool { false }
        } else {
            fn works3() -> bool { true }
        }
    }

    cfg_if! {
        if #[cfg(test)] {
            use core::option::Option as Option3;
            fn works4() -> Option3<u32> { Some(1) }
        }
    }

    cfg_if! {
        if #[cfg(foo)] {
            fn works5() -> bool { false }
        } else if #[cfg(test)] {
            fn works5() -> bool { true }
        }
    }

    #[test]
    fn it_works() {
        assert!(works1().is_some());
        assert!(works2());
        assert!(works3());
        assert!(works4().is_some());
        assert!(works5());
    }

    extern crate std;
    #[test]
    fn test_usage_within_a_function() {
        cfg_if! {if #[cfg(debug_assertions)] {
            std::println!("debug");
            std::println!("debug2");
        } else {
            std::println!("release");
            std::println!("release2");
        }}
    }
}
<|endoftext|>


<commit_before>//! A minimal implementation of SHA1 for rust.
//!
//! Example:
//!
//! ```rust
//! extern crate sha1;
//! # fn main() {
//!
//! let mut m = sha1::Sha1::new();
//! m.update("Hello World!".as_bytes());
//! assert_eq!(m.hexdigest().as_slice(),
//!            "2ef7bde608ce5404e97d5f042f95f89f1c232871");
//! # }
//! ```

#![feature(slicing_syntax)]

#![experimental]

use std::io::BufWriter;

/// Represents a Sha1 hash object in memory.
#[derive(Clone)]
pub struct Sha1 {
    state: [u32; 5],
    data: Vec<u8>,
    len: u64,
}

const DEFAULT_STATE : [u32; 5] =
    [0x67452301, 0xefcdab89, 0x98badcfe, 0x10325476, 0xc3d2e1f0];

fn to_hex(input: &[u8]) -> String {
    let mut s = String::new();
    for b in input.iter() {
        s.push_str(format!("{:02x}", *b)[]);
    }
    return s;
}

impl Sha1 {

    /// Creates an fresh sha1 hash object.
    pub fn new() -> Sha1 {
        Sha1 {
            state: DEFAULT_STATE,
            data: Vec::new(),
            len: 0,
        }
    }

    fn process_block(&mut self, block: &[u8]) {
        assert_eq!(block.len(), 64);

        let mut words = [0u32; 80];
        for (i, chunk) in block.chunks(4).enumerate() {
            words[i] = (chunk[3] as u32) |
                       ((chunk[2] as u32) << 8) |
                       ((chunk[1] as u32) << 16) |
                       ((chunk[0] as u32) << 24);
        }

        fn ff(b: u32, c: u32, d: u32) -> u32 { d ^ (b & (c ^ d)) }
        fn gg(b: u32, c: u32, d: u32) -> u32 { b ^ c ^ d }
        fn hh(b: u32, c: u32, d: u32) -> u32 { (b & c) | (d & (b | c)) }
        fn ii(b: u32, c: u32, d: u32) -> u32 { b ^ c ^ d }

        fn left_rotate(x: u32, n: u32) -> u32 { (x << n as uint) | (x >> (32 - n) as uint) }

        for i in range(16u, 80u) {
            let n = words[i - 3] ^ words[i - 8] ^ words[i - 14] ^ words[i - 16];
            words[i] = left_rotate(n, 1);
        }

        let mut a = self.state[0];
        let mut b = self.state[1];
        let mut c = self.state[2];
        let mut d = self.state[3];
        let mut e = self.state[4];

        for i in range(0u, 80u) {
            let (f, k) = match i {
                0 ... 19 => (ff(b, c, d), 0x5a827999),
                20 ... 39 => (gg(b, c, d), 0x6ed9eba1),
                40 ... 59 => (hh(b, c, d), 0x8f1bbcdc),
                60 ... 79 => (ii(b, c, d), 0xca62c1d6),
                _ => (0, 0),
            };

            let tmp = left_rotate(a, 5) + f + e + k + words[i];
            e = d;
            d = c;
            c = left_rotate(b, 30);
            b = a;
            a = tmp;
        }

        self.state[0] += a;
        self.state[1] += b;
        self.state[2] += c;
        self.state[3] += d;
        self.state[4] += e;
    }

    /// Resets the hash object to it's initial state.
    pub fn reset(&mut self) {
        self.state = DEFAULT_STATE;
        self.data.clear();
        self.len = 0;
    }

    /// Update hash with input data.
    pub fn update(&mut self, data: &[u8]) {
        let mut d = self.data.clone();
        self.data.clear();

        d.push_all(data);

        for chunk in d[].chunks(64) {
            self.len += chunk.len() as u64;

            if chunk.len() == 64 {
                self.process_block(chunk);
            } else {
                self.data.push_all(chunk);
            }
        }
    }

    /// Retrieve digest result.  The output must be large enough to
    /// contain result (20 bytes).
    pub fn output(&self, out: &mut [u8]) {
        // these are unlikely to fail, since we're writing to memory
        #![allow(unused_must_use)]

        let mut m = Sha1 {
            state: self.state,
            data: Vec::new(),
            len: 0,
        };

        let mut w = Vec::<u8>::new();
        w.write(self.data[]);
        w.write_u8(0x80 as u8);
        let padding = (((56 - self.len as int - 1) % 64) + 64) % 64;
        for _ in range(0, padding) {
            w.write_u8(0u8);
        }
        w.write_be_u64(self.len * 8);
        for chunk in w[].chunks(64) {
            m.process_block(chunk);
        }

        let mut w = BufWriter::new(out);
        for &n in m.state.iter() {
            w.write_be_u32(n);
        }
    }

    /// Shortcut for getting `output` into a new vector.
    pub fn digest(&self) -> Vec<u8> {
        let mut buf = [0u8; 20].to_vec();
        self.output(&mut *buf[]);
        buf
    }

    /// Shortcut for getting a hex output of the vector.
    pub fn hexdigest(&self) -> String {
        to_hex(self.digest()[])
    }
}

#[test]
fn test_simple() {
    let mut m = Sha1::new();

    let tests = [
        ("The quick brown fox jumps over the lazy dog",
         "2fd4e1c67a2d28fced849ee1bb76e7391b93eb12"),
        ("The quick brown fox jumps over the lazy cog",
         "de9f2c7fd25e1b3afad3e85a0bd17d9b100db4b3"),
        ("", "da39a3ee5e6b4b0d3255bfef95601890afd80709"),
        ("testing\n", "9801739daae44ec5293d4e1f53d3f4d2d426d91c"),
        ("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
         "025ecbd5d70f8fb3c5457cd96bab13fda305dc59"),
    ];

    for &(s, ref h) in tests.iter() {
        let data = s.as_bytes();

        m.reset();
        m.update(data);
        let hh = m.hexdigest();

        assert_eq!(hh.len(), h.len());
        assert_eq!(hh[], *h);
    }
}
<commit_msg>fix invalid message length calculation<commit_after>//! A minimal implementation of SHA1 for rust.
//!
//! Example:
//!
//! ```rust
//! extern crate sha1;
//! # fn main() {
//!
//! let mut m = sha1::Sha1::new();
//! m.update("Hello World!".as_bytes());
//! assert_eq!(m.hexdigest().as_slice(),
//!            "2ef7bde608ce5404e97d5f042f95f89f1c232871");
//! # }
//! ```

#![feature(slicing_syntax)]

#![experimental]

#[cfg(test)] extern crate test;

use std::io::BufWriter;

/// Represents a Sha1 hash object in memory.
#[derive(Clone)]
pub struct Sha1 {
    state: [u32; 5],
    data: Vec<u8>,
    len: u64,
}

const DEFAULT_STATE : [u32; 5] =
    [0x67452301, 0xefcdab89, 0x98badcfe, 0x10325476, 0xc3d2e1f0];

fn to_hex(input: &[u8]) -> String {
    let mut s = String::new();
    for b in input.iter() {
        s.push_str(format!("{:02x}", *b)[]);
    }
    return s;
}

impl Sha1 {

    /// Creates an fresh sha1 hash object.
    pub fn new() -> Sha1 {
        Sha1 {
            state: DEFAULT_STATE,
            data: Vec::new(),
            len: 0,
        }
    }

    fn process_block(&mut self, block: &[u8]) {
        assert_eq!(block.len(), 64);

        let mut words = [0u32; 80];
        for (i, chunk) in block.chunks(4).enumerate() {
            words[i] = (chunk[3] as u32) |
                       ((chunk[2] as u32) << 8) |
                       ((chunk[1] as u32) << 16) |
                       ((chunk[0] as u32) << 24);
        }

        fn ff(b: u32, c: u32, d: u32) -> u32 { d ^ (b & (c ^ d)) }
        fn gg(b: u32, c: u32, d: u32) -> u32 { b ^ c ^ d }
        fn hh(b: u32, c: u32, d: u32) -> u32 { (b & c) | (d & (b | c)) }
        fn ii(b: u32, c: u32, d: u32) -> u32 { b ^ c ^ d }

        fn left_rotate(x: u32, n: u32) -> u32 { (x << n as uint) | (x >> (32 - n) as uint) }

        for i in range(16u, 80u) {
            let n = words[i - 3] ^ words[i - 8] ^ words[i - 14] ^ words[i - 16];
            words[i] = left_rotate(n, 1);
        }

        let mut a = self.state[0];
        let mut b = self.state[1];
        let mut c = self.state[2];
        let mut d = self.state[3];
        let mut e = self.state[4];

        for i in range(0u, 80u) {
            let (f, k) = match i {
                0 ... 19 => (ff(b, c, d), 0x5a827999),
                20 ... 39 => (gg(b, c, d), 0x6ed9eba1),
                40 ... 59 => (hh(b, c, d), 0x8f1bbcdc),
                60 ... 79 => (ii(b, c, d), 0xca62c1d6),
                _ => (0, 0),
            };

            let tmp = left_rotate(a, 5) + f + e + k + words[i];
            e = d;
            d = c;
            c = left_rotate(b, 30);
            b = a;
            a = tmp;
        }

        self.state[0] += a;
        self.state[1] += b;
        self.state[2] += c;
        self.state[3] += d;
        self.state[4] += e;
    }

    /// Resets the hash object to it's initial state.
    pub fn reset(&mut self) {
        self.state = DEFAULT_STATE;
        self.data.clear();
        self.len = 0;
    }

    /// Update hash with input data.
    pub fn update(&mut self, data: &[u8]) {
        let mut d = self.data.clone();
        self.data.clear();

        d.push_all(data);

        for chunk in d[].chunks(64) {
            if chunk.len() == 64 {
                self.len += 64;
                self.process_block(chunk);
            } else {
                self.data.push_all(chunk);
            }
        }
    }

    /// Retrieve digest result.  The output must be large enough to
    /// contain result (20 bytes).
    pub fn output(&self, out: &mut [u8]) {
        // these are unlikely to fail, since we're writing to memory
        #![allow(unused_must_use)]

        let mut m = Sha1 {
            state: self.state.clone(),
            data: Vec::new(),
            len: 0,
        };

        let mut w = Vec::<u8>::new();
        w.write(self.data[]);
        w.write_u8(0x80u8);
        let padding = 64 - ((self.data.len() + 9) % 64);
        for _ in range(0, padding) {
            w.write_u8(0u8);
        }
        w.write_be_u64((self.data.len() as u64 + self.len) * 8);
        for chunk in w[].chunks(64) {
            m.process_block(chunk);
        }

        let mut w = BufWriter::new(out);
        for &n in m.state.iter() {
            w.write_be_u32(n);
        }
    }

    /// Shortcut for getting `output` into a new vector.
    pub fn digest(&self) -> Vec<u8> {
        let mut buf = [0u8; 20].to_vec();
        self.output(&mut *buf[]);
        buf
    }

    /// Shortcut for getting a hex output of the vector.
    pub fn hexdigest(&self) -> String {
        to_hex(self.digest()[])
    }
}

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

    #[test]
    fn test_simple() {
        let mut m = Sha1::new();

        let tests = [
            ("The quick brown fox jumps over the lazy dog",
            "2fd4e1c67a2d28fced849ee1bb76e7391b93eb12"),
            ("The quick brown fox jumps over the lazy cog",
            "de9f2c7fd25e1b3afad3e85a0bd17d9b100db4b3"),
            ("", "da39a3ee5e6b4b0d3255bfef95601890afd80709"),
            ("testing\n", "9801739daae44ec5293d4e1f53d3f4d2d426d91c"),
            ("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
            "025ecbd5d70f8fb3c5457cd96bab13fda305dc59"),
        ];

        for &(s, ref h) in tests.iter() {
            let data = s.as_bytes();

            m.reset();
            m.update(data);
            let hh = m.hexdigest();

            assert_eq!(hh.len(), h.len());
            assert_eq!(hh[], *h);
        }
    }

    #[bench]
    fn test_sha1_speed(b: &mut Bencher) {
        let mut m = Sha1::new();
        let s = "The quick brown fox jumps over the lazy dog.";
        let n = 1000u64;

        b.bytes = n * s.len() as u64;
        b.iter(|| {
            m.reset();
            for _ in range(0, n) {
                m.update(s.as_bytes());
            }
            assert_eq!(m.hexdigest()[], "7ca27655f67fceaa78ed2e645a81c7f1d6e249d2");
        });
    }
}
<|endoftext|>