<commit_before>//! Manager that is called when an output is created or destroyed.
//! Pass a struct that implements this trait to the `Compositor` during
//! initialization.

use compositor::{Compositor, COMPOSITOR_PTR};
use libc;
use manager::{OutputHandler, UserOutput};
use types::Output;

use wayland_sys::server::WAYLAND_SERVER_HANDLE;
use wayland_sys::server::signal::wl_signal_add;
use wlroots_sys::wlr_output;

/// Used to ensure the output sets the mode before doing any other
/// operation on the Output.
pub struct OutputBuilder<'output> {
    output: &'output mut Output
}

/// Used to ensure that the builder is used to construct
/// the OutputHandler instance.
pub struct OutputBuilderResult<'output> {
    pub output: &'output mut Output,
    result: Box<OutputHandler>
}

/// Wrapper around Output destruction so that you can't call
/// unsafe methods (e.g anything like setting the mode).
pub struct OutputDestruction<'output>(&'output mut Output);

/// Handles output addition and removal.
pub trait OutputManagerHandler {
    /// Called whenever an output is added.
    fn output_added<'output>(&mut self,
                             &mut Compositor,
                             _: OutputBuilder<'output>)
                             -> Option<OutputBuilderResult<'output>> {
        None
    }

    /// Called whenever an output is removed.
    fn output_removed(&mut self, &mut Compositor, OutputDestruction) {
        // TODO
    }
}

impl<'output> OutputBuilder<'output> {
    pub fn build_best_mode<T: OutputHandler + 'static>(self,
                                                       data: T)
                                                       -> OutputBuilderResult<'output> {
        self.output.choose_best_mode();
        OutputBuilderResult { output: self.output,
                              result: Box::new(data) }
    }
}

impl<'output> OutputDestruction<'output> {
    // TODO Functions which are safe to use
}

wayland_listener!(OutputManager, (Vec<Box<UserOutput>>, Box<OutputManagerHandler>), [
    add_listener => add_notify: |this: &mut OutputManager, data: *mut libc::c_void,| unsafe {
        let (ref mut outputs, ref mut manager) = this.data;
        let data = data as *mut wlr_output;
        let mut output = Output::new(data as *mut wlr_output);
        // NOTE
        // This clone is required because we pass it mutably to the output builder,
        // but due to lack of NLL there's no way to tell Rust it's safe to use it in
        // in the if branch.
        //
        // Thus, we need to clone it here and then drop the original once at the end.
        //
        // This is not a real clone, but an pub(crate) unsafe one we added, so it doesn't
        // break safety concerns in user code. Just an unfortunate hack we have to put here.
        let output_clone = output.clone();
        let builder = OutputBuilder { output: &mut output };
        let compositor = &mut *COMPOSITOR_PTR;
        output_clone.set_lock(true);
        let build_result = manager.output_added(compositor, builder);
        if let Some(OutputBuilderResult {result: output_ptr, .. }) = build_result {
            output_clone.set_lock(false);
            let mut output = UserOutput::new((output_clone, output_ptr));
            // Add the output frame event to this manager
            wl_signal_add(&mut (*data).events.frame as *mut _ as _,
                        output.frame_listener() as _);
            // Add the output resolution event to this manager
            wl_signal_add(&mut (*data).events.resolution as *mut _ as _,
                        output.resolution_listener() as _);
            // Store the user UserOutput, free later in remove listener
            outputs.push(output);
        }
    };
    remove_listener => remove_notify: |this: &mut OutputManager, data: *mut libc::c_void,| unsafe {
        let (ref mut outputs, ref mut manager) = this.data;
        let data = data as *mut wlr_output;
        // NOTE
        // We get it from the list so that we can get the Rc'd `Output`, because there's
        // no way to re-construct that using just the raw pointer.
        if let Some(output) = outputs.iter_mut().find(|output| output.output_ptr() == data) {
            let output = output.output_mut();
            let compositor = &mut *COMPOSITOR_PTR;
            output.set_lock(true);
            manager.output_removed(compositor, OutputDestruction(output));
            // NOTE We don't remove the lock because we are removing it
            if let Some(layout) = output.layout() {
                layout.borrow_mut().remove(output);
            }
        }
        // Remove user output data
        if let Some(index) = outputs.iter().position(|output| output.output_ptr() == data) {
            let mut removed_output = outputs.remove(index);
            ffi_dispatch!(WAYLAND_SERVER_HANDLE,
                          wl_list_remove,
                          &mut (*removed_output.frame_listener()).link as *mut _ as _);
            ffi_dispatch!(WAYLAND_SERVER_HANDLE,
                          wl_list_remove,
                          &mut (*removed_output.resolution_listener()).link as *mut _ as _);

        }
    };
]);
<commit_msg>A panic in output setup now causes an abort<commit_after>//! Manager that is called when an output is created or destroyed.
//! Pass a struct that implements this trait to the `Compositor` during
//! initialization.

use compositor::{Compositor, COMPOSITOR_PTR};
use libc;
use manager::{OutputHandler, UserOutput};
use types::Output;

use wayland_sys::server::WAYLAND_SERVER_HANDLE;
use wayland_sys::server::signal::wl_signal_add;
use wlroots_sys::wlr_output;

use std::panic;

/// Used to ensure the output sets the mode before doing any other
/// operation on the Output.
pub struct OutputBuilder<'output> {
    output: &'output mut Output
}

/// Used to ensure that the builder is used to construct
/// the OutputHandler instance.
pub struct OutputBuilderResult<'output> {
    pub output: &'output mut Output,
    result: Box<OutputHandler>
}

/// Wrapper around Output destruction so that you can't call
/// unsafe methods (e.g anything like setting the mode).
pub struct OutputDestruction<'output>(&'output mut Output);

/// Handles output addition and removal.
pub trait OutputManagerHandler {
    /// Called whenever an output is added.
    ///
    /// # Panics
    /// Any panic in this function will cause the process to abort.
    fn output_added<'output>(&mut self,
                             &mut Compositor,
                             _: OutputBuilder<'output>)
                             -> Option<OutputBuilderResult<'output>> {
        None
    }

    /// Called whenever an output is removed.
    fn output_removed(&mut self, &mut Compositor, OutputDestruction) {
        // TODO
    }
}

impl<'output> OutputBuilder<'output> {
    pub fn build_best_mode<T: OutputHandler + 'static>(self,
                                                       data: T)
                                                       -> OutputBuilderResult<'output> {
        self.output.choose_best_mode();
        OutputBuilderResult { output: self.output,
                              result: Box::new(data) }
    }
}

impl<'output> OutputDestruction<'output> {
    // TODO Functions which are safe to use
}

wayland_listener!(OutputManager, (Vec<Box<UserOutput>>, Box<OutputManagerHandler>), [
    add_listener => add_notify: |this: &mut OutputManager, data: *mut libc::c_void,| unsafe {
        let (ref mut outputs, ref mut manager) = this.data;
        let data = data as *mut wlr_output;
        let mut output = Output::new(data as *mut wlr_output);
        // NOTE
        // This clone is required because we pass it mutably to the output builder,
        // but due to lack of NLL there's no way to tell Rust it's safe to use it in
        // in the if branch.
        //
        // Thus, we need to clone it here and then drop the original once at the end.
        //
        // This is not a real clone, but an pub(crate) unsafe one we added, so it doesn't
        // break safety concerns in user code. Just an unfortunate hack we have to put here.
        let output_clone = output.clone();
        let builder = OutputBuilder { output: &mut output };
        let compositor = &mut *COMPOSITOR_PTR;
        output_clone.set_lock(true);
        let res = panic::catch_unwind(
            panic::AssertUnwindSafe(||manager.output_added(compositor, builder)));
        let build_result = match res {
            Ok(res) => res,
            // NOTE
            // Either Wayland or wlroots does not handle failure to set up output correctly.
            // Calling wl_display_terminate does not work if output is incorrectly set up.
            //
            // Instead, execution keeps going with an eventual segfault (if lucky).
            //
            // To fix this, we abort the process if there was a panic in output setup.
            Err(_) => ::std::process::abort()
        };
        if let Some(OutputBuilderResult {result: output_ptr, .. }) = build_result {
            output_clone.set_lock(false);
            let mut output = UserOutput::new((output_clone, output_ptr));
            // Add the output frame event to this manager
            wl_signal_add(&mut (*data).events.frame as *mut _ as _,
                        output.frame_listener() as _);
            // Add the output resolution event to this manager
            wl_signal_add(&mut (*data).events.resolution as *mut _ as _,
                        output.resolution_listener() as _);
            // Store the user UserOutput, free later in remove listener
            outputs.push(output);
        }
    };
    remove_listener => remove_notify: |this: &mut OutputManager, data: *mut libc::c_void,| unsafe {
        let (ref mut outputs, ref mut manager) = this.data;
        let data = data as *mut wlr_output;
        // NOTE
        // We get it from the list so that we can get the Rc'd `Output`, because there's
        // no way to re-construct that using just the raw pointer.
        if let Some(output) = outputs.iter_mut().find(|output| output.output_ptr() == data) {
            let output = output.output_mut();
            let compositor = &mut *COMPOSITOR_PTR;
            output.set_lock(true);
            manager.output_removed(compositor, OutputDestruction(output));
            // NOTE We don't remove the lock because we are removing it
            if let Some(layout) = output.layout() {
                layout.borrow_mut().remove(output);
            }
        }
        // Remove user output data
        if let Some(index) = outputs.iter().position(|output| output.output_ptr() == data) {
            let mut removed_output = outputs.remove(index);
            ffi_dispatch!(WAYLAND_SERVER_HANDLE,
                          wl_list_remove,
                          &mut (*removed_output.frame_listener()).link as *mut _ as _);
            ffi_dispatch!(WAYLAND_SERVER_HANDLE,
                          wl_list_remove,
                          &mut (*removed_output.resolution_listener()).link as *mut _ as _);

        }
    };
]);
<|endoftext|>


<commit_before>//! A symbolic memory model.
//!
//! Each cell of `SymbolicMemory` is a valid `il::Expr`. Concrete values may be stored as
//! `il::Constant`, and symbolic values stored as valid expressions.
//!
//! `SymbolicMemory` is paged under-the-hood with reference-counted pages. When these pages
//! are written to, we use the copy-on-write functionality of rust's `std::sync::Arc` type,
//! giving us copy-on-write paging. This allows for very fast forks of `SymbolicMemory`.

use engine;
use error::*;
use il;
use std::collections::BTreeMap;
use std::sync::Arc;

const PAGE_SIZE: usize = 1024;

