<commit_before>//! Hosting the CLR from unix with Rust

extern crate libloading as libl;

use std::collections::HashMap;
use std::ffi::CString;
use std::io;
use std::os::raw::{c_char, c_int, c_uint, c_void};

// -- Private utility functions -- //
fn to_c_str<T: Into<Vec<u8>>>(t: T) -> *const c_char {
    CString::new(t).unwrap().as_ptr() as *const c_char
}

/// A function pointer type for the init function
type CoreClrInitializeFn = unsafe extern fn(
    *const c_char,
    *const c_char,
    c_int,
    *const *const c_char,
    *const *const c_char,
    *const *const c_void,
    *const c_uint) -> c_int;

/// A function pointer type for the shutdown function
type CoreClrShutdownFn = unsafe extern fn(*const c_void, c_uint) -> c_int;

/// A function pointer type for the second shutdown function
type CoreClrShutdown2Fn = unsafe extern fn(*const c_void, c_uint, *const c_int) -> c_int;

/// A function pointer type for the create delegate function
type CoreClrCreateDelegateFn = unsafe extern fn(
    *const c_void,
    c_uint,
    *const c_char,
    *const c_char,
    *const c_char,
    *const *const c_void) -> c_int;

/// The CoreClr object represents a binding to the CLR maintained by a private handle and domain id
pub struct CoreClr {
    host_handle: *const c_void,
    domain_id: c_uint
}

impl CoreClr {
    /// Private helper function to grab a reference to the library in the current context
    fn library() -> libl::Result<libl::Library> {
        libl::Library::new("/usr/local/share/dotnet/shared/Microsoft.NETCore.App/2.0.0")
    }

    /// Creates a new CLR object
    pub fn init(
        exe_path: &str,
        app_domain_friendly_name: &str,
        properties_option: Option<HashMap<&str, &str>>) -> libl::Result<CoreClr>
    {
        // Create the host handle and its ref
        let host_handle = 0 as *const c_void;
        let host_handle_ref = &host_handle as *const *const c_void;

        // Create the domain id and its ref
        let domain_id = 0 as c_uint;
        let domain_id_ref = &domain_id as *const c_uint;

        // Raw C string refs from exe_path and app_domain_friendly_name slices
        let exe_path_raw = to_c_str(exe_path);
        let app_domain_friendly_name_raw = to_c_str(app_domain_friendly_name);

        // Either use the provided option or a blank hashmap
        let properties = properties_option.unwrap_or_else(HashMap::new);

        // Count for upcoming vecs
        let properties_count = properties.len() as c_int;

        // Collect property keys into a vec
        let properties_keys: Vec<*const c_char> = properties.keys()
                .map(|&k| to_c_str(k))
                .collect();

        // Collect equivalent property values into another vec
        let properties_values: Vec<*const c_char> = properties.values()
                .map(|&v| to_c_str(v))
                .collect();

        // Grab refs of the two to pass into the actual function call
        let properties_keys_ref = properties_keys.as_ptr();
        let properties_values_ref = properties_values.as_ptr();

        // Open up an unsafe block for actually loading functions from the CLR libs
        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_initialize: libl::Symbol<CoreClrInitializeFn> = coreclr_library.get(b"coreclr_initialize")?;

            // Initialize the CLR
            match coreclr_initialize(
                exe_path_raw,
                app_domain_friendly_name_raw,
                properties_count,
                properties_keys_ref,
                properties_values_ref,
                host_handle_ref,
                domain_id_ref)
            {
                // If healthy exit code, return a model of the CLR
                0 => Ok(CoreClr {
                    host_handle: host_handle,
                    domain_id: domain_id
                }),
                // Else panic
                _ => panic!("Failed to initialize")
            }
        }
    }

    /// Shuts down the CLR
    pub fn shutdown(self: Self) -> io::Result<()> {
        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_shutdown: libl::Symbol<CoreClrShutdownFn> = coreclr_library.get(b"coreclr_shutdown")?;

            // Shutdown the CLR
            match coreclr_shutdown(self.host_handle, self.domain_id) {
                // If healthy exit code, return unit
                0 => Ok(()),
                // Else panic
                _ => panic!("Failed to shutdown")
            }
        }
    }

    /// Shuts down the CLR, returning a latched exit code?  Should follow up on what that is.
    pub fn shutdown_2(self: Self) -> io::Result<c_int> {
        let latched_exit_code = -1 as c_int;
        let latched_exit_code_ref = &latched_exit_code as *const c_int;

        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_shutdown_2: libl::Symbol<CoreClrShutdown2Fn> = coreclr_library.get(b"coreclr_shutdown_2")?;

            // Shutdown the CLR
            match coreclr_shutdown_2(self.host_handle, self.domain_id, latched_exit_code_ref) {
                // If healthy exit code, return the resulting exit code
                0 => Ok(latched_exit_code),
                // Else panic
                _ => panic!("Failed to shutdown")
            }
        }
    }

    /// Spins a raw delegate pointer out of the CLR
    pub fn create_raw_delegate(
        self: &Self,
        entry_point_assembly_name: &str,
        entry_point_type_name: &str,
        entry_point_method_name: &str) -> io::Result<*const c_void>
    {
        let coreclr_delegate = 0 as *const c_void;
        let coreclr_delegate_ref = &coreclr_delegate as *const *const c_void;

        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_create_delegate: libl::Symbol<CoreClrCreateDelegateFn> = coreclr_library.get(b"coreclr_create_delegate")?;

            // Create the delegate
            match coreclr_create_delegate(
                self.host_handle,
                self.domain_id,
                to_c_str(entry_point_assembly_name),
                to_c_str(entry_point_type_name),
                to_c_str(entry_point_method_name),
                coreclr_delegate_ref)
            {
                // If healthy exit code, return the resulting exit code
                0 => Ok(coreclr_delegate),
                // Else panic
                _ => panic!("Failed to shutdown")
            }
        }
    }
}<commit_msg>Added (failing) test to init and shutdown the CLR on unix hosts<commit_after>//! Hosting the CLR from unix with Rust

extern crate libloading as libl;

use std::collections::HashMap;
use std::ffi::CString;
use std::io;
use std::os::raw::{c_char, c_int, c_uint, c_void};

// -- Private utility functions -- //
fn to_c_str<T: Into<Vec<u8>>>(t: T) -> *const c_char {
    CString::new(t).unwrap().as_ptr() as *const c_char
}

/// A function pointer type for the init function
type CoreClrInitializeFn = unsafe extern fn(
    *const c_char,
    *const c_char,
    c_int,
    *const *const c_char,
    *const *const c_char,
    *const *const c_void,
    *const c_uint) -> c_int;

/// A function pointer type for the shutdown function
type CoreClrShutdownFn = unsafe extern fn(*const c_void, c_uint) -> c_int;

/// A function pointer type for the second shutdown function
type CoreClrShutdown2Fn = unsafe extern fn(*const c_void, c_uint, *const c_int) -> c_int;

/// A function pointer type for the create delegate function
type CoreClrCreateDelegateFn = unsafe extern fn(
    *const c_void,
    c_uint,
    *const c_char,
    *const c_char,
    *const c_char,
    *const *const c_void) -> c_int;

/// The CoreClr object represents a binding to the CLR maintained by a private handle and domain id
pub struct CoreClr {
    host_handle: *const c_void,
    domain_id: c_uint
}

impl CoreClr {
    /// Private helper function to grab a reference to the library in the current context
    fn library() -> libl::Result<libl::Library> {
        libl::Library::new("/usr/local/share/dotnet/shared/Microsoft.NETCore.App/2.0.0/libcoreclr.dylib")
    }

    /// Creates a new CLR object
    pub fn init(
        exe_path: &str,
        app_domain_friendly_name: &str,
        properties_option: Option<HashMap<&str, &str>>) -> libl::Result<CoreClr>
    {
        // Create the host handle and its ref
        let host_handle = 0 as *const c_void;
        let host_handle_ref = &host_handle as *const *const c_void;

        // Create the domain id and its ref
        let domain_id = 0 as c_uint;
        let domain_id_ref = &domain_id as *const c_uint;

        // Raw C string refs from exe_path and app_domain_friendly_name slices
        let exe_path_raw = to_c_str(exe_path);
        let app_domain_friendly_name_raw = to_c_str(app_domain_friendly_name);

        // Either use the provided option or a blank hashmap
        let properties = properties_option.unwrap_or_else(HashMap::new);

        // Count for upcoming vecs
        let properties_count = properties.len() as c_int;

        // Collect property keys into a vec
        let properties_keys: Vec<*const c_char> = properties.keys()
                .map(|&k| to_c_str(k))
                .collect();

        // Collect equivalent property values into another vec
        let properties_values: Vec<*const c_char> = properties.values()
                .map(|&v| to_c_str(v))
                .collect();

        // Grab refs of the two to pass into the actual function call
        let properties_keys_ref = properties_keys.as_ptr();
        let properties_values_ref = properties_values.as_ptr();

        // Open up an unsafe block for actually loading functions from the CLR libs
        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_initialize: libl::Symbol<CoreClrInitializeFn> = coreclr_library.get(b"coreclr_initialize")?;

            // Initialize the CLR
            match coreclr_initialize(
                exe_path_raw,
                app_domain_friendly_name_raw,
                properties_count,
                properties_keys_ref,
                properties_values_ref,
                host_handle_ref,
                domain_id_ref)
            {
                // If healthy exit code, return a model of the CLR
                0 => Ok(CoreClr {
                    host_handle: host_handle,
                    domain_id: domain_id
                }),
                // Else panic
                _ => panic!("Failed to initialize, {:?} {:?}", host_handle, domain_id)
            }
        }
    }

    /// Shuts down the CLR
    pub fn shutdown(self: Self) -> io::Result<()> {
        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_shutdown: libl::Symbol<CoreClrShutdownFn> = coreclr_library.get(b"coreclr_shutdown")?;

            // Shutdown the CLR
            match coreclr_shutdown(self.host_handle, self.domain_id) {
                // If healthy exit code, return unit
                0 => Ok(()),
                // Else panic
                _ => panic!("Failed to shutdown")
            }
        }
    }

    /// Shuts down the CLR, returning a latched exit code?  Should follow up on what that is.
    pub fn shutdown_2(self: Self) -> io::Result<c_int> {
        let latched_exit_code = -1 as c_int;
        let latched_exit_code_ref = &latched_exit_code as *const c_int;

        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_shutdown_2: libl::Symbol<CoreClrShutdown2Fn> = coreclr_library.get(b"coreclr_shutdown_2")?;

            // Shutdown the CLR
            match coreclr_shutdown_2(self.host_handle, self.domain_id, latched_exit_code_ref) {
                // If healthy exit code, return the resulting exit code
                0 => Ok(latched_exit_code),
                // Else panic
                _ => panic!("Failed to shutdown")
            }
        }
    }

    /// Spins a raw delegate pointer out of the CLR
    pub fn create_raw_delegate(
        self: &Self,
        entry_point_assembly_name: &str,
        entry_point_type_name: &str,
        entry_point_method_name: &str) -> io::Result<*const c_void>
    {
        let coreclr_delegate = 0 as *const c_void;
        let coreclr_delegate_ref = &coreclr_delegate as *const *const c_void;

        unsafe {
            let coreclr_library = CoreClr::library()?;
            let coreclr_create_delegate: libl::Symbol<CoreClrCreateDelegateFn> = coreclr_library.get(b"coreclr_create_delegate")?;

            // Create the delegate
            match coreclr_create_delegate(
                self.host_handle,
                self.domain_id,
                to_c_str(entry_point_assembly_name),
                to_c_str(entry_point_type_name),
                to_c_str(entry_point_method_name),
                coreclr_delegate_ref)
            {
                // If healthy exit code, return the resulting exit code
                0 => Ok(coreclr_delegate),
                // Else panic
                _ => panic!("Failed to shutdown")
            }
        }
    }
}

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

    #[test]
    fn init_and_shutdown() {
        let coreclr = CoreClr::init("", "", None).unwrap();
        coreclr.shutdown();
    }
}<|endoftext|>


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

