<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,
            _ => 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>Add 3A+ revision code<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[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
    }
}
<|endoftext|>


<commit_before>extern crate serialize;

use serialize::{Decodable,json};
use std::io::Process;
use std::io::fs::File;
use std::io::process::ProcessOutput;
use std::io::{fs,IoResult,Truncate,UserRWX,Write};
use std::str;

#[deriving(Decodable,Show)]
struct Example {
    id: ~str,
    title: ~str,
}

fn read(file: &Path) -> IoResult<~str> {
    let mut file = try!(File::open(file));

    Ok(try!(file.read_to_str()))
}

fn write(file: &Path, contents: &str) -> IoResult<()> {
    let mut file = try!(File::open_mode(file, Truncate, Write));

    Ok(try!(file.write_str(contents)))
}

fn update(example: &Example) -> bool {
    print!("{}... ", example.id);

    let src_dir = Path::new(format!("src/{}", example.id));
    let out_dir = Path::new(format!("output/examples/{}", example.id));

    let sources: Vec<Path> = match fs::readdir(&src_dir) {
        Err(_) => {
            println!("couldn't read {}", src_dir.display());
            return false
        } Ok(contents) => contents.move_iter().filter(|path| match path.extension_str() {
            None => false,
            Some(extension) => extension == "rs",
        }).collect(),
    };

    let mut template = match read(&src_dir.join("input.md")) {
        Err(_) => {
            println!("couldn't read input.md");
            return false
        } Ok(contents) => format!("\\# {}\n\n{}", example.title, contents),
    };

    // insert source code in markdown
    for source in sources.iter() {
        let filename = source.filename_str().unwrap();
        let code = match read(source) {
            Err(_) => {
                println!("couldn't read {}", filename);
                return false;
            }, Ok(contents) => format!("``` rust\n// {}\n{}```", filename, contents),
        };

        template = template.replace(format!("\\{{}\\}", filename), code);
    }

    // insert program output in markdown
    for source in sources.iter() {
        let filename = source.filename_str().unwrap();
        let token = format!("\\{{}\\}", source.with_extension("out").filename_str().unwrap());

        if template.contains(token) {
            match compile_run(source) {
                Err(action) => {
                    println!("couldn't {} {}", action, filename);
                    return false;
                }, Ok(output) => {
                    let output = format!("```\n$ rustc {} && ./{}\n{}```",
                                         filename,
                                         filename.split('.').nth(0).unwrap(),
                                         output);

                    template = template.replace(token, output);
                }
            }
        }
    }

    if !out_dir.exists() {
        fs::mkdir(&out_dir, UserRWX).unwrap();
    }

    match write(&out_dir.join("README.md"), template) {
        Err(_) => {
            println!("couldn't write README.md");
            false
        } Ok(_) => {
            println!("DONE");
            true
        }
    }
}

fn compile_run(path: &Path) -> Result<~str, &'static str> {
    match Process::output("rustc", [path.as_str().unwrap().to_owned(),
                                    "-o".to_owned(),
                                    "executable".to_owned()]) {
        Err(_) => return Err("compile"),
        Ok(out) => {
            if !out.status.success() {
                return Ok(str::from_utf8_owned(out.error.move_iter().collect::<~[u8]>()).unwrap());
            }
        }
    }

    match Process::output("./executable", []) {
        Err(_) => Err("run"),
        Ok(out) => {
            fs::unlink(&Path::new("./executable")).unwrap();
            let ProcessOutput { status: _, output: out, error: err } = out;
            let stdout = str::from_utf8_owned(out.move_iter().collect::<~[u8]>()).unwrap();
            let stderr = str::from_utf8_owned(err.move_iter().collect::<~[u8]>()).unwrap();

            Ok(vec!(stdout, stderr).concat())
        }
    }
}

fn main() {
    let src_dir = Path::new("src");

    let examples: Vec<Example> = match read(&src_dir.join("order.json")) {
        Err(err) => fail!("couldn't read order.json: {}", err),
        Ok(string) => match json::from_str(string) {
            Err(err) => fail!("order.json is invalid json: {}", err),
            Ok(json) => match Decodable::decode(&mut json::Decoder::new(json)) {
                Err(err) => fail!("error decoding order.json: {}", err),
                Ok(data) => data,
            }
        }
    };

    let mut indent = false;
    let summary = examples.move_iter().map(|example| {
        if example.id.as_slice() == "staging" {
            indent = true;
        }

        let chapter = if update(&example) {
            format!("* [{}](examples/{}/README.md)", example.title, example.id)
        } else {
            format!("* {}", example.title)
        };

        if indent &&
            example.id.as_slice() != "staging" &&
            example.id.as_slice() != "todo" {
            format!("  {}", chapter)
        } else {
            chapter
        }
    }).collect::<Vec<~str>>().connect("\n");

    if write(&Path::new("output/SUMMARY.md"), summary).is_err() {
        fail!("failed to write SUMMARY.md");
    }
}
<commit_msg>update.rs: update to nightly Rust<commit_after>extern crate serialize;

use serialize::{Decodable,json};
use std::io::Process;
use std::io::fs::File;
use std::io::process::ProcessOutput;
use std::io::{fs,IoResult,Truncate,UserRWX,Write};
use std::str;

#[deriving(Decodable,Show)]
struct Example {
    id: ~str,
    title: ~str,
}

fn read(file: &Path) -> IoResult<~str> {
    let mut file = try!(File::open(file));

    Ok(try!(file.read_to_str()))
}

fn write(file: &Path, contents: &str) -> IoResult<()> {
    let mut file = try!(File::open_mode(file, Truncate, Write));

    Ok(try!(file.write_str(contents)))
}

fn update(example: &Example) -> bool {
    print!("{}... ", example.id);

    let src_dir = Path::new(format!("src/{}", example.id));
    let out_dir = Path::new(format!("output/examples/{}", example.id));

    let sources: Vec<Path> = match fs::readdir(&src_dir) {
        Err(_) => {
            println!("couldn't read {}", src_dir.display());
            return false
        } Ok(contents) => contents.move_iter().filter(|path| match path.extension_str() {
            None => false,
            Some(extension) => extension == "rs",
        }).collect(),
    };

    let mut template = match read(&src_dir.join("input.md")) {
        Err(_) => {
            println!("couldn't read input.md");
            return false
        } Ok(contents) => format!("\\# {}\n\n{}", example.title, contents),
    };

    // insert source code in markdown
    for source in sources.iter() {
        let filename = source.filename_str().unwrap();
        let code = match read(source) {
            Err(_) => {
                println!("couldn't read {}", filename);
                return false;
            }, Ok(contents) => format!("``` rust\n// {}\n{}```", filename, contents),
        };

        template = template.replace(format!("\\{{}\\}", filename), code);
    }

    // insert program output in markdown
    for source in sources.iter() {
        let filename = source.filename_str().unwrap();
        let token = format!("\\{{}\\}", source.with_extension("out").filename_str().unwrap());

        if template.contains(token) {
            match compile_run(source) {
                Err(action) => {
                    println!("couldn't {} {}", action, filename);
                    return false;
                }, Ok(output) => {
                    let output = format!("```\n$ rustc {} && ./{}\n{}```",
                                         filename,
                                         filename.split('.').nth(0).unwrap(),
                                         output);

                    template = template.replace(token, output);
                }
            }
        }
    }

    if !out_dir.exists() {
        fs::mkdir(&out_dir, UserRWX).unwrap();
    }

    match write(&out_dir.join("README.md"), template) {
        Err(_) => {
            println!("couldn't write README.md");
            false
        } Ok(_) => {
            println!("DONE");
            true
        }
    }
}