/// The underlying endianness of this memory model.
#[derive(Clone, Debug)]
pub enum Endian {
    Big,
    Little
}

#[derive(Clone)]
struct SymbolicPage {
    size: usize,
    cells: Vec<il::Expression>
}

impl SymbolicPage {
    fn new(size: usize) -> SymbolicPage {
        let mut v = Vec::new();
        for _ in 0..size {
            v.push(il::expr_const(0, 8));
        }

        SymbolicPage {
            size: size,
            cells: v
        }
    }

    fn store(&mut self, offset: usize, value: il::Expression) -> Result<()> {
        if value.bits() != 8 {
            bail!("SymbolicPage tried to store value with bits={}", value.bits());
        }

        if offset >= self.size {
            bail!("Out of bounds offset {} for SymbolicPage with size {}", offset, self.size);
        }

        self.cells.as_mut_slice()[offset] = value;

        Ok(())
    }

    fn load(&self, offset: usize) -> Result<il::Expression> {
        if offset >= self.size {
            bail!("Out of bounds offset {} for SymbolicPage with size {}", offset, self.size);
        }

        Ok(self.cells[offset].clone())
    }
}

/// A symbolic memory model for Falcon IL expressions.
#[derive(Clone)]
pub struct SymbolicMemory {
    endian: Endian,
    pages: BTreeMap<u64, Arc<SymbolicPage>>
}

impl SymbolicMemory {
    /// Create a new `SymbolicMemory`.
    pub fn new(endian: Endian) -> SymbolicMemory {
        SymbolicMemory {
            endian: endian,
            pages: BTreeMap::new()
        }
    }

    /// Get the endianness of this `SymbolicMemory`.
    pub fn endian(&self) -> &Endian {
        &self.endian
    }

    fn store_byte(&mut self, address: u64, value: il::Expression) -> Result<()> {
        let page_address = address & !(PAGE_SIZE as u64 - 1);
        let offset = (address & (PAGE_SIZE as u64 - 1)) as usize;

        if let Some(mut page) = self.pages.get_mut(&page_address) {
            Arc::make_mut(&mut page).store(offset, value)?;
            return Ok(())
        }

        let mut page = SymbolicPage::new(PAGE_SIZE);
        page.store(offset, value)?;
        self.pages.insert(page_address, Arc::new(page));

        Ok(())
    }

    fn load_byte(&self, address: u64) -> Result<Option<il::Expression>> {
        let page_address = address & !(PAGE_SIZE as u64 - 1);
        let offset = (address & (PAGE_SIZE as u64 - 1)) as usize;
        match self.pages.get(&page_address) {
            Some(page) => Ok(Some(page.load(offset)?)),
            None => Ok(None)
        }
    }

    /// Store an expression at the given address.
    ///
    /// The value must have a bit-width >= 8, and the bit-width must be evenly divisible
    /// by 8.
    pub fn store(&mut self, address: u64, value: il::Expression) -> Result<()> {
        if value.bits() % 8 != 0 {
            return Err(format!("Storing value in symbolic with bit width not divisible by 8 {}",
                value.bits()).into());
        }
        if value.bits() > 8 {
            let bytes = value.bits() / 8;
            for offset in 0..bytes {
                let offset = offset as u64;
                let shift = match self.endian {
                    Endian::Big => (bytes as u64 - offset - 1) * 8,
                    Endian::Little => offset * 8
                };
                let shift = il::expr_const(shift, value.bits());
                let value = il::Expression::shr(value.clone(), shift)?;
                let value = il::Expression::trun(8, value)?;
                // trace!("STORE [{:x}]={}", address + offset, value);
                self.store_byte(address + offset, value)?;
            }
        }
        else if value.bits() == 8 {
            self.store_byte(address, value)?;
        }
        else {
            return Err(format!("Invalid bit width in symbolic memory store: {}",
                value.bits()).into());
        }
        Ok(())
    }

    /// Loads an expression from the given address.
    ///
    /// `bits` must be >= 8, and evenly divisible by 8.
    /// If a value exists, it will be returned in `Some(expr)`. If no value exists for all
    /// bits at the given address, `None` will be returned.
    pub fn load(&self, address: u64, bits: usize) -> Result<Option<il::Expression>> {
        if bits % 8 != 0 {
            Err(format!("Loading symbolic memory with non-8 bit-width {}", bits).into())
        }
        else if bits == 0 {
            Err("Loading symbolic memory with 0 bit-width".into())
        }
        else if bits == 8 {
            match self.load_byte(address)? {
                Some(ref expr) => Ok(Some(engine::simplify_expression(expr)?)),
                None => Ok(None)
            }
        }
        else {
            let mut result = None;
            let bytes = (bits / 8) as u64;
            for offset in 0..bytes {
                let expr = match self.load_byte(address + offset)? {
                    Some(expr) => expr,
                    None => return Ok(None)
                };
                // trace!("LOAD [{:x}]={}", address + offset, expr);
                let expr = il::Expression::zext(bits, expr.clone())?;
                let shift = match self.endian {
                    Endian::Big => (bytes - offset - 1) * 8,
                    Endian::Little => offset * 8
                };
                let shift = il::expr_const(shift, bits);
                let expr = il::Expression::shl(expr, shift)?;
                result = match result {
                    Some(result) => Some(il::Expression::or(result, expr)?),
                    None => Some(expr)
                };
            }

            if let Some(expr) = result {
                result = Some(engine::simplify_expression(&expr)?);
            }

            Ok(result)
        }
    }
}<commit_msg>no split memory seems stable.<commit_after>//! A symbolic memory model.
//!
//! Each cell of `SymbolicMemory` is a valid `il::Expr`. Concrete values may be stored as
//! `il::Constant`, and symbolic values stored as valid expressions.
//!
//! `SymbolicMemory` is paged under-the-hood with reference-counted pages. When these pages
//! are written to, we use the copy-on-write functionality of rust's `std::sync::Arc` type,
//! giving us copy-on-write paging. This allows for very fast forks of `SymbolicMemory`.

use error::*;
use il;
use std::collections::BTreeMap;
use std::sync::Arc;

const PAGE_SIZE: usize = 1024;

/// The underlying endianness of this memory model.
#[derive(Clone, Debug)]
pub enum Endian {
    Big,
    Little
}

/// We would prefer to avoid splitting values when reading/writing memory values
/// > 1 byte to memory. One way we can do this is by enabling memory to hold values
/// > 1 byte in length. Yeah, no, I'm sure this will work out well.
/// Every memory location will hold either an expression, or a Backref, which is
/// an address of the beginning of an expression which extends to this address.
///
/// If we get this right, it should be pretty awesome. _If_we_get_this_right_.
#[derive(Clone)]
enum MemoryCell {
    Expression(il::Expression),
    Backref(u64)
}

impl MemoryCell {
    fn expression(&self) -> Option<&il::Expression> {
        match self {
            &MemoryCell::Expression(ref expr) => Some(expr),
            &MemoryCell::Backref(_) => None
        }
    }
}

#[derive(Clone)]
struct SymbolicPage {
    size: usize,
    cells: Vec<MemoryCell>
}

impl SymbolicPage {
    fn new(size: usize) -> SymbolicPage {
        let mut v = Vec::new();
        for _ in 0..size {
            v.push(MemoryCell::Expression(il::expr_const(0, 8)));
        }

        SymbolicPage {
            size: size,
            cells: v
        }
    }

    fn store(&mut self, offset: usize, cell: MemoryCell) -> Result<()> {
        if offset >= self.size {
            bail!("Out of bounds offset {} for SymbolicPage with size {}", offset, self.size);
        }

        self.cells.as_mut_slice()[offset] = cell;

        Ok(())
    }

    fn load(&self, offset: usize) -> Result<&MemoryCell> {
        if offset >= self.size {
            bail!("Out of bounds offset {} for SymbolicPage with size {}", offset, self.size);
        }

        Ok(&self.cells[offset])
    }
}

/// A symbolic memory model for Falcon IL expressions.
#[derive(Clone)]
pub struct SymbolicMemory {
    endian: Endian,
    pages: BTreeMap<u64, Arc<SymbolicPage>>
}

impl SymbolicMemory {
    /// Create a new `SymbolicMemory`.
    pub fn new(endian: Endian) -> SymbolicMemory {
        SymbolicMemory {
            endian: endian,
            pages: BTreeMap::new()
        }
    }

    /// Get the endianness of this `SymbolicMemory`.
    pub fn endian(&self) -> &Endian {
        &self.endian
    }

    fn store_cell(&mut self, address: u64, cell: MemoryCell) -> Result<()> {
        let page_address = address & !(PAGE_SIZE as u64 - 1);
        let offset = (address & (PAGE_SIZE as u64 - 1)) as usize;

        if let Some(mut page) = self.pages.get_mut(&page_address) {
            Arc::make_mut(&mut page).store(offset, cell)?;
            return Ok(())
        }

        let mut page = SymbolicPage::new(PAGE_SIZE);
        page.store(offset, cell)?;
        self.pages.insert(page_address, Arc::new(page));

        Ok(())
    }

    fn load_cell(&self, address: u64) -> Result<Option<&MemoryCell>> {
        let page_address = address & !(PAGE_SIZE as u64 - 1);
        let offset = (address & (PAGE_SIZE as u64 - 1)) as usize;
        match self.pages.get(&page_address) {
            Some(page) => Ok(Some(page.load(offset)?)),
            None => Ok(None)
        }
    }