//! Low-level devicemapper configuration of the running kernel.

#[allow(dead_code, non_camel_case_types)]
mod dm_ioctl;

use std::fs::File;
use std::io::{Result, Error, BufReader, BufRead};
use std::os::unix::io::AsRawFd;
use std::mem;
use std::slice;
use std::slice::bytes::copy_memory;
use std::collections::BTreeSet;

use byteorder::{NativeEndian, ByteOrder};

use nix::sys::ioctl;

use dm::dm_ioctl as dmi;
use lv::LV;
use vg::VG;
use util::align_to;
use ::Device;

const DM_IOCTL: u8 = 0xfd;
const DM_CTL_PATH: &'static str= "/dev/mapper/control";

const DM_VERSION_MAJOR: u32 = 4;
const DM_VERSION_MINOR: u32 = 30;
const DM_VERSION_PATCHLEVEL: u32 = 0;

// Status bits
//const DM_READONLY_FLAG: u32 = 1;
const DM_SUSPEND_FLAG: u32 = 2;
//const DM_PERSISTENT_DEV_FLAG: u32 = 8;

/// Major numbers used by DM.
pub fn dev_majors() -> BTreeSet<u32> {
    let mut set = BTreeSet::new();

    let f = File::open("/proc/devices")
        .ok().expect("Could not open /proc/devices");

    let reader = BufReader::new(f);

    for line in reader.lines()
        .filter_map(|x| x.ok())
        .skip_while(|x| x != "Block devices:")
        .skip(1) {
            let spl: Vec<_> = line.split_whitespace().collect();

            if spl[1] == "device-mapper" {
                set.insert(spl[0].parse::<u32>().unwrap());
            }
        }

    set
}

/// Context needed for communicating with devicemapper.
pub struct DM<'a> {
    file: File,
    vg: &'a VG,
}

impl <'a> DM<'a> {
    /// Create a new context for communicating about a given VG with DM.
    pub fn new(vg: &'a VG) -> Result<Self> {
        Ok(DM {
            file: try!(File::open(DM_CTL_PATH)),
            vg: vg,
        })
    }

    fn initialize_hdr(hdr: &mut dmi::Struct_dm_ioctl) -> () {
        hdr.version[0] = DM_VERSION_MAJOR;
        hdr.version[1] = DM_VERSION_MINOR;
        hdr.version[2] = DM_VERSION_PATCHLEVEL;

        hdr.data_start = mem::size_of::<dmi::Struct_dm_ioctl>() as u32;
    }

    fn hdr_set_name(hdr: &mut dmi::Struct_dm_ioctl, vg_name: &str, lv_name: &str) -> () {
        let name = format!("{}-{}", vg_name.replace("-", "--"),
                           lv_name.replace("-", "--"));
        let name_dest: &mut [u8; 128] = unsafe { mem::transmute(&mut hdr.name) };
        copy_memory(name.as_bytes(), &mut name_dest[..]);
    }

    fn hdr_set_uuid(hdr: &mut dmi::Struct_dm_ioctl, vg_uuid: &str, lv_uuid: &str) -> () {
        let uuid = format!("LVM-{}{}", vg_uuid.replace("-", ""), lv_uuid.replace("-", ""));
        let uuid_dest: &mut [u8; 129] = unsafe { mem::transmute(&mut hdr.uuid) };
        copy_memory(uuid.as_bytes(), &mut uuid_dest[..]);
    }

    /// Devicemapper version information: Major, Minor, and patchlevel versions.
    pub fn get_version(&self) -> Result<(u32, u32, u32)> {

        let mut hdr: dmi::Struct_dm_ioctl = Default::default();
        hdr.version[0] = DM_VERSION_MAJOR;
        hdr.version[1] = DM_VERSION_MINOR;
        hdr.version[2] = DM_VERSION_PATCHLEVEL;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_VERSION_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => {},
        };

        Ok((hdr.version[0], hdr.version[1], hdr.version[2]))
    }

    /// Returns a list of tuples containing DM device names within this VG,
    /// and their major/minor device numbers.
    pub fn list_devices(&self) -> Result<Vec<(String, Device)>> {
        let mut buf = [0u8; 16 * 1024];
        let mut hdr: &mut dmi::Struct_dm_ioctl = unsafe {mem::transmute(&mut buf)};

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = buf.len() as u32;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_LIST_DEVICES_CMD as u8, buf.len());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut buf) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => {},
        };

        let mut devs = Vec::new();
        if (hdr.data_size - hdr.data_start as u32) != 0 {
            let mut result = &buf[hdr.data_start as usize..];

            loop {
                let slc = slice_to_null(&result[12..]).expect("Bad data from ioctl");
                let devno = NativeEndian::read_u64(&result[..8]);
                let dm_name = String::from_utf8_lossy(slc);
                let mut vg_name = self.vg.name().replace("-", "--");
                vg_name.push('-');

                // Return only devices within this VG
                if dm_name.starts_with(&vg_name) {
                    let lv_name = dm_name.trim_left_matches(&vg_name).replace("--", "-");
                    devs.push((lv_name, devno.into()));
                }

                let next = NativeEndian::read_u32(&result[8..12]);
                if next == 0 { break }

                result = &result[next as usize..];
            }
        }

        Ok(devs)
    }

    /// Query DM for which devices depend on this device.
    pub fn list_deps(&self, dev: Device) -> Result<Vec<Device>> {
        let mut buf = [0u8; 16 * 1024];
        let mut hdr: &mut dmi::Struct_dm_ioctl = unsafe {mem::transmute(&mut buf)};

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = buf.len() as u32;
        hdr.dev = dev.into();

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_TABLE_DEPS_CMD as u8, buf.len());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut buf) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => {},
        };

        // TODO: Check DM_BUFFER_FULL_FLAG for:
        // DM_DEVICE_LIST_VERSIONS, DM_DEVICE_LIST, DM_DEVICE_DEPS,
        // DM_DEVICE_STATUS, DM_DEVICE_TABLE, DM_DEVICE_WAITEVENT,
        // DM_DEVICE_TARGET_MSG

        let mut devs = Vec::new();
        if (hdr.data_size - hdr.data_start as u32) != 0 {
            let result = &buf[hdr.data_start as usize..];
            let entries = NativeEndian::read_u32(&result[..4]) as usize;

            for entry in 0..entries {
                let dev = &result[(8*entry)+8..(8*entry)+16];
                devs.push(Device::from(NativeEndian::read_u64(&dev)));
            }
        }

        Ok(devs)
    }

    /// Recursively walk DM deps to see if device is present
    pub fn depends_on(&self, dev: Device, dm_majors: &BTreeSet<u32>) -> bool {
        if !dm_majors.contains(&dev.major) {
            return false;
        }

        if let Ok(dep_list) = self.list_deps(dev) {
            for d in dep_list {
                if d == dev {
                    return true;
                } else if self.depends_on(d, dm_majors) {
                    return true;
                }
            }
        }

        false
    }

    fn create_device(&self, lv: &mut LV) -> Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        Self::hdr_set_name(&mut hdr, &self.vg.name(), &lv.name);
        Self::hdr_set_uuid(&mut hdr, &self.vg.id(), &lv.id);
        hdr.data_size = hdr.data_start;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_CREATE_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => { }
        };

        lv.device = Some(Device::from(hdr.dev));

        Ok(())
    }

    fn remove_device(&self, lv: &LV) -> Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = hdr.data_start;
        Self::hdr_set_name(&mut hdr, &self.vg.name(), &lv.name);

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_REMOVE_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => Ok(())
        }
    }

    fn load_device(&self, lv: &LV) -> Result<()> {
        let sectors_per_extent = self.vg.extent_size();
        let mut targs = Vec::new();

        // Construct targets first, since we need to know how many & size
        // before initializing the header.
        for seg in &lv.segments {
            let mut targ: dmi::Struct_dm_target_spec = Default::default();
            targ.sector_start = seg.start_extent() * sectors_per_extent;
            targ.length = seg.extent_count() * sectors_per_extent;
            targ.status = 0;

            let mut dst: &mut [u8] = unsafe {
                mem::transmute(&mut targ.target_type[..])
            };
            copy_memory(seg.dm_type().as_bytes(), &mut dst);

            let mut params = seg.dm_params(self.vg);

            let pad_bytes = align_to(
                params.len() + 1usize, 8usize) - params.len();
            params.extend(vec!["\0"; pad_bytes]);

            targ.next = (mem::size_of::<dmi::Struct_dm_target_spec>()
                         + params.len()) as u32;

            targs.push((targ, params));
        }

        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        Self::hdr_set_name(&mut hdr, &self.vg.name(), &lv.name);

        hdr.data_start = mem::size_of::<dmi::Struct_dm_ioctl>() as u32;
        hdr.data_size = hdr.data_start + targs.iter()
            .map(|&(t, _)| t.next)
            .sum::<u32>();
        hdr.target_count = targs.len() as u32;

        // Flatten into buf
        let mut buf: Vec<u8> = Vec::with_capacity(hdr.data_size as usize);
        unsafe {
            let ptr: *mut u8 = mem::transmute(&mut hdr);
            let slc = slice::from_raw_parts(ptr, hdr.data_start as usize);
            buf.extend(slc);
        }

        for (targ, param) in targs {
            unsafe {
                let ptr: *mut u8 = mem::transmute(&targ);
                let slc = slice::from_raw_parts(
                    ptr, mem::size_of::<dmi::Struct_dm_target_spec>());
                buf.extend(slc);
            }

            buf.extend(param.as_bytes());
        }

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_TABLE_LOAD_CMD as u8, buf.len());

        match unsafe { ioctl::read_into_ptr(self.file.as_raw_fd(), op, buf.as_mut_ptr()) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => Ok(())
        }
    }

    fn suspend_device(&self, lv: &LV) -> Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = hdr.data_start;
        Self::hdr_set_name(&mut hdr, &self.vg.name(), &lv.name);
        hdr.flags = DM_SUSPEND_FLAG;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_SUSPEND_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => Ok(())
        }
    }

    fn resume_device(&self, lv: &LV) -> Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = hdr.data_start;
        Self::hdr_set_name(&mut hdr, &self.vg.name(), &lv.name);
        // DM_SUSPEND_FLAG not set = resume

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_SUSPEND_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => Ok(())
        }
    }

    /// Activate a Logical Volume.
    ///
    /// Also populates the LV's device field.
    pub fn activate_device(&self, lv: &mut LV) -> Result<()> {
        try!(self.create_device(lv));
        try!(self.load_device(lv));
        self.resume_device(lv)
    }

    /// Deactivate a Logical Volume.
    pub fn deactivate_device(&self, lv: &LV) -> Result<()> {
        try!(self.suspend_device(lv));
        self.remove_device(lv)
    }
}