fn compile_run(path: &Path) -> Result<~str, &'static str> {
    match Process::output("rustc", [path.as_str().unwrap().to_owned(),
                                    "-o".to_owned(),
                                    "executable".to_owned()]) {
        Err(_) => return Err("compile"),
        Ok(out) => {
            if !out.status.success() {
                return Ok(str::from_utf8_owned(out.error.as_slice().to_owned()).unwrap());
            }
        }
    }

    match Process::output("./executable", []) {
        Err(_) => Err("run"),
        Ok(out) => {
            fs::unlink(&Path::new("./executable")).unwrap();
            let ProcessOutput { status: _, output: out, error: err } = out;
            let stdout = str::from_utf8_owned(out.as_slice().to_owned()).unwrap();
            let stderr = str::from_utf8_owned(err.as_slice().to_owned()).unwrap();

            Ok(vec!(stdout, stderr).concat())
        }
    }
}

fn main() {
    let src_dir = Path::new("src");

    let examples: Vec<Example> = match read(&src_dir.join("order.json")) {
        Err(err) => fail!("couldn't read order.json: {}", err),
        Ok(string) => match json::from_str(string) {
            Err(err) => fail!("order.json is invalid json: {}", err),
            Ok(json) => match Decodable::decode(&mut json::Decoder::new(json)) {
                Err(err) => fail!("error decoding order.json: {}", err),
                Ok(data) => data,
            }
        }
    };

    let mut indent = false;
    let summary = examples.move_iter().map(|example| {
        if example.id.as_slice() == "staging" {
            indent = true;
        }

        let chapter = if update(&example) {
            format!("* [{}](examples/{}/README.md)", example.title, example.id)
        } else {
            format!("* {}", example.title)
        };

        if indent &&
            example.id.as_slice() != "staging" &&
            example.id.as_slice() != "todo" {
            format!("  {}", chapter)
        } else {
            chapter
        }
    }).collect::<Vec<~str>>().connect("\n");

    if write(&Path::new("output/SUMMARY.md"), summary).is_err() {
        fail!("failed to write SUMMARY.md");
    }
}
<|endoftext|>


<commit_before>use utils::address::Address;
use utils::gas::Gas;

use vm::errors::MachineError;
use vm::{Memory, Machine, MachineStatus};

use std::str::FromStr;

use tiny_keccak::Keccak;
use sha2::Sha256;
use ripemd160::Ripemd160;
use secp256k1::{Secp256k1, RecoverableSignature, Message, RecoveryId, Error};
use digest::{Input, FixedOutput};

pub fn gas_div_ceil(a: Gas, b: Gas) -> Gas {
    if a % b == Gas::zero() {
        a / b
    } else {
        a / b + Gas::from(1u64)
    }
}

pub fn is_precompiled(address: Address) -> bool {
    let ecrec_address = Address::from_str("0x0000000000000000000000000000000000000001").unwrap();
    let sha256_address = Address::from_str("0x0000000000000000000000000000000000000002").unwrap();
    let rip160_address = Address::from_str("0x0000000000000000000000000000000000000003").unwrap();
    let id_address = Address::from_str("0x0000000000000000000000000000000000000004").unwrap();

    address == ecrec_address ||
        address == sha256_address ||
        address == rip160_address ||
        address == id_address
}

impl<M: Memory + Default> Machine<M> {
    #[allow(unused_variables)]
    pub fn step_precompiled(&mut self) -> bool {
        let ecrec_address = Address::from_str("0x0000000000000000000000000000000000000001").unwrap();
        let sha256_address = Address::from_str("0x0000000000000000000000000000000000000002").unwrap();
        let rip160_address = Address::from_str("0x0000000000000000000000000000000000000003").unwrap();
        let id_address = Address::from_str("0x0000000000000000000000000000000000000004").unwrap();

        if self.state.context.address == id_address {
            self.step_precompiled_id();
            return true;
        } else if self.state.context.address == rip160_address {
            self.step_precompiled_rip160();
            return true;
        } else if self.state.context.address == sha256_address {
            self.step_precompiled_sha256();
            return true;
        } else if self.state.context.address == ecrec_address {
            self.step_precompiled_ecrec();
            return true;
        } else {
            return false;
        }
    }

    fn step_precompiled_id(&mut self) {
        let gas = Gas::from(15u64) +
            Gas::from(3u64) * gas_div_ceil(Gas::from(self.state.context.data.len()),
                                           Gas::from(32u64));
        if gas > self.state.context.gas_limit {
            self.state.used_gas = self.state.context.gas_limit;
            self.status = MachineStatus::ExitedErr(MachineError::EmptyGas);
        } else {
            self.state.used_gas = gas;
            self.state.out = self.state.context.data.clone();
            self.status = MachineStatus::ExitedOk;
        }
    }

    fn step_precompiled_rip160(&mut self) {
        let gas = Gas::from(600u64) +
            Gas::from(120u64) * gas_div_ceil(Gas::from(self.state.context.data.len()),
                                             Gas::from(32u64));
        if gas > self.state.context.gas_limit {
            self.state.used_gas = self.state.context.gas_limit;
            self.status = MachineStatus::ExitedErr(MachineError::EmptyGas);
        } else {
            self.state.used_gas = gas;
            let mut ripemd = Ripemd160::default();
            ripemd.digest(&self.state.context.data);
            let fixed = ripemd.fixed_result();
            let mut result: [u8; 32] = [0u8; 32];
            for i in 0..20 {
                result[i + 12] = fixed[i];
            }
            self.state.out = result.as_ref().into();
            self.status = MachineStatus::ExitedOk;
        }
    }

    fn step_precompiled_sha256(&mut self) {
        let gas = Gas::from(60u64) +
            Gas::from(12u64) * gas_div_ceil(Gas::from(self.state.context.data.len()),
                                            Gas::from(32u64));
        if gas > self.state.context.gas_limit {
            self.state.used_gas = self.state.context.gas_limit;
            self.status = MachineStatus::ExitedErr(MachineError::EmptyGas);
        } else {
            self.state.used_gas = gas;
            let mut sha2 = Sha256::default();
            sha2.digest(&self.state.context.data);
            let fixed = sha2.fixed_result();
            let mut result: [u8; 32] = [0u8; 32];
            for i in 0..32 {
                result[i] = fixed[i];
            }
            self.state.out = result.as_ref().into();
            self.status = MachineStatus::ExitedOk;
        }
    }

    fn step_precompiled_ecrec(&mut self) {
        let gas = Gas::from(3000u64);
        if gas > self.state.context.gas_limit {
            self.state.used_gas = self.state.context.gas_limit;
            self.status = MachineStatus::ExitedErr(MachineError::EmptyGas);
        } else {
            self.state.used_gas = gas;
            let mut data = [0u8; 128];
            for i in 0..self.state.context.data.len() {
                data[i] = self.state.context.data[i];
            }
            match kececrec(&data) {
                Ok(mut ret) => {
                    for i in 0..12 {
                        ret[i] = 0u8;
                    }
                    self.state.out = ret.as_ref().into();
                },
                Err(_) => (),
            }
            self.status = MachineStatus::ExitedOk;
        }
    }
}

fn kececrec(data: &[u8]) -> Result<[u8; 32], Error> {
    let ecrec = Secp256k1::new();
    let message = Message::from_slice(&data[0..32])?;
    let recid = RecoveryId::from_i32(data[63] as i32)?;
    let sig = RecoverableSignature::from_compact(&ecrec, &data[64..128], recid)?;
    let recovered = ecrec.recover(&message, &sig)?;
    let key = recovered.serialize_vec(&ecrec, false);
    let mut kec = Keccak::new_keccak256();
    kec.update(&key);
    let mut ret = [0u8; 32];
    kec.finalize(&mut ret);
    Ok(ret)
}
<commit_msg>Problem: gas_div_ceil should not be pub<commit_after>use utils::address::Address;
use utils::gas::Gas;

use vm::errors::MachineError;
use vm::{Memory, Machine, MachineStatus};

use std::str::FromStr;