    /// Store an expression at the given address.
    ///
    /// The value must have a bit-width >= 8, and the bit-width must be evenly divisible
    /// by 8.
    pub fn store(&mut self, address: u64, value: il::Expression) -> Result<()> {
        if value.bits() % 8 != 0 || value.bits() == 0 {
            return Err(format!("Storing value in symbolic with bit width not divisible by 8 and > 0 {}",
                value.bits()).into());
        }

        // There are a few scenarios here we need to account for
        // E is for Expression, B is for Backref. Consider a 4-byte write, with
        // the original memory on top and our write immediately underneath that.
        //
        // Case 0
        // EEEBEE   The easiest scenario, we just replace expressions in place
        //  WWWW
        //
        // Case 1
        // EBBEEE   We overwrite some backrefs that refer to before our write.
        //  WWWW    We need to truncate the expression before.
        //          First byte we overwrite is a backref.
        //
        // Case 2
        // EEEBBB   We overwrite an expression that starts in the middle of our
        //  WWWW    write.
        //          Byte after last byte is a backref.
        //
        // Case 3
        // EBBBBB   We overwrite an expression that starts before our expression,
        //  WWWW    and continues after out expression.
        //          Handle case 2, then case 1, and case 3 will be fine.

        // If the byte after the last byte is a Backref
        let address_after_write = address + (value.bits() / 8) as u64;

        let write = if let Some(cell) = self.load_cell(address_after_write)? {
            if let MemoryCell::Backref(backref_address) = *cell {
                let expr = self.load_cell(backref_address)?
                               .unwrap()
                               .expression()
                               .unwrap();
                let expr_bits = expr.bits();
                let shift_bits = (address_after_write - backref_address) * 8;
                let final_bits = expr_bits - shift_bits as usize;
                match self.endian {
                    Endian::Little => {
                        let expr = il::Expression::shr(
                            expr.clone(),
                            il::expr_const(shift_bits, expr_bits)
                        )?;
                        let expr = il::Expression::trun(final_bits, expr)?;
                        Some((address_after_write, expr))
                    },
                    Endian::Big => {
                        let expr = il::Expression::trun(final_bits, expr.clone())?;
                        Some((address_after_write, expr))
                    }
                }
            }
            else {
                None
            }
        }
        else {
            None
        };

        if let Some((address, expr)) = write {
            let expr_bytes = (expr.bits() / 8) as u64;
            self.store_cell(address, MemoryCell::Expression(expr))?;
            for i in 1..expr_bytes {
                self.store_cell(address + i, MemoryCell::Backref(address))?;
            }
        }

        // If the first byte of the write is a Backref
        let write = if let Some(cell) = self.load_cell(address)? {
            if let MemoryCell::Backref(backref_address) = *cell {
                let expr = self.load_cell(backref_address)?
                               .unwrap()
                               .expression()
                               .unwrap();
                let expr_bits = expr.bits();
                let final_bits = (address - backref_address) as usize * 8;
                let shift_bits = (expr_bits - final_bits) as u64;
                match self.endian {
                    Endian::Little => {
                        let expr = il::Expression::trun(final_bits, expr.clone())?;
                        Some((backref_address, expr))
                    },
                    Endian::Big => {
                        let expr = il::Expression::shr(expr.clone(), il::expr_const(shift_bits, expr_bits))?;
                        let expr = il::Expression::trun(final_bits, expr)?;
                        Some((backref_address, expr))
                    }
                }
            }
            else {
                None
            }
        }
        else {
            None
        };

        if let Some((address, expr)) = write {
            self.store(address, expr)?;
        }

        // Now store this value and set its backrefs
        let bits = value.bits();
        self.store_cell(address, MemoryCell::Expression(value))?;

        for i in 1..(bits / 8) {
            self.store_cell(address + i as u64, MemoryCell::Backref(address))?;
        }

        Ok(())
    }

    /// Loads an expression from the given address.
    ///
    /// `bits` must be >= 8, and evenly divisible by 8.
    /// If a value exists, it will be returned in `Some(expr)`. If no value exists for all
    /// bits at the given address, `None` will be returned.
    pub fn load(&self, address: u64, bits: usize) -> Result<Option<il::Expression>> {
        if bits % 8 != 0 {
            return Err(format!("Loading symbolic memory with non-8 bit-width {}", bits).into());
        }
        else if bits == 0 {
            return Err("Loading symbolic memory with 0 bit-width".into());
        }

        // The scenarios we need to account for
        // E is Expression, B is Backref, L is for Load
        //
        // Case 0
        // EEBBBE   A perfect match. No adjustments required.
        //  LLLL
        //
        // Case 1
        // EEBBBB   An expression extends beyond the load, truncate
        //  LLLL
        //
        // Case 2
        // EBBBBB   An expression overlaps on both sides. Shift and truncate.
        //  LLLL
        //
        // Case 3
        // EEBEBE   Multiple sub-expressions. Shift and or them together.
        //  LLLL
        //
        // Case 4
        // EBEBEB   Overlapping expression before, in the middle, and after
        //  LLLL    Handle case 2, then case 3
        //

        // This will be a nightmare to deal with endian-wise, so we're going to
        // attempt to load everything the first time, and if this doesn't work
        // then we'll do single-byte loads and patch everything together.

        // Get started with our first load
        let load_expr = if let Some(cell) = self.load_cell(address)? {
            match *cell {
                MemoryCell::Expression(ref expr) => {
                    if expr.bits() <= bits {
                        expr.clone()
                    }
                    else {
                        il::Expression::trun(bits, expr.clone())?
                    }
                },
                MemoryCell::Backref(backref_address) => {
                    let expr = self.load_cell(backref_address)?
                                   .unwrap()
                                   .expression()
                                   .unwrap();
                    let expr_bits = expr.bits();
                    let shift_bits = (address - backref_address) * 8;
                    let final_bits = expr_bits - shift_bits as usize;
                    let expr = match self.endian {
                        Endian::Little => {
                            let expr = il::Expression::shl(expr.clone(),
                                    il::expr_const(shift_bits, expr_bits))?;
                            let expr = il::Expression::trun(final_bits, expr)?;
                            expr
                        },
                        Endian::Big => {
                            let expr = il::Expression::trun(final_bits, expr.clone())?;
                            expr
                        }
                    };
                    if expr.bits() > bits {
                        il::Expression::trun(bits, expr)?
                    }
                    else {
                        expr
                    }
                }
            }
        }
        else {
            return Ok(None);
        };

        // if we're done, finish
        if load_expr.bits() == bits {
            return Ok(Some(load_expr));
        }

        // Fall back to single-byte loads
        let mut result = None;
        let bytes = (bits / 8) as u64;
        for offset in 0..bytes {
            let expr = match self.load(address + offset, 8)? {
                Some(expr) => expr,
                None => return Ok(None)
            };
            // trace!("LOAD [{:x}]={}", address + offset, expr);
            let expr = il::Expression::zext(bits, expr.clone())?;
            let shift = match self.endian {
                Endian::Big => (bytes - offset - 1) * 8,
                Endian::Little => offset * 8
            };
            let shift = il::expr_const(shift, bits);
            let expr = il::Expression::shl(expr, shift)?;
            result = match result {
                Some(result) => Some(il::Expression::or(result, expr)?),
                None => Some(expr)
            };
        }

        Ok(result)
    }
}<|endoftext|>


<commit_before>// Rust Bitcoin Library
// Written in 2014 by
//     Andrew Poelstra <apoelstra@wpsoftware.net>
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # Network message
//!
//! This module defines the `Message` traits which are used
//! for (de)serializing Bitcoin objects for transmission on the network. It
//! also defines (de)serialization routines for many primitives.
//!

use std::iter;
use std::io::Cursor;
use std::sync::mpsc::Sender;

use blockdata::block;
use blockdata::transaction;
use network::address::Address;
use network::message_network;
use network::message_blockdata;
use network::encodable::{ConsensusDecodable, ConsensusEncodable};
use network::encodable::CheckedData;
use network::serialize::{serialize, RawDecoder, SimpleEncoder, SimpleDecoder};
use util::{self, propagate_err};

/// Serializer for command string
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct CommandString(pub String);

impl<S: SimpleEncoder> ConsensusEncodable<S> for CommandString {
    #[inline]
    fn consensus_encode(&self, s: &mut S) -> Result<(), S::Error> {
        let &CommandString(ref inner_str) = self;
        let mut rawbytes = [0u8; 12];
        let strbytes = inner_str.as_bytes();
        if strbytes.len() > 12 {
            panic!("Command string longer than 12 bytes");
        }
        for x in 0..strbytes.len() {
            rawbytes[x] = strbytes[x];
        }
        rawbytes.consensus_encode(s)
    }
}

impl<D: SimpleDecoder> ConsensusDecodable<D> for CommandString {
    #[inline]
    fn consensus_decode(d: &mut D) -> Result<CommandString, D::Error> {
        let rawbytes: [u8; 12] = try!(ConsensusDecodable::consensus_decode(d));
        let rv = iter::FromIterator::from_iter(rawbytes.iter().filter_map(|&u| if u > 0 { Some(u as char) } else { None }));
        Ok(CommandString(rv))
    }
}

/// A Network message
pub struct RawNetworkMessage {
    /// Magic bytes to identify the network these messages are meant for
    pub magic: u32,
    /// The actual message data
    pub payload: NetworkMessage
}

/// A response from the peer-connected socket
pub enum SocketResponse {
    /// A message was received
    MessageReceived(NetworkMessage),
    /// An error occured and the socket needs to close
    ConnectionFailed(util::Error, Sender<()>)
}

#[derive(Clone, PartialEq, Eq, Debug)]
/// A Network message payload. Proper documentation is available on the Bitcoin
/// wiki https://en.bitcoin.it/wiki/Protocol_specification
pub enum NetworkMessage {
    /// `version`
    Version(message_network::VersionMessage),
    /// `verack`
    Verack,
    /// `addr`
    Addr(Vec<(u32, Address)>),
    /// `inv`
    Inv(Vec<message_blockdata::Inventory>),
    /// `getdata`
    GetData(Vec<message_blockdata::Inventory>),
    /// `notfound`
    NotFound(Vec<message_blockdata::Inventory>),
    /// `getblocks`
    GetBlocks(message_blockdata::GetBlocksMessage),
    /// `getheaders`
    GetHeaders(message_blockdata::GetHeadersMessage),
    /// `mempool`
    MemPool,
    /// tx
    Tx(transaction::Transaction),
    /// `block`
    Block(block::Block),
    /// `headers`
    Headers(Vec<block::LoneBlockHeader>),
    // TODO: getaddr,
    // TODO: checkorder,
    // TODO: submitorder,
    // TODO: reply,
    /// `ping`
    Ping(u64),
    /// `pong`
    Pong(u64)
    // TODO: reject,
    // TODO: bloom filtering
    // TODO: alert
}

impl RawNetworkMessage {
    fn command(&self) -> String {
        match self.payload {
            NetworkMessage::Version(_) => "version",
            NetworkMessage::Verack     => "verack",
            NetworkMessage::Addr(_)    => "addr",
            NetworkMessage::Inv(_)     => "inv",
            NetworkMessage::GetData(_) => "getdata",
            NetworkMessage::NotFound(_) => "notfound",
            NetworkMessage::GetBlocks(_) => "getblocks",
            NetworkMessage::GetHeaders(_) => "getheaders",
            NetworkMessage::MemPool    => "mempool",
            NetworkMessage::Tx(_)      => "tx",
            NetworkMessage::Block(_)   => "block",
            NetworkMessage::Headers(_) => "headers",
            NetworkMessage::Ping(_)    => "ping",
            NetworkMessage::Pong(_)    => "pong",
        }.to_owned()
    }
}

