<commit_before>use alloc::boxed::Box;

use collections::Vec;
use collections::string::ToString;

use common::random::rand;

use core::{cmp, mem, ptr, slice, str};

use fs::{KScheme, Resource, Url};

use network::common::{n16, Checksum, Ipv4Addr, IP_ADDR, FromBytes, ToBytes};

use system::error::{Error, Result, ENOENT};

#[derive(Copy, Clone)]
#[repr(packed)]
pub struct UdpHeader {
    pub src: n16,
    pub dst: n16,
    pub len: n16,
    pub checksum: Checksum,
}

pub struct Udp {
    pub header: UdpHeader,
    pub data: Vec<u8>,
}

impl FromBytes for Udp {
    fn from_bytes(bytes: Vec<u8>) -> Option<Self> {
        if bytes.len() >= mem::size_of::<UdpHeader>() {
            unsafe {
                Option::Some(Udp {
                    header: ptr::read(bytes.as_ptr() as *const UdpHeader),
                    data: bytes[mem::size_of::<UdpHeader>()..bytes.len()].to_vec(),
                })
            }
        } else {
            Option::None
        }
    }
}

impl ToBytes for Udp {
    fn to_bytes(&self) -> Vec<u8> {
        unsafe {
            let header_ptr: *const UdpHeader = &self.header;
            let mut ret = Vec::from(slice::from_raw_parts(header_ptr as *const u8,
                                                          mem::size_of::<UdpHeader>()));
            ret.extend_from_slice(&self.data);
            ret
        }
    }
}

/// UDP resource
pub struct UdpResource {
    ip: Box<Resource>,
    data: Vec<u8>,
    peer_addr: Ipv4Addr,
    peer_port: u16,
    host_port: u16,
}

impl Resource for UdpResource {
    fn dup(&self) -> Result<Box<Resource>> {
        match self.ip.dup() {
            Ok(ip) => {
                Ok(Box::new(UdpResource {
                    ip: ip,
                    data: self.data.clone(),
                    peer_addr: self.peer_addr,
                    peer_port: self.peer_port,
                    host_port: self.host_port,
                }))
            }
            Err(err) => Err(err),
        }
    }

    fn path(&self, buf: &mut [u8]) -> Result<usize> {
        let path_string = format!("udp:{}:{}/{}", self.peer_addr.to_string(), self.peer_port, self.host_port);
        let path = path_string.as_bytes();

        for (b, p) in buf.iter_mut().zip(path.iter()) {
            *b = *p;
        }

        Ok(cmp::min(buf.len(), path.len()))
    }

    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        if ! self.data.is_empty() {
            let mut bytes: Vec<u8> = Vec::new();
            mem::swap(&mut self.data, &mut bytes);

            // TODO: Allow splitting
            let i = 0;
            while i < buf.len() && i < bytes.len() {
                buf[i] = bytes[i];
            }
            return Ok(i);
        }

        loop {
            let mut bytes = [0; 8192];
            match self.ip.read(&mut bytes) {
                Ok(count) => {
                    if let Some(datagram) = Udp::from_bytes(bytes[.. count].to_vec()) {
                        if datagram.header.dst.get() == self.host_port &&
                           datagram.header.src.get() == self.peer_port {
                            // TODO: Allow splitting
                            let i = 0;
                            while i < buf.len() && i < datagram.data.len() {
                                buf[i] = datagram.data[i];
                            }
                            return Ok(i);
                        }
                    }
                }
                Err(err) => return Err(err),
            }
        }
    }

    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        let udp_data = Vec::from(buf);

        let mut udp = Udp {
            header: UdpHeader {
                src: n16::new(self.host_port),
                dst: n16::new(self.peer_port),
                len: n16::new((mem::size_of::<UdpHeader>() + udp_data.len()) as u16),
                checksum: Checksum { data: 0 },
            },
            data: udp_data,
        };

        unsafe {
            let proto = n16::new(0x11);
            let datagram_len = n16::new((mem::size_of::<UdpHeader>() + udp.data.len()) as u16);
            udp.header.checksum.data =
                Checksum::compile(Checksum::sum((&IP_ADDR as *const Ipv4Addr) as usize,
                                                mem::size_of::<Ipv4Addr>()) +
                                  Checksum::sum((&self.peer_addr as *const Ipv4Addr) as usize,
                                                mem::size_of::<Ipv4Addr>()) +
                                  Checksum::sum((&proto as *const n16) as usize,
                                                mem::size_of::<n16>()) +
                                  Checksum::sum((&datagram_len as *const n16) as usize,
                                                mem::size_of::<n16>()) +
                                  Checksum::sum((&udp.header as *const UdpHeader) as usize,
                                                mem::size_of::<UdpHeader>()) +
                                  Checksum::sum(udp.data.as_ptr() as usize, udp.data.len()));
        }

        match self.ip.write(&udp.to_bytes()) {
            Ok(_) => Ok(buf.len()),
            Err(err) => Err(err),
        }
    }

    fn sync(&mut self) -> Result<()> {
        self.ip.sync()
    }
}

/// UDP UdpScheme
pub struct UdpScheme;

impl KScheme for UdpScheme {
    fn scheme(&self) -> &str {
        "udp"
    }

    fn open(&mut self, url: Url, _: usize) -> Result<Box<Resource>> {
        let mut parts = url.reference().split('/');
        let remote = parts.next().unwrap_or("");
        let path = parts.next().unwrap_or("");

        debugln!("UDP OPEN R {} P {}", remote, path);

        // Check host and port vs path
        if ! path.is_empty() {
            let host_port = path.parse::<u16>().unwrap_or(0);
            if host_port > 0 {
                while let Ok(mut ip) = Url::from_str("ip:/11").unwrap().open() {
                    let mut bytes = [0; 8192];
                    if let Ok(count) = ip.read(&mut bytes) {
                        if let Some(datagram) = Udp::from_bytes(bytes[.. count].to_vec()) {
                            debugln!("UDP {} to {}", datagram.header.src.get(), datagram.header.dst.get());
                            if datagram.header.dst.get() == host_port {
                                debugln!("GOT PACKET");
                                let mut path = [0; 256];
                                if let Ok(path_count) = ip.path(&mut path) {
                                    let ip_reference = unsafe { str::from_utf8_unchecked(&path[.. path_count]) }.split(':').nth(1).unwrap_or("");
                                    let peer_addr = ip_reference.split('/').next().unwrap_or("").split(':').next().unwrap_or("");

                                    debugln!("FROM {} = {}: {} = {:?}", ip_reference, peer_addr, datagram.data.len(), datagram.data);
                                    return Ok(Box::new(UdpResource {
                                        ip: ip,
                                        data: datagram.data,
                                        peer_addr: Ipv4Addr::from_string(&peer_addr.to_string()),
                                        peer_port: datagram.header.src.get(),
                                        host_port: host_port,
                                    }));
                                } else {
                                    debugln!("FAILED TO GET PATH");
                                }
                            } else {
                                debugln!("WRONG DST {}", datagram.header.dst.get());
                            }
                        } else {
                            debugln!("Not UDP {}", count);
                        }
                    } else {
                        debugln!("NO DATA");
                    }
                }
            }
        } else {
            let mut remote_parts = remote.split(':');
            let peer_addr = remote_parts.next().unwrap_or("");
            let peer_port = remote_parts.next().unwrap_or("").parse::<usize>().unwrap_or(0);
            if peer_port > 0 && peer_port < 65536 {
                let host_port = (rand() % 32768 + 32768) as u16;

                if let Ok(ip) = Url::from_str(&format!("ip:{}/11", peer_addr)).unwrap().open() {
                    return Ok(Box::new(UdpResource {
                        ip: ip,
                        data: Vec::new(),
                        peer_addr: Ipv4Addr::from_string(&peer_addr.to_string()),
                        peer_port: peer_port as u16,
                        host_port: host_port,
                    }));
                }
            }
        }

        Err(Error::new(ENOENT))
    }
}
<commit_msg>Fix udp scheme<commit_after>use alloc::boxed::Box;

use collections::Vec;
use collections::string::ToString;

use common::random::rand;

use core::{cmp, mem, ptr, slice, str};

use fs::{KScheme, Resource, Url};

use network::common::{n16, Checksum, Ipv4Addr, IP_ADDR, FromBytes, ToBytes};

use system::error::{Error, Result, ENOENT};

#[derive(Copy, Clone)]
#[repr(packed)]
pub struct UdpHeader {
    pub src: n16,
    pub dst: n16,
    pub len: n16,
    pub checksum: Checksum,
}

pub struct Udp {
    pub header: UdpHeader,
    pub data: Vec<u8>,
}

impl FromBytes for Udp {
    fn from_bytes(bytes: Vec<u8>) -> Option<Self> {
        if bytes.len() >= mem::size_of::<UdpHeader>() {
            unsafe {
                Option::Some(Udp {
                    header: ptr::read(bytes.as_ptr() as *const UdpHeader),
                    data: bytes[mem::size_of::<UdpHeader>()..bytes.len()].to_vec(),
                })
            }
        } else {
            Option::None
        }
    }
}

impl ToBytes for Udp {
    fn to_bytes(&self) -> Vec<u8> {
        unsafe {
            let header_ptr: *const UdpHeader = &self.header;
            let mut ret = Vec::from(slice::from_raw_parts(header_ptr as *const u8,
                                                          mem::size_of::<UdpHeader>()));
            ret.extend_from_slice(&self.data);
            ret
        }
    }
}

/// UDP resource
pub struct UdpResource {
    ip: Box<Resource>,
    data: Vec<u8>,
    peer_addr: Ipv4Addr,
    peer_port: u16,
    host_port: u16,
}