use tiny_keccak::Keccak;
use sha2::Sha256;
use ripemd160::Ripemd160;
use secp256k1::{Secp256k1, RecoverableSignature, Message, RecoveryId, Error};
use digest::{Input, FixedOutput};

fn gas_div_ceil(a: Gas, b: Gas) -> Gas {
    if a % b == Gas::zero() {
        a / b
    } else {
        a / b + Gas::from(1u64)
    }
}

pub fn is_precompiled(address: Address) -> bool {
    let ecrec_address = Address::from_str("0x0000000000000000000000000000000000000001").unwrap();
    let sha256_address = Address::from_str("0x0000000000000000000000000000000000000002").unwrap();
    let rip160_address = Address::from_str("0x0000000000000000000000000000000000000003").unwrap();
    let id_address = Address::from_str("0x0000000000000000000000000000000000000004").unwrap();

    address == ecrec_address ||
        address == sha256_address ||
        address == rip160_address ||
        address == id_address
}

impl<M: Memory + Default> Machine<M> {
    #[allow(unused_variables)]
    pub fn step_precompiled(&mut self) -> bool {
        let ecrec_address = Address::from_str("0x0000000000000000000000000000000000000001").unwrap();
        let sha256_address = Address::from_str("0x0000000000000000000000000000000000000002").unwrap();
        let rip160_address = Address::from_str("0x0000000000000000000000000000000000000003").unwrap();
        let id_address = Address::from_str("0x0000000000000000000000000000000000000004").unwrap();

        if self.state.context.address == id_address {
            self.step_precompiled_id();
            return true;
        } else if self.state.context.address == rip160_address {
            self.step_precompiled_rip160();
            return true;
        } else if self.state.context.address == sha256_address {
            self.step_precompiled_sha256();
            return true;
        } else if self.state.context.address == ecrec_address {
            self.step_precompiled_ecrec();
            return true;
        } else {
            return false;
        }
    }

    fn step_precompiled_id(&mut self) {
        let gas = Gas::from(15u64) +
            Gas::from(3u64) * gas_div_ceil(Gas::from(self.state.context.data.len()),
                                           Gas::from(32u64));
        if gas > self.state.context.gas_limit {
            self.state.used_gas = self.state.context.gas_limit;
            self.status = MachineStatus::ExitedErr(MachineError::EmptyGas);
        } else {
            self.state.used_gas = gas;
            self.state.out = self.state.context.data.clone();
            self.status = MachineStatus::ExitedOk;
        }
    }

    fn step_precompiled_rip160(&mut self) {
        let gas = Gas::from(600u64) +
            Gas::from(120u64) * gas_div_ceil(Gas::from(self.state.context.data.len()),
                                             Gas::from(32u64));
        if gas > self.state.context.gas_limit {
            self.state.used_gas = self.state.context.gas_limit;
            self.status = MachineStatus::ExitedErr(MachineError::EmptyGas);
        } else {
            self.state.used_gas = gas;
            let mut ripemd = Ripemd160::default();
            ripemd.digest(&self.state.context.data);
            let fixed = ripemd.fixed_result();
            let mut result: [u8; 32] = [0u8; 32];
            for i in 0..20 {
                result[i + 12] = fixed[i];
            }
            self.state.out = result.as_ref().into();
            self.status = MachineStatus::ExitedOk;
        }
    }

    fn step_precompiled_sha256(&mut self) {
        let gas = Gas::from(60u64) +
            Gas::from(12u64) * gas_div_ceil(Gas::from(self.state.context.data.len()),
                                            Gas::from(32u64));
        if gas > self.state.context.gas_limit {
            self.state.used_gas = self.state.context.gas_limit;
            self.status = MachineStatus::ExitedErr(MachineError::EmptyGas);
        } else {
            self.state.used_gas = gas;
            let mut sha2 = Sha256::default();
            sha2.digest(&self.state.context.data);
            let fixed = sha2.fixed_result();
            let mut result: [u8; 32] = [0u8; 32];
            for i in 0..32 {
                result[i] = fixed[i];
            }
            self.state.out = result.as_ref().into();
            self.status = MachineStatus::ExitedOk;
        }
    }

    fn step_precompiled_ecrec(&mut self) {
        let gas = Gas::from(3000u64);
        if gas > self.state.context.gas_limit {
            self.state.used_gas = self.state.context.gas_limit;
            self.status = MachineStatus::ExitedErr(MachineError::EmptyGas);
        } else {
            self.state.used_gas = gas;
            let mut data = [0u8; 128];
            for i in 0..self.state.context.data.len() {
                data[i] = self.state.context.data[i];
            }
            match kececrec(&data) {
                Ok(mut ret) => {
                    for i in 0..12 {
                        ret[i] = 0u8;
                    }
                    self.state.out = ret.as_ref().into();
                },
                Err(_) => (),
            }
            self.status = MachineStatus::ExitedOk;
        }
    }
}

fn kececrec(data: &[u8]) -> Result<[u8; 32], Error> {
    let ecrec = Secp256k1::new();
    let message = Message::from_slice(&data[0..32])?;
    let recid = RecoveryId::from_i32(data[63] as i32)?;
    let sig = RecoverableSignature::from_compact(&ecrec, &data[64..128], recid)?;
    let recovered = ecrec.recover(&message, &sig)?;
    let key = recovered.serialize_vec(&ecrec, false);
    let mut kec = Keccak::new_keccak256();
    kec.update(&key);
    let mut ret = [0u8; 32];
    kec.finalize(&mut ret);
    Ok(ret)
}
<|endoftext|>


<commit_before>extern crate libc;
extern crate std;

use super::liblua;
use super::Index;
use super::Lua;
use super::Pushable;
use super::Readable;

macro_rules! integer_impl(
    ($t:ident) => (
        impl Pushable for $t {
            fn push_to_lua(&self, lua: &Lua) {
                unsafe { liblua::lua_pushinteger(lua.lua, *self as liblua::lua_Integer) }
            }
        }
        impl Readable for $t {
            fn read_from_lua(lua: &Lua, index: i32) -> Option<$t> {
                unsafe {
                    let success: libc::c_int = std::mem::uninitialized();
                    let val = liblua::lua_tointegerx(lua.lua, index, std::mem::transmute(&success));
                    match success {
                        0 => None,
                        _ => Some(val as $t)
                    }
                }
            }
        }
        impl Index for $t {
        }
    );
)

integer_impl!(int)
integer_impl!(i8)
integer_impl!(i16)
integer_impl!(i32)
//integer_impl!(i64)   // data loss

macro_rules! unsigned_impl(
    ($t:ident) => (
        impl Pushable for $t {
            fn push_to_lua(&self, lua: &Lua) {
                unsafe { liblua::lua_pushunsigned(lua.lua, *self as liblua::lua_Unsigned) }
            }
        }
        impl Readable for $t {
            fn read_from_lua(lua: &Lua, index: i32) -> Option<$t> {
                unsafe {
                    let success: libc::c_int = std::mem::uninitialized();
                    let val = liblua::lua_tounsignedx(lua.lua, index, std::mem::transmute(&success));
                    match success {
                        0 => None,
                        _ => Some(val as $t)
                    }
                }
            }
        }
        impl Index for $t {
        }
    );
)

unsigned_impl!(uint)
unsigned_impl!(u8)
unsigned_impl!(u16)
unsigned_impl!(u32)
//unsigned_impl!(u64)   // data loss

macro_rules! numeric_impl(
    ($t:ident) => (
        impl Pushable for $t {
            fn push_to_lua(&self, lua: &Lua) {
                unsafe { liblua::lua_pushnumber(lua.lua, *self as f64) }
            }
        }
        impl Readable for $t {
            fn read_from_lua(lua: &Lua, index: i32) -> Option<$t> {
                unsafe {
                    let success: libc::c_int = std::mem::uninitialized();
                    let val = liblua::lua_tonumberx(lua.lua, index, std::mem::transmute(&success));
                    match success {
                        0 => None,
                        _ => Some(val as $t)
                    }
                }
            }
        }
        impl Index for $t {
        }
    );
)