impl<S: SimpleEncoder> ConsensusEncodable<S> for RawNetworkMessage {
    fn consensus_encode(&self, s: &mut S) -> Result<(), S::Error> {
        try!(self.magic.consensus_encode(s));
        try!(CommandString(self.command()).consensus_encode(s));
        try!(CheckedData(match self.payload {
            NetworkMessage::Version(ref dat) => serialize(dat),
            NetworkMessage::Verack           => Ok(vec![]),
            NetworkMessage::Addr(ref dat)    => serialize(dat),
            NetworkMessage::Inv(ref dat)     => serialize(dat),
            NetworkMessage::GetData(ref dat) => serialize(dat),
            NetworkMessage::NotFound(ref dat) => serialize(dat),
            NetworkMessage::GetBlocks(ref dat) => serialize(dat),
            NetworkMessage::GetHeaders(ref dat) => serialize(dat),
            NetworkMessage::MemPool          => Ok(vec![]),
            NetworkMessage::Tx(ref dat)      => serialize(dat),
            NetworkMessage::Block(ref dat)   => serialize(dat),
            NetworkMessage::Headers(ref dat) => serialize(dat),
            NetworkMessage::Ping(ref dat)    => serialize(dat),
            NetworkMessage::Pong(ref dat)    => serialize(dat),
        }.unwrap()).consensus_encode(s));
        Ok(())
    }
}