impl Resource for UdpResource {
    fn dup(&self) -> Result<Box<Resource>> {
        match self.ip.dup() {
            Ok(ip) => {
                Ok(Box::new(UdpResource {
                    ip: ip,
                    data: self.data.clone(),
                    peer_addr: self.peer_addr,
                    peer_port: self.peer_port,
                    host_port: self.host_port,
                }))
            }
            Err(err) => Err(err),
        }
    }

    fn path(&self, buf: &mut [u8]) -> Result<usize> {
        let path_string = format!("udp:{}:{}/{}", self.peer_addr.to_string(), self.peer_port, self.host_port);
        let path = path_string.as_bytes();

        for (b, p) in buf.iter_mut().zip(path.iter()) {
            *b = *p;
        }

        Ok(cmp::min(buf.len(), path.len()))
    }

    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        if ! self.data.is_empty() {
            let mut bytes: Vec<u8> = Vec::new();
            mem::swap(&mut self.data, &mut bytes);

            // TODO: Allow splitting
            let mut i = 0;
            while i < buf.len() && i < bytes.len() {
                buf[i] = bytes[i];
                i += 1;
            }
            return Ok(i);
        }

        loop {
            let mut bytes = [0; 8192];
            let count = try!(self.ip.read(&mut bytes));
            if let Some(datagram) = Udp::from_bytes(bytes[.. count].to_vec()) {
                if datagram.header.dst.get() == self.host_port &&
                   datagram.header.src.get() == self.peer_port {
                    // TODO: Allow splitting
                    let mut i = 0;
                    while i < buf.len() && i < datagram.data.len() {
                        buf[i] = datagram.data[i];
                        i += 1;
                    }
                    return Ok(i);
                }
            }
        }
    }

    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        let mut udp = Udp {
            header: UdpHeader {
                src: n16::new(self.host_port),
                dst: n16::new(self.peer_port),
                len: n16::new((mem::size_of::<UdpHeader>() + buf.len()) as u16),
                checksum: Checksum { data: 0 },
            },
            data: Vec::from(buf),
        };

        unsafe {
            let proto = n16::new(0x11);
            let datagram_len = n16::new((mem::size_of::<UdpHeader>() + udp.data.len()) as u16);
            udp.header.checksum.data =
                Checksum::compile(Checksum::sum((&IP_ADDR as *const Ipv4Addr) as usize,
                                                mem::size_of::<Ipv4Addr>()) +
                                  Checksum::sum((&self.peer_addr as *const Ipv4Addr) as usize,
                                                mem::size_of::<Ipv4Addr>()) +
                                  Checksum::sum((&proto as *const n16) as usize,
                                                mem::size_of::<n16>()) +
                                  Checksum::sum((&datagram_len as *const n16) as usize,
                                                mem::size_of::<n16>()) +
                                  Checksum::sum((&udp.header as *const UdpHeader) as usize,
                                                mem::size_of::<UdpHeader>()) +
                                  Checksum::sum(udp.data.as_ptr() as usize, udp.data.len()));
        }

        self.ip.write(&udp.to_bytes()).and(Ok(buf.len()))
    }

    fn sync(&mut self) -> Result<()> {
        self.ip.sync()
    }
}

/// UDP UdpScheme
pub struct UdpScheme;

impl KScheme for UdpScheme {
    fn scheme(&self) -> &str {
        "udp"
    }

    fn open(&mut self, url: Url, _: usize) -> Result<Box<Resource>> {
        let mut parts = url.reference().split('/');
        let remote = parts.next().unwrap_or("");
        let path = parts.next().unwrap_or("");

        // Check host and port vs path
        if ! path.is_empty() {
            let host_port = path.parse::<u16>().unwrap_or(0);
            if host_port > 0 {
                while let Ok(mut ip) = Url::from_str("ip:/11").unwrap().open() {
                    let mut bytes = [0; 8192];
                    if let Ok(count) = ip.read(&mut bytes) {
                        if let Some(datagram) = Udp::from_bytes(bytes[.. count].to_vec()) {
                            if datagram.header.dst.get() == host_port {
                                let mut path = [0; 256];
                                if let Ok(path_count) = ip.path(&mut path) {
                                    let ip_reference = unsafe { str::from_utf8_unchecked(&path[.. path_count]) }.split(':').nth(1).unwrap_or("");
                                    let peer_addr = ip_reference.split('/').next().unwrap_or("").split(':').next().unwrap_or("");

                                    return Ok(Box::new(UdpResource {
                                        ip: ip,
                                        data: datagram.data,
                                        peer_addr: Ipv4Addr::from_string(&peer_addr.to_string()),
                                        peer_port: datagram.header.src.get(),
                                        host_port: host_port,
                                    }));
                                }
                            }
                        }
                    }
                }
            }
        } else {
            let mut remote_parts = remote.split(':');
            let peer_addr = remote_parts.next().unwrap_or("");
            let peer_port = remote_parts.next().unwrap_or("").parse::<usize>().unwrap_or(0);
            if peer_port > 0 && peer_port < 65536 {
                let host_port = (rand() % 32768 + 32768) as u16;

                if let Ok(ip) = Url::from_str(&format!("ip:{}/11", peer_addr)).unwrap().open() {
                    return Ok(Box::new(UdpResource {
                        ip: ip,
                        data: Vec::new(),
                        peer_addr: Ipv4Addr::from_string(&peer_addr.to_string()),
                        peer_port: peer_port as u16,
                        host_port: host_port,
                    }));
                }
            }
        }

        Err(Error::new(ENOENT))
    }
}
<|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/.

#![cfg_attr(not(feature = "clippy"), allow(unknown_lints))]
#![allow(doc_markdown)]

extern crate devicemapper;
extern crate libstratis;
#[macro_use]
extern crate log;
extern crate clap;
#[cfg(feature = "dbus_enabled")]
extern crate dbus;
extern crate env_logger;
extern crate libc;
extern crate libudev;
extern crate nix;

use std::cell::RefCell;
use std::env;
use std::fs::{File, OpenOptions};
use std::io::{ErrorKind, Read, Write};
use std::os::unix::io::AsRawFd;
use std::path::PathBuf;
use std::process::exit;
use std::rc::Rc;

use clap::{App, Arg, ArgMatches};
use env_logger::LogBuilder;
use libc::pid_t;
use log::{LogLevelFilter, SetLoggerError};
use nix::fcntl::{flock, FlockArg};
use nix::unistd::getpid;

#[cfg(feature = "dbus_enabled")]
use dbus::WatchEvent;

use devicemapper::Device;

use libstratis::engine::{Engine, SimEngine, StratEngine};
use libstratis::stratis::{StratisError, StratisResult, VERSION};

const STRATISD_PID_PATH: &str = "/var/run/stratisd.pid";

/// If writing a program error to stderr fails, panic.
fn print_err(err: &StratisError) -> () {
    eprintln!("{}", err);
}

/// Configure and initialize the logger.
/// If debug is true, log at debug level. Otherwise read log configuration
/// parameters from the environment if RUST_LOG is set. Otherwise, just
/// accept the default configuration.
fn initialize_log(debug: bool) -> Result<(), SetLoggerError> {
    let mut builder = LogBuilder::new();
    if debug {
        builder.filter(Some("stratisd"), LogLevelFilter::Debug);
        builder.filter(Some("libstratis"), LogLevelFilter::Debug);
    } else if let Ok(s) = env::var("RUST_LOG") {
        builder.parse(&s);
    };

    builder.init()
}

/// Given a udev event check to see if it's an add and if it is return the device node and
/// devicemapper::Device.
fn handle_udev_add(event: &libudev::Event) -> Option<(Device, PathBuf)> {
    if event.event_type() == libudev::EventType::Add {
        let device = event.device();
        return device.devnode().and_then(|devnode| {
            device
                .devnum()
                .and_then(|devnum| Some((Device::from(devnum), PathBuf::from(devnode))))
        });
    }
    None
}

/// To ensure only one instance of stratisd runs at a time, acquire an
/// exclusive lock. Return an error if lock attempt fails.
fn trylock_pid_file() -> StratisResult<File> {
    let mut f = match OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .open(STRATISD_PID_PATH)
    {
        Ok(f) => f,
        Err(e) => {
            if e.kind() == ErrorKind::PermissionDenied {
                return Err(StratisError::Error(
                    "Must be running as root in order to start daemon.".to_string(),
                ));
            }
            return Err(e.into());
        }
    };
    match flock(f.as_raw_fd(), FlockArg::LockExclusiveNonblock) {
        Ok(_) => {
            f.write_all(format!("{}\n", getpid()).as_bytes())?;
            Ok(f)
        }
        Err(_) => {
            let mut buf = String::new();
            f.read_to_string(&mut buf)?;
            // pidfile is supposed to contain pid of holder. But you never
            // know so be paranoid.
            let pid_str = buf.split_whitespace()
                .next()
                .and_then(|s| s.parse::<pid_t>().ok())
                .map(|pid| format!("{}", pid))
                .unwrap_or_else(|| "<unknown>".into());
            Err(StratisError::Error(format!(
                "Daemon already running with pid: {}",
                pid_str
            )))
        }
    }
}