numeric_impl!(f32)
numeric_impl!(f64)

impl Pushable for std::string::String {
    fn push_to_lua(&self, lua: &Lua) {
        unsafe { liblua::lua_pushstring(lua.lua, self.to_c_str().unwrap()) }
    }
}

impl Readable for String {
    fn read_from_lua(lua: &Lua, index: i32) -> Option<std::string::String> {
        unsafe {
            let size: libc::size_t = std::mem::uninitialized();
            let cStr = liblua::lua_tolstring(lua.lua, index, std::mem::transmute(&size));
            if cStr.is_null() {
                return None;
            }

            // TODO: check this transmute, why is from_raw_parts taking a mut ptr?
            let val = std::string::String::from_raw_parts(size as uint, size as uint, std::mem::transmute(cStr));

            Some(val)
        }
    }
}

impl Index for String {
}

impl Pushable for ~str {
    fn push_to_lua(&self, lua: &Lua) {
        String::from_str(*self).push_to_lua(lua)
    }
}

impl Pushable for &'static str {
    fn push_to_lua(&self, lua: &Lua) {
        String::from_str(*self).push_to_lua(lua)
    }
}

impl Readable for ~str {
    fn read_from_lua(lua: &Lua, index: i32) -> Option<~str> {
        unimplemented!()
    }
}

impl Readable for &'static str {
    fn read_from_lua(lua: &Lua, index: i32) -> Option<&'static str> {
        unimplemented!()
    }
}

impl Index for ~str {
}

impl Index for &'static str {
    fn lua_set_global<T: Pushable>(&self, lua: &Lua, value: T) {
        value.push_to_lua(lua);
        unsafe { liblua::lua_setglobal(lua.lua, self.to_c_str().unwrap()); }
    }

    fn lua_get_global<T: Readable>(&self, lua: &Lua) -> Option<T> {
        unsafe { liblua::lua_getglobal(lua.lua, self.to_c_str().unwrap()); }
        Readable::read_from_lua(lua, -1)
    }
}

/*impl GlobalsIndex for &'static str {
    fn lua_set_global(&self, lua: &Lua) {
        unsafe { liblua::lua_setglobal(lua.lua, self.to_c_str().unwrap()) }
    }

    fn lua_get_global(&self, lua: &Lua) {
        unsafe { liblua::lua_getglobal(lua.lua, self.to_c_str().unwrap()) }
    }
}

impl GlobalsIndex for ~str {
    fn lua_set_global(&self, lua: &Lua) {
        unsafe { liblua::lua_setglobal(lua.lua, self.to_c_str().unwrap()) }
    }

    fn lua_get_global(&self, lua: &Lua) {
        unsafe { liblua::lua_getglobal(lua.lua, self.to_c_str().unwrap()) }
    }
}*/

#[cfg(test)]
mod tests {
    #[test]
    fn readwrite_ints() {
        let mut lua = super::super::Lua::new();

        lua.set("a", 2);
        let x: int = lua.get("a").unwrap();
        assert_eq!(x, 2)
    }
}
<commit_msg>Fixed string reading<commit_after>extern crate libc;
extern crate std;

use super::liblua;
use super::Index;
use super::Lua;
use super::Pushable;
use super::Readable;

macro_rules! integer_impl(
    ($t:ident) => (
        impl Pushable for $t {
            fn push_to_lua(&self, lua: &Lua) {
                unsafe { liblua::lua_pushinteger(lua.lua, *self as liblua::lua_Integer) }
            }
        }
        impl Readable for $t {
            fn read_from_lua(lua: &Lua, index: i32) -> Option<$t> {
                unsafe {
                    let success: libc::c_int = std::mem::uninitialized();
                    let val = liblua::lua_tointegerx(lua.lua, index, std::mem::transmute(&success));
                    match success {
                        0 => None,
                        _ => Some(val as $t)
                    }
                }
            }
        }
        impl Index for $t {
        }
    );
)

integer_impl!(int)
integer_impl!(i8)
integer_impl!(i16)
integer_impl!(i32)
//integer_impl!(i64)   // data loss

macro_rules! unsigned_impl(
    ($t:ident) => (
        impl Pushable for $t {
            fn push_to_lua(&self, lua: &Lua) {
                unsafe { liblua::lua_pushunsigned(lua.lua, *self as liblua::lua_Unsigned) }
            }
        }
        impl Readable for $t {
            fn read_from_lua(lua: &Lua, index: i32) -> Option<$t> {
                unsafe {
                    let success: libc::c_int = std::mem::uninitialized();
                    let val = liblua::lua_tounsignedx(lua.lua, index, std::mem::transmute(&success));
                    match success {
                        0 => None,
                        _ => Some(val as $t)
                    }
                }
            }
        }
        impl Index for $t {
        }
    );
)

unsigned_impl!(uint)
unsigned_impl!(u8)
unsigned_impl!(u16)
unsigned_impl!(u32)
//unsigned_impl!(u64)   // data loss

macro_rules! numeric_impl(
    ($t:ident) => (
        impl Pushable for $t {
            fn push_to_lua(&self, lua: &Lua) {
                unsafe { liblua::lua_pushnumber(lua.lua, *self as f64) }
            }
        }
        impl Readable for $t {
            fn read_from_lua(lua: &Lua, index: i32) -> Option<$t> {
                unsafe {
                    let success: libc::c_int = std::mem::uninitialized();
                    let val = liblua::lua_tonumberx(lua.lua, index, std::mem::transmute(&success));
                    match success {
                        0 => None,
                        _ => Some(val as $t)
                    }
                }
            }
        }
        impl Index for $t {
        }
    );
)

numeric_impl!(f32)
numeric_impl!(f64)

impl Pushable for std::string::String {
    fn push_to_lua(&self, lua: &Lua) {
        unsafe { liblua::lua_pushstring(lua.lua, self.to_c_str().unwrap()) }
    }
}

impl Readable for String {
    fn read_from_lua(lua: &Lua, index: i32) -> Option<std::string::String> {
        unsafe {
            let size: libc::size_t = std::mem::uninitialized();
            let cStr = liblua::lua_tolstring(lua.lua, index, std::mem::transmute(&size));
            if cStr.is_null() {
                return None;
            }

            // I can't manage to make this compile properly, so we just transmute into a static slice
            let slice = std::slice::raw::buf_as_slice(cStr as *u8, size as uint, |buf| { let b:&'static [u8] = std::mem::transmute(buf); b });
            String::from_utf8(Vec::from_slice(slice)).ok()
        }
    }
}

impl Index for String {
}

impl Pushable for ~str {
    fn push_to_lua(&self, lua: &Lua) {
        String::from_str(*self).push_to_lua(lua)
    }
}

impl Pushable for &'static str {
    fn push_to_lua(&self, lua: &Lua) {
        String::from_str(*self).push_to_lua(lua)
    }
}

impl Readable for ~str {
    fn read_from_lua(lua: &Lua, index: i32) -> Option<~str> {
        unimplemented!()
    }
}

impl Readable for &'static str {
    fn read_from_lua(lua: &Lua, index: i32) -> Option<&'static str> {
        unimplemented!()
    }
}

impl Index for ~str {
}

impl Index for &'static str {
    fn lua_set_global<T: Pushable>(&self, lua: &Lua, value: T) {
        value.push_to_lua(lua);
        unsafe { liblua::lua_setglobal(lua.lua, self.to_c_str().unwrap()); }
    }

    fn lua_get_global<T: Readable>(&self, lua: &Lua) -> Option<T> {
        unsafe { liblua::lua_getglobal(lua.lua, self.to_c_str().unwrap()); }
        Readable::read_from_lua(lua, -1)
    }
}