// TODO: restriction on D::Error is so that `propagate_err` will work;
//       is there a more generic way to handle this?
impl<D: SimpleDecoder<Error=util::Error>> ConsensusDecodable<D> for RawNetworkMessage {
    fn consensus_decode(d: &mut D) -> Result<RawNetworkMessage, D::Error> {
        let magic = try!(ConsensusDecodable::consensus_decode(d));
        let CommandString(cmd): CommandString= try!(ConsensusDecodable::consensus_decode(d));
        let CheckedData(raw_payload): CheckedData = try!(ConsensusDecodable::consensus_decode(d));

        let mut mem_d = RawDecoder::new(Cursor::new(raw_payload));
        let payload = match &cmd[..] {
            "version" => NetworkMessage::Version(try!(propagate_err("version".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "verack"  => NetworkMessage::Verack,
            "addr"    => NetworkMessage::Addr(try!(propagate_err("addr".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "inv"     => NetworkMessage::Inv(try!(propagate_err("inv".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "getdata" => NetworkMessage::GetData(try!(propagate_err("getdata".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "notfound" => NetworkMessage::NotFound(try!(propagate_err("notfound".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "getblocks" => NetworkMessage::GetBlocks(try!(propagate_err("getblocks".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "getheaders" => NetworkMessage::GetHeaders(try!(propagate_err("getheaders".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "mempool" => NetworkMessage::MemPool,
            "block"   => NetworkMessage::Block(try!(propagate_err("block".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "headers" => NetworkMessage::Headers(try!(propagate_err("headers".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "ping"    => NetworkMessage::Ping(try!(propagate_err("ping".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "pong"    => NetworkMessage::Ping(try!(propagate_err("pong".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "tx"      => NetworkMessage::Tx(try!(propagate_err("tx".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            cmd => return Err(d.error(format!("unrecognized network command `{}`", cmd)))
        };
        Ok(RawNetworkMessage {
            magic: magic,
            payload: payload
        })
    }
}

#[cfg(test)]
mod test {
    use super::{RawNetworkMessage, NetworkMessage, CommandString};

    use network::serialize::{deserialize, serialize};

    #[test]
    fn serialize_commandstring_test() {
        let cs = CommandString("Andrew".to_owned());
        assert_eq!(serialize(&cs).ok(), Some(vec![0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0, 0]));
    }

    #[test]
    fn deserialize_commandstring_test() {
        let cs: Result<CommandString, _> = deserialize(&[0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0, 0]);
        assert!(cs.is_ok());
        assert_eq!(cs.unwrap(), CommandString("Andrew".to_owned()));

        let short_cs: Result<CommandString, _> = deserialize(&[0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0]);
        assert!(short_cs.is_err());
    }

    #[test]
    fn serialize_verack_test() {
        assert_eq!(serialize(&RawNetworkMessage { magic: 0xd9b4bef9, payload: NetworkMessage::Verack }).ok(),
                             Some(vec![0xf9, 0xbe, 0xb4, 0xd9, 0x76, 0x65, 0x72, 0x61,
                                       0x63, 0x6B, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                       0x00, 0x00, 0x00, 0x00, 0x5d, 0xf6, 0xe0, 0xe2]));
    }

    #[test]
    fn serialize_ping_test() {
        assert_eq!(serialize(&RawNetworkMessage { magic: 0xd9b4bef9, payload: NetworkMessage::Ping(100) }).ok(),
                             Some(vec![0xf9, 0xbe, 0xb4, 0xd9, 0x70, 0x69, 0x6e, 0x67,
                                       0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                       0x08, 0x00, 0x00, 0x00, 0x24, 0x67, 0xf1, 0x1d,
                                       0x64, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]));
    }

    #[test]
    fn serialize_mempool_test() {
        assert_eq!(serialize(&RawNetworkMessage { magic: 0xd9b4bef9, payload: NetworkMessage::MemPool }).ok(),
                             Some(vec![0xf9, 0xbe, 0xb4, 0xd9, 0x6d, 0x65, 0x6d, 0x70,
                                       0x6f, 0x6f, 0x6c, 0x00, 0x00, 0x00, 0x00, 0x00,
                                       0x00, 0x00, 0x00, 0x00, 0x5d, 0xf6, 0xe0, 0xe2]));
    }

}

<commit_msg>Fix decoding of Pong message<commit_after>// Rust Bitcoin Library
// Written in 2014 by
//     Andrew Poelstra <apoelstra@wpsoftware.net>
//
// To the extent possible under law, the author(s) have dedicated all
// copyright and related and neighboring rights to this software to
// the public domain worldwide. This software is distributed without
// any warranty.
//
// You should have received a copy of the CC0 Public Domain Dedication
// along with this software.
// If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
//

//! # Network message
//!
//! This module defines the `Message` traits which are used
//! for (de)serializing Bitcoin objects for transmission on the network. It
//! also defines (de)serialization routines for many primitives.
//!

use std::iter;
use std::io::Cursor;
use std::sync::mpsc::Sender;

use blockdata::block;
use blockdata::transaction;
use network::address::Address;
use network::message_network;
use network::message_blockdata;
use network::encodable::{ConsensusDecodable, ConsensusEncodable};
use network::encodable::CheckedData;
use network::serialize::{serialize, RawDecoder, SimpleEncoder, SimpleDecoder};
use util::{self, propagate_err};

/// Serializer for command string
#[derive(PartialEq, Eq, Clone, Debug)]
pub struct CommandString(pub String);

impl<S: SimpleEncoder> ConsensusEncodable<S> for CommandString {
    #[inline]
    fn consensus_encode(&self, s: &mut S) -> Result<(), S::Error> {
        let &CommandString(ref inner_str) = self;
        let mut rawbytes = [0u8; 12];
        let strbytes = inner_str.as_bytes();
        if strbytes.len() > 12 {
            panic!("Command string longer than 12 bytes");
        }
        for x in 0..strbytes.len() {
            rawbytes[x] = strbytes[x];
        }
        rawbytes.consensus_encode(s)
    }
}

impl<D: SimpleDecoder> ConsensusDecodable<D> for CommandString {
    #[inline]
    fn consensus_decode(d: &mut D) -> Result<CommandString, D::Error> {
        let rawbytes: [u8; 12] = try!(ConsensusDecodable::consensus_decode(d));
        let rv = iter::FromIterator::from_iter(rawbytes.iter().filter_map(|&u| if u > 0 { Some(u as char) } else { None }));
        Ok(CommandString(rv))
    }
}

/// A Network message
pub struct RawNetworkMessage {
    /// Magic bytes to identify the network these messages are meant for
    pub magic: u32,
    /// The actual message data
    pub payload: NetworkMessage
}

/// A response from the peer-connected socket
pub enum SocketResponse {
    /// A message was received
    MessageReceived(NetworkMessage),
    /// An error occured and the socket needs to close
    ConnectionFailed(util::Error, Sender<()>)
}

#[derive(Clone, PartialEq, Eq, Debug)]
/// A Network message payload. Proper documentation is available on the Bitcoin
/// wiki https://en.bitcoin.it/wiki/Protocol_specification
pub enum NetworkMessage {
    /// `version`
    Version(message_network::VersionMessage),
    /// `verack`
    Verack,
    /// `addr`
    Addr(Vec<(u32, Address)>),
    /// `inv`
    Inv(Vec<message_blockdata::Inventory>),
    /// `getdata`
    GetData(Vec<message_blockdata::Inventory>),
    /// `notfound`
    NotFound(Vec<message_blockdata::Inventory>),
    /// `getblocks`
    GetBlocks(message_blockdata::GetBlocksMessage),
    /// `getheaders`
    GetHeaders(message_blockdata::GetHeadersMessage),
    /// `mempool`
    MemPool,
    /// tx
    Tx(transaction::Transaction),
    /// `block`
    Block(block::Block),
    /// `headers`
    Headers(Vec<block::LoneBlockHeader>),
    // TODO: getaddr,
    // TODO: checkorder,
    // TODO: submitorder,
    // TODO: reply,
    /// `ping`
    Ping(u64),
    /// `pong`
    Pong(u64)
    // TODO: reject,
    // TODO: bloom filtering
    // TODO: alert
}

impl RawNetworkMessage {
    fn command(&self) -> String {
        match self.payload {
            NetworkMessage::Version(_) => "version",
            NetworkMessage::Verack     => "verack",
            NetworkMessage::Addr(_)    => "addr",
            NetworkMessage::Inv(_)     => "inv",
            NetworkMessage::GetData(_) => "getdata",
            NetworkMessage::NotFound(_) => "notfound",
            NetworkMessage::GetBlocks(_) => "getblocks",
            NetworkMessage::GetHeaders(_) => "getheaders",
            NetworkMessage::MemPool    => "mempool",
            NetworkMessage::Tx(_)      => "tx",
            NetworkMessage::Block(_)   => "block",
            NetworkMessage::Headers(_) => "headers",
            NetworkMessage::Ping(_)    => "ping",
            NetworkMessage::Pong(_)    => "pong",
        }.to_owned()
    }
}

impl<S: SimpleEncoder> ConsensusEncodable<S> for RawNetworkMessage {
    fn consensus_encode(&self, s: &mut S) -> Result<(), S::Error> {
        try!(self.magic.consensus_encode(s));
        try!(CommandString(self.command()).consensus_encode(s));
        try!(CheckedData(match self.payload {
            NetworkMessage::Version(ref dat) => serialize(dat),
            NetworkMessage::Verack           => Ok(vec![]),
            NetworkMessage::Addr(ref dat)    => serialize(dat),
            NetworkMessage::Inv(ref dat)     => serialize(dat),
            NetworkMessage::GetData(ref dat) => serialize(dat),
            NetworkMessage::NotFound(ref dat) => serialize(dat),
            NetworkMessage::GetBlocks(ref dat) => serialize(dat),
            NetworkMessage::GetHeaders(ref dat) => serialize(dat),
            NetworkMessage::MemPool          => Ok(vec![]),
            NetworkMessage::Tx(ref dat)      => serialize(dat),
            NetworkMessage::Block(ref dat)   => serialize(dat),
            NetworkMessage::Headers(ref dat) => serialize(dat),
            NetworkMessage::Ping(ref dat)    => serialize(dat),
            NetworkMessage::Pong(ref dat)    => serialize(dat),
        }.unwrap()).consensus_encode(s));
        Ok(())
    }
}

// TODO: restriction on D::Error is so that `propagate_err` will work;
//       is there a more generic way to handle this?
impl<D: SimpleDecoder<Error=util::Error>> ConsensusDecodable<D> for RawNetworkMessage {
    fn consensus_decode(d: &mut D) -> Result<RawNetworkMessage, D::Error> {
        let magic = try!(ConsensusDecodable::consensus_decode(d));
        let CommandString(cmd): CommandString= try!(ConsensusDecodable::consensus_decode(d));
        let CheckedData(raw_payload): CheckedData = try!(ConsensusDecodable::consensus_decode(d));

        let mut mem_d = RawDecoder::new(Cursor::new(raw_payload));
        let payload = match &cmd[..] {
            "version" => NetworkMessage::Version(try!(propagate_err("version".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "verack"  => NetworkMessage::Verack,
            "addr"    => NetworkMessage::Addr(try!(propagate_err("addr".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "inv"     => NetworkMessage::Inv(try!(propagate_err("inv".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "getdata" => NetworkMessage::GetData(try!(propagate_err("getdata".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "notfound" => NetworkMessage::NotFound(try!(propagate_err("notfound".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "getblocks" => NetworkMessage::GetBlocks(try!(propagate_err("getblocks".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "getheaders" => NetworkMessage::GetHeaders(try!(propagate_err("getheaders".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "mempool" => NetworkMessage::MemPool,
            "block"   => NetworkMessage::Block(try!(propagate_err("block".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "headers" => NetworkMessage::Headers(try!(propagate_err("headers".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "ping"    => NetworkMessage::Ping(try!(propagate_err("ping".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "pong"    => NetworkMessage::Pong(try!(propagate_err("pong".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            "tx"      => NetworkMessage::Tx(try!(propagate_err("tx".to_owned(), ConsensusDecodable::consensus_decode(&mut mem_d)))),
            cmd => return Err(d.error(format!("unrecognized network command `{}`", cmd)))
        };
        Ok(RawNetworkMessage {
            magic: magic,
            payload: payload
        })
    }
}

#[cfg(test)]
mod test {
    use super::{RawNetworkMessage, NetworkMessage, CommandString};

    use network::serialize::{deserialize, serialize};

    #[test]
    fn serialize_commandstring_test() {
        let cs = CommandString("Andrew".to_owned());
        assert_eq!(serialize(&cs).ok(), Some(vec![0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0, 0]));
    }

    #[test]
    fn deserialize_commandstring_test() {
        let cs: Result<CommandString, _> = deserialize(&[0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0, 0]);
        assert!(cs.is_ok());
        assert_eq!(cs.unwrap(), CommandString("Andrew".to_owned()));

        let short_cs: Result<CommandString, _> = deserialize(&[0x41u8, 0x6e, 0x64, 0x72, 0x65, 0x77, 0, 0, 0, 0, 0]);
        assert!(short_cs.is_err());
    }

    #[test]
    fn serialize_verack_test() {
        assert_eq!(serialize(&RawNetworkMessage { magic: 0xd9b4bef9, payload: NetworkMessage::Verack }).ok(),
                             Some(vec![0xf9, 0xbe, 0xb4, 0xd9, 0x76, 0x65, 0x72, 0x61,
                                       0x63, 0x6B, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                       0x00, 0x00, 0x00, 0x00, 0x5d, 0xf6, 0xe0, 0xe2]));
    }

    #[test]
    fn serialize_ping_test() {
        assert_eq!(serialize(&RawNetworkMessage { magic: 0xd9b4bef9, payload: NetworkMessage::Ping(100) }).ok(),
                             Some(vec![0xf9, 0xbe, 0xb4, 0xd9, 0x70, 0x69, 0x6e, 0x67,
                                       0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                                       0x08, 0x00, 0x00, 0x00, 0x24, 0x67, 0xf1, 0x1d,
                                       0x64, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]));
    }

    #[test]
    fn serialize_mempool_test() {
        assert_eq!(serialize(&RawNetworkMessage { magic: 0xd9b4bef9, payload: NetworkMessage::MemPool }).ok(),
                             Some(vec![0xf9, 0xbe, 0xb4, 0xd9, 0x6d, 0x65, 0x6d, 0x70,
                                       0x6f, 0x6f, 0x6c, 0x00, 0x00, 0x00, 0x00, 0x00,
                                       0x00, 0x00, 0x00, 0x00, 0x5d, 0xf6, 0xe0, 0xe2]));
    }

}
<|endoftext|>


<commit_before>extern crate open189;

fn main() {
    let args: Vec<_> = std::env::args().collect();
    if args.len() < 6 {
        println!("usage: {} <app id> <secret> <access token> <phone> <code> <expire time>", args[0]);
        std::process::exit(1);
    }

    let ref app_id = args[1];
    let ref secret = args[2];
    let ref access_token = args[3];
    let ref phone = args[4];
    let ref code = args[5];
    let expire_time: Option<usize> = if args.len() < 7 {
        None
    } else {
        Some(args[6].parse().unwrap())
    };

    let app = open189::Open189App::new(app_id, secret);
    let sms_token = app.sms_get_token(access_token);
    println!("sms token = {:?}", sms_token);
    let sms_token = sms_token.unwrap();

    let result = app.sms_send_verification_code(access_token, &sms_token, phone, code, expire_time);
    println!("send result = {:?}", result);
}
<commit_msg>style: rustfmt<commit_after>extern crate open189;

fn main() {
    let args: Vec<_> = std::env::args().collect();
    if args.len() < 6 {
        println!("usage: {} <app id> <secret> <access token> <phone> <code> <expire time>",
                 args[0]);
        std::process::exit(1);
    }

    let ref app_id = args[1];
    let ref secret = args[2];
    let ref access_token = args[3];
    let ref phone = args[4];
    let ref code = args[5];
    let expire_time: Option<usize> = if args.len() < 7 {
        None
    } else {
        Some(args[6].parse().unwrap())
    };

    let app = open189::Open189App::new(app_id, secret);
    let sms_token = app.sms_get_token(access_token);
    println!("sms token = {:?}", sms_token);
    let sms_token = sms_token.unwrap();

    let result = app.sms_send_verification_code(access_token, &sms_token, phone, code, expire_time);
    println!("send result = {:?}", result);
}
<|endoftext|>


<commit_before>use {io, EventSet, PollOpt, Token};
use event::{self, Event};
use nix::unistd::close;
use nix::sys::event::{EventFilter, EventFlag, FilterFlag, KEvent, kqueue, kevent, kevent_ts};
use nix::sys::event::{EV_ADD, EV_CLEAR, EV_DELETE, EV_DISABLE, EV_ENABLE, EV_EOF, EV_ERROR, EV_ONESHOT};
use libc::{timespec, time_t};
use std::{cmp, fmt, slice};
use std::cell::UnsafeCell;
use std::os::unix::io::RawFd;
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
use std::time::Duration;

/// Each Selector has a globally unique(ish) ID associated with it. This ID
/// gets tracked by `TcpStream`, `TcpListener`, etc... when they are first
/// registered with the `Selector`. If a type that is previously associatd with
/// a `Selector` attempts to register itself with a different `Selector`, the
/// operation will return with an error. This matches windows behavior.
static NEXT_ID: AtomicUsize = ATOMIC_USIZE_INIT;

pub struct Selector {
    id: usize,
    kq: RawFd,
    changes: UnsafeCell<Vec<KEvent>>,
}

#[cfg(not(target_os = "netbsd"))]
type UData = usize;

#[cfg(target_os = "netbsd")]
type UData = i64;

impl Selector {
    pub fn new() -> io::Result<Selector> {
        // offset by 1 to avoid choosing 0 as the id of a selector
        let id = NEXT_ID.fetch_add(1, Ordering::Relaxed) + 1;
        let kq = try!(kqueue().map_err(super::from_nix_error));

        Ok(Selector {
            id: id,
            kq: kq,
            changes: UnsafeCell::new(vec![]),
        })
    }

    pub fn id(&self) -> usize {
        self.id
    }

    pub fn select(&self, evts: &mut Events, awakener: Token, timeout: Option<Duration>) -> io::Result<bool> {
        let timeout = timeout.map(|to| timespec {
            tv_sec: cmp::min(to.as_secs(), time_t::max_value() as u64) as time_t,
            tv_nsec: to.subsec_nanos() as i64,
        });

        let cnt = try!(kevent_ts(self.kq, &[], evts.as_mut_slice(), timeout)
                                  .map_err(super::from_nix_error));

        self.mut_changes().clear();

        unsafe {
            evts.sys_events.set_len(cnt);
        }

        Ok(evts.coalesce(awakener))
    }

    pub fn register(&self, fd: RawFd, token: Token, interests: EventSet, opts: PollOpt) -> io::Result<()> {
        trace!("registering; token={:?}; interests={:?}", token, interests);

        self.ev_register(fd, usize::from(token), EventFilter::EVFILT_READ, interests.contains(EventSet::readable()), opts);
        self.ev_register(fd, usize::from(token), EventFilter::EVFILT_WRITE, interests.contains(EventSet::writable()), opts);

        self.flush_changes()
    }

    pub fn reregister(&self, fd: RawFd, token: Token, interests: EventSet, opts: PollOpt) -> io::Result<()> {
        // Just need to call register here since EV_ADD is a mod if already
        // registered
        self.register(fd, token, interests, opts)
    }

    pub fn deregister(&self, fd: RawFd) -> io::Result<()> {
        self.ev_push(fd, 0, EventFilter::EVFILT_READ, EV_DELETE);
        self.ev_push(fd, 0, EventFilter::EVFILT_WRITE, EV_DELETE);

        self.flush_changes()
    }

    fn ev_register(&self, fd: RawFd, token: usize, filter: EventFilter, enable: bool, opts: PollOpt) {
        let mut flags = EV_ADD;

        if enable {
            flags = flags | EV_ENABLE;
        } else {
            flags = flags | EV_DISABLE;
        }

        if opts.contains(PollOpt::edge()) {
            flags = flags | EV_CLEAR;
        }

        if opts.contains(PollOpt::oneshot()) {
            flags = flags | EV_ONESHOT;
        }

        self.ev_push(fd, token, filter, flags);
    }

    fn ev_push(&self, fd: RawFd, token: usize, filter: EventFilter, flags: EventFlag) {
        self.mut_changes().push(
            KEvent {
                ident: fd as ::libc::uintptr_t,
                filter: filter,
                flags: flags,
                fflags: FilterFlag::empty(),
                data: 0,
                udata: token as UData,
            });
    }

    fn flush_changes(&self) -> io::Result<()> {
        let result = kevent(self.kq, self.changes(), &mut [], 0).map(|_| ())
            .map_err(super::from_nix_error).map(|_| ());

        self.mut_changes().clear();
        result
    }

    fn changes(&self) -> &[KEvent] {
        unsafe { &(*self.changes.get())[..] }
    }

    fn mut_changes(&self) -> &mut Vec<KEvent> {
        unsafe { &mut *self.changes.get() }
    }
}

impl fmt::Debug for Selector {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Selector")
            .field("id", &self.id)
            .field("kq", &self.kq)
            .field("changes", &self.changes().len())
            .finish()
    }
}

impl Drop for Selector {
    fn drop(&mut self) {
        let _ = close(self.kq);
    }
}

pub struct Events {
    sys_events: Vec<KEvent>,
    events: Vec<Event>,
    event_map: HashMap<Token, usize>,
}

impl Events {
    pub fn new() -> Events {
        Events {
            sys_events: Vec::with_capacity(1024),
            events: Vec::with_capacity(1024),
            event_map: HashMap::with_capacity(1024)
        }
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.events.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.events.is_empty()
    }

    pub fn get(&self, idx: usize) -> Option<Event> {
        self.events.get(idx).map(|e| *e)
    }

    fn coalesce(&mut self, awakener: Token) -> bool {
        let mut ret = false;
        self.events.clear();
        self.event_map.clear();

        for e in &self.sys_events {
            let token = Token(e.udata as usize);
            let len = self.events.len();

            if token == awakener {
                // TODO: Should this return an error if event is an error. It
                // is not critical as spurious wakeups are permitted.
                ret = true;
                continue;
            }

            let idx = *self.event_map.entry(token)
                .or_insert(len);

            if idx == len {
                // New entry, insert the default
                self.events.push(Event::new(EventSet::none(), token));

            }

            if e.flags.contains(EV_ERROR) {
                event::kind_mut(&mut self.events[idx]).insert(EventSet::error());
            }

            if e.filter == EventFilter::EVFILT_READ {
                event::kind_mut(&mut self.events[idx]).insert(EventSet::readable());
            } else if e.filter == EventFilter::EVFILT_WRITE {
                event::kind_mut(&mut self.events[idx]).insert(EventSet::writable());
            }

            if e.flags.contains(EV_EOF) {
                event::kind_mut(&mut self.events[idx]).insert(EventSet::hup());

                // When the read end of the socket is closed, EV_EOF is set on
                // flags, and fflags contains the error if there is one.
                if !e.fflags.is_empty() {
                    event::kind_mut(&mut self.events[idx]).insert(EventSet::error());
                }
            }
        }

        ret
    }

    pub fn push_event(&mut self, event: Event) {
        self.events.push(event);
    }

    fn as_mut_slice(&mut self) -> &mut [KEvent] {
        unsafe {
            let ptr = (&mut self.sys_events[..]).as_mut_ptr();
            slice::from_raw_parts_mut(ptr, self.sys_events.capacity())
        }
    }
}

impl fmt::Debug for Events {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Events {{ len: {} }}", self.sys_events.len())
    }
}
<commit_msg>Fix NetBSD build<commit_after>use {io, EventSet, PollOpt, Token};
use event::{self, Event};
use nix::unistd::close;
use nix::sys::event::{EventFilter, EventFlag, FilterFlag, KEvent, kqueue, kevent, kevent_ts};
use nix::sys::event::{EV_ADD, EV_CLEAR, EV_DELETE, EV_DISABLE, EV_ENABLE, EV_EOF, EV_ERROR, EV_ONESHOT};
use libc::{timespec, time_t};
use std::{cmp, fmt, slice};
use std::cell::UnsafeCell;
use std::os::unix::io::RawFd;
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
use std::time::Duration;

/// Each Selector has a globally unique(ish) ID associated with it. This ID
/// gets tracked by `TcpStream`, `TcpListener`, etc... when they are first
/// registered with the `Selector`. If a type that is previously associatd with
/// a `Selector` attempts to register itself with a different `Selector`, the
/// operation will return with an error. This matches windows behavior.
static NEXT_ID: AtomicUsize = ATOMIC_USIZE_INIT;

pub struct Selector {
    id: usize,
    kq: RawFd,
    changes: UnsafeCell<Vec<KEvent>>,
}

#[cfg(not(target_os = "netbsd"))]
type UData = usize;

#[cfg(target_os = "netbsd")]
type UData = isize;

impl Selector {
    pub fn new() -> io::Result<Selector> {
        // offset by 1 to avoid choosing 0 as the id of a selector
        let id = NEXT_ID.fetch_add(1, Ordering::Relaxed) + 1;
        let kq = try!(kqueue().map_err(super::from_nix_error));

        Ok(Selector {
            id: id,
            kq: kq,
            changes: UnsafeCell::new(vec![]),
        })
    }

    pub fn id(&self) -> usize {
        self.id
    }

    pub fn select(&self, evts: &mut Events, awakener: Token, timeout: Option<Duration>) -> io::Result<bool> {
        let timeout = timeout.map(|to| timespec {
            tv_sec: cmp::min(to.as_secs(), time_t::max_value() as u64) as time_t,
            tv_nsec: to.subsec_nanos() as i64,
        });

        let cnt = try!(kevent_ts(self.kq, &[], evts.as_mut_slice(), timeout)
                                  .map_err(super::from_nix_error));

        self.mut_changes().clear();

        unsafe {
            evts.sys_events.set_len(cnt);
        }

        Ok(evts.coalesce(awakener))
    }

    pub fn register(&self, fd: RawFd, token: Token, interests: EventSet, opts: PollOpt) -> io::Result<()> {
        trace!("registering; token={:?}; interests={:?}", token, interests);

        self.ev_register(fd, usize::from(token), EventFilter::EVFILT_READ, interests.contains(EventSet::readable()), opts);
        self.ev_register(fd, usize::from(token), EventFilter::EVFILT_WRITE, interests.contains(EventSet::writable()), opts);

        self.flush_changes()
    }

    pub fn reregister(&self, fd: RawFd, token: Token, interests: EventSet, opts: PollOpt) -> io::Result<()> {
        // Just need to call register here since EV_ADD is a mod if already
        // registered
        self.register(fd, token, interests, opts)
    }

    pub fn deregister(&self, fd: RawFd) -> io::Result<()> {
        self.ev_push(fd, 0, EventFilter::EVFILT_READ, EV_DELETE);
        self.ev_push(fd, 0, EventFilter::EVFILT_WRITE, EV_DELETE);

        self.flush_changes()
    }

    fn ev_register(&self, fd: RawFd, token: usize, filter: EventFilter, enable: bool, opts: PollOpt) {
        let mut flags = EV_ADD;

        if enable {
            flags = flags | EV_ENABLE;
        } else {
            flags = flags | EV_DISABLE;
        }

        if opts.contains(PollOpt::edge()) {
            flags = flags | EV_CLEAR;
        }

        if opts.contains(PollOpt::oneshot()) {
            flags = flags | EV_ONESHOT;
        }

        self.ev_push(fd, token, filter, flags);
    }

    fn ev_push(&self, fd: RawFd, token: usize, filter: EventFilter, flags: EventFlag) {
        self.mut_changes().push(
            KEvent {
                ident: fd as ::libc::uintptr_t,
                filter: filter,
                flags: flags,
                fflags: FilterFlag::empty(),
                data: 0,
                udata: token as UData,
            });
    }

    fn flush_changes(&self) -> io::Result<()> {
        let result = kevent(self.kq, self.changes(), &mut [], 0).map(|_| ())
            .map_err(super::from_nix_error).map(|_| ());

        self.mut_changes().clear();
        result
    }

    fn changes(&self) -> &[KEvent] {
        unsafe { &(*self.changes.get())[..] }
    }

    fn mut_changes(&self) -> &mut Vec<KEvent> {
        unsafe { &mut *self.changes.get() }
    }
}

impl fmt::Debug for Selector {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Selector")
            .field("id", &self.id)
            .field("kq", &self.kq)
            .field("changes", &self.changes().len())
            .finish()
    }
}

impl Drop for Selector {
    fn drop(&mut self) {
        let _ = close(self.kq);
    }
}

pub struct Events {
    sys_events: Vec<KEvent>,
    events: Vec<Event>,
    event_map: HashMap<Token, usize>,
}

impl Events {
    pub fn new() -> Events {
        Events {
            sys_events: Vec::with_capacity(1024),
            events: Vec::with_capacity(1024),
            event_map: HashMap::with_capacity(1024)
        }
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.events.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.events.is_empty()
    }

    pub fn get(&self, idx: usize) -> Option<Event> {
        self.events.get(idx).map(|e| *e)
    }

    fn coalesce(&mut self, awakener: Token) -> bool {
        let mut ret = false;
        self.events.clear();
        self.event_map.clear();

        for e in &self.sys_events {
            let token = Token(e.udata as usize);
            let len = self.events.len();

            if token == awakener {
                // TODO: Should this return an error if event is an error. It
                // is not critical as spurious wakeups are permitted.
                ret = true;
                continue;
            }

            let idx = *self.event_map.entry(token)
                .or_insert(len);

            if idx == len {
                // New entry, insert the default
                self.events.push(Event::new(EventSet::none(), token));

            }

            if e.flags.contains(EV_ERROR) {
                event::kind_mut(&mut self.events[idx]).insert(EventSet::error());
            }

            if e.filter == EventFilter::EVFILT_READ {
                event::kind_mut(&mut self.events[idx]).insert(EventSet::readable());
            } else if e.filter == EventFilter::EVFILT_WRITE {
                event::kind_mut(&mut self.events[idx]).insert(EventSet::writable());
            }

            if e.flags.contains(EV_EOF) {
                event::kind_mut(&mut self.events[idx]).insert(EventSet::hup());

                // When the read end of the socket is closed, EV_EOF is set on
                // flags, and fflags contains the error if there is one.
                if !e.fflags.is_empty() {
                    event::kind_mut(&mut self.events[idx]).insert(EventSet::error());
                }
            }
        }

        ret
    }

    pub fn push_event(&mut self, event: Event) {
        self.events.push(event);
    }

    fn as_mut_slice(&mut self) -> &mut [KEvent] {
        unsafe {
            let ptr = (&mut self.sys_events[..]).as_mut_ptr();
            slice::from_raw_parts_mut(ptr, self.sys_events.capacity())
        }
    }
}

impl fmt::Debug for Events {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "Events {{ len: {} }}", self.sys_events.len())
    }
}
<|endoftext|>


<commit_before>use regex::Regex;
use std::path::PathBuf;
use std::result::Result;
use subprocess::{Exec, Redirection};

const LOOP_BACKING_DEVICE_BASENAME: &'static str = "flaglock6-home";
const LOOP_DEVICE_PARTITION_LABEL: &'static str = "fl6";
const RSYNC_DOT_NET_REMOTE_DIR: &'static str = "usw-s007.rsync.net:./";

#[derive(Debug)]
pub enum AuricError {
    Invocation(String),
    Io(String),
    Subprocess(String),
    NotFound,
}

enum AuricOperationMode {
    Mount,
    Unmount,
}

impl From<subprocess::PopenError> for AuricError {
    fn from(err: subprocess::PopenError) -> AuricError {
        AuricError::Subprocess(err.to_string())
    }
}

struct SshfsManager {
    mountpoint: PathBuf,
}

struct LoopManager {
    target: Regex,
    backing: PathBuf,
}

struct LuksManager {
    device: PathBuf,
}

struct AuricImpl {
    sshfs_manager: SshfsManager,
    loop_manager: LoopManager,
    luks_manager: LuksManager,
}

fn exec_with_stdout(command: &str, args: &[&str], error_hint: &str) -> Result<String, AuricError> {
    let result = Exec::cmd(command)
        .args(args)
        .stdout(Redirection::Pipe)
        .stderr(Redirection::Pipe)
        .capture()?;
    if !result.success() {
        return Err(AuricError::Subprocess(format!(
            "{} failed: ``{}''",
            error_hint,
            result.stderr_str(),
        )));
    }
    Ok(result.stdout_str())
}

fn exec(command: &str, args: &[&str], error_hint: &str) -> Result<(), AuricError> {
    exec_with_stdout(command, args, error_hint).map(|_stdout_str| ())
}

impl SshfsManager {
    pub fn new() -> Result<SshfsManager, AuricError> {
        let xdg_runtime_dir = match std::env::var("XDG_RUNTIME_DIR") {
            Ok(val) => val,
            Err(e) => {
                return Err(AuricError::Invocation(format!(
                    "fatal error getting ${{XDG_RUNTIME_DIR}}: ``{}''",
                    e
                )))
            }
        };
        let mut mountpoint = PathBuf::new();
        mountpoint.push(xdg_runtime_dir);
        mountpoint.push("rsync.net-mountpoint");

        Ok(SshfsManager {
            mountpoint: mountpoint,
        })
    }

    fn ensure_mountpoint_is_dir(&self) -> Result<(), AuricError> {
        if self.mountpoint.is_dir() {
            return Ok(());
        }
        std::fs::create_dir(self.mountpoint.as_path())
            .map_err(|e| AuricError::Io(format!("bad mountpoint: ``{}''", e)))
    }

    // Does _not_ ensure that backing device exists.
    pub fn loop_backing_device_path(&self) -> PathBuf {
        let mut result: PathBuf = self.mountpoint.clone();
        result.push(LOOP_BACKING_DEVICE_BASENAME);
        result
    }

    fn loop_backing_device_present(&self) -> bool {
        self.loop_backing_device_path().is_file()
    }

    pub fn mount(&self) -> Result<(), AuricError> {
        self.ensure_mountpoint_is_dir()?;
        if self.loop_backing_device_present() {
            return Ok(());
        }

        // Invoke `sshfs`
        // *    with compression,
        // *    allowing other users to access the mounted filesystem
        //      (necessary to allow loop-mounts), and
        // *    attempting to reconnect if we lose the connection.
        exec(
            "sshfs",
            &[
                "-C",
                "-o",
                "allow_other,reconnect",
                RSYNC_DOT_NET_REMOTE_DIR,
                self.mountpoint.to_str().unwrap(),
            ],
            "sshfs mount",
        )
    }

    pub fn unmount(&self) -> Result<(), AuricError> {
        if !self.loop_backing_device_present() {
            return Ok(());
        }
        exec(
            "fusermount",
            &["-u", self.mountpoint.to_str().unwrap()],
            "fusermount -u",
        )
    }
}

impl LoopManager {
    pub fn new(backing: &PathBuf) -> LoopManager {
        let target = format!(
            r"^(/dev/loop\d)(\s+\d+){{4}}\s+({})\s+",
            backing.to_str().unwrap()
        );
        LoopManager {
            target: Regex::new(&target).unwrap(),
            backing: backing.clone(),
        }
    }

    pub fn find(&self) -> Result<PathBuf, AuricError> {
        let losetup_stdout = exec_with_stdout("losetup", &["-l"], "loop device enumeration")?;
        for line in losetup_stdout.lines() {
            if let Some(captures) = self.target.captures(line) {
                return Ok(PathBuf::from(captures.get(1).unwrap().as_str()));
            }
        }
        Err(AuricError::NotFound)
    }

    pub fn attach(&self) -> Result<(), AuricError> {
        if let Ok(_attached_loop) = self.find() {
            return Ok(());
        }

        exec(
            "udisksctl",
            &["loop-setup", "-f", self.backing.to_str().unwrap()],
            "loop device attach",
        )
    }

    // Note that the loop device will linger if the volume remains
    // unlocked. Prefer not to call this before locking.
    pub fn detach(&self) -> Result<(), AuricError> {
        let attached_loop = match self.find() {
            Ok(path) => path,
            Err(e) => {
                if let AuricError::NotFound = e {
                    return Ok(());
                }
                return Err(e);
            }
        };
        exec(
            "udisksctl",
            &["loop-delete", "-b", attached_loop.to_str().unwrap()],
            "loop device detach",
        )
    }
}

impl LuksManager {
    pub fn new() -> LuksManager {
        let mut device = PathBuf::from("/dev/disk/by-label");
        device.push(LOOP_DEVICE_PARTITION_LABEL);
        LuksManager { device: device }
    }

    pub fn unlock(&self, locked_device: &PathBuf) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["unlock", "-b", locked_device.to_str().unwrap()],
            format!("unlocking {}", locked_device.to_str().unwrap()).as_str(),
        )
    }

    pub fn lock(&self, underlying_device: &PathBuf) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["lock", "-b", underlying_device.to_str().unwrap()],
            format!("locking {}", underlying_device.to_str().unwrap()).as_str(),
        )
    }

    pub fn mount(&self) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["mount", "-b", self.device.to_str().unwrap()],
            format!("mounting {}", self.device.to_str().unwrap()).as_str(),
        )
    }

    pub fn unmount(&self) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["unmount", "-b", self.device.to_str().unwrap()],
            format!("unmounting {}", self.device.to_str().unwrap()).as_str(),
        )
    }
}

fn get_action() -> Result<AuricOperationMode, AuricError> {
    let mut args_iter = std::env::args();
    args_iter.next();
    if let Some(action) = args_iter.next() {
        match action.as_str() {
            "mount" => return Ok(AuricOperationMode::Mount),
            "unmount" => return Ok(AuricOperationMode::Unmount),
            _ => {
                return Err(AuricError::Invocation(format!(
                    "unknown action: ``{}''",
                    action
                )))
            }
        }
    }
    Err(AuricError::Invocation("no action given".to_string()))
}

impl AuricImpl {
    pub fn new() -> Result<AuricImpl, AuricError> {
        let sshfs_manager = SshfsManager::new()?;
        let loop_backing_device_path = sshfs_manager.loop_backing_device_path();
        Ok(AuricImpl {
            sshfs_manager: sshfs_manager,
            loop_manager: LoopManager::new(&loop_backing_device_path),
            luks_manager: LuksManager::new(),
        })
    }

    fn mount(&self) -> Result<(), AuricError> {
        println!("{}", "Mounting sshfs...");
        self.sshfs_manager.mount()?;
        println!("{}", "Attaching loop device...");
        self.loop_manager.attach()?;
        println!("{}", "Finding newly attached loop device...");
        let loop_device = self.loop_manager.find()?;
        println!("{}", "Unlocking volume...");
        self.luks_manager.unlock(&loop_device)?;
        println!("{}", "Mounting volume...");
        self.luks_manager.mount()
    }

    fn unmount(&self) -> Result<(), AuricError> {
        println!("{}", "Unmounting volume...");
        self.luks_manager.unmount()?;
        println!("{}", "Finding backing loop device...");
        let loop_device = self.loop_manager.find()?;
        println!("{}", "Locking volume...");
        self.luks_manager.lock(&loop_device)?;
        println!("{}", "Detaching loop device...");
        self.loop_manager.detach()?;
        println!("{}", "Unmounting sshfs...");
        self.sshfs_manager.unmount()
    }

    pub fn act(&self, mode: AuricOperationMode) -> Result<(), AuricError> {
        match mode {
            AuricOperationMode::Mount => return self.mount(),
            AuricOperationMode::Unmount => return self.unmount(),
        }
    }
}

fn main_impl() -> Result<(), AuricError> {
    let action = get_action()?;
    let auric_impl = AuricImpl::new()?;
    auric_impl.act(action)
}

fn main() {
    std::process::exit(match main_impl() {
        Ok(_) => 0,
        Err(e) => {
            eprintln!("{:?}", e);
            1
        }
    });
}
<commit_msg>apricot-auric: Implement `rsync` command<commit_after>use regex::Regex;
use std::path::PathBuf;
use std::result::Result;
use subprocess::{Exec, ExitStatus, Redirection};

const LOOP_BACKING_DEVICE_BASENAME: &'static str = "flaglock6-home";
const LOOP_DEVICE_PARTITION_LABEL: &'static str = "fl6";
const RSYNC_DOT_NET_REMOTE_DIR: &'static str = "usw-s007.rsync.net:./";

#[derive(Debug)]
pub enum AuricError {
    Invocation(String),
    Io(String),
    Subprocess(String),
    NotFound,
}

enum AuricOperationMode {
    Mount,
    Unmount,
    Rsync,
}

impl From<subprocess::PopenError> for AuricError {
    fn from(err: subprocess::PopenError) -> AuricError {
        AuricError::Subprocess(err.to_string())
    }
}

struct SshfsManager {
    mountpoint: PathBuf,
}

struct LoopManager {
    target: Regex,
    backing: PathBuf,
}

struct LuksManager {
    device: PathBuf,
}

struct AuricImpl {
    sshfs_manager: SshfsManager,
    loop_manager: LoopManager,
    luks_manager: LuksManager,
}

fn exec_with_stdout(command: &str, args: &[&str], error_hint: &str) -> Result<String, AuricError> {
    let result = Exec::cmd(command)
        .args(args)
        .stdout(Redirection::Pipe)
        .stderr(Redirection::Pipe)
        .capture()?;
    if !result.success() {
        return Err(AuricError::Subprocess(format!(
            "{} failed: ``{}''",
            error_hint,
            result.stderr_str(),
        )));
    }
    Ok(result.stdout_str())
}

fn exec(command: &str, args: &[&str], error_hint: &str) -> Result<(), AuricError> {
    exec_with_stdout(command, args, error_hint).map(|_stdout_str| ())
}

impl SshfsManager {
    pub fn new() -> Result<SshfsManager, AuricError> {
        let xdg_runtime_dir = match std::env::var("XDG_RUNTIME_DIR") {
            Ok(val) => val,
            Err(e) => {
                return Err(AuricError::Invocation(format!(
                    "fatal error getting ${{XDG_RUNTIME_DIR}}: ``{}''",
                    e
                )))
            }
        };
        let mut mountpoint = PathBuf::new();
        mountpoint.push(xdg_runtime_dir);
        mountpoint.push("rsync.net-mountpoint");

        Ok(SshfsManager {
            mountpoint: mountpoint,
        })
    }

    fn ensure_mountpoint_is_dir(&self) -> Result<(), AuricError> {
        if self.mountpoint.is_dir() {
            return Ok(());
        }
        std::fs::create_dir(self.mountpoint.as_path())
            .map_err(|e| AuricError::Io(format!("bad mountpoint: ``{}''", e)))
    }

    // Does _not_ ensure that backing device exists.
    pub fn loop_backing_device_path(&self) -> PathBuf {
        let mut result: PathBuf = self.mountpoint.clone();
        result.push(LOOP_BACKING_DEVICE_BASENAME);
        result
    }

    fn loop_backing_device_present(&self) -> bool {
        self.loop_backing_device_path().is_file()
    }

    pub fn mount(&self) -> Result<(), AuricError> {
        self.ensure_mountpoint_is_dir()?;
        if self.loop_backing_device_present() {
            return Ok(());
        }

        // Invoke `sshfs`
        // *    with compression,
        // *    allowing other users to access the mounted filesystem
        //      (necessary to allow loop-mounts), and
        // *    attempting to reconnect if we lose the connection.
        exec(
            "sshfs",
            &[
                "-C",
                "-o",
                "allow_other,reconnect",
                RSYNC_DOT_NET_REMOTE_DIR,
                self.mountpoint.to_str().unwrap(),
            ],
            "sshfs mount",
        )
    }

    pub fn unmount(&self) -> Result<(), AuricError> {
        if !self.loop_backing_device_present() {
            return Ok(());
        }
        exec(
            "fusermount",
            &["-u", self.mountpoint.to_str().unwrap()],
            "fusermount -u",
        )
    }
}

impl LoopManager {
    pub fn new(backing: &PathBuf) -> LoopManager {
        let target = format!(
            r"^(/dev/loop\d)(\s+\d+){{4}}\s+({})\s+",
            backing.to_str().unwrap()
        );
        LoopManager {
            target: Regex::new(&target).unwrap(),
            backing: backing.clone(),
        }
    }

    pub fn find(&self) -> Result<PathBuf, AuricError> {
        let losetup_stdout = exec_with_stdout("losetup", &["-l"], "loop device enumeration")?;
        for line in losetup_stdout.lines() {
            if let Some(captures) = self.target.captures(line) {
                return Ok(PathBuf::from(captures.get(1).unwrap().as_str()));
            }
        }
        Err(AuricError::NotFound)
    }

    pub fn attach(&self) -> Result<(), AuricError> {
        if let Ok(_attached_loop) = self.find() {
            return Ok(());
        }

        exec(
            "udisksctl",
            &["loop-setup", "-f", self.backing.to_str().unwrap()],
            "loop device attach",
        )
    }

    // Note that the loop device will linger if the volume remains
    // unlocked. Prefer not to call this before locking.
    pub fn detach(&self) -> Result<(), AuricError> {
        let attached_loop = match self.find() {
            Ok(path) => path,
            Err(e) => {
                if let AuricError::NotFound = e {
                    return Ok(());
                }
                return Err(e);
            }
        };
        exec(
            "udisksctl",
            &["loop-delete", "-b", attached_loop.to_str().unwrap()],
            "loop device detach",
        )
    }
}

impl LuksManager {
    pub fn new() -> LuksManager {
        let mut device = PathBuf::from("/dev/disk/by-label");
        device.push(LOOP_DEVICE_PARTITION_LABEL);
        LuksManager { device: device }
    }

    pub fn unlock(&self, locked_device: &PathBuf) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["unlock", "-b", locked_device.to_str().unwrap()],
            format!("unlocking {}", locked_device.to_str().unwrap()).as_str(),
        )
    }

    pub fn lock(&self, underlying_device: &PathBuf) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["lock", "-b", underlying_device.to_str().unwrap()],
            format!("locking {}", underlying_device.to_str().unwrap()).as_str(),
        )
    }

    pub fn mount(&self) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["mount", "-b", self.device.to_str().unwrap()],
            format!("mounting {}", self.device.to_str().unwrap()).as_str(),
        )
    }

    pub fn unmount(&self) -> Result<(), AuricError> {
        exec(
            "udisksctl",
            &["unmount", "-b", self.device.to_str().unwrap()],
            format!("unmounting {}", self.device.to_str().unwrap()).as_str(),
        )
    }
}

fn get_action() -> Result<AuricOperationMode, AuricError> {
    let mut args_iter = std::env::args();
    args_iter.next();
    if let Some(action) = args_iter.next() {
        match action.as_str() {
            "mount" => return Ok(AuricOperationMode::Mount),
            "unmount" => return Ok(AuricOperationMode::Unmount),
            "rsync" => return Ok(AuricOperationMode::Rsync),
            _ => {
                return Err(AuricError::Invocation(format!(
                    "unknown action: ``{}''",
                    action
                )))
            }
        }
    }
    Err(AuricError::Invocation("no action given".to_string()))
}

impl AuricImpl {
    pub fn new() -> Result<AuricImpl, AuricError> {
        let sshfs_manager = SshfsManager::new()?;
        let loop_backing_device_path = sshfs_manager.loop_backing_device_path();
        Ok(AuricImpl {
            sshfs_manager: sshfs_manager,
            loop_manager: LoopManager::new(&loop_backing_device_path),
            luks_manager: LuksManager::new(),
        })
    }

    fn mount(&self) -> Result<(), AuricError> {
        println!("{}", "Mounting sshfs...");
        self.sshfs_manager.mount()?;
        println!("{}", "Attaching loop device...");
        self.loop_manager.attach()?;
        println!("{}", "Finding newly attached loop device...");
        let loop_device = self.loop_manager.find()?;
        println!("{}", "Unlocking volume...");
        self.luks_manager.unlock(&loop_device)?;
        println!("{}", "Mounting volume...");
        self.luks_manager.mount()
    }

    fn unmount(&self) -> Result<(), AuricError> {
        println!("{}", "Unmounting volume...");
        self.luks_manager.unmount()?;
        println!("{}", "Finding backing loop device...");
        let loop_device = self.loop_manager.find()?;
        println!("{}", "Locking volume...");
        self.luks_manager.lock(&loop_device)?;
        println!("{}", "Detaching loop device...");
        self.loop_manager.detach()?;
        println!("{}", "Unmounting sshfs...");
        self.sshfs_manager.unmount()
    }

    fn rsync(&self) -> Result<(), AuricError> {
        // I'm not pulling in the `home_dir` crate just to grab a value
        // that I suspect will be immutable for many years to come.
        // Likewise, the user mount directory is lazily hardcoded.
        let mut remote_directory = PathBuf::from(r"/run/media/kalvin/");
        remote_directory.push(LOOP_DEVICE_PARTITION_LABEL);
        if !remote_directory.is_dir() {
            return Err(AuricError::Invocation(format!(
                "remote directory ``{}'' not mounted",
                remote_directory.display()
            )));
        }

        let cmd = Exec::cmd("rsync").args(&[
            "--delete",
            "-avPS",
            "--exclude=/.cache/",
            "--exclude=/.cargo/",
            "--exclude=/.local/share/",
            "--exclude=/.mozilla/",
            "--exclude=/.thunderbird/",
            "--exclude=/Downloads/.firefox-nightly/",
            "--exclude=/Downloads/.thunderbird-beta/",
            "--delete-excluded",
            "/home/kalvin/",
            remote_directory.to_str().unwrap(),
        ]);
        let result = cmd.join()?;
        if !result.success() {
            match result {
                ExitStatus::Exited(code) => {
                    return Err(AuricError::Subprocess(format!(
                        "rsync exited with code {}",
                        code
                    )))
                }
                ExitStatus::Signaled(signum) => {
                    return Err(AuricError::Subprocess(format!(
                        "rsync exited with signal {}",
                        signum
                    )))
                }
                ExitStatus::Other(code) => {
                    return Err(AuricError::Subprocess(format!(
                        "unknown rsync error: {}",
                        code
                    )))
                }
                ExitStatus::Undetermined => {
                    return Err(AuricError::Subprocess("unknown rsync error!".to_string()))
                }
            }
        }
        Ok(())
    }

    pub fn act(&self, mode: AuricOperationMode) -> Result<(), AuricError> {
        match mode {
            AuricOperationMode::Mount => return self.mount(),
            AuricOperationMode::Unmount => return self.unmount(),
            AuricOperationMode::Rsync => return self.rsync(),
        }
    }
}

fn main_impl() -> Result<(), AuricError> {
    let action = get_action()?;
    let auric_impl = AuricImpl::new()?;
    auric_impl.act(action)
}

fn main() {
    std::process::exit(match main_impl() {
        Ok(_) => 0,
        Err(e) => {
            eprintln!("{:?}", e);
            1
        }
    });
}
<|endoftext|>