fn run(matches: &ArgMatches) -> StratisResult<()> {
    // Setup a udev listener before initializing the engine. A device may
    // appear after the engine has read the /dev directory but before it has
    // completed initialization. Unless the udev event has been recorded, the
    // engine will miss the device.
    // This is especially important since stratisd must run during early boot.
    let context = libudev::Context::new()?;
    let mut monitor = libudev::Monitor::new(&context)?;
    monitor.match_subsystem_devtype("block", "disk")?;
    let mut udev = monitor.listen()?;

    let engine: Rc<RefCell<Engine>> = {
        if matches.is_present("sim") {
            info!("Using SimEngine");
            Rc::new(RefCell::new(SimEngine::default()))
        } else {
            info!("Using StratEngine");
            Rc::new(RefCell::new(StratEngine::initialize()?))
        }
    };

    /*
    The file descriptor array indexes are laid out in the following:

    0   == Always udev fd index
    1   == engine index if eventable
    1/2 == Start of dbus client file descriptor(s), 1 if engine is not eventable, else 2
    */
    const FD_INDEX_UDEV: usize = 0;
    const FD_INDEX_ENGINE: usize = 1;

    let mut fds = Vec::new();

    fds.push(libc::pollfd {
        fd: udev.as_raw_fd(),
        revents: 0,
        events: libc::POLLIN,
    });

    let eventable = engine.borrow().get_eventable();

    // The variable _dbus_client_index_start is only used when dbus support is compiled in, thus
    // we denote the value as not needed to compile when dbus support is not included.
    let (poll_timeout, _dbus_client_index_start) = match eventable {
        Some(ref evt) => {
            fds.push(libc::pollfd {
                fd: evt.get_pollable_fd(),
                revents: 0,
                events: libc::POLLIN,
            });

            (-1, FD_INDEX_ENGINE + 1)
        }
        None => (10000, FD_INDEX_ENGINE),
    };

    #[cfg(feature = "dbus_enabled")]
    let (dbus_conn, mut tree, base_object_path, dbus_context) =
        libstratis::dbus_api::connect(Rc::clone(&engine))?;

    #[cfg(feature = "dbus_enabled")]
    for (_, pool_uuid, pool) in engine.borrow().pools() {
        libstratis::dbus_api::register_pool(
            &dbus_conn,
            &dbus_context,
            &mut tree,
            pool_uuid,
            pool,
            &base_object_path,
        )?;
    }

    loop {
        // Process any udev block events
        if fds[FD_INDEX_UDEV].revents != 0 {
            while let Some(event) = udev.receive_event() {
                if let Some((device, devnode)) = handle_udev_add(&event) {
                    // If block evaluate returns an error we are going to ignore it as
                    // there is nothing we can do for a device we are getting errors with.
                    let pool_uuid = engine
                        .borrow_mut()
                        .block_evaluate(device, devnode)
                        .unwrap_or(None);

                    // We need to pretend that we aren't using the variable _pool_uuid so
                    // that we can conditionally compile out the register_pool when dbus
                    // is not enabled.
                    if let Some(_pool_uuid) = pool_uuid {
                        #[cfg(feature = "dbus_enabled")]
                        libstratis::dbus_api::register_pool(
                            &dbus_conn,
                            &dbus_context,
                            &mut tree,
                            _pool_uuid,
                            engine
                                .borrow()
                                .get_pool(_pool_uuid)
                                .expect(
                                    "block_evaluate() returned a pool UUID, pool must be available",
                                )
                                .1,
                            &base_object_path,
                        )?;
                    }
                }
            }
        }

        // Handle engine events, if the engine is eventable
        match eventable {
            Some(ref evt) => {
                if fds[FD_INDEX_ENGINE].revents != 0 {
                    fds[FD_INDEX_ENGINE].revents = 0;
                    evt.clear_event()?;
                    engine.borrow_mut().evented()?;
                }
            }
            None => {
                // Unconditionally call evented() if engine has no eventable.
                // This looks like a bad idea, but the only engine that has
                // no eventable is the sim engine, and for that engine,
                // evented() is essentially a no-op.
                engine.borrow_mut().evented()?;
            }
        }

        // Iterate through D-Bus file descriptors (if enabled)
        #[cfg(feature = "dbus_enabled")]
        {
            for pfd in fds[_dbus_client_index_start..]
                .iter()
                .filter(|pfd| pfd.revents != 0)
            {
                for item in dbus_conn.watch_handle(pfd.fd, WatchEvent::from_revents(pfd.revents)) {
                    if let Err(r) =
                        libstratis::dbus_api::handle(&dbus_conn, &item, &mut tree, &dbus_context)
                    {
                        print_err(&From::from(r));
                    }
                }
            }

            // Refresh list of dbus fds to poll for every time. This can change as
            // D-Bus clients come and go.
            fds.truncate(_dbus_client_index_start);

            fds.extend(dbus_conn.watch_fds().iter().map(|w| w.to_pollfd()));
        }

        let r = unsafe { libc::poll(fds.as_mut_ptr(), fds.len() as libc::c_ulong, poll_timeout) };
        assert!(r >= 0);
    }
}

fn main() {
    let matches = App::new("stratis")
        .version(VERSION)
        .about("Stratis storage management")
        .arg(
            Arg::with_name("debug")
                .long("debug")
                .help("Print additional output for debugging"),
        )
        .arg(
            Arg::with_name("sim")
                .long("sim")
                .help("Use simulator engine"),
        )
        .get_matches();

    let result = trylock_pid_file()
        .and_then(|_pidfile| {
            Ok(initialize_log(matches.is_present("debug"))
                .expect("This is the first and only invocation of this method; it must succeed."))
        })
        .and_then(|_| run(&matches));
    if let Err(err) = result {
        print_err(&err);
        exit(1);
    } else {
        exit(0);
    }
}
<commit_msg>Add a comment explaining why polling is done with libc<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/.

#![cfg_attr(not(feature = "clippy"), allow(unknown_lints))]
#![allow(doc_markdown)]

extern crate devicemapper;
extern crate libstratis;
#[macro_use]
extern crate log;
extern crate clap;
#[cfg(feature = "dbus_enabled")]
extern crate dbus;
extern crate env_logger;
extern crate libc;
extern crate libudev;
extern crate nix;

use std::cell::RefCell;
use std::env;
use std::fs::{File, OpenOptions};
use std::io::{ErrorKind, Read, Write};
use std::os::unix::io::AsRawFd;
use std::path::PathBuf;
use std::process::exit;
use std::rc::Rc;

use clap::{App, Arg, ArgMatches};
use env_logger::LogBuilder;
use libc::pid_t;
use log::{LogLevelFilter, SetLoggerError};
use nix::fcntl::{flock, FlockArg};
use nix::unistd::getpid;

#[cfg(feature = "dbus_enabled")]
use dbus::WatchEvent;

use devicemapper::Device;

use libstratis::engine::{Engine, SimEngine, StratEngine};
use libstratis::stratis::{StratisError, StratisResult, VERSION};

const STRATISD_PID_PATH: &str = "/var/run/stratisd.pid";

/// If writing a program error to stderr fails, panic.
fn print_err(err: &StratisError) -> () {
    eprintln!("{}", err);
}

/// Configure and initialize the logger.
/// If debug is true, log at debug level. Otherwise read log configuration
/// parameters from the environment if RUST_LOG is set. Otherwise, just
/// accept the default configuration.
fn initialize_log(debug: bool) -> Result<(), SetLoggerError> {
    let mut builder = LogBuilder::new();
    if debug {
        builder.filter(Some("stratisd"), LogLevelFilter::Debug);
        builder.filter(Some("libstratis"), LogLevelFilter::Debug);
    } else if let Ok(s) = env::var("RUST_LOG") {
        builder.parse(&s);
    };

    builder.init()
}

/// Given a udev event check to see if it's an add and if it is return the device node and
/// devicemapper::Device.
fn handle_udev_add(event: &libudev::Event) -> Option<(Device, PathBuf)> {
    if event.event_type() == libudev::EventType::Add {
        let device = event.device();
        return device.devnode().and_then(|devnode| {
            device
                .devnum()
                .and_then(|devnum| Some((Device::from(devnum), PathBuf::from(devnode))))
        });
    }
    None
}

/// To ensure only one instance of stratisd runs at a time, acquire an
/// exclusive lock. Return an error if lock attempt fails.
fn trylock_pid_file() -> StratisResult<File> {
    let mut f = match OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .open(STRATISD_PID_PATH)
    {
        Ok(f) => f,
        Err(e) => {
            if e.kind() == ErrorKind::PermissionDenied {
                return Err(StratisError::Error(
                    "Must be running as root in order to start daemon.".to_string(),
                ));
            }
            return Err(e.into());
        }
    };
    match flock(f.as_raw_fd(), FlockArg::LockExclusiveNonblock) {
        Ok(_) => {
            f.write_all(format!("{}\n", getpid()).as_bytes())?;
            Ok(f)
        }
        Err(_) => {
            let mut buf = String::new();
            f.read_to_string(&mut buf)?;
            // pidfile is supposed to contain pid of holder. But you never
            // know so be paranoid.
            let pid_str = buf.split_whitespace()
                .next()
                .and_then(|s| s.parse::<pid_t>().ok())
                .map(|pid| format!("{}", pid))
                .unwrap_or_else(|| "<unknown>".into());
            Err(StratisError::Error(format!(
                "Daemon already running with pid: {}",
                pid_str
            )))
        }
    }
}