/*impl GlobalsIndex for &'static str {
    fn lua_set_global(&self, lua: &Lua) {
        unsafe { liblua::lua_setglobal(lua.lua, self.to_c_str().unwrap()) }
    }

    fn lua_get_global(&self, lua: &Lua) {
        unsafe { liblua::lua_getglobal(lua.lua, self.to_c_str().unwrap()) }
    }
}

impl GlobalsIndex for ~str {
    fn lua_set_global(&self, lua: &Lua) {
        unsafe { liblua::lua_setglobal(lua.lua, self.to_c_str().unwrap()) }
    }

    fn lua_get_global(&self, lua: &Lua) {
        unsafe { liblua::lua_getglobal(lua.lua, self.to_c_str().unwrap()) }
    }
}*/

#[cfg(test)]
mod tests {
    #[test]
    fn readwrite_ints() {
        let mut lua = super::super::Lua::new();

        lua.set("a", 2);
        let x: int = lua.get("a").unwrap();
        assert_eq!(x, 2)
    }
}
<|endoftext|>


<commit_before>use byteorder::{ByteOrder, LittleEndian};
use log::{debug, trace, warn};
use goblin::{Object};
use failure::{Error, Fail};
use std::ops::Range;

use super::super::super::arch::{RVA};
use super::super::super::loader::Permissions;
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};
use std::collections::{HashSet};

#[derive(Debug, Fail)]
pub enum RelocAnalyzerError {
    #[fail(display = "Invalid relocation type")]
    InvalidRelocType,
    #[fail(display = "Relocation target not found in image")]
    InvalidTargetAddress,
}

pub struct RelocAnalyzer {}

impl RelocAnalyzer {
    pub fn new() -> RelocAnalyzer {
        RelocAnalyzer {}
    }
}

// TODO: make this much faster
fn is_in_insn(ws: &Workspace, rva: RVA) -> bool {
    let start: usize = rva.into();
    // TODO: underflow
    // TODO: remove harded max insn length
    let end: usize = start - 0x10;

    for i in (start..end).rev() {
        let i = RVA::from(i);
        if let Some(meta) = ws.get_meta(i) {
            if !meta.is_insn() {
                continue;
            }

            if let Ok(len) = meta.get_insn_length() {
                if i + len > rva {
                    return true;
                }
            }
        }
    }
    return false;
}

fn is_ptr(ws: &Workspace, rva: RVA) -> bool {
    if let Ok(ptr) = ws.read_va(rva) {
        if let Some(ptr) = ws.rva(ptr) {
            return ws.probe(ptr, 1, Permissions::R);
        }
    }
    return false;
}

fn is_zero(ws: &Workspace, rva: RVA) -> bool {
    if let Ok(v) = ws.read_u32(rva) {
        return v == 0;
    }
    return false;
}

#[derive(Debug)]
pub enum RelocationType {
    ImageRelBasedAbsolute,
    ImageRelBasedHigh,
    ImageRelBasedLow,
    ImageRelBasedHighLow,
    ImageRelBasedHighAdj,

    // ImageRelBasedMIPS_JmpAddr,
    // ImageRelBasedARM_MOV32,
    // ImageRelBasedRiscV_High20,
    ImageRelArch1,

    ImageRelReserved,

    // ImageRelBasedTHUMB_MOV32,
    // ImageRelBasedRiscV_Low12I,
    ImageRelArch2,

    ImageRelBasedRiscVLow12S,
    ImageRelBasedMIPSJmpAddr16,
    ImageRelBasedDir64,
}

pub struct Reloc {
    pub typ: RelocationType,
    pub offset: RVA,
}

fn parse_reloc(base: RVA, entry: u16) -> Result<Reloc, Error> {
    let reloc_type   = (entry & 0b1111000000000000) >> 12;
    let reloc_offset =  entry & 0b0000111111111111;

    // TODO: this should probably be on `RelocationType` as a `From` trait.
    let reloc_type = match reloc_type {
        0 => RelocationType::ImageRelBasedAbsolute,
        1 => RelocationType::ImageRelBasedHigh,
        2 => RelocationType::ImageRelBasedLow,
        3 => RelocationType::ImageRelBasedHighLow,
        4 => RelocationType::ImageRelBasedHighAdj,
        5 => RelocationType::ImageRelArch1,
        6 => RelocationType::ImageRelReserved,
        7 => RelocationType::ImageRelArch2,
        8 => RelocationType::ImageRelBasedRiscVLow12S,
        9 => RelocationType::ImageRelBasedMIPSJmpAddr16,
        10 => RelocationType::ImageRelBasedDir64,
        _ => return Err(RelocAnalyzerError::InvalidRelocType.into()),
    };

    Ok(Reloc {
        typ: reloc_type,
        offset: base + RVA::from(reloc_offset as i64),
    })
}

fn split_u32(e: u32) -> (u16, u16) {
    let m = ((e & 0x0000FFFF) >>  0) as u16;
    let n = ((e & 0xFFFF0000) >> 16) as u16;
    (m, n)
}

/// ```
/// use lancelot::rsrc::*;
/// use lancelot::arch::*;
/// use lancelot::workspace::Workspace;
/// use lancelot::analysis::pe::relocs;
///
/// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
///    .disable_analysis()
///    .load().unwrap();
/// let relocs = relocs::get_relocs(&ws).unwrap();
/// assert_eq!(relocs[0].offset,   RVA(0x76008));
/// assert_eq!(relocs[277].offset, RVA(0xA8000));
/// ```
pub fn get_relocs(ws: &Workspace) -> Result<Vec<Reloc>, Error> {
    let pe = match Object::parse(&ws.buf) {
        Ok(Object::PE(pe)) => pe,
        _ => return Ok(vec![]),
    };

    let opt_header = match pe.header.optional_header {
        Some(opt_header) => opt_header,
        _ => return Ok(vec![]),
    };

    let reloc_directory = match opt_header.data_directories.get_base_relocation_table() {
        Some(reloc_directory) => reloc_directory,
        _ => return Ok(vec![]),
    };

    let dir_start = RVA::from(reloc_directory.virtual_address as i64);
    let buf = ws.read_bytes(dir_start, reloc_directory.size as usize)?;

    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    let mut ret = vec![];
    let mut index = 0;
    loop {
        // parse base relocation chunk
        // until there's no data left

        let page_rva = entries.get(index);
        let block_size = entries.get(index+1);
        let (page_rva, block_size) = match (page_rva, block_size) {
            (Some(0),        _               ) => break,
            (_,              Some(0)         ) => break,
            (Some(page_rva), Some(block_size)) => (RVA::from(*page_rva as i64), *block_size),
            _ => break,
        };

        // cast from u32 to usize
        let chunk_entries_count = ((block_size) / 4) as usize;

        for i in 2..chunk_entries_count {
            if let Some(&entry) = entries.get(index + i) {
                let (m, n) = split_u32(entry);

                let reloc1 = parse_reloc(page_rva, m)?;
                if !ws.probe(reloc1.offset, 4, Permissions::R) {
                    break
                }
                ret.push(reloc1);

                let reloc2 = parse_reloc(page_rva, n)?;
                if !ws.probe(reloc2.offset, 4, Permissions::R) {
                    break
                }
                ret.push(reloc2);
            } else {
                break
            }
        };

        index += chunk_entries_count;
    }

    Ok(ret)
}

fn sections_contain(sections: &[Range<RVA>], rva: RVA) -> bool {
    sections.iter()
        .find(|section| section.contains(&rva))
        .is_some()
}