//
// Return up to the first \0, or None
//
fn slice_to_null(slc: &[u8]) -> Option<&[u8]> {
    for (i, c) in slc.iter().enumerate() {
        if *c == b'\0' { return Some(&slc[..i]) };
    }
    None
}
<commit_msg>Add consts for UUID_NAME_LEN and DM_NAME_LEN<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! Low-level devicemapper configuration of the running kernel.

#[allow(dead_code, non_camel_case_types)]
mod dm_ioctl;

use std::fs::File;
use std::io::{Result, Error, BufReader, BufRead};
use std::os::unix::io::AsRawFd;
use std::mem;
use std::slice;
use std::slice::bytes::copy_memory;
use std::collections::BTreeSet;

use byteorder::{NativeEndian, ByteOrder};

use nix::sys::ioctl;

use dm::dm_ioctl as dmi;
use lv::LV;
use vg::VG;
use util::align_to;
use ::Device;

const DM_IOCTL: u8 = 0xfd;
const DM_CTL_PATH: &'static str= "/dev/mapper/control";

const DM_VERSION_MAJOR: u32 = 4;
const DM_VERSION_MINOR: u32 = 30;
const DM_VERSION_PATCHLEVEL: u32 = 0;

const DM_NAME_LEN: usize = 128;
const DM_UUID_LEN: usize = 129;

// Status bits
//const DM_READONLY_FLAG: u32 = 1;
const DM_SUSPEND_FLAG: u32 = 2;
//const DM_PERSISTENT_DEV_FLAG: u32 = 8;

/// Major numbers used by DM.
pub fn dev_majors() -> BTreeSet<u32> {
    let mut set = BTreeSet::new();

    let f = File::open("/proc/devices")
        .ok().expect("Could not open /proc/devices");

    let reader = BufReader::new(f);

    for line in reader.lines()
        .filter_map(|x| x.ok())
        .skip_while(|x| x != "Block devices:")
        .skip(1) {
            let spl: Vec<_> = line.split_whitespace().collect();

            if spl[1] == "device-mapper" {
                set.insert(spl[0].parse::<u32>().unwrap());
            }
        }

    set
}

/// Context needed for communicating with devicemapper.
pub struct DM<'a> {
    file: File,
    vg: &'a VG,
}

impl <'a> DM<'a> {
    /// Create a new context for communicating about a given VG with DM.
    pub fn new(vg: &'a VG) -> Result<Self> {
        Ok(DM {
            file: try!(File::open(DM_CTL_PATH)),
            vg: vg,
        })
    }

    fn initialize_hdr(hdr: &mut dmi::Struct_dm_ioctl) -> () {
        hdr.version[0] = DM_VERSION_MAJOR;
        hdr.version[1] = DM_VERSION_MINOR;
        hdr.version[2] = DM_VERSION_PATCHLEVEL;

        hdr.data_start = mem::size_of::<dmi::Struct_dm_ioctl>() as u32;
    }

    fn hdr_set_name(hdr: &mut dmi::Struct_dm_ioctl, vg_name: &str, lv_name: &str) -> () {
        let name = format!("{}-{}", vg_name.replace("-", "--"),
                           lv_name.replace("-", "--"));
        let name_dest: &mut [u8; DM_NAME_LEN] = unsafe { mem::transmute(&mut hdr.name) };
        copy_memory(name.as_bytes(), &mut name_dest[..]);
    }

    fn hdr_set_uuid(hdr: &mut dmi::Struct_dm_ioctl, vg_uuid: &str, lv_uuid: &str) -> () {
        let uuid = format!("LVM-{}{}", vg_uuid.replace("-", ""), lv_uuid.replace("-", ""));
        let uuid_dest: &mut [u8; DM_UUID_LEN] = unsafe { mem::transmute(&mut hdr.uuid) };
        copy_memory(uuid.as_bytes(), &mut uuid_dest[..]);
    }

    /// Devicemapper version information: Major, Minor, and patchlevel versions.
    pub fn get_version(&self) -> Result<(u32, u32, u32)> {

        let mut hdr: dmi::Struct_dm_ioctl = Default::default();
        hdr.version[0] = DM_VERSION_MAJOR;
        hdr.version[1] = DM_VERSION_MINOR;
        hdr.version[2] = DM_VERSION_PATCHLEVEL;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_VERSION_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => {},
        };