fn run(matches: &ArgMatches) -> StratisResult<()> {
    // Setup a udev listener before initializing the engine. A device may
    // appear after the engine has read the /dev directory but before it has
    // completed initialization. Unless the udev event has been recorded, the
    // engine will miss the device.
    // This is especially important since stratisd must run during early boot.
    let context = libudev::Context::new()?;
    let mut monitor = libudev::Monitor::new(&context)?;
    monitor.match_subsystem_devtype("block", "disk")?;
    let mut udev = monitor.listen()?;

    let engine: Rc<RefCell<Engine>> = {
        if matches.is_present("sim") {
            info!("Using SimEngine");
            Rc::new(RefCell::new(SimEngine::default()))
        } else {
            info!("Using StratEngine");
            Rc::new(RefCell::new(StratEngine::initialize()?))
        }
    };

    /*
    The file descriptor array indexes are laid out in the following:

    0   == Always udev fd index
    1   == engine index if eventable
    1/2 == Start of dbus client file descriptor(s), 1 if engine is not eventable, else 2
    */
    const FD_INDEX_UDEV: usize = 0;
    const FD_INDEX_ENGINE: usize = 1;

    /*
    fds is a Vec of libc::pollfd structs. Ideally, it would be possible
    to use the higher level nix crate to handle polling. If this were possible,
    then the Vec would be one of nix::poll::PollFds and this would be more
    rustic. Unfortunately, the rust D-Bus library requires an explicit file
    descriptor to be passed as an argument to Connection::watch_handle(),
    and the explicit file descriptor can not be extracted from the PollFd
    struct. So, at this time, sticking with libc is less complex than
    converting to using nix, because if using nix, the file descriptor would
    have to be maintained in the Vec as well as the PollFd struct.
    */
    let mut fds = Vec::new();

    fds.push(libc::pollfd {
        fd: udev.as_raw_fd(),
        revents: 0,
        events: libc::POLLIN,
    });

    let eventable = engine.borrow().get_eventable();

    // The variable _dbus_client_index_start is only used when dbus support is compiled in, thus
    // we denote the value as not needed to compile when dbus support is not included.
    let (poll_timeout, _dbus_client_index_start) = match eventable {
        Some(ref evt) => {
            fds.push(libc::pollfd {
                fd: evt.get_pollable_fd(),
                revents: 0,
                events: libc::POLLIN,
            });

            (-1, FD_INDEX_ENGINE + 1)
        }
        None => (10000, FD_INDEX_ENGINE),
    };

    #[cfg(feature = "dbus_enabled")]
    let (dbus_conn, mut tree, base_object_path, dbus_context) =
        libstratis::dbus_api::connect(Rc::clone(&engine))?;

    #[cfg(feature = "dbus_enabled")]
    for (_, pool_uuid, pool) in engine.borrow().pools() {
        libstratis::dbus_api::register_pool(
            &dbus_conn,
            &dbus_context,
            &mut tree,
            pool_uuid,
            pool,
            &base_object_path,
        )?;
    }

    loop {
        // Process any udev block events
        if fds[FD_INDEX_UDEV].revents != 0 {
            while let Some(event) = udev.receive_event() {
                if let Some((device, devnode)) = handle_udev_add(&event) {
                    // If block evaluate returns an error we are going to ignore it as
                    // there is nothing we can do for a device we are getting errors with.
                    let pool_uuid = engine
                        .borrow_mut()
                        .block_evaluate(device, devnode)
                        .unwrap_or(None);

                    // We need to pretend that we aren't using the variable _pool_uuid so
                    // that we can conditionally compile out the register_pool when dbus
                    // is not enabled.
                    if let Some(_pool_uuid) = pool_uuid {
                        #[cfg(feature = "dbus_enabled")]
                        libstratis::dbus_api::register_pool(
                            &dbus_conn,
                            &dbus_context,
                            &mut tree,
                            _pool_uuid,
                            engine
                                .borrow()
                                .get_pool(_pool_uuid)
                                .expect(
                                    "block_evaluate() returned a pool UUID, pool must be available",
                                )
                                .1,
                            &base_object_path,
                        )?;
                    }
                }
            }
        }

        // Handle engine events, if the engine is eventable
        match eventable {
            Some(ref evt) => {
                if fds[FD_INDEX_ENGINE].revents != 0 {
                    fds[FD_INDEX_ENGINE].revents = 0;
                    evt.clear_event()?;
                    engine.borrow_mut().evented()?;
                }
            }
            None => {
                // Unconditionally call evented() if engine has no eventable.
                // This looks like a bad idea, but the only engine that has
                // no eventable is the sim engine, and for that engine,
                // evented() is essentially a no-op.
                engine.borrow_mut().evented()?;
            }
        }

        // Iterate through D-Bus file descriptors (if enabled)
        #[cfg(feature = "dbus_enabled")]
        {
            for pfd in fds[_dbus_client_index_start..]
                .iter()
                .filter(|pfd| pfd.revents != 0)
            {
                for item in dbus_conn.watch_handle(pfd.fd, WatchEvent::from_revents(pfd.revents)) {
                    if let Err(r) =
                        libstratis::dbus_api::handle(&dbus_conn, &item, &mut tree, &dbus_context)
                    {
                        print_err(&From::from(r));
                    }
                }
            }

            // Refresh list of dbus fds to poll for every time. This can change as
            // D-Bus clients come and go.
            fds.truncate(_dbus_client_index_start);

            fds.extend(dbus_conn.watch_fds().iter().map(|w| w.to_pollfd()));
        }

        let r = unsafe { libc::poll(fds.as_mut_ptr(), fds.len() as libc::c_ulong, poll_timeout) };
        assert!(r >= 0);
    }
}

fn main() {
    let matches = App::new("stratis")
        .version(VERSION)
        .about("Stratis storage management")
        .arg(
            Arg::with_name("debug")
                .long("debug")
                .help("Print additional output for debugging"),
        )
        .arg(
            Arg::with_name("sim")
                .long("sim")
                .help("Use simulator engine"),
        )
        .get_matches();

    let result = trylock_pid_file()
        .and_then(|_pidfile| {
            Ok(initialize_log(matches.is_present("debug"))
                .expect("This is the first and only invocation of this method; it must succeed."))
        })
        .and_then(|_| run(&matches));
    if let Err(err) = result {
        print_err(&err);
        exit(1);
    } else {
        exit(0);
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use common::*;
use message;

pub type SegmentId = u32;

pub struct SegmentReader<'self> {
    messageReader : * message::MessageReader<'self>,
    segment : &'self [Word]
}

impl <'self> SegmentReader<'self> {

    pub unsafe fn getStartPtr(&self) -> *Word {
        self.segment.unsafe_ref(0)
    }

    pub unsafe fn containsInterval(&self, from : *Word, to : *Word) -> bool {
        let fromAddr : uint = std::cast::transmute(from);
        let toAddr : uint = std::cast::transmute(to);
        let thisBegin : uint = std::cast::transmute(self.segment.unsafe_ref(0));
        let thisEnd : uint = std::cast::transmute(
            self.segment.unsafe_ref(self.segment.len()));
        return (fromAddr >= thisBegin && toAddr <= thisEnd);
        // TODO readLimiter
    }
}

pub struct SegmentBuilder {
    messageBuilder : *mut message::MessageBuilder,
    id : SegmentId,
    ptr : *mut Word,
    pos : WordCount,
    size : WordCount
}

impl SegmentBuilder {

    pub fn new(messageBuilder : *mut message::MessageBuilder,
               size : WordCount) -> SegmentBuilder {
        let idx = unsafe {((*messageBuilder).segments.len() - 1) as SegmentId};
        SegmentBuilder {
            messageBuilder : messageBuilder,
            ptr : unsafe {(*messageBuilder).segments[idx].unsafe_mut_ref(0)},
            id : idx,
            pos : 0,
            size : size
        }
    }

    pub fn getWordOffsetTo(&mut self, ptr : *mut Word) -> WordCount {
        let thisAddr : uint = self.ptr.to_uint();
        let ptrAddr : uint = ptr.to_uint();
        assert!(ptrAddr >= thisAddr);
        let result = (ptrAddr - thisAddr) / BYTES_PER_WORD;
        return result;
    }

    pub fn allocate(&mut self, amount : WordCount) -> Option<*mut Word> {
        if (amount > self.size - self.pos) {
            return None;
        } else {
            let result = unsafe { self.ptr.offset(self.pos as int) };
            self.pos += amount;
            return Some(result);
        }
    }

    pub fn available(&self) -> WordCount {
        self.size - self.pos
    }

    #[inline]
    pub unsafe fn getPtrUnchecked(&mut self, offset : WordCount) -> *mut Word {
        self.ptr.offset(offset as int)
    }

    pub fn asReader<T>(&mut self, f : &fn(&SegmentReader) -> T) -> T {
        unsafe {
            do (*self.messageBuilder).asReader |messageReader| {
                f(&SegmentReader {
                        messageReader : std::ptr::to_unsafe_ptr(messageReader),
                        segment : messageReader.segments[self.id]
                    })
            }
        }
    }
}

pub trait Arena<'a> {
    fn tryGetSegment(&self, id : SegmentId) -> *SegmentReader<'a>;
}

pub struct ReaderArena<'a> {
    message : message::MessageReader<'a>,
    segment0 : SegmentReader<'a>,

    moreSegments : Option<~[SegmentReader<'a>]>
    //XXX should this be a map as in capnproto-c++?
}

impl <'a> Arena<'a> for ReaderArena<'a> {
    fn tryGetSegment(&self, id : SegmentId) -> *SegmentReader<'a> {
        if (id == 0) {
            return std::ptr::to_unsafe_ptr(&self.segment0);
        } else {
            match self.moreSegments {
                None => {fail!("no segments!")}
                Some(ref segs) => {
                    unsafe {segs.unsafe_ref(id as uint - 1)}
                }
            }
        }
    }
}

pub struct BuilderArena<'a> {
    message : *message::MessageBuilder,
    segment0 : SegmentBuilder
}

impl <'a> Arena<'a> for BuilderArena<'a> {
    fn tryGetSegment(&self, _id : SegmentId) -> *SegmentReader<'a> {
        fail!()
    }
}
<commit_msg>this may be a more viable strategy<commit_after>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use common::*;
use message;