impl Analyzer for RelocAnalyzer {
    fn get_name(&self) -> String {
        "PE relocation analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::RelocAnalyzer;
    /// lancelot::test::init_logging();
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    ///
    /// let anal = RelocAnalyzer::new();
    /// anal.analyze(&mut ws).unwrap();
    ///
    /// assert!(ws.get_meta(RVA(0xC7F0)).unwrap().is_insn());
    /// ```
    fn analyze(&self, ws: &mut Workspace)-> Result<(), Error> {
        let x_sections: Vec<Range<RVA>> = ws.module.sections.iter()
            .filter(|section| section.perms.intersects(Permissions::X))
            .map(|section| Range {
                start: section.addr,
                end: section.end(),
            })
            .collect();

        let relocs: Vec<Reloc> = get_relocs(ws)?
            .into_iter()
            .filter(|r| match &r.typ {
                RelocationType::ImageRelBasedHighLow => true,
                RelocationType::ImageRelBasedDir64 => true,
                reloc_type @ _ => {
                    // all other reloc types are currently unsupported (uncommon)
                    warn!("ignoring relocation with unsupported type: {:?}", reloc_type);
                    false
                },
            })
            .collect();
        debug!("found {} total relocs", relocs.len());

        for reloc in relocs.iter() {
            let target = ws.rva(ws.read_va(reloc.offset)?).ok_or(RelocAnalyzerError::InvalidTargetAddress)?;
            trace!("found reloc {} -> {} {}",
                   reloc.offset,
                   target,
                if sections_contain(&x_sections, target) {"(code)"} else {""}
            );
        }

        // scan for relocations to code.
        //
        // a relocation is a hardcoded offset that must be fixed up if the desired base address
        //  cannot be used.
        // for example, the function pointer passed to CreateThread will be a hardcoded address
        //  of the start of the function.
        //
        // relocated pointers may point to the .data section, e.g. strings or other constants.
        // we want to ignore these.
        // we are only interested in targets in executable sections.
        // we assume these are pointers to instructions/code.
        //
        // looking for pointers into the .text section
        // to things that
        //   1. are not already in an instruction
        //   2. don't appear to be a pointer
        // and assume this is code.

        let mut unique_targets: HashSet<RVA> = HashSet::new();
        unique_targets.extend(relocs
            .iter()
            .map(|reloc| reloc.offset)
            .map(|rva| ws.read_va(rva))
            .filter_map(Result::ok)
            .filter_map(|va| ws.rva(va)));

        debug!("reduced to {} unique reloc targets", unique_targets.len());

       let o: Vec<RVA> = unique_targets
            .iter()
            .filter(|&&rva| sections_contain(&x_sections, rva))
            .filter(|&&rva| !is_in_insn(ws, rva))
            .filter(|&&rva| !is_ptr(ws, rva))
            .filter(|&&rva| !is_zero(ws, rva))
            .map(|&rva| rva)
            // TODO: maybe ensure that the insn decodes.
            .collect();

        debug!("found {} relocs that point to instructions", o.len());

        o.iter().for_each(|&rva| {
            debug!("found pointer via relocations from executable section to {} (code)", rva);

            // TODO: consume result
            ws.make_insn(rva).unwrap();
            // TODO: consume result
            ws.analyze().unwrap();
        });
        Ok(())
    }
}
<commit_msg>reloc: don't terminate early when encountering an invalid relocation target.<commit_after>use byteorder::{ByteOrder, LittleEndian};
use log::{debug, trace, warn};
use goblin::{Object};
use failure::{Error, Fail};
use std::ops::Range;

use super::super::super::arch::{RVA};
use super::super::super::loader::Permissions;
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};
use std::collections::{HashSet};

#[derive(Debug, Fail)]
pub enum RelocAnalyzerError {
    #[fail(display = "Invalid relocation type")]
    InvalidRelocType,
    #[fail(display = "Relocation target not found in image")]
    InvalidTargetAddress,
}

pub struct RelocAnalyzer {}

impl RelocAnalyzer {
    pub fn new() -> RelocAnalyzer {
        RelocAnalyzer {}
    }
}

// TODO: make this much faster
fn is_in_insn(ws: &Workspace, rva: RVA) -> bool {
    let start: usize = rva.into();
    // TODO: underflow
    // TODO: remove harded max insn length
    let end: usize = start - 0x10;

    for i in (start..end).rev() {
        let i = RVA::from(i);
        if let Some(meta) = ws.get_meta(i) {
            if !meta.is_insn() {
                continue;
            }

            if let Ok(len) = meta.get_insn_length() {
                if i + len > rva {
                    return true;
                }
            }
        }
    }
    return false;
}

fn is_ptr(ws: &Workspace, rva: RVA) -> bool {
    if let Ok(ptr) = ws.read_va(rva) {
        if let Some(ptr) = ws.rva(ptr) {
            return ws.probe(ptr, 1, Permissions::R);
        }
    }
    return false;
}

fn is_zero(ws: &Workspace, rva: RVA) -> bool {
    if let Ok(v) = ws.read_u32(rva) {
        return v == 0;
    }
    return false;
}

#[derive(Debug)]
pub enum RelocationType {
    ImageRelBasedAbsolute,
    ImageRelBasedHigh,
    ImageRelBasedLow,
    ImageRelBasedHighLow,
    ImageRelBasedHighAdj,

    // ImageRelBasedMIPS_JmpAddr,
    // ImageRelBasedARM_MOV32,
    // ImageRelBasedRiscV_High20,
    ImageRelArch1,

    ImageRelReserved,

    // ImageRelBasedTHUMB_MOV32,
    // ImageRelBasedRiscV_Low12I,
    ImageRelArch2,

    ImageRelBasedRiscVLow12S,
    ImageRelBasedMIPSJmpAddr16,
    ImageRelBasedDir64,
}

pub struct Reloc {
    pub typ: RelocationType,
    pub offset: RVA,
}

fn parse_reloc(base: RVA, entry: u16) -> Result<Reloc, Error> {
    let reloc_type   = (entry & 0b1111000000000000) >> 12;
    let reloc_offset =  entry & 0b0000111111111111;

    // TODO: this should probably be on `RelocationType` as a `From` trait.
    let reloc_type = match reloc_type {
        0 => RelocationType::ImageRelBasedAbsolute,
        1 => RelocationType::ImageRelBasedHigh,
        2 => RelocationType::ImageRelBasedLow,
        3 => RelocationType::ImageRelBasedHighLow,
        4 => RelocationType::ImageRelBasedHighAdj,
        5 => RelocationType::ImageRelArch1,
        6 => RelocationType::ImageRelReserved,
        7 => RelocationType::ImageRelArch2,
        8 => RelocationType::ImageRelBasedRiscVLow12S,
        9 => RelocationType::ImageRelBasedMIPSJmpAddr16,
        10 => RelocationType::ImageRelBasedDir64,
        _ => return Err(RelocAnalyzerError::InvalidRelocType.into()),
    };

    Ok(Reloc {
        typ: reloc_type,
        offset: base + RVA::from(reloc_offset as i64),
    })
}

fn split_u32(e: u32) -> (u16, u16) {
    let m = ((e & 0x0000FFFF) >>  0) as u16;
    let n = ((e & 0xFFFF0000) >> 16) as u16;
    (m, n)
}