        Ok((hdr.version[0], hdr.version[1], hdr.version[2]))
    }

    /// Returns a list of tuples containing DM device names within this VG,
    /// and their major/minor device numbers.
    pub fn list_devices(&self) -> Result<Vec<(String, Device)>> {
        let mut buf = [0u8; 16 * 1024];
        let mut hdr: &mut dmi::Struct_dm_ioctl = unsafe {mem::transmute(&mut buf)};

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = buf.len() as u32;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_LIST_DEVICES_CMD as u8, buf.len());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut buf) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => {},
        };

        let mut devs = Vec::new();
        if (hdr.data_size - hdr.data_start as u32) != 0 {
            let mut result = &buf[hdr.data_start as usize..];

            loop {
                let slc = slice_to_null(&result[12..]).expect("Bad data from ioctl");
                let devno = NativeEndian::read_u64(&result[..8]);
                let dm_name = String::from_utf8_lossy(slc);
                let mut vg_name = self.vg.name().replace("-", "--");
                vg_name.push('-');

                // Return only devices within this VG
                if dm_name.starts_with(&vg_name) {
                    let lv_name = dm_name.trim_left_matches(&vg_name).replace("--", "-");
                    devs.push((lv_name, devno.into()));
                }

                let next = NativeEndian::read_u32(&result[8..12]);
                if next == 0 { break }

                result = &result[next as usize..];
            }
        }

        Ok(devs)
    }

    /// Query DM for which devices depend on this device.
    pub fn list_deps(&self, dev: Device) -> Result<Vec<Device>> {
        let mut buf = [0u8; 16 * 1024];
        let mut hdr: &mut dmi::Struct_dm_ioctl = unsafe {mem::transmute(&mut buf)};

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = buf.len() as u32;
        hdr.dev = dev.into();

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_TABLE_DEPS_CMD as u8, buf.len());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut buf) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => {},
        };

        // TODO: Check DM_BUFFER_FULL_FLAG for:
        // DM_DEVICE_LIST_VERSIONS, DM_DEVICE_LIST, DM_DEVICE_DEPS,
        // DM_DEVICE_STATUS, DM_DEVICE_TABLE, DM_DEVICE_WAITEVENT,
        // DM_DEVICE_TARGET_MSG

        let mut devs = Vec::new();
        if (hdr.data_size - hdr.data_start as u32) != 0 {
            let result = &buf[hdr.data_start as usize..];
            let entries = NativeEndian::read_u32(&result[..4]) as usize;

            for entry in 0..entries {
                let dev = &result[(8*entry)+8..(8*entry)+16];
                devs.push(Device::from(NativeEndian::read_u64(&dev)));
            }
        }

        Ok(devs)
    }

    /// Recursively walk DM deps to see if device is present
    pub fn depends_on(&self, dev: Device, dm_majors: &BTreeSet<u32>) -> bool {
        if !dm_majors.contains(&dev.major) {
            return false;
        }

        if let Ok(dep_list) = self.list_deps(dev) {
            for d in dep_list {
                if d == dev {
                    return true;
                } else if self.depends_on(d, dm_majors) {
                    return true;
                }
            }
        }

        false
    }

    fn create_device(&self, lv: &mut LV) -> Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        Self::hdr_set_name(&mut hdr, &self.vg.name(), &lv.name);
        Self::hdr_set_uuid(&mut hdr, &self.vg.id(), &lv.id);
        hdr.data_size = hdr.data_start;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_CREATE_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => { }
        };

        lv.device = Some(Device::from(hdr.dev));

        Ok(())
    }

    fn remove_device(&self, lv: &LV) -> Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = hdr.data_start;
        Self::hdr_set_name(&mut hdr, &self.vg.name(), &lv.name);

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_REMOVE_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => Ok(())
        }
    }

    fn load_device(&self, lv: &LV) -> Result<()> {
        let sectors_per_extent = self.vg.extent_size();
        let mut targs = Vec::new();

        // Construct targets first, since we need to know how many & size
        // before initializing the header.
        for seg in &lv.segments {
            let mut targ: dmi::Struct_dm_target_spec = Default::default();
            targ.sector_start = seg.start_extent() * sectors_per_extent;
            targ.length = seg.extent_count() * sectors_per_extent;
            targ.status = 0;

            let mut dst: &mut [u8] = unsafe {
                mem::transmute(&mut targ.target_type[..])
            };
            copy_memory(seg.dm_type().as_bytes(), &mut dst);

            let mut params = seg.dm_params(self.vg);

            let pad_bytes = align_to(
                params.len() + 1usize, 8usize) - params.len();
            params.extend(vec!["\0"; pad_bytes]);

            targ.next = (mem::size_of::<dmi::Struct_dm_target_spec>()
                         + params.len()) as u32;

            targs.push((targ, params));
        }

        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        Self::hdr_set_name(&mut hdr, &self.vg.name(), &lv.name);

        hdr.data_start = mem::size_of::<dmi::Struct_dm_ioctl>() as u32;
        hdr.data_size = hdr.data_start + targs.iter()
            .map(|&(t, _)| t.next)
            .sum::<u32>();
        hdr.target_count = targs.len() as u32;

        // Flatten into buf
        let mut buf: Vec<u8> = Vec::with_capacity(hdr.data_size as usize);
        unsafe {
            let ptr: *mut u8 = mem::transmute(&mut hdr);
            let slc = slice::from_raw_parts(ptr, hdr.data_start as usize);
            buf.extend(slc);
        }

        for (targ, param) in targs {
            unsafe {
                let ptr: *mut u8 = mem::transmute(&targ);
                let slc = slice::from_raw_parts(
                    ptr, mem::size_of::<dmi::Struct_dm_target_spec>());
                buf.extend(slc);
            }

            buf.extend(param.as_bytes());
        }

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_TABLE_LOAD_CMD as u8, buf.len());

        match unsafe { ioctl::read_into_ptr(self.file.as_raw_fd(), op, buf.as_mut_ptr()) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => Ok(())
        }
    }

    fn suspend_device(&self, lv: &LV) -> Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = hdr.data_start;
        Self::hdr_set_name(&mut hdr, &self.vg.name(), &lv.name);
        hdr.flags = DM_SUSPEND_FLAG;

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_SUSPEND_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => Ok(())
        }
    }

    fn resume_device(&self, lv: &LV) -> Result<()> {
        let mut hdr: dmi::Struct_dm_ioctl = Default::default();

        Self::initialize_hdr(&mut hdr);
        hdr.data_size = hdr.data_start;
        Self::hdr_set_name(&mut hdr, &self.vg.name(), &lv.name);
        // DM_SUSPEND_FLAG not set = resume

        let op = ioctl::op_read_write(DM_IOCTL, dmi::DM_DEV_SUSPEND_CMD as u8,
                                      mem::size_of::<dmi::Struct_dm_ioctl>());

        match unsafe { ioctl::read_into(self.file.as_raw_fd(), op, &mut hdr) } {
            Err(_) => return Err((Error::last_os_error())),
            _ => Ok(())
        }
    }

    /// Activate a Logical Volume.
    ///
    /// Also populates the LV's device field.
    pub fn activate_device(&self, lv: &mut LV) -> Result<()> {
        try!(self.create_device(lv));
        try!(self.load_device(lv));
        self.resume_device(lv)
    }

    /// Deactivate a Logical Volume.
    pub fn deactivate_device(&self, lv: &LV) -> Result<()> {
        try!(self.suspend_device(lv));
        self.remove_device(lv)
    }
}

//
// Return up to the first \0, or None
//
fn slice_to_null(slc: &[u8]) -> Option<&[u8]> {
    for (i, c) in slc.iter().enumerate() {
        if *c == b'\0' { return Some(&slc[..i]) };
    }
    None
}
<|endoftext|>


<commit_before>// error-pattern:yummy
#![feature(box_syntax)]
#![feature(rustc_private)]
#![allow(unknown_lints, missing_docs_in_private_items)]

extern crate clippy_lints;
extern crate getopts;
extern crate rustc;
extern crate rustc_driver;
extern crate rustc_errors;
extern crate rustc_plugin;
extern crate rustc_trans_utils;
extern crate syntax;

use rustc_driver::{driver, Compilation, CompilerCalls, RustcDefaultCalls};
use rustc_trans_utils::trans_crate::TransCrate;
use rustc::session::{config, Session};
use rustc::session::config::{ErrorOutputType, Input};
use std::path::PathBuf;
use std::process::Command;
use syntax::ast;

struct ClippyCompilerCalls {
    default: RustcDefaultCalls,
    run_lints: bool,
}

impl ClippyCompilerCalls {
    fn new(run_lints: bool) -> Self {
        Self {
            default: RustcDefaultCalls,
            run_lints,
        }
    }
}

impl<'a> CompilerCalls<'a> for ClippyCompilerCalls {
    fn early_callback(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        descriptions: &rustc_errors::registry::Registry,
        output: ErrorOutputType,
    ) -> Compilation {
        self.default
            .early_callback(matches, sopts, cfg, descriptions, output)
    }
    fn no_input(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
        descriptions: &rustc_errors::registry::Registry,
    ) -> Option<(Input, Option<PathBuf>)> {
        self.default
            .no_input(matches, sopts, cfg, odir, ofile, descriptions)
    }
    fn late_callback(
        &mut self,
        trans_crate: &TransCrate,
        matches: &getopts::Matches,
        sess: &Session,
        crate_stores: &rustc::middle::cstore::CrateStore,
        input: &Input,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
    ) -> Compilation {
        self.default
            .late_callback(trans_crate, matches, sess, crate_stores, input, odir, ofile)
    }
    fn build_controller(&mut self, sess: &Session, matches: &getopts::Matches) -> driver::CompileController<'a> {
        let mut control = self.default.build_controller(sess, matches);

        if self.run_lints {
            let old = std::mem::replace(&mut control.after_parse.callback, box |_| {});
            control.after_parse.callback = Box::new(move |state| {
                {
                    let mut registry = rustc_plugin::registry::Registry::new(
                        state.session,
                        state
                            .krate
                            .as_ref()
                            .expect(
                                "at this compilation stage \
                                 the crate must be parsed",
                            )
                            .span,
                    );
                    registry.args_hidden = Some(Vec::new());
                    clippy_lints::register_plugins(&mut registry);

                    let rustc_plugin::registry::Registry {
                        early_lint_passes,
                        late_lint_passes,
                        lint_groups,
                        llvm_passes,
                        attributes,
                        ..
                    } = registry;
                    let sess = &state.session;
                    let mut ls = sess.lint_store.borrow_mut();
                    for pass in early_lint_passes {
                        ls.register_early_pass(Some(sess), true, pass);
                    }
                    for pass in late_lint_passes {
                        ls.register_late_pass(Some(sess), true, pass);
                    }

                    for (name, to) in lint_groups {
                        ls.register_group(Some(sess), true, name, to);
                    }

                    sess.plugin_llvm_passes.borrow_mut().extend(llvm_passes);
                    sess.plugin_attributes.borrow_mut().extend(attributes);
                }
                old(state);
            });
        }

        control
    }
}

#[allow(print_stdout)]
fn show_version() {
    println!(env!("CARGO_PKG_VERSION"));
}