pub type SegmentId = u32;

pub struct SegmentReader<'self> {
    messageReader : * message::MessageReader<'self>,
    segment : &'self [Word]
}

impl <'self> SegmentReader<'self> {

    pub unsafe fn getStartPtr(&self) -> *Word {
        self.segment.unsafe_ref(0)
    }

    pub unsafe fn containsInterval(&self, from : *Word, to : *Word) -> bool {
        let fromAddr : uint = std::cast::transmute(from);
        let toAddr : uint = std::cast::transmute(to);
        let thisBegin : uint = std::cast::transmute(self.segment.unsafe_ref(0));
        let thisEnd : uint = std::cast::transmute(
            self.segment.unsafe_ref(self.segment.len()));
        return (fromAddr >= thisBegin && toAddr <= thisEnd);
        // TODO readLimiter
    }
}

pub struct SegmentBuilder {
    messageBuilder : *mut message::MessageBuilder,
    id : SegmentId,
    ptr : *mut Word,
    pos : WordCount,
    size : WordCount
}

impl SegmentBuilder {

    pub fn new(messageBuilder : *mut message::MessageBuilder,
               size : WordCount) -> SegmentBuilder {
        let idx = unsafe {((*messageBuilder).segments.len() - 1) as SegmentId};
        SegmentBuilder {
            messageBuilder : messageBuilder,
            ptr : unsafe {(*messageBuilder).segments[idx].unsafe_mut_ref(0)},
            id : idx,
            pos : 0,
            size : size
        }
    }

    pub fn getWordOffsetTo(&mut self, ptr : *mut Word) -> WordCount {
        let thisAddr : uint = self.ptr.to_uint();
        let ptrAddr : uint = ptr.to_uint();
        assert!(ptrAddr >= thisAddr);
        let result = (ptrAddr - thisAddr) / BYTES_PER_WORD;
        return result;
    }

    pub fn allocate(&mut self, amount : WordCount) -> Option<*mut Word> {
        if (amount > self.size - self.pos) {
            return None;
        } else {
            let result = unsafe { self.ptr.offset(self.pos as int) };
            self.pos += amount;
            return Some(result);
        }
    }

    pub fn available(&self) -> WordCount {
        self.size - self.pos
    }

    #[inline]
    pub unsafe fn getPtrUnchecked(&mut self, offset : WordCount) -> *mut Word {
        self.ptr.offset(offset as int)
    }

    pub fn asReader<T>(&mut self, f : &fn(&SegmentReader) -> T) -> T {
        unsafe {
            do (*self.messageBuilder).asReader |messageReader| {
                f(&SegmentReader {
                        messageReader : std::ptr::to_unsafe_ptr(messageReader),
                        segment : messageReader.segments[self.id]
                    })
            }
        }
    }
}

pub struct ReaderArena<'a> {
    message : message::MessageReader<'a>,
    segment0 : SegmentReader<'a>,

    moreSegments : Option<~[SegmentReader<'a>]>
    //XXX should this be a map as in capnproto-c++?
}

pub struct BuilderArena {
    message : *message::MessageBuilder,
    segment0 : SegmentBuilder
}

pub enum Arena<'a> {
    Reader_(ReaderArena<'a>),
    Builder_(BuilderArena)
}

impl <'a> Arena<'a>  {
    pub fn tryGetSegment(&self, id : SegmentId) -> *SegmentReader<'a> {
        match self {
            &Reader_(ref reader) => {
                if (id == 0) {
                    return std::ptr::to_unsafe_ptr(&reader.segment0);
                } else {
                    match reader.moreSegments {
                        None => {fail!("no segments!")}
                        Some(ref segs) => {
                            unsafe {segs.unsafe_ref(id as uint - 1)}
                        }
                    }
                }
            }
            &Builder_(ref _builder) => {
                fail!()
            }
        }
    }
}
<|endoftext|>


<commit_before>//! A button that was pressed on a keyboard by a user.
//!
//! This is mostly used to define keybindings.

use rlua::{self, Lua, Table, ToLua, UserData, UserDataMethods, Value};
use wlroots::{self, xkbcommon::xkb};

use common::{class::{self, Class, ClassBuilder},
             object::{self, Object},
             property::Property};
use lua::mods_to_num;

#[derive(Clone, Debug, Default)]
pub struct KeyState {
    modifiers: u32,
    keysym: wlroots::Key,
    keycode: xkb::Keycode
}

pub type Key<'lua> = Object<'lua, KeyState>;

impl<'lua> Key<'lua> {
    fn new(lua: &'lua Lua, args: Table) -> rlua::Result<Key<'lua>> {
        // TODO FIXME
        let class = class::class_setup(lua, "key")?;
        Ok(Key::allocate(lua, class)?.handle_constructor_argument(args)?
                                     .build())
    }

    pub fn set_modifiers(&mut self, modifiers: u32) -> rlua::Result<()> {
        let mut state = self.state_mut()?;
        state.modifiers = modifiers;
        Ok(())
    }

    pub fn modifiers(&self) -> rlua::Result<u32> {
        let state = self.state()?;
        Ok(state.modifiers)
    }

    pub fn set_keysym(&mut self, keysym: wlroots::Key) -> rlua::Result<()> {
        let mut state = self.state_mut()?;
        state.keysym = keysym;
        Ok(())
    }

    pub fn keysym(&self) -> rlua::Result<wlroots::Key> {
        let state = self.state()?;
        Ok(state.keysym)
    }

    pub fn set_keycode(&mut self, keycode: xkb::Keycode) -> rlua::Result<()> {
        let mut state = self.state_mut()?;
        state.keycode = keycode;
        Ok(())
    }

    #[allow(dead_code)]
    pub fn keycode(&self) -> rlua::Result<xkb::Keycode> {
        let state = self.state()?;
        Ok(state.keycode)
    }
}

impl UserData for KeyState {
    fn add_methods(methods: &mut UserDataMethods<Self>) {
        object::default_add_methods(methods);
    }
}

pub fn init(lua: &Lua) -> rlua::Result<Class<KeyState>> {
    property_setup(lua, method_setup(lua, Class::builder(lua, "key", None)?)?)?.save_class("key")?
                                                                               .build()
}

fn method_setup<'lua>(lua: &'lua Lua,
                      builder: ClassBuilder<'lua, KeyState>)
                      -> rlua::Result<ClassBuilder<'lua, KeyState>> {
    // TODO Do properly
    builder.method("__call".into(),
                   lua.create_function(|lua, args: Table| Key::new(lua, args))?)
}

fn property_setup<'lua>(lua: &'lua Lua,
                        builder: ClassBuilder<'lua, KeyState>)
                        -> rlua::Result<ClassBuilder<'lua, KeyState>> {
    // TODO Do properly
    builder.property(Property::new("key".into(),
                                   Some(lua.create_function(set_key)?),
                                   Some(lua.create_function(get_key)?),
                                   Some(lua.create_function(set_key)?)))?
           .property(Property::new("keysym".into(),
                                   None,
                                   Some(lua.create_function(get_keysym)?),
                                   None))?
           .property(Property::new("modifiers".into(),
                                   Some(lua.create_function(set_modifiers)?),
                                   Some(lua.create_function(get_modifiers)?),
                                   Some(lua.create_function(set_modifiers)?)))
}

fn get_modifiers<'lua>(_: &'lua Lua, key: Key<'lua>) -> rlua::Result<u32> {
    key.modifiers()
}

fn set_modifiers<'lua>(_: &'lua Lua,
                       (mut key, mods): (Key<'lua>, Table<'lua>))
                       -> rlua::Result<()> {
    key.set_modifiers(mods_to_num(mods)?.bits())
}

fn get_keysym<'lua>(lua: &'lua Lua, key: Key<'lua>) -> rlua::Result<Value<'lua>> {
    // TODO Shouldn't this be able to fail?
    xkb::keysym_get_name(key.keysym()?).to_lua(lua)
}

fn get_key<'lua>(lua: &'lua Lua, key: Key<'lua>) -> rlua::Result<Value<'lua>> {
    key.keysym()?.to_lua(lua)
}

fn set_key<'lua>(_: &'lua Lua,
                 (mut key, key_name): (Key<'lua>, String))
                 -> rlua::Result<Value<'lua>> {
    if key_name.starts_with('#') && key_name.len() >= 2 {
        let number = key_name[1..].parse::<xkb::Keycode>()
            .map_err(|err| rlua::Error::RuntimeError(format!("Parse error: {:?}", err)))?;
        // the - 8 is because of xcb conventions, where "#10" is the keysim for 1,
        // and the keycode of 1 is 0x02 (obviously)
        key.set_keycode(number - 8 as xkb::Keycode)?;
    } else {
        let keysym = xkb::keysym_from_name(key_name.as_str(), 0);
        key.set_keysym(keysym)?;
    }
    Ok(rlua::Value::Nil)
}
<commit_msg>Fix a `trivial interger cast` which slipped through the verification<commit_after>//! A button that was pressed on a keyboard by a user.
//!
//! This is mostly used to define keybindings.

use rlua::{self, Lua, Table, ToLua, UserData, UserDataMethods, Value};
use wlroots::{self, xkbcommon::xkb};

use common::{class::{self, Class, ClassBuilder},
             object::{self, Object},
             property::Property};
use lua::mods_to_num;

#[derive(Clone, Debug, Default)]
pub struct KeyState {
    modifiers: u32,
    keysym: wlroots::Key,
    keycode: xkb::Keycode
}

pub type Key<'lua> = Object<'lua, KeyState>;