/// ```
/// use lancelot::rsrc::*;
/// use lancelot::arch::*;
/// use lancelot::workspace::Workspace;
/// use lancelot::analysis::pe::relocs;
///
/// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
///    .disable_analysis()
///    .load().unwrap();
/// let relocs = relocs::get_relocs(&ws).unwrap();
/// assert_eq!(relocs[0].offset,   RVA(0x76008));
/// assert_eq!(relocs[277].offset, RVA(0xA8000));
/// ```
pub fn get_relocs(ws: &Workspace) -> Result<Vec<Reloc>, Error> {
    let pe = match Object::parse(&ws.buf) {
        Ok(Object::PE(pe)) => pe,
        _ => return Ok(vec![]),
    };

    let opt_header = match pe.header.optional_header {
        Some(opt_header) => opt_header,
        _ => return Ok(vec![]),
    };

    let reloc_directory = match opt_header.data_directories.get_base_relocation_table() {
        Some(reloc_directory) => reloc_directory,
        _ => return Ok(vec![]),
    };

    let dir_start = RVA::from(reloc_directory.virtual_address as i64);
    let buf = ws.read_bytes(dir_start, reloc_directory.size as usize)?;

    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    let mut ret = vec![];
    let mut index = 0;
    loop {
        // parse base relocation chunk
        // until there's no data left

        let page_rva = entries.get(index);
        let block_size = entries.get(index+1);
        let (page_rva, block_size) = match (page_rva, block_size) {
            (Some(0),        _               ) => break,
            (_,              Some(0)         ) => break,
            (Some(page_rva), Some(block_size)) => (RVA::from(*page_rva as i64), *block_size),
            _ => break,
        };

        // cast from u32 to usize
        let chunk_entries_count = ((block_size) / 4) as usize;

        for i in 2..chunk_entries_count {
            if let Some(&entry) = entries.get(index + i) {
                let (m, n) = split_u32(entry);

                let reloc1 = parse_reloc(page_rva, m)?;
                if !ws.probe(reloc1.offset, 4, Permissions::R) {
                    break
                }
                ret.push(reloc1);

                let reloc2 = parse_reloc(page_rva, n)?;
                if !ws.probe(reloc2.offset, 4, Permissions::R) {
                    break
                }
                ret.push(reloc2);
            } else {
                break
            }
        };

        index += chunk_entries_count;
    }

    Ok(ret)
}

fn sections_contain(sections: &[Range<RVA>], rva: RVA) -> bool {
    sections.iter()
        .find(|section| section.contains(&rva))
        .is_some()
}

impl Analyzer for RelocAnalyzer {
    fn get_name(&self) -> String {
        "PE relocation analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::RelocAnalyzer;
    /// lancelot::test::init_logging();
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    ///
    /// let anal = RelocAnalyzer::new();
    /// anal.analyze(&mut ws).unwrap();
    ///
    /// assert!(ws.get_meta(RVA(0xC7F0)).unwrap().is_insn());
    /// ```
    fn analyze(&self, ws: &mut Workspace)-> Result<(), Error> {
        let x_sections: Vec<Range<RVA>> = ws.module.sections.iter()
            .filter(|section| section.perms.intersects(Permissions::X))
            .map(|section| Range {
                start: section.addr,
                end: section.end(),
            })
            .collect();

        let relocs: Vec<Reloc> = get_relocs(ws)?
            .into_iter()
            .filter(|r| match &r.typ {
                RelocationType::ImageRelBasedHighLow => true,
                RelocationType::ImageRelBasedDir64 => true,
                reloc_type @ _ => {
                    // all other reloc types are currently unsupported (uncommon)
                    warn!("ignoring relocation with unsupported type: {:?}", reloc_type);
                    false
                },
            })
            .collect();
        debug!("found {} total relocs", relocs.len());

        for reloc in relocs.iter() {
            // a relocation fixes up an existing pointer (hardcoded virtual address).
            // so we'll validate that the pointer actually points to something in the image.
            //
            // we might want to simply assume an invalid pointer means the reloc table is invalid;
            // however,
            // legit `Windows.Media.Protection.PlayReady.dll` has some invalid relocs (not sure way).
            let ptr = match ws.read_va(reloc.offset) {
                Ok(ptr) => ptr,
                Err(_) => {
                    warn!("reloc fixes up an invalid pointer: {}", reloc.offset);
                    continue;
                }
            };

            let target = match ws.rva(ptr) {
                Some(target) => target,
                None => {
                    warn!("reloc fixes up a pointer to unmapped data: {} {}", reloc.offset, ptr);
                    continue
                }
            };

            trace!("found reloc {} -> {} {}",
                   reloc.offset,
                   target,
                if sections_contain(&x_sections, target) {"(code)"} else {""}
            );
        }

        // scan for relocations to code.
        //
        // a relocation is a hardcoded offset that must be fixed up if the desired base address
        //  cannot be used.
        // for example, the function pointer passed to CreateThread will be a hardcoded address
        //  of the start of the function.
        //
        // relocated pointers may point to the .data section, e.g. strings or other constants.
        // we want to ignore these.
        // we are only interested in targets in executable sections.
        // we assume these are pointers to instructions/code.
        //
        // looking for pointers into the .text section
        // to things that
        //   1. are not already in an instruction
        //   2. don't appear to be a pointer
        // and assume this is code.

        let mut unique_targets: HashSet<RVA> = HashSet::new();
        unique_targets.extend(relocs
            .iter()
            .map(|reloc| reloc.offset)
            .map(|rva| ws.read_va(rva))
            .filter_map(Result::ok)
            .filter_map(|va| ws.rva(va)));

        debug!("reduced to {} unique reloc targets", unique_targets.len());

       let o: Vec<RVA> = unique_targets
            .iter()
            .filter(|&&rva| sections_contain(&x_sections, rva))
            .filter(|&&rva| !is_in_insn(ws, rva))
            .filter(|&&rva| !is_ptr(ws, rva))
            .filter(|&&rva| !is_zero(ws, rva))
            .map(|&rva| rva)
            // TODO: maybe ensure that the insn decodes.
            .collect();

        debug!("found {} relocs that point to instructions", o.len());

        o.iter().for_each(|&rva| {
            debug!("found pointer via relocations from executable section to {} (code)", rva);

            // TODO: consume result
            ws.make_insn(rva).unwrap();
            // TODO: consume result
            ws.analyze().unwrap();
        });
        Ok(())
    }
}
<|endoftext|>


<commit_before>//
// nfa.rs
// Nondeterministic Finite Automata
//

type StateID = usize;

struct State {
    id: StateID,
    accept: bool,
    trans: Vec<Transition>
}

impl State {
    pub fn find_transition(&self, label: Label) -> Option<&Transition> {
        for t in self.trans.iter() {
            if t.label == label {
                return Some(&t);
            }
        }

        None
    }
}

#[derive(PartialEq, Eq)]
enum Label {
    Epsilon,
    Symbol(char)
}

struct Transition {
    label: Label,
    target: StateID
}

struct NFASet {
    states: Vec<State>
}

impl NFASet {
    pub fn new() -> NFASet {
        let st : Vec<State> = Vec::with_capacity(10);
        NFASet { states: st }
    }

    pub fn create_state(&mut self, acc: bool) -> StateID {
        let id = self.states.len();
        let trs : Vec<Transition> = Vec::with_capacity(2);
        let st = State { id: id, accept: acc, trans: trs };
        self.states.push(st);
        id
    }

    pub fn check_id(&self, id: StateID) -> bool {
        id < self.states.len()
    }

    fn get_state_mut(&mut self, id: StateID) -> &mut State {
        // TODO: change to safe, checked access
        unsafe {
            let res : &mut State = self.states.get_unchecked_mut(id);
            res
        }
    }

    pub fn get_state(&self, id: StateID) -> Option<&State> {
        self.states.get(id)
    }

    pub fn add_transition(&mut self, src_id: StateID, dst_id: StateID, label: Label) -> bool {
        if self.check_id(src_id) && self.check_id(dst_id) {
            let src_st = self.get_state_mut(src_id);
            let trans = Transition { label: label, target: dst_id };
            src_st.trans.push(trans);
            true
        }
        else {
            false
        }
    }

    pub fn create_nfa(&self, start: StateID) -> Option<NFA> {
        if self.check_id(start) {
            Some(NFA { start: start })
        }
        else {
            None
        }
    }

    pub fn nfa_union(&mut self, n1: NFA, n2: NFA) -> NFA {
        NFA { start: 0 }  // TODO: implement
    }

    pub fn nfa_concat(&mut self, n1: NFA, n2: NFA) -> NFA {
        NFA { start: 0 }  // TODO: implement
    }