pub fn main() {
    use std::env;

    if std::env::args().any(|a| a == "--version" || a == "-V") {
        show_version();
        return;
    }

    let sys_root = option_env!("SYSROOT")
        .map(String::from)
        .or_else(|| std::env::var("SYSROOT").ok())
        .or_else(|| {
            let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
            let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
            home.and_then(|home| toolchain.map(|toolchain| format!("{}/toolchains/{}", home, toolchain)))
        })
        .or_else(|| {
            Command::new("rustc")
                .arg("--print")
                .arg("sysroot")
                .output()
                .ok()
                .and_then(|out| String::from_utf8(out.stdout).ok())
                .map(|s| s.trim().to_owned())
        })
        .expect("need to specify SYSROOT env var during clippy compilation, or use rustup or multirust");

    // Setting RUSTC_WRAPPER causes Cargo to pass 'rustc' as the first argument.
    // We're invoking the compiler programmatically, so we ignore this/
    let mut orig_args: Vec<String> = env::args().collect();
    if orig_args.len() <= 1 {
        std::process::exit(1);
    }
    if orig_args[1] == "rustc" {
        // we still want to be able to invoke it normally though
        orig_args.remove(1);
    }
    // this conditional check for the --sysroot flag is there so users can call
    // `clippy_driver` directly
    // without having to pass --sysroot or anything
    let mut args: Vec<String> = if orig_args.iter().any(|s| s == "--sysroot") {
        orig_args.clone()
    } else {
        orig_args
            .clone()
            .into_iter()
            .chain(Some("--sysroot".to_owned()))
            .chain(Some(sys_root))
            .collect()
    };

    // this check ensures that dependencies are built but not linted and the final
    // crate is
    // linted but not built
    let clippy_enabled = env::var("CLIPPY_TESTS")
        .ok()
        .map_or(false, |val| val == "true")
        || orig_args.iter().any(|s| s == "--emit=dep-info,metadata");

    if clippy_enabled {
        args.extend_from_slice(&["--cfg".to_owned(), r#"feature="cargo-clippy""#.to_owned()]);
        if let Ok(extra_args) = env::var("CLIPPY_ARGS") {
            args.extend(extra_args.split("__CLIPPY_HACKERY__").filter(|s| !s.is_empty()).map(str::to_owned));
        }
    }

    let mut ccc = ClippyCompilerCalls::new(clippy_enabled);
    rustc_driver::run(move || rustc_driver::run_compiler(&args, &mut ccc, None, None));
}
<commit_msg>Rename trans to codegen<commit_after>// error-pattern:yummy
#![feature(box_syntax)]
#![feature(rustc_private)]
#![allow(unknown_lints, missing_docs_in_private_items)]

extern crate clippy_lints;
extern crate getopts;
extern crate rustc;
extern crate rustc_driver;
extern crate rustc_errors;
extern crate rustc_plugin;
extern crate rustc_codegen_utils;
extern crate syntax;

use rustc_driver::{driver, Compilation, CompilerCalls, RustcDefaultCalls};
use rustc_codegen_utils::codegen_backend::CodegenBackend;
use rustc::session::{config, Session};
use rustc::session::config::{ErrorOutputType, Input};
use std::path::PathBuf;
use std::process::Command;
use syntax::ast;

struct ClippyCompilerCalls {
    default: RustcDefaultCalls,
    run_lints: bool,
}

impl ClippyCompilerCalls {
    fn new(run_lints: bool) -> Self {
        Self {
            default: RustcDefaultCalls,
            run_lints,
        }
    }
}

impl<'a> CompilerCalls<'a> for ClippyCompilerCalls {
    fn early_callback(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        descriptions: &rustc_errors::registry::Registry,
        output: ErrorOutputType,
    ) -> Compilation {
        self.default
            .early_callback(matches, sopts, cfg, descriptions, output)
    }
    fn no_input(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
        descriptions: &rustc_errors::registry::Registry,
    ) -> Option<(Input, Option<PathBuf>)> {
        self.default
            .no_input(matches, sopts, cfg, odir, ofile, descriptions)
    }
    fn late_callback(
        &mut self,
        trans_crate: &CodegenBackend,
        matches: &getopts::Matches,
        sess: &Session,
        crate_stores: &rustc::middle::cstore::CrateStore,
        input: &Input,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
    ) -> Compilation {
        self.default
            .late_callback(trans_crate, matches, sess, crate_stores, input, odir, ofile)
    }
    fn build_controller(&mut self, sess: &Session, matches: &getopts::Matches) -> driver::CompileController<'a> {
        let mut control = self.default.build_controller(sess, matches);

        if self.run_lints {
            let old = std::mem::replace(&mut control.after_parse.callback, box |_| {});
            control.after_parse.callback = Box::new(move |state| {
                {
                    let mut registry = rustc_plugin::registry::Registry::new(
                        state.session,
                        state
                            .krate
                            .as_ref()
                            .expect(
                                "at this compilation stage \
                                 the crate must be parsed",
                            )
                            .span,
                    );
                    registry.args_hidden = Some(Vec::new());
                    clippy_lints::register_plugins(&mut registry);

                    let rustc_plugin::registry::Registry {
                        early_lint_passes,
                        late_lint_passes,
                        lint_groups,
                        llvm_passes,
                        attributes,
                        ..
                    } = registry;
                    let sess = &state.session;
                    let mut ls = sess.lint_store.borrow_mut();
                    for pass in early_lint_passes {
                        ls.register_early_pass(Some(sess), true, pass);
                    }
                    for pass in late_lint_passes {
                        ls.register_late_pass(Some(sess), true, pass);
                    }

                    for (name, to) in lint_groups {
                        ls.register_group(Some(sess), true, name, to);
                    }

                    sess.plugin_llvm_passes.borrow_mut().extend(llvm_passes);
                    sess.plugin_attributes.borrow_mut().extend(attributes);
                }
                old(state);
            });
        }

        control
    }
}

#[allow(print_stdout)]
fn show_version() {
    println!(env!("CARGO_PKG_VERSION"));
}

pub fn main() {
    use std::env;

    if std::env::args().any(|a| a == "--version" || a == "-V") {
        show_version();
        return;
    }

    let sys_root = option_env!("SYSROOT")
        .map(String::from)
        .or_else(|| std::env::var("SYSROOT").ok())
        .or_else(|| {
            let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
            let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
            home.and_then(|home| toolchain.map(|toolchain| format!("{}/toolchains/{}", home, toolchain)))
        })
        .or_else(|| {
            Command::new("rustc")
                .arg("--print")
                .arg("sysroot")
                .output()
                .ok()
                .and_then(|out| String::from_utf8(out.stdout).ok())
                .map(|s| s.trim().to_owned())
        })
        .expect("need to specify SYSROOT env var during clippy compilation, or use rustup or multirust");

    // Setting RUSTC_WRAPPER causes Cargo to pass 'rustc' as the first argument.
    // We're invoking the compiler programmatically, so we ignore this/
    let mut orig_args: Vec<String> = env::args().collect();
    if orig_args.len() <= 1 {
        std::process::exit(1);
    }
    if orig_args[1] == "rustc" {
        // we still want to be able to invoke it normally though
        orig_args.remove(1);
    }
    // this conditional check for the --sysroot flag is there so users can call
    // `clippy_driver` directly
    // without having to pass --sysroot or anything
    let mut args: Vec<String> = if orig_args.iter().any(|s| s == "--sysroot") {
        orig_args.clone()
    } else {
        orig_args
            .clone()
            .into_iter()
            .chain(Some("--sysroot".to_owned()))
            .chain(Some(sys_root))
            .collect()
    };

    // this check ensures that dependencies are built but not linted and the final
    // crate is
    // linted but not built
    let clippy_enabled = env::var("CLIPPY_TESTS")
        .ok()
        .map_or(false, |val| val == "true")
        || orig_args.iter().any(|s| s == "--emit=dep-info,metadata");

    if clippy_enabled {
        args.extend_from_slice(&["--cfg".to_owned(), r#"feature="cargo-clippy""#.to_owned()]);
        if let Ok(extra_args) = env::var("CLIPPY_ARGS") {
            args.extend(extra_args.split("__CLIPPY_HACKERY__").filter(|s| !s.is_empty()).map(str::to_owned));
        }
    }

    let mut ccc = ClippyCompilerCalls::new(clippy_enabled);
    rustc_driver::run(move || rustc_driver::run_compiler(&args, &mut ccc, None, None));
}
<|endoftext|>


<commit_before>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
use self::WhichLine::*;

use std::fs::File;
use std::io::BufReader;
use std::io::prelude::*;
use std::path::Path;

pub struct ExpectedError {
    pub line: usize,
    pub kind: String,
    pub msg: String,
}

#[derive(PartialEq, Debug)]
enum WhichLine { ThisLine, FollowPrevious(usize), AdjustBackward(usize) }

/// Looks for either "//~| KIND MESSAGE" or "//~^^... KIND MESSAGE"
/// The former is a "follow" that inherits its target from the preceding line;
/// the latter is an "adjusts" that goes that many lines up.
///
/// Goal is to enable tests both like: //~^^^ ERROR go up three
/// and also //~^ ERROR message one for the preceding line, and
///          //~| ERROR message two for that same line.
///
/// If cfg is not None (i.e., in an incremental test), then we look
/// for `//[X]~` instead, where `X` is the current `cfg`.
pub fn load_errors(testfile: &Path, cfg: Option<&str>) -> Vec<ExpectedError> {
    let rdr = BufReader::new(File::open(testfile).unwrap());

    // `last_nonfollow_error` tracks the most recently seen
    // line with an error template that did not use the
    // follow-syntax, "//~| ...".
    //
    // (pnkfelix could not find an easy way to compose Iterator::scan
    // and Iterator::filter_map to pass along this information into
    // `parse_expected`. So instead I am storing that state here and
    // updating it in the map callback below.)
    let mut last_nonfollow_error = None;

    let tag = match cfg {
        Some(rev) => format!("//[{}]~", rev),
        None => format!("//~")
    };

    rdr.lines()
       .enumerate()
       .filter_map(|(line_no, ln)| {
           parse_expected(last_nonfollow_error,
                          line_no + 1,
                          &ln.unwrap(),
                          &tag)
               .map(|(which, error)| {
                   match which {
                       FollowPrevious(_) => {}
                       _ => last_nonfollow_error = Some(error.line),
                   }
                   error
               })
       })
       .collect()
}

fn parse_expected(last_nonfollow_error: Option<usize>,
                  line_num: usize,
                  line: &str,
                  tag: &str)
                  -> Option<(WhichLine, ExpectedError)> {
    let start = match line.find(tag) { Some(i) => i, None => return None };
    let (follow, adjusts) = if line.char_at(start + tag.len()) == '|' {
        (true, 0)
    } else {
        (false, line[start + tag.len()..].chars().take_while(|c| *c == '^').count())
    };
    let kind_start = start + tag.len() + adjusts + (follow as usize);
    let letters = line[kind_start..].chars();
    let kind = letters.skip_while(|c| c.is_whitespace())
                      .take_while(|c| !c.is_whitespace())
                      .flat_map(|c| c.to_lowercase())
                      .collect::<String>();
    let letters = line[kind_start..].chars();
    let msg = letters.skip_while(|c| c.is_whitespace())
                     .skip_while(|c| !c.is_whitespace())
                     .collect::<String>().trim().to_owned();

    let (which, line) = if follow {
        assert!(adjusts == 0, "use either //~| or //~^, not both.");
        let line = last_nonfollow_error.unwrap_or_else(|| {
            panic!("encountered //~| without preceding //~^ line.")
        });
        (FollowPrevious(line), line)
    } else {
        let which =
            if adjusts > 0 { AdjustBackward(adjusts) } else { ThisLine };
        let line = line_num - adjusts;
        (which, line)
    };

    debug!("line={} tag={:?} which={:?} kind={:?} msg={:?}",
           line_num, tag, which, kind, msg);

    Some((which, ExpectedError { line: line,
                                 kind: kind,
                                 msg: msg, }))
}
<commit_msg>Prefer `Option::expect` over explicit unwrapping.<commit_after>// Copyright 2012-2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
use self::WhichLine::*;

use std::fs::File;
use std::io::BufReader;
use std::io::prelude::*;
use std::path::Path;

pub struct ExpectedError {
    pub line: usize,
    pub kind: String,
    pub msg: String,
}

#[derive(PartialEq, Debug)]
enum WhichLine { ThisLine, FollowPrevious(usize), AdjustBackward(usize) }

/// Looks for either "//~| KIND MESSAGE" or "//~^^... KIND MESSAGE"
/// The former is a "follow" that inherits its target from the preceding line;
/// the latter is an "adjusts" that goes that many lines up.
///
/// Goal is to enable tests both like: //~^^^ ERROR go up three
/// and also //~^ ERROR message one for the preceding line, and
///          //~| ERROR message two for that same line.
///
/// If cfg is not None (i.e., in an incremental test), then we look
/// for `//[X]~` instead, where `X` is the current `cfg`.
pub fn load_errors(testfile: &Path, cfg: Option<&str>) -> Vec<ExpectedError> {
    let rdr = BufReader::new(File::open(testfile).unwrap());

    // `last_nonfollow_error` tracks the most recently seen
    // line with an error template that did not use the
    // follow-syntax, "//~| ...".
    //
    // (pnkfelix could not find an easy way to compose Iterator::scan
    // and Iterator::filter_map to pass along this information into
    // `parse_expected`. So instead I am storing that state here and
    // updating it in the map callback below.)
    let mut last_nonfollow_error = None;

    let tag = match cfg {
        Some(rev) => format!("//[{}]~", rev),
        None => format!("//~")
    };

    rdr.lines()
       .enumerate()
       .filter_map(|(line_no, ln)| {
           parse_expected(last_nonfollow_error,
                          line_no + 1,
                          &ln.unwrap(),
                          &tag)
               .map(|(which, error)| {
                   match which {
                       FollowPrevious(_) => {}
                       _ => last_nonfollow_error = Some(error.line),
                   }
                   error
               })
       })
       .collect()
}

fn parse_expected(last_nonfollow_error: Option<usize>,
                  line_num: usize,
                  line: &str,
                  tag: &str)
                  -> Option<(WhichLine, ExpectedError)> {
    let start = match line.find(tag) { Some(i) => i, None => return None };
    let (follow, adjusts) = if line.char_at(start + tag.len()) == '|' {
        (true, 0)
    } else {
        (false, line[start + tag.len()..].chars().take_while(|c| *c == '^').count())
    };
    let kind_start = start + tag.len() + adjusts + (follow as usize);
    let letters = line[kind_start..].chars();
    let kind = letters.skip_while(|c| c.is_whitespace())
                      .take_while(|c| !c.is_whitespace())
                      .flat_map(|c| c.to_lowercase())
                      .collect::<String>();
    let letters = line[kind_start..].chars();
    let msg = letters.skip_while(|c| c.is_whitespace())
                     .skip_while(|c| !c.is_whitespace())
                     .collect::<String>().trim().to_owned();

    let (which, line) = if follow {
        assert!(adjusts == 0, "use either //~| or //~^, not both.");
        let line = last_nonfollow_error.expect("encountered //~| without \
                                                preceding //~^ line.");
        (FollowPrevious(line), line)
    } else {
        let which =
            if adjusts > 0 { AdjustBackward(adjusts) } else { ThisLine };
        let line = line_num - adjusts;
        (which, line)
    };

    debug!("line={} tag={:?} which={:?} kind={:?} msg={:?}",
           line_num, tag, which, kind, msg);

    Some((which, ExpectedError { line: line,
                                 kind: kind,
                                 msg: msg, }))
}
<|endoftext|>


<commit_before>use decoders::*;
use decoders::tiff::*;
use decoders::basics::*;
use std::mem::transmute;
use std::f32;
use std::cmp;

#[derive(Debug, Clone)]
pub struct ArwDecoder<'a> {
  buffer: &'a [u8],
  rawloader: &'a RawLoader,
  tiff: TiffIFD<'a>,
}

impl<'a> ArwDecoder<'a> {
  pub fn new(buf: &'a [u8], tiff: TiffIFD<'a>, rawloader: &'a RawLoader) -> ArwDecoder<'a> {
    ArwDecoder {
      buffer: buf,
      tiff: tiff,
      rawloader: rawloader,
    }
  }
}

impl<'a> Decoder for ArwDecoder<'a> {
  fn image(&self) -> Result<Image,String> {
    let camera = try!(self.rawloader.check_supported(&self.tiff));
    let data = self.tiff.find_ifds_with_tag(Tag::StripOffsets);
    if data.len() == 0 {
      if camera.model == "DSLR-A100" {
        return self.image_a100(camera)
      } else { // try decoding as SRF
        return self.image_srf(camera)
      }
    }
    let raw = data[0];
    let width = fetch_tag!(raw, Tag::ImageWidth).get_u32(0);
    let mut height = fetch_tag!(raw, Tag::ImageLength).get_u32(0);
    let offset = fetch_tag!(raw, Tag::StripOffsets).get_u32(0) as usize;
    let count = fetch_tag!(raw, Tag::StripByteCounts).get_u32(0) as usize;
    let compression = fetch_tag!(raw, Tag::Compression).get_u32(0);
    let bps = if camera.bps != 0 {
      camera.bps
    } else {
      fetch_tag!(raw, Tag::BitsPerSample).get_u32(0)
    };
    let src = &self.buffer[offset .. self.buffer.len()];

    let image = match compression {
      1 => {
        if camera.model == "DSC-R1" {
          decode_14be_unpacked(src, width as usize, height as usize)
        } else {
          decode_16le(src, width as usize, height as usize)
        }
      }
      32767 => {
        if ((width*height*bps) as usize) != count*8 {
          height += 8;
          ArwDecoder::decode_arw1(src, width as usize, height as usize)
        } else {
          match bps {
            8 => {let curve = try!(ArwDecoder::get_curve(raw)); ArwDecoder::decode_arw2(src, width as usize, height as usize, &curve)},
            12 => decode_12le(src, width as usize, height as usize),
            _ => return Err(format!("ARW2: Don't know how to decode images with {} bps", bps)),
          }
        }
      },
      _ => return Err(format!("ARW: Don't know how to decode type {}", compression).to_string()),
    };

    ok_image(camera, width, height, try!(self.get_wb()), image)
  }
}

impl<'a> ArwDecoder<'a> {
  fn image_a100(&self, camera: &Camera) -> Result<Image,String> {
    // We've caught the elusive A100 in the wild, a transitional format
    // between the simple sanity of the MRW custom format and the wordly
    // wonderfullness of the Tiff-based ARW format, let's shoot from the hip
    let data = self.tiff.find_ifds_with_tag(Tag::SubIFDs);
    if data.len() == 0 {
      return Err("ARW: Couldn't find the data IFD!".to_string())
    }
    let raw = data[0];
    let width = 3881;
    let height = 2608;
    let offset = fetch_tag!(raw, Tag::SubIFDs).get_u32(0) as usize;

    let src = &self.buffer[offset .. self.buffer.len()];
    let image = ArwDecoder::decode_arw1(src, width as usize, height as usize);

    // Get the WB the MRW way
    let priv_offset = fetch_tag!(self.tiff, Tag::DNGPrivateArea).get_u32(0);
    let buf = &self.buffer[priv_offset as usize..];
    let mut currpos: usize = 8;
    let mut wb_vals: [u16;4] = [0;4];
    // At most we read 20 bytes from currpos so check we don't step outside that
    while currpos+20 < buf.len() {
      let tag: u32 = BEu32(buf,currpos);
      let len: u32 = LEu32(buf,currpos+4);
      if tag == 0x574247 { // WBG
        for i in 0..4 {
          wb_vals[i] = LEu16(buf, currpos+12+i*2);
        }
        break;
      }
      currpos += (len+8) as usize;
    }
    let wb_coeffs = [wb_vals[0] as f32, wb_vals[1] as f32, wb_vals[3] as f32, f32::NAN];

    ok_image(camera, width, height, wb_coeffs, image)
  }

  fn image_srf(&self, camera: &Camera) -> Result<Image,String> {
    let data = self.tiff.find_ifds_with_tag(Tag::ImageWidth);
    if data.len() == 0 {
      return Err("ARW: Couldn't find the data IFD!".to_string())
    }
    let raw = data[0];

    let width = fetch_tag!(raw, Tag::ImageWidth).get_u32(0);
    let height = fetch_tag!(raw, Tag::ImageLength).get_u32(0);
    let len = (width*height*2) as usize;

    // Constants taken from dcraw
    let off: usize = 862144;
    let key_off: usize = 200896;
    let head_off: usize = 164600;

    // Replicate the dcraw contortions to get the "decryption" key
    let offset = (self.buffer[key_off as usize] as usize)*4;
    let first_key = BEu32(self.buffer, key_off+offset);
    let head = ArwDecoder::sony_decrypt(self.buffer, head_off, 40, first_key);
    let second_key = LEu32(&head, 22);

    // "Decrypt" the whole image buffer
    let image_data = ArwDecoder::sony_decrypt(self.buffer, off, len, second_key);
    let image = decode_16be(&image_data, width as usize, height as usize);

    ok_image(camera, width, height, [f32::NAN,f32::NAN,f32::NAN,f32::NAN], image)
  }

  fn decode_arw1(buf: &[u8], width: usize, height: usize) -> Vec<u16> {
    let mut pump = BitPumpMSB::new(buf);
    let mut out: Vec<u16> = vec![0; (width*height) as usize];

    let mut sum: i32 = 0;
    for x in 0..width {
      let col = width-1-x;
      let mut row = 0;
      while row <= height {
        if row == height {
          row = 1;
        }

        let mut len: u32 = 4 - pump.get_bits(2);
        if len == 3 && pump.get_bits(1) != 0 {
          len = 0;
        } else if len == 4 {
          let zeros = pump.peek_bits(13).leading_zeros() - 19;
          len += zeros;
          pump.get_bits(cmp::min(13, zeros+1));
        }
        let diff: i32 = pump.get_ibits(len);
        sum += diff;
        if len > 0 && (diff & (1 << (len - 1))) == 0 {
          sum -= (1 << len) - 1;
        }
        out[row*width+col] = sum as u16;
        row += 2
      }
    }
    out
  }

  fn decode_arw2(buf: &[u8], width: usize, height: usize, curve: &LookupTable) -> Vec<u16> {
    decode_threaded(width, height, &(|out: &mut [u16], row| {
      let mut pump = BitPumpLSB::new(&buf[(row*width)..]);

      let mut random = pump.peek_bits(24);
      for out in out.chunks_mut(32) {
        // Process 32 pixels at a time in interleaved fashion
        for j in 0..2 {
          let max = pump.get_bits(11);
          let min = pump.get_bits(11);
          let delta = max-min;
          // Calculate the size of the data shift needed by how large the delta is
          // A delta with 11 bits requires a shift of 4, 10 bits of 3, etc
          let delta_shift: u32 = cmp::max(0, (32-(delta.leading_zeros() as i32)) - 7) as u32;
          let imax = pump.get_bits(4);
          let imin = pump.get_bits(4);

          for i in 0..16 {
            let val = if i == imax {
              max
            } else if i == imin {
              min
            } else {
              cmp::min(0x7ff,(pump.get_bits(7) << delta_shift) + min)
            };
            out[j+(i*2) as usize] = curve.dither((val<<1) as u16, &mut random);
          }
        }
      }
    }))
  }

  fn get_wb(&self) -> Result<[f32;4], String> {
    let priv_offset = fetch_tag!(self.tiff, Tag::DNGPrivateArea).get_force_u32(0);
    let priv_tiff = TiffIFD::new(self.buffer, priv_offset as usize, 0, 0, 0, LITTLE_ENDIAN).unwrap();
    let sony_offset = fetch_tag!(priv_tiff, Tag::SonyOffset).get_u32(0) as usize;
    let sony_length = fetch_tag!(priv_tiff, Tag::SonyLength).get_u32(0) as usize;
    let sony_key = fetch_tag!(priv_tiff, Tag::SonyKey).get_u32(0);
    let decrypted_buf = ArwDecoder::sony_decrypt(self.buffer, sony_offset, sony_length, sony_key);
    let decrypted_tiff = TiffIFD::new(&decrypted_buf, 0, sony_offset, 0, 0, LITTLE_ENDIAN).unwrap();
    let grgb_levels = decrypted_tiff.find_entry(Tag::SonyGRBG);
    let rggb_levels = decrypted_tiff.find_entry(Tag::SonyRGGB);
    if grgb_levels.is_some() {
      let levels = grgb_levels.unwrap();
      Ok([levels.get_u32(1) as f32, levels.get_u32(0) as f32, levels.get_u32(2) as f32, f32::NAN])
    } else if rggb_levels.is_some() {
      let levels = rggb_levels.unwrap();
      Ok([levels.get_u32(0) as f32, levels.get_u32(1) as f32, levels.get_u32(3) as f32, f32::NAN])
    } else {
      Err("ARW: Couldn't find GRGB or RGGB levels".to_string())
    }
  }

  fn get_curve(raw: &TiffIFD) -> Result<LookupTable, String> {
    let centry = fetch_tag!(raw, Tag::SonyCurve);
    let mut curve: [u32;6] = [ 0, 0, 0, 0, 0, 4095 ];

    for i in 0..4 {
      curve[i+1] = ((centry.get_u32(i) >> 2) & 0xfff) as u32;
    }

    let mut out = vec![0 as u16; (curve[5]+1) as usize];
    for i in 0..5 {
      for j in (curve[i]+1)..(curve[i+1]+1) {
        out[j as usize] = out[(j-1) as usize] + (1<<i);
      }
    }

    Ok(LookupTable::new(&out))
  }

  fn sony_decrypt(buf: &[u8], offset: usize, length: usize, key: u32) -> Vec<u8>{
    let mut pad: [u32; 128] = [0 as u32; 128];
    let mut mkey = key;
    // Initialize the decryption pad from the key
    for p in 0..4 {
      mkey = mkey.wrapping_mul(48828125).wrapping_add(1);
      pad[p] = mkey;
    }
    pad[3] = pad[3] << 1 | (pad[0]^pad[2]) >> 31;
    for p in 4..127 {
      pad[p] = (pad[p-4]^pad[p-2]) << 1 | (pad[p-3]^pad[p-1]) >> 31;
    }
    for p in 0..127 {
      pad[p] = u32::from_be(pad[p]);
    }

    let mut out = Vec::with_capacity(length+4);
    for i in 0..(length/4+1) {
      let p = i + 127;
      pad[p & 127] = pad[(p+1) & 127] ^ pad[(p+1+64) & 127];
      let output = LEu32(buf, offset+i*4) ^ pad[p & 127];
      let bytes: [u8;4] = unsafe { transmute(output.to_le()) };
      out.push(bytes[0]);
      out.push(bytes[1]);
      out.push(bytes[2]);
      out.push(bytes[3]);
    }
    out
  }
}
<commit_msg>ARW: Fix the Sony A100 WB and cleanup a bit<commit_after>use decoders::*;
use decoders::tiff::*;
use decoders::basics::*;
use std::mem::transmute;
use std::f32::NAN;
use std::cmp;

#[derive(Debug, Clone)]
pub struct ArwDecoder<'a> {
  buffer: &'a [u8],
  rawloader: &'a RawLoader,
  tiff: TiffIFD<'a>,
}

impl<'a> ArwDecoder<'a> {
  pub fn new(buf: &'a [u8], tiff: TiffIFD<'a>, rawloader: &'a RawLoader) -> ArwDecoder<'a> {
    ArwDecoder {
      buffer: buf,
      tiff: tiff,
      rawloader: rawloader,
    }
  }
}

impl<'a> Decoder for ArwDecoder<'a> {
  fn image(&self) -> Result<Image,String> {
    let camera = try!(self.rawloader.check_supported(&self.tiff));
    let data = self.tiff.find_ifds_with_tag(Tag::StripOffsets);
    if data.len() == 0 {
      if camera.model == "DSLR-A100" {
        return self.image_a100(camera)
      } else { // try decoding as SRF
        return self.image_srf(camera)
      }
    }
    let raw = data[0];
    let width = fetch_tag!(raw, Tag::ImageWidth).get_u32(0);
    let mut height = fetch_tag!(raw, Tag::ImageLength).get_u32(0);
    let offset = fetch_tag!(raw, Tag::StripOffsets).get_u32(0) as usize;
    let count = fetch_tag!(raw, Tag::StripByteCounts).get_u32(0) as usize;
    let compression = fetch_tag!(raw, Tag::Compression).get_u32(0);
    let bps = if camera.bps != 0 {
      camera.bps
    } else {
      fetch_tag!(raw, Tag::BitsPerSample).get_u32(0)
    };
    let src = &self.buffer[offset .. self.buffer.len()];

    let image = match compression {
      1 => {
        if camera.model == "DSC-R1" {
          decode_14be_unpacked(src, width as usize, height as usize)
        } else {
          decode_16le(src, width as usize, height as usize)
        }
      }
      32767 => {
        if ((width*height*bps) as usize) != count*8 {
          height += 8;
          ArwDecoder::decode_arw1(src, width as usize, height as usize)
        } else {
          match bps {
            8 => {let curve = try!(ArwDecoder::get_curve(raw)); ArwDecoder::decode_arw2(src, width as usize, height as usize, &curve)},
            12 => decode_12le(src, width as usize, height as usize),
            _ => return Err(format!("ARW2: Don't know how to decode images with {} bps", bps)),
          }
        }
      },
      _ => return Err(format!("ARW: Don't know how to decode type {}", compression).to_string()),
    };

    ok_image(camera, width, height, try!(self.get_wb()), image)
  }
}

impl<'a> ArwDecoder<'a> {
  fn image_a100(&self, camera: &Camera) -> Result<Image,String> {
    // We've caught the elusive A100 in the wild, a transitional format
    // between the simple sanity of the MRW custom format and the wordly
    // wonderfullness of the Tiff-based ARW format, let's shoot from the hip
    let data = self.tiff.find_ifds_with_tag(Tag::SubIFDs);
    if data.len() == 0 {
      return Err("ARW: Couldn't find the data IFD!".to_string())
    }
    let raw = data[0];
    let width = 3881;
    let height = 2608;
    let offset = fetch_tag!(raw, Tag::SubIFDs).get_u32(0) as usize;

    let src = &self.buffer[offset .. self.buffer.len()];
    let image = ArwDecoder::decode_arw1(src, width as usize, height as usize);

    // Get the WB the MRW way
    let priv_offset = fetch_tag!(self.tiff, Tag::DNGPrivateArea).get_force_u32(0);
    let buf = &self.buffer[priv_offset as usize..];
    let mut currpos: usize = 8;
    let mut wb_coeffs: [f32;4] = [0.0, 0.0, 0.0, NAN];
    // At most we read 20 bytes from currpos so check we don't step outside that
    while currpos+20 < buf.len() {
      let tag: u32 = BEu32(buf,currpos);
      let len: u32 = LEu32(buf,currpos+4);
      if tag == 0x574247 { // WBG
        wb_coeffs[0] = LEu16(buf, currpos+12) as f32;
        wb_coeffs[1] = LEu16(buf, currpos+14) as f32;
        wb_coeffs[2] = LEu16(buf, currpos+18) as f32;
        break;
      }
      currpos += (len+8) as usize;
    }

    ok_image(camera, width, height, wb_coeffs, image)
  }

  fn image_srf(&self, camera: &Camera) -> Result<Image,String> {
    let data = self.tiff.find_ifds_with_tag(Tag::ImageWidth);
    if data.len() == 0 {
      return Err("ARW: Couldn't find the data IFD!".to_string())
    }
    let raw = data[0];

    let width = fetch_tag!(raw, Tag::ImageWidth).get_u32(0);
    let height = fetch_tag!(raw, Tag::ImageLength).get_u32(0);
    let len = (width*height*2) as usize;

    // Constants taken from dcraw
    let off: usize = 862144;
    let key_off: usize = 200896;
    let head_off: usize = 164600;

    // Replicate the dcraw contortions to get the "decryption" key
    let offset = (self.buffer[key_off as usize] as usize)*4;
    let first_key = BEu32(self.buffer, key_off+offset);
    let head = ArwDecoder::sony_decrypt(self.buffer, head_off, 40, first_key);
    let second_key = LEu32(&head, 22);

    // "Decrypt" the whole image buffer
    let image_data = ArwDecoder::sony_decrypt(self.buffer, off, len, second_key);
    let image = decode_16be(&image_data, width as usize, height as usize);

    ok_image(camera, width, height, [NAN,NAN,NAN,NAN], image)
  }

  fn decode_arw1(buf: &[u8], width: usize, height: usize) -> Vec<u16> {
    let mut pump = BitPumpMSB::new(buf);
    let mut out: Vec<u16> = vec![0; (width*height) as usize];

    let mut sum: i32 = 0;
    for x in 0..width {
      let col = width-1-x;
      let mut row = 0;
      while row <= height {
        if row == height {
          row = 1;
        }

        let mut len: u32 = 4 - pump.get_bits(2);
        if len == 3 && pump.get_bits(1) != 0 {
          len = 0;
        } else if len == 4 {
          let zeros = pump.peek_bits(13).leading_zeros() - 19;
          len += zeros;
          pump.get_bits(cmp::min(13, zeros+1));
        }
        let diff: i32 = pump.get_ibits(len);
        sum += diff;
        if len > 0 && (diff & (1 << (len - 1))) == 0 {
          sum -= (1 << len) - 1;
        }
        out[row*width+col] = sum as u16;
        row += 2
      }
    }
    out
  }

  fn decode_arw2(buf: &[u8], width: usize, height: usize, curve: &LookupTable) -> Vec<u16> {
    decode_threaded(width, height, &(|out: &mut [u16], row| {
      let mut pump = BitPumpLSB::new(&buf[(row*width)..]);

      let mut random = pump.peek_bits(24);
      for out in out.chunks_mut(32) {
        // Process 32 pixels at a time in interleaved fashion
        for j in 0..2 {
          let max = pump.get_bits(11);
          let min = pump.get_bits(11);
          let delta = max-min;
          // Calculate the size of the data shift needed by how large the delta is
          // A delta with 11 bits requires a shift of 4, 10 bits of 3, etc
          let delta_shift: u32 = cmp::max(0, (32-(delta.leading_zeros() as i32)) - 7) as u32;
          let imax = pump.get_bits(4);
          let imin = pump.get_bits(4);

          for i in 0..16 {
            let val = if i == imax {
              max
            } else if i == imin {
              min
            } else {
              cmp::min(0x7ff,(pump.get_bits(7) << delta_shift) + min)
            };
            out[j+(i*2) as usize] = curve.dither((val<<1) as u16, &mut random);
          }
        }
      }
    }))
  }

  fn get_wb(&self) -> Result<[f32;4], String> {
    let priv_offset = fetch_tag!(self.tiff, Tag::DNGPrivateArea).get_force_u32(0);
    let priv_tiff = TiffIFD::new(self.buffer, priv_offset as usize, 0, 0, 0, LITTLE_ENDIAN).unwrap();
    let sony_offset = fetch_tag!(priv_tiff, Tag::SonyOffset).get_u32(0) as usize;
    let sony_length = fetch_tag!(priv_tiff, Tag::SonyLength).get_u32(0) as usize;
    let sony_key = fetch_tag!(priv_tiff, Tag::SonyKey).get_u32(0);
    let decrypted_buf = ArwDecoder::sony_decrypt(self.buffer, sony_offset, sony_length, sony_key);
    let decrypted_tiff = TiffIFD::new(&decrypted_buf, 0, sony_offset, 0, 0, LITTLE_ENDIAN).unwrap();
    let grgb_levels = decrypted_tiff.find_entry(Tag::SonyGRBG);
    let rggb_levels = decrypted_tiff.find_entry(Tag::SonyRGGB);
    if grgb_levels.is_some() {
      let levels = grgb_levels.unwrap();
      Ok([levels.get_u32(1) as f32, levels.get_u32(0) as f32, levels.get_u32(2) as f32, NAN])
    } else if rggb_levels.is_some() {
      let levels = rggb_levels.unwrap();
      Ok([levels.get_u32(0) as f32, levels.get_u32(1) as f32, levels.get_u32(3) as f32, NAN])
    } else {
      Err("ARW: Couldn't find GRGB or RGGB levels".to_string())
    }
  }

  fn get_curve(raw: &TiffIFD) -> Result<LookupTable, String> {
    let centry = fetch_tag!(raw, Tag::SonyCurve);
    let mut curve: [u32;6] = [ 0, 0, 0, 0, 0, 4095 ];

    for i in 0..4 {
      curve[i+1] = ((centry.get_u32(i) >> 2) & 0xfff) as u32;
    }

    let mut out = vec![0 as u16; (curve[5]+1) as usize];
    for i in 0..5 {
      for j in (curve[i]+1)..(curve[i+1]+1) {
        out[j as usize] = out[(j-1) as usize] + (1<<i);
      }
    }

    Ok(LookupTable::new(&out))
  }

  fn sony_decrypt(buf: &[u8], offset: usize, length: usize, key: u32) -> Vec<u8>{
    let mut pad: [u32; 128] = [0 as u32; 128];
    let mut mkey = key;
    // Initialize the decryption pad from the key
    for p in 0..4 {
      mkey = mkey.wrapping_mul(48828125).wrapping_add(1);
      pad[p] = mkey;
    }
    pad[3] = pad[3] << 1 | (pad[0]^pad[2]) >> 31;
    for p in 4..127 {
      pad[p] = (pad[p-4]^pad[p-2]) << 1 | (pad[p-3]^pad[p-1]) >> 31;
    }
    for p in 0..127 {
      pad[p] = u32::from_be(pad[p]);
    }

    let mut out = Vec::with_capacity(length+4);
    for i in 0..(length/4+1) {
      let p = i + 127;
      pad[p & 127] = pad[(p+1) & 127] ^ pad[(p+1+64) & 127];
      let output = LEu32(buf, offset+i*4) ^ pad[p & 127];
      let bytes: [u8;4] = unsafe { transmute(output.to_le()) };
      out.push(bytes[0]);
      out.push(bytes[1]);
      out.push(bytes[2]);
      out.push(bytes[3]);
    }
    out
  }
}
<|endoftext|>


<commit_before>use byteorder::{ByteOrder, LittleEndian};

use opcodes::{AddressingMode, OpCode};

pub fn disassemble(raw: &[u8]) -> String {
    let mut result = String::new();

    let mut i: usize = 0;
    while i < raw.len() {
        let opcode = OpCode::from_raw_byte(raw[i]);
        let val = match opcode.mode {
            AddressingMode::Implied |
            AddressingMode::Accumulator => String::from(""),
            AddressingMode::Immediate => format!(" #${:X}", raw[i + 1]),
            AddressingMode::Absolute => format!(" ${:X}", LittleEndian::read_u16(&raw[i + 1..])),
            _ => String::from(""),
        };
        let opcode_text = format!("{}{}\n", opcode.mnemonic, val);
        result.push_str(&opcode_text);
        i += opcode.length as usize;
    }

    result
}

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

    #[test]
    fn can_disassemble_basic_instructions() {
        let code: Vec<u8> = vec![0xA9, 0x20, 0x8D, 0x00, 0x44];
        let asm = disassemble(&code);

        assert_eq!(clean_asm("

            LDA #$20
            STA $4400

        "),
                   clean_asm(asm));
    }

    /// Returns a Vector of Strings where each entry
    /// is a non-empty line from some text input, with
    /// all leading and trailing whitespace removed.
    fn clean_asm<I>(input: I) -> Vec<String>
        where I: Into<String>
    {
        input.into()
            .lines()
            .map(|line| line.trim())
            .map(String::from)
            .filter(|line| line.len() > 0)
            .collect()
    }
}
<commit_msg>Add more disassembler addressing modes and tests<commit_after>use byteorder::{ByteOrder, LittleEndian};

use opcodes::{AddressingMode, OpCode};

pub fn disassemble(raw: &[u8]) -> String {
    let mut result = String::new();

    let mut i: usize = 0;
    while i < raw.len() {
        let opcode = OpCode::from_raw_byte(raw[i]);
        let val = match opcode.mode {
            AddressingMode::Implied |
            AddressingMode::Accumulator => String::from(""),
            AddressingMode::Immediate => format!(" #${:X}", raw[i + 0x01]),
            AddressingMode::Indirect => {
                format!(" (${:X})", LittleEndian::read_u16(&raw[i + 0x01..]))
            }
            AddressingMode::Relative => format!(" #${:X}", raw[i + 0x01]),
            AddressingMode::Absolute => format!(" ${:X}", LittleEndian::read_u16(&raw[i + 0x01..])),
            AddressingMode::ZeroPage => format!(" ${:X}", raw[i + 0x01]),
            AddressingMode::ZeroPageX => format!(" ${:X},X", raw[i + 0x01]),
            AddressingMode::ZeroPageY => format!(" ${:X},Y", raw[i + 0x01]),
            _ => String::from(""),
        };
        let opcode_text = format!("{}{}\n", opcode.mnemonic, val);
        result.push_str(&opcode_text);
        i += opcode.length as usize;
    }

    result
}

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

    #[test]
    fn can_disassemble_basic_instructions() {
        let code: Vec<u8> = vec![0xA9, 0x20, 0x8D, 0x00, 0x44];
        let asm = disassemble(&code);

        assert_eq!(clean_asm("

            LDA #$20
            STA $4400

        "),
                   clean_asm(asm));
    }

    #[test]
    fn can_disassemble_indirect_jmp() {
        let code: Vec<u8> = vec![0x6C, 0x00, 0x44];
        let asm = disassemble(&code);

        assert_eq!(clean_asm("

            JMP ($4400)

        "),
                   clean_asm(asm));
    }

    #[test]
    fn can_disassemble_relative_addressing() {
        let code: Vec<u8> = vec![0xD0, 0xFF];
        let asm = disassemble(&code);

        assert_eq!(clean_asm("

            BNE #$FF

        "),
                   clean_asm(asm));
    }

    #[test]
    fn can_disassemble_zero_page_addressing() {
        let code: Vec<u8> = vec![0xA5, 0x35];
        let asm = disassemble(&code);

        assert_eq!(clean_asm("

            LDA $35

        "),
                   clean_asm(asm));
    }

    #[test]
    fn can_disassemble_zero_page_indexed_addressing() {
        let code: Vec<u8> = vec![0x95, 0x44, 0x96, 0xFE];
        let asm = disassemble(&code);

        assert_eq!(clean_asm("

            STA $44,X
            STX $FE,Y

        "),
                   clean_asm(asm));
    }

    /// Returns a Vector of Strings where each entry
    /// is a non-empty line from some text input, with
    /// all leading and trailing whitespace removed.
    fn clean_asm<I>(input: I) -> Vec<String>
        where I: Into<String>
    {
        input.into()
            .lines()
            .map(|line| line.trim())
            .map(String::from)
            .filter(|line| line.len() > 0)
            .collect()
    }
}
<|endoftext|>