impl<'lua> Key<'lua> {
    fn new(lua: &'lua Lua, args: Table) -> rlua::Result<Key<'lua>> {
        // TODO FIXME
        let class = class::class_setup(lua, "key")?;
        Ok(Key::allocate(lua, class)?.handle_constructor_argument(args)?
                                     .build())
    }

    pub fn set_modifiers(&mut self, modifiers: u32) -> rlua::Result<()> {
        let mut state = self.state_mut()?;
        state.modifiers = modifiers;
        Ok(())
    }

    pub fn modifiers(&self) -> rlua::Result<u32> {
        let state = self.state()?;
        Ok(state.modifiers)
    }

    pub fn set_keysym(&mut self, keysym: wlroots::Key) -> rlua::Result<()> {
        let mut state = self.state_mut()?;
        state.keysym = keysym;
        Ok(())
    }

    pub fn keysym(&self) -> rlua::Result<wlroots::Key> {
        let state = self.state()?;
        Ok(state.keysym)
    }

    pub fn set_keycode(&mut self, keycode: xkb::Keycode) -> rlua::Result<()> {
        let mut state = self.state_mut()?;
        state.keycode = keycode;
        Ok(())
    }

    #[allow(dead_code)]
    pub fn keycode(&self) -> rlua::Result<xkb::Keycode> {
        let state = self.state()?;
        Ok(state.keycode)
    }
}

impl UserData for KeyState {
    fn add_methods(methods: &mut UserDataMethods<Self>) {
        object::default_add_methods(methods);
    }
}

pub fn init(lua: &Lua) -> rlua::Result<Class<KeyState>> {
    property_setup(lua, method_setup(lua, Class::builder(lua, "key", None)?)?)?.save_class("key")?
                                                                               .build()
}

fn method_setup<'lua>(lua: &'lua Lua,
                      builder: ClassBuilder<'lua, KeyState>)
                      -> rlua::Result<ClassBuilder<'lua, KeyState>> {
    // TODO Do properly
    builder.method("__call".into(),
                   lua.create_function(|lua, args: Table| Key::new(lua, args))?)
}

fn property_setup<'lua>(lua: &'lua Lua,
                        builder: ClassBuilder<'lua, KeyState>)
                        -> rlua::Result<ClassBuilder<'lua, KeyState>> {
    // TODO Do properly
    builder.property(Property::new("key".into(),
                                   Some(lua.create_function(set_key)?),
                                   Some(lua.create_function(get_key)?),
                                   Some(lua.create_function(set_key)?)))?
           .property(Property::new("keysym".into(),
                                   None,
                                   Some(lua.create_function(get_keysym)?),
                                   None))?
           .property(Property::new("modifiers".into(),
                                   Some(lua.create_function(set_modifiers)?),
                                   Some(lua.create_function(get_modifiers)?),
                                   Some(lua.create_function(set_modifiers)?)))
}

fn get_modifiers<'lua>(_: &'lua Lua, key: Key<'lua>) -> rlua::Result<u32> {
    key.modifiers()
}

fn set_modifiers<'lua>(_: &'lua Lua,
                       (mut key, mods): (Key<'lua>, Table<'lua>))
                       -> rlua::Result<()> {
    key.set_modifiers(mods_to_num(mods)?.bits())
}

fn get_keysym<'lua>(lua: &'lua Lua, key: Key<'lua>) -> rlua::Result<Value<'lua>> {
    // TODO Shouldn't this be able to fail?
    xkb::keysym_get_name(key.keysym()?).to_lua(lua)
}

fn get_key<'lua>(lua: &'lua Lua, key: Key<'lua>) -> rlua::Result<Value<'lua>> {
    key.keysym()?.to_lua(lua)
}

fn set_key<'lua>(_: &'lua Lua,
                 (mut key, key_name): (Key<'lua>, String))
                 -> rlua::Result<Value<'lua>> {
    if key_name.starts_with('#') && key_name.len() >= 2 {
        let number = key_name[1..].parse::<xkb::Keycode>()
            .map_err(|err| rlua::Error::RuntimeError(format!("Parse error: {:?}", err)))?;
        // the - 8 is because of xcb conventions, where "#10" is the keysim for 1,
        // and the keycode of 1 is 0x02 (obviously)
        key.set_keycode(number - 8)?;
    } else {
        let keysym = xkb::keysym_from_name(key_name.as_str(), 0);
        key.set_keysym(keysym)?;
    }
    Ok(rlua::Value::Nil)
}
<|endoftext|>


<commit_before>use failure::Error;
use std::str::FromStr;
use std::net::IpAddr;
use std::collections::BTreeMap;

use quire::validate::{Structure, Scalar, Numeric, Mapping};
use quire::{Options, parse_string};

#[derive(Serialize, Deserialize)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum ChildKind {
    Daemon,
    Command,
}

// Note everything here should be stable-serializable
#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
pub struct ChildConfig {
    pub instances: usize,
    pub image: String,
    pub config: String,
    #[serde(skip_serializing_if="BTreeMap::is_empty", default)]
    pub variables: BTreeMap<String, String>,
    #[serde(skip_serializing_if="Vec::is_empty", default)]
    pub ip_addresses: Vec<IpAddr>,
    pub kind: ChildKind,
}

impl ChildConfig {
    pub fn instantiate(&self, instance: usize) -> Result<ChildConfig, Error>
    {
        let mut cfg = self.clone();
        if self.ip_addresses.len() > 0 {
            if let Some(addr) = self.ip_addresses.get(instance) {
                cfg.ip_addresses = vec![*addr];
            } else {
                bail!("Instance no {}, but there's only {} ip addresses",
                    instance, cfg.ip_addresses.len());
            }
        }
        return Ok(cfg);
    }
    pub fn mapping_validator<'x>() -> Mapping<'x> {
        return Mapping::new(
            Scalar::new(),
            ChildConfig::validator());
    }
    pub fn validator<'x>() -> Structure<'x> {
        Structure::new()
        .member("instances", Numeric::new().default(1))
        .member("image", Scalar::new())
        .member("config", Scalar::new())
        .member("variables", Mapping::new(Scalar::new(), Scalar::new()))
        .member("kind", Scalar::new().default("Daemon"))
    }
}

impl FromStr for ChildConfig {
    type Err = ();
    fn from_str(body: &str) -> Result<ChildConfig, ()> {
        parse_string("<command-line>", body,
            &ChildConfig::validator(), &Options::default())
            .map_err(|_| ())
    }
}

#[cfg(test)]
mod test {
    use std::collections::BTreeMap;
    use std::str::FromStr;
    use super::ChildConfig;
    use super::ChildKind::Daemon;
    use serde_json::{to_string, from_str};

    #[test]
    fn deserialize_compat() {
        let data = r#"{
            "instances":1,
            "image":"myproj.4a20772b",
            "config":"/config/staging/myproj.yaml",
            "kind":"Daemon"}"#;
        let cc = ChildConfig::from_str(data).unwrap();
        assert_eq!(cc, ChildConfig {
            instances: 1,
            image: String::from("myproj.4a20772b"),
            config: String::from("/config/staging/myproj.yaml"),
            variables: BTreeMap::new(),
            kind: Daemon,
        });

        let cc: ChildConfig = from_str(&data).unwrap();
        assert_eq!(cc, ChildConfig {
            instances: 1,
            image: String::from("myproj.4a20772b"),
            config: String::from("/config/staging/myproj.yaml"),
            variables: BTreeMap::new(),
            kind: Daemon,
        });
    }

    #[test]
    fn dedeserialize_vars() {
        let data = r#"{
            "instances":1,
            "image":"myproj.4a20772b",
            "config":"/config/staging/myproj.yaml",
            "variables": {"a": "b"},
            "kind":"Daemon"}"#;
        let cc = ChildConfig::from_str(data).unwrap();
        assert_eq!(cc, ChildConfig {
            instances: 1,
            image: String::from("myproj.4a20772b"),
            config: String::from("/config/staging/myproj.yaml"),
            variables: vec![
                (String::from("a"), String::from("b")),
            ].into_iter().collect(),
            kind: Daemon,
        })
    }

    #[test]
    fn serialize_compat() {
        let data = to_string(&ChildConfig {
            instances: 1,
            image: String::from("myproj.4a20772b"),
            config: String::from("/config/staging/myproj.yaml"),
            variables: BTreeMap::new(),
            kind: Daemon,
        }).unwrap();
        assert_eq!(data, "{\
            \"instances\":1,\
            \"image\":\"myproj.4a20772b\",\
            \"config\":\"/config/staging/myproj.yaml\",\
            \"kind\":\"Daemon\"}");
    }

    #[test]
    fn serialize_vars() {
        let data = to_string(&ChildConfig {
            instances: 1,
            image: String::from("myproj.4a20772b"),
            config: String::from("/config/staging/myproj.yaml"),
            variables: vec![
                (String::from("a"), String::from("b")),
                (String::from("c"), String::from("d")),
            ].into_iter().collect(),
            kind: Daemon,
        }).unwrap();
        assert_eq!(data, "{\
            \"instances\":1,\
            \"image\":\"myproj.4a20772b\",\
            \"config\":\"/config/staging/myproj.yaml\",\
            \"variables\":{\"a\":\"b\",\"c\":\"d\"},\
            \"kind\":\"Daemon\"}");
    }
}
<commit_msg>Fix tests<commit_after>use failure::Error;
use std::str::FromStr;
use std::net::IpAddr;
use std::collections::BTreeMap;

use quire::validate::{Structure, Scalar, Numeric, Mapping};
use quire::{Options, parse_string};

#[derive(Serialize, Deserialize)]
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum ChildKind {
    Daemon,
    Command,
}