    pub fn nfa_star(&mut self, n: NFA) -> NFA {
        NFA { start: 0 }  // TODO: implement
    }
}

struct NFA {
    start: StateID,
}

#[test]
fn states() {
    let mut ns = NFASet::new();
    let id1 = ns.create_state(false);
    let id2 = ns.create_state(false);
    assert_eq!(id1, 0);
    assert_eq!(id2, 1);

    assert!(ns.check_id(id1));
    assert!(ns.check_id(id2));
    assert!(!ns.check_id(100));
}

#[test]
fn add_transition() {
    let mut ns = NFASet::new();
    let id1 = ns.create_state(false);
    let id2 = ns.create_state(false);
    let id3 = ns.create_state(false);
    assert!(ns.add_transition(id1, id2, Label::Epsilon));
    assert!(ns.add_transition(id2, id3, Label::Symbol('a')));

    assert!(!ns.add_transition(id1, 1200, Label::Epsilon));
    assert!(!ns.add_transition(1020, id2, Label::Epsilon));
    assert!(!ns.add_transition(1320, 1200, Label::Epsilon));
}

#[test]
fn transitions() {
    let mut ns = NFASet::new();
    let id1 = ns.create_state(false);
    let id2 = ns.create_state(false);
    let id3 = ns.create_state(false);
    assert!(ns.add_transition(id1, id2, Label::Epsilon));
    assert!(ns.add_transition(id2, id3, Label::Symbol('a')));

    let st1 = ns.get_state(id1).unwrap();
    let t1 = st1.find_transition(Label::Epsilon).unwrap();

    let st2 = ns.get_state(t1.target).unwrap();
    let t2 = st2.find_transition(Label::Symbol('a')).unwrap();

    assert_eq!(t2.target, id3);
}

<commit_msg>Implemented NFA operations<commit_after>//
// nfa.rs
// Nondeterministic Finite Automata
//

type StateID = usize;

struct State {
    id: StateID,
    trans: Vec<Transition>
}

impl State {
    // finds the outgoing transitions for a given label
    pub fn find_transition(&self, label: Label) -> Vec<&Transition> {
        let mut res : Vec<&Transition> = Vec::new();
        for t in self.trans.iter() {
            if t.label == label {
                res.push(&t);
            }
        }

        res
    }
}

#[derive(PartialEq, Eq)]
enum Label {
    Epsilon,
    Symbol(char)
}

struct Transition {
    label: Label,
    target: StateID
}

struct NFASet {
    states: Vec<State>
}

impl NFASet {
    pub fn new() -> NFASet {
        let st : Vec<State> = Vec::with_capacity(10);
        NFASet { states: st }
    }

    pub fn create_state(&mut self) -> StateID {
        let id = self.states.len();
        let trs : Vec<Transition> = Vec::with_capacity(2);
        let st = State { id: id, trans: trs };
        self.states.push(st);
        id
    }

    pub fn check_id(&self, id: StateID) -> bool {
        id < self.states.len()
    }

    fn get_state_mut(&mut self, id: StateID) -> &mut State {
        // TODO: change to safe, checked access
        unsafe {
            let res : &mut State = self.states.get_unchecked_mut(id);
            res
        }
    }

    pub fn get_state(&self, id: StateID) -> Option<&State> {
        self.states.get(id)
    }

    pub fn add_transition(&mut self, src_id: StateID, dst_id: StateID, label: Label) -> bool {
        if self.check_id(src_id) && self.check_id(dst_id) {
            let src_st = self.get_state_mut(src_id);
            let trans = Transition { label: label, target: dst_id };
            src_st.trans.push(trans);
            true
        }
        else {
            false
        }
    }

    pub fn create_nfa(&self, start: StateID, acc: StateID) -> Option<NFA> {
        if self.check_id(start) && self.check_id(acc) {
            Some(NFA { start: start, accept: acc })
        }
        else {
            None
        }
    }

    pub fn nfa_epsilon(&mut self) -> NFA {
        let start = self.create_state();
        let acc = self.create_state();

        self.add_transition(start, acc, Label::Epsilon);

        NFA { start: start, accept: acc }
    }

    pub fn nfa_symbol(&mut self, symbol: char) -> NFA {
        let start = self.create_state();
        let acc = self.create_state();

        self.add_transition(start, acc, Label::Symbol(symbol));

        NFA { start: start, accept: acc }
    }

    pub fn nfa_union(&mut self, n1: NFA, n2: NFA) -> NFA {
        let start = self.create_state();
        self.add_transition(start, n1.start, Label::Epsilon);
        self.add_transition(start, n2.start, Label::Epsilon);

        let acc = self.create_state();
        self.add_transition(n1.accept, acc, Label::Epsilon);
        self.add_transition(n2.accept, acc, Label::Epsilon);

        NFA { start: start, accept: acc }
    }

    pub fn nfa_concat(&mut self, n1: NFA, n2: NFA) -> NFA {
        self.add_transition(n1.accept, n2.start, Label::Epsilon);
        NFA { start: n1.start, accept: n2.accept }
    }

    pub fn nfa_star(&mut self, n: NFA) -> NFA {
        let start = self.create_state();
        let acc = self.create_state();

        // transitions connecting old and new start/accept states
        self.add_transition(start, n.start, Label::Epsilon);
        self.add_transition(n.accept, acc, Label::Epsilon);

        // transition connecting start and accept for 0 occurrences
        self.add_transition(start, acc, Label::Epsilon);

        // loop transition for any number of occurrences
        self.add_transition(n.accept, n.start, Label::Epsilon);

        NFA { start: start, accept: acc }
    }
}

// NFAs with a single accepting state, which is sufficient
// for automata generated from regular expressions
struct NFA {
    start: StateID,
    accept: StateID
}

#[test]
fn states() {
    let mut ns = NFASet::new();
    let id1 = ns.create_state();
    let id2 = ns.create_state();
    assert_eq!(id1, 0);
    assert_eq!(id2, 1);

    assert!(ns.check_id(id1));
    assert!(ns.check_id(id2));
    assert!(!ns.check_id(100));
}

#[test]
fn add_transition() {
    let mut ns = NFASet::new();
    let id1 = ns.create_state();
    let id2 = ns.create_state();
    let id3 = ns.create_state();
    assert!(ns.add_transition(id1, id2, Label::Epsilon));
    assert!(ns.add_transition(id2, id3, Label::Symbol('a')));

    assert!(!ns.add_transition(id1, 1200, Label::Epsilon));
    assert!(!ns.add_transition(1020, id2, Label::Epsilon));
    assert!(!ns.add_transition(1320, 1200, Label::Epsilon));
}

#[test]
fn transitions() {
    let mut ns = NFASet::new();
    let id1 = ns.create_state();
    let id2 = ns.create_state();
    let id3 = ns.create_state();
    assert!(ns.add_transition(id1, id2, Label::Epsilon));
    assert!(ns.add_transition(id1, id3, Label::Epsilon));
    assert!(ns.add_transition(id2, id3, Label::Symbol('a')));

    let st1 = ns.get_state(id1).unwrap();
    let t0 = st1.find_transition(Label::Symbol('b'));
    assert_eq!(t0.len(), 0);
    let t1 = st1.find_transition(Label::Epsilon);
    assert_eq!(t1.len(), 2);

    let st2 = ns.get_state(t1[0].target).unwrap();
    let t2 = st2.find_transition(Label::Symbol('a'));
    assert_eq!(t2.len(), 1);

    assert_eq!(t2[0].target, id3);
    assert_eq!(t1[1].target, id3);
}

#[test]
fn epsilon() {
    let mut ns = NFASet::new();
    let n1 = ns.nfa_epsilon();

    let st1 = ns.get_state(n1.start).unwrap();
    let t1 = st1.find_transition(Label::Epsilon);
    assert_eq!(t1.len(), 1);

    let st2 = ns.get_state(t1[0].target).unwrap();
    let t2 = st2.find_transition(Label::Epsilon);
    assert_eq!(t2.len(), 0);
}
<|endoftext|>