// Note everything here should be stable-serializable
#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
pub struct ChildConfig {
    pub instances: usize,
    pub image: String,
    pub config: String,
    #[serde(skip_serializing_if="BTreeMap::is_empty", default)]
    pub variables: BTreeMap<String, String>,
    #[serde(skip_serializing_if="Vec::is_empty", default)]
    pub ip_addresses: Vec<IpAddr>,
    pub kind: ChildKind,
}

impl ChildConfig {
    pub fn instantiate(&self, instance: usize) -> Result<ChildConfig, Error>
    {
        let mut cfg = self.clone();
        if self.ip_addresses.len() > 0 {
            if let Some(addr) = self.ip_addresses.get(instance) {
                cfg.ip_addresses = vec![*addr];
            } else {
                bail!("Instance no {}, but there's only {} ip addresses",
                    instance, cfg.ip_addresses.len());
            }
        }
        return Ok(cfg);
    }
    pub fn mapping_validator<'x>() -> Mapping<'x> {
        return Mapping::new(
            Scalar::new(),
            ChildConfig::validator());
    }
    pub fn validator<'x>() -> Structure<'x> {
        Structure::new()
        .member("instances", Numeric::new().default(1))
        .member("image", Scalar::new())
        .member("config", Scalar::new())
        .member("variables", Mapping::new(Scalar::new(), Scalar::new()))
        .member("kind", Scalar::new().default("Daemon"))
    }
}

impl FromStr for ChildConfig {
    type Err = ();
    fn from_str(body: &str) -> Result<ChildConfig, ()> {
        parse_string("<command-line>", body,
            &ChildConfig::validator(), &Options::default())
            .map_err(|_| ())
    }
}

#[cfg(test)]
mod test {
    use std::collections::BTreeMap;
    use std::str::FromStr;
    use super::ChildConfig;
    use super::ChildKind::Daemon;
    use serde_json::{to_string, from_str};

    #[test]
    fn deserialize_compat() {
        let data = r#"{
            "instances":1,
            "image":"myproj.4a20772b",
            "config":"/config/staging/myproj.yaml",
            "kind":"Daemon"}"#;
        let cc = ChildConfig::from_str(data).unwrap();
        assert_eq!(cc, ChildConfig {
            instances: 1,
            image: String::from("myproj.4a20772b"),
            config: String::from("/config/staging/myproj.yaml"),
            variables: BTreeMap::new(),
            ip_addresses: Vec::new(),
            kind: Daemon,
        });

        let cc: ChildConfig = from_str(&data).unwrap();
        assert_eq!(cc, ChildConfig {
            instances: 1,
            image: String::from("myproj.4a20772b"),
            config: String::from("/config/staging/myproj.yaml"),
            variables: BTreeMap::new(),
            ip_addresses: Vec::new(),
            kind: Daemon,
        });
    }

    #[test]
    fn dedeserialize_vars() {
        let data = r#"{
            "instances":1,
            "image":"myproj.4a20772b",
            "config":"/config/staging/myproj.yaml",
            "variables": {"a": "b"},
            "kind":"Daemon"}"#;
        let cc = ChildConfig::from_str(data).unwrap();
        assert_eq!(cc, ChildConfig {
            instances: 1,
            image: String::from("myproj.4a20772b"),
            config: String::from("/config/staging/myproj.yaml"),
            variables: vec![
                (String::from("a"), String::from("b")),
            ].into_iter().collect(),
            ip_addresses: Vec::new(),
            kind: Daemon,
        })
    }

    #[test]
    fn serialize_compat() {
        let data = to_string(&ChildConfig {
            instances: 1,
            image: String::from("myproj.4a20772b"),
            config: String::from("/config/staging/myproj.yaml"),
            variables: BTreeMap::new(),
            ip_addresses: Vec::new(),
            kind: Daemon,
        }).unwrap();
        assert_eq!(data, "{\
            \"instances\":1,\
            \"image\":\"myproj.4a20772b\",\
            \"config\":\"/config/staging/myproj.yaml\",\
            \"kind\":\"Daemon\"}");
    }

    #[test]
    fn serialize_vars() {
        let data = to_string(&ChildConfig {
            instances: 1,
            image: String::from("myproj.4a20772b"),
            config: String::from("/config/staging/myproj.yaml"),
            variables: vec![
                (String::from("a"), String::from("b")),
                (String::from("c"), String::from("d")),
            ].into_iter().collect(),
            ip_addresses: Vec::new(),
            kind: Daemon,
        }).unwrap();
        assert_eq!(data, "{\
            \"instances\":1,\
            \"image\":\"myproj.4a20772b\",\
            \"config\":\"/config/staging/myproj.yaml\",\
            \"variables\":{\"a\":\"b\",\"c\":\"d\"},\
            \"kind\":\"Daemon\"}");
    }
}
<|endoftext|>


<commit_before>extern crate glob;
extern crate clap;
extern crate elfkit;

use std::fs::{File};
use elfkit::Elf;
use std::path::Path;
use std::collections::HashSet;
use std::io::{self};
use std::io::BufReader;
use std::io::BufRead;
use glob::glob;

fn recursive_ldd(lpaths: &Vec<String>, path: &str, visited: &mut HashSet<String>) {
    let mut f = File::open(path).unwrap();
    let mut elf = Elf::from_reader(&mut f).unwrap();

    let mut deps = Vec::new();
    for shndx in 0..elf.sections.len() {
        if elf.sections[shndx].header.shtype == elfkit::types::SectionType::DYNAMIC {
            elf.load(shndx, &mut f).unwrap();
            let dynamic = elf.sections[shndx].content.as_dynamic().unwrap();
            for dyn in dynamic.iter() {
                if dyn.dhtype == elfkit::types::DynamicType::NEEDED {
                    if let elfkit::dynamic::DynamicContent::String(ref name) = dyn.content {
                        deps.push(String::from_utf8_lossy(&name.0).into_owned());
                    }
                }
            }
        }
    }

    for dep in &mut deps {
        let mut found = false;
        for lpath in lpaths.iter().map(|p| Path::new(p)) {
            let joined = lpath.join(&dep);
            if joined.exists() {
                *dep = joined.to_string_lossy().into_owned();
                found = true;
                break;
            }
        }
        if found {
            if visited.insert(dep.clone()) {
                println!("{}", dep);
                recursive_ldd(lpaths, dep, visited);
            }
        } else {
            panic!("unable to find dependcy {} in {:?}", dep, lpaths);
        }
    }

}

fn parse_ld_so_conf(path: &str) -> io::Result<Vec<String>> {
    let mut paths = Vec::new();

    let f = File::open(path)?;
    let f = BufReader::new(&f);
    for line in f.lines() {
        let line = line?;
        let line = line.trim();
        if line.starts_with("#") {
            continue;
        }
        if line == "" {
            continue;
        }

        if line.contains(" ") {
            if line.starts_with("include ") {
                for entry in glob(line.split(" ").last().unwrap()).expect("Failed to read glob pattern") {
                    paths.extend(parse_ld_so_conf(&entry.unwrap().to_string_lossy().into_owned())?);
                }

            }
        } else {
            paths.push(line.to_owned());
        }
    }
    Ok(paths)
}

fn main() {
    let matches = clap::App::new("elfkit-ldd")
        .setting(clap::AppSettings::ArgRequiredElseHelp)
        .setting(clap::AppSettings::UnifiedHelpMessage)
        .setting(clap::AppSettings::DisableHelpSubcommand)
        .version("0.5")
        .arg(clap::Arg::with_name("file")
             .required(true)
             .help("path to binary to inspect")
             .takes_value(true)
             .index(1)
            )
        .arg(clap::Arg::with_name("library-path")
             .short("L")
             .long("library-path")
             .takes_value(true)
             .multiple(true)
             .help("specify library lookup path, ignoring sysroot/etc/ld.so.conf")
             )
        .arg(clap::Arg::with_name("sysroot")
             .short("R")
             .long("sysroot")
             .takes_value(true)
             .help("specify sysroot to look up dependencies in, instead of /")
             )
        .get_matches();

    let sysroot = matches.value_of("sysroot").unwrap_or("/").to_owned();

    let lpaths = match matches.values_of("library-path") {
        None => {
            parse_ld_so_conf(&(sysroot.clone() + "/etc/ld.so.conf")).unwrap_or(
                vec![sysroot.clone() + "/lib", sysroot.clone() + "/usr/lib"])
        },
        Some(vals) => {
            vals.map(|v|v.to_owned()).collect()
        }
    };

    let path  = matches.value_of("file").unwrap();
    let mut visited = HashSet::new();
    recursive_ldd(&lpaths, &path, &mut visited);
}

<commit_msg>ldd: pass sysroot correctly<commit_after>extern crate glob;
extern crate clap;
extern crate elfkit;

use std::fs::{File};
use elfkit::Elf;
use std::path::Path;
use std::collections::HashSet;
use std::io::{self};
use std::io::BufReader;
use std::io::BufRead;
use glob::glob;

fn recursive_ldd(lpaths: &Vec<String>, path: &str, visited: &mut HashSet<String>) {
    let mut f = File::open(path).unwrap();
    let mut elf = Elf::from_reader(&mut f).unwrap();

    let mut deps = Vec::new();
    for shndx in 0..elf.sections.len() {
        if elf.sections[shndx].header.shtype == elfkit::types::SectionType::DYNAMIC {
            elf.load(shndx, &mut f).unwrap();
            let dynamic = elf.sections[shndx].content.as_dynamic().unwrap();
            for dyn in dynamic.iter() {
                if dyn.dhtype == elfkit::types::DynamicType::NEEDED {
                    if let elfkit::dynamic::DynamicContent::String(ref name) = dyn.content {
                        deps.push(String::from_utf8_lossy(&name.0).into_owned());
                    }
                }
            }
        }
    }

    for dep in &mut deps {
        let mut found = false;
        for lpath in lpaths.iter().map(|p| Path::new(p)) {
            let joined = lpath.join(&dep);
            if joined.exists() {
                *dep = joined.to_string_lossy().into_owned();
                found = true;
                break;
            }
        }
        if found {
            if visited.insert(dep.clone()) {
                println!("{}", dep);
                recursive_ldd(lpaths, dep, visited);
            }
        } else {
            panic!("unable to find dependcy {} in {:?}", dep, lpaths);
        }
    }

}

fn parse_ld_so_conf(path: &str) -> io::Result<Vec<String>> {
    let mut paths = Vec::new();

    let f = File::open(path)?;
    let f = BufReader::new(&f);
    for line in f.lines() {
        let line = line?;
        let line = line.trim();
        if line.starts_with("#") {
            continue;
        }
        if line == "" {
            continue;
        }

        if line.contains(" ") {
            if line.starts_with("include ") {
                for entry in glob(line.split(" ").last().unwrap()).expect("Failed to read glob pattern") {
                    paths.extend(parse_ld_so_conf(&entry.unwrap().to_string_lossy().into_owned())?);
                }

            }
        } else {
            paths.push(line.to_owned());
        }
    }
    Ok(paths)
}

fn main() {
    let matches = clap::App::new("elfkit-ldd")
        .setting(clap::AppSettings::ArgRequiredElseHelp)
        .setting(clap::AppSettings::UnifiedHelpMessage)
        .setting(clap::AppSettings::DisableHelpSubcommand)
        .version("0.5")
        .arg(clap::Arg::with_name("file")
             .required(true)
             .help("path to binary to inspect")
             .takes_value(true)
             .index(1)
            )
        .arg(clap::Arg::with_name("library-path")
             .short("L")
             .long("library-path")
             .takes_value(true)
             .multiple(true)
             .help("specify library lookup path, ignoring sysroot/etc/ld.so.conf")
             )
        .arg(clap::Arg::with_name("sysroot")
             .short("R")
             .long("sysroot")
             .takes_value(true)
             .help("specify sysroot to look up dependencies in, instead of /")
             )
        .get_matches();

    let sysroot = matches.value_of("sysroot").unwrap_or("/").to_owned();

    let lpaths = match matches.values_of("library-path") {
        None => {
            match parse_ld_so_conf(&(sysroot.clone() + "/etc/ld.so.conf")) {
                Ok(l) => l.into_iter().map(|l| sysroot.clone() + &l).collect(),
                Err(_) => vec![sysroot.clone() + "/lib", sysroot.clone() + "/usr/lib"],
            }
        },
        Some(vals) => {
            vals.map(|v|v.to_owned()).collect()
        }
    };

    let path  = matches.value_of("file").unwrap();
    let mut visited = HashSet::new();
    recursive_ldd(&lpaths, &path, &mut visited);
}

<|endoftext|>


<commit_before>mod command;

use n64::*;

use self::command::*;

use std::io::*;

pub struct Debugger {
    n64: N64
}

impl Debugger {
    pub fn new(n64: N64) -> Debugger {
        Debugger {
            n64: n64
        }
    }

    pub fn run(&mut self) {
        loop {
            print!("r64> ");
            stdout().flush().unwrap();

            let command = read_stdin().parse();
            match command {
                Ok(Command::Step) => self.step(),
                _ => println!("Invalid input")
            }
        }
    }

    pub fn step(&mut self) {
        self.n64.step();
    }
}

fn read_stdin() -> String {
    let mut input = String::new();
    stdin().read_line(&mut input).unwrap();
    return input.trim().into();
}
<commit_msg>Fixed exit/quit command<commit_after>mod command;

use n64::*;

use self::command::*;

use std::io::*;

pub struct Debugger {
    n64: N64
}

impl Debugger {
    pub fn new(n64: N64) -> Debugger {
        Debugger {
            n64: n64
        }
    }

    pub fn run(&mut self) {
        loop {
            print!("r64> ");
            stdout().flush().unwrap();

            let command = read_stdin().parse();
            match command {
                Ok(Command::Step) => self.step(),
                Ok(Command::Exit) => break,
                Err(_) => println!("Invalid input")
            }
        }
    }

    pub fn step(&mut self) {
        self.n64.step();
    }
}

fn read_stdin() -> String {
    let mut input = String::new();
    stdin().read_line(&mut input).unwrap();
    return input.trim().into();
}
<|endoftext|>


<commit_before>use decoders::*;
use decoders::tiff::*;
use decoders::basics::*;
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<RawImage,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_usize(0);
    let mut height = fetch_tag!(raw, Tag::ImageLength).get_usize(0);
    let offset = fetch_tag!(raw, Tag::StripOffsets).get_usize(0);
    let count = fetch_tag!(raw, Tag::StripByteCounts).get_usize(0);
    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_usize(0)
    };
    let src = &self.buffer[offset..];

    let image = match compression {
      1 => {
        if camera.model == "DSC-R1" {
          decode_14be_unpacked(src, width, height)
        } else {
          decode_16le(src, width, height)
        }
      }
      32767 => {
        if (width*height*bps) != count*8 {
          height += 8;
          ArwDecoder::decode_arw1(src, width, height)
        } else {
          match bps {
            8 => {let curve = try!(ArwDecoder::get_curve(raw)); ArwDecoder::decode_arw2(src, width, height, &curve)},
            12 => decode_12le(src, width, height),
            _ => 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<RawImage,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_usize(0);

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

    // Get the WB the MRW way
    let priv_offset = fetch_tag!(self.tiff, Tag::DNGPrivateArea).get_force_u32(0) as usize;
    let buf = &self.buffer[priv_offset..];
    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: usize = LEu32(buf,currpos+4) as usize;
      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;
    }

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

  fn image_srf(&self, camera: Camera) -> Result<RawImage,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_usize(0);
    let height = fetch_tag!(raw, Tag::ImageLength).get_usize(0);
    let len = width*height*2;

    // 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)*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, height);

    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];

    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) as usize;
          let imin = pump.get_bits(4) as usize;

          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)] = 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) as usize;
    let priv_tiff = try!(TiffIFD::new(self.buffer, priv_offset, 0, 0, 0, LITTLE_ENDIAN));
    let sony_offset = fetch_tag!(priv_tiff, Tag::SonyOffset).get_usize(0);
    let sony_length = fetch_tag!(priv_tiff, Tag::SonyLength).get_usize(0);
    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: [usize;6] = [ 0, 0, 0, 0, 0, 4095 ];

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

    let mut out = vec![0 as u16; curve[5]+1];
    for i in 0..5 {
      for j in (curve[i]+1)..(curve[i+1]+1) {
        out[j] = out[(j-1)] + (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];
      out.push(((output >>  0) & 0xff) as u8);
      out.push(((output >>  8) & 0xff) as u8);
      out.push(((output >> 16) & 0xff) as u8);
      out.push(((output >> 24) & 0xff) as u8);
    }
    out
  }
}
<commit_msg>arw2: avoid underflow with the new pump size<commit_after>use decoders::*;
use decoders::tiff::*;
use decoders::basics::*;
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<RawImage,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_usize(0);
    let mut height = fetch_tag!(raw, Tag::ImageLength).get_usize(0);
    let offset = fetch_tag!(raw, Tag::StripOffsets).get_usize(0);
    let count = fetch_tag!(raw, Tag::StripByteCounts).get_usize(0);
    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_usize(0)
    };
    let src = &self.buffer[offset..];

    let image = match compression {
      1 => {
        if camera.model == "DSC-R1" {
          decode_14be_unpacked(src, width, height)
        } else {
          decode_16le(src, width, height)
        }
      }
      32767 => {
        if (width*height*bps) != count*8 {
          height += 8;
          ArwDecoder::decode_arw1(src, width, height)
        } else {
          match bps {
            8 => {let curve = try!(ArwDecoder::get_curve(raw)); ArwDecoder::decode_arw2(src, width, height, &curve)},
            12 => decode_12le(src, width, height),
            _ => 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<RawImage,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_usize(0);

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

    // Get the WB the MRW way
    let priv_offset = fetch_tag!(self.tiff, Tag::DNGPrivateArea).get_force_u32(0) as usize;
    let buf = &self.buffer[priv_offset..];
    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: usize = LEu32(buf,currpos+4) as usize;
      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;
    }

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

  fn image_srf(&self, camera: Camera) -> Result<RawImage,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_usize(0);
    let height = fetch_tag!(raw, Tag::ImageLength).get_usize(0);
    let len = width*height*2;

    // 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)*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, height);

    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];

    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(16);
      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) as usize;
          let imin = pump.get_bits(4) as usize;

          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)] = 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) as usize;
    let priv_tiff = try!(TiffIFD::new(self.buffer, priv_offset, 0, 0, 0, LITTLE_ENDIAN));
    let sony_offset = fetch_tag!(priv_tiff, Tag::SonyOffset).get_usize(0);
    let sony_length = fetch_tag!(priv_tiff, Tag::SonyLength).get_usize(0);
    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: [usize;6] = [ 0, 0, 0, 0, 0, 4095 ];

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

    let mut out = vec![0 as u16; curve[5]+1];
    for i in 0..5 {
      for j in (curve[i]+1)..(curve[i+1]+1) {
        out[j] = out[(j-1)] + (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];
      out.push(((output >>  0) & 0xff) as u8);
      out.push(((output >>  8) & 0xff) as u8);
      out.push(((output >> 16) & 0xff) as u8);
      out.push(((output >> 24) & 0xff) as u8);
    }
    out
  }
}
<|endoftext|>