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

//! Untyped root container for a Cap'n Proto value.

use any_pointer;
use private::capability::ClientHook;
use private::units::*;
use private::arena::{BuilderArena, ReaderArena, SegmentBuilder, SegmentReader};
use private::layout;
use traits::{FromPointerReader, FromPointerBuilder, SetPointerBuilder};
use {OutputSegments, Result, Word};

/// Options controlling how data is read.
#[derive(Clone, Copy)]
pub struct ReaderOptions {

    /// Limits how many total words of data are allowed to be traversed. Traversal is counted when
    /// a new struct or list builder is obtained, e.g. from a get() accessor. This means that calling
    /// the getter for the same sub-struct multiple times will cause it to be double-counted. Once
    /// the traversal limit is reached, an error will be reported.
    ///
    /// This limit exists for security reasons. It is possible for an attacker to construct a message
    /// in which multiple pointers point at the same location. This is technically invalid, but hard
    /// to detect. Using such a message, an attacker could cause a message which is small on the wire
    /// to appear much larger when actually traversed, possibly exhausting server resources leading to
    /// denial-of-service.
    ///
    /// It makes sense to set a traversal limit that is much larger than the underlying message.
    /// Together with sensible coding practices (e.g. trying to avoid calling sub-object getters
    /// multiple times, which is expensive anyway), this should provide adequate protection without
    /// inconvenience.
    pub traversal_limit_in_words : u64,

    /// Limits how deeply nested a message structure can be, e.g. structs containing other structs or
    /// lists of structs.
    ///
    /// Like the traversal limit, this limit exists for security reasons. Since it is common to use
    /// recursive code to traverse recursive data structures, an attacker could easily cause a stack
    /// overflow by sending a very-depply-nested (or even cyclic) message, without the message even
    /// being very large. The default limit of 64 is probably low enough to prevent any chance of
    /// stack overflow, yet high enough that it is never a problem in practice.
    pub nesting_limit : i32,
}

pub const DEFAULT_READER_OPTIONS : ReaderOptions =
    ReaderOptions { traversal_limit_in_words : 8 * 1024 * 1024, nesting_limit : 64 };

impl ReaderOptions {
    pub fn new() -> ReaderOptions { DEFAULT_READER_OPTIONS }

    pub fn nesting_limit<'a>(&'a mut self, value : i32) -> &'a mut ReaderOptions {
        self.nesting_limit = value;
        return self;
    }

    pub fn traversal_limit_in_words<'a>(&'a mut self, value : u64) -> &'a mut ReaderOptions {
        self.traversal_limit_in_words = value;
        return self;
    }
}

type SegmentId = u32;

pub unsafe trait ReaderSegments {
    // UNSAFETY ALERT: The callee is responsible for ensuring that the returned memory is valid
    // for the lifetime of the object. This could be much longer than `'a`.
    fn get_segment<'a>(&'a self, id: u32) -> Option<&'a [Word]>;
}

pub struct SegmentArray<'a> {
    segments: &'a [&'a [Word]],
}

impl <'a> SegmentArray<'a> {
    pub fn new(segments: &'a [&'a [Word]]) -> SegmentArray<'a> {
        SegmentArray { segments: segments }
    }
}

unsafe impl <'b> ReaderSegments for SegmentArray<'b> {
    fn get_segment<'a>(&'a self, id: u32) -> Option<&'a [Word]> {
        if id < self.segments.len() as u32 {
            Some(self.segments[id as usize])
        } else {
            None
        }
    }
}

/// An abstract container used to read a message.
pub struct Reader<S> where S: ReaderSegments {
    arena: Box<ReaderArena>,
    segments: Box<S>,
    options: ReaderOptions,
}

unsafe impl <S> Send for Reader<S> where S: Send + ReaderSegments {}

impl <S> Reader<S> where S: ReaderSegments {
    pub fn new(segments: S, options: ReaderOptions) -> Reader<S> {
        let boxed_segments = Box::new(segments);
        let boxed_segments_ref = {
            let r :&S = &*boxed_segments;
            unsafe { ::std::mem::transmute(r as &ReaderSegments) }
        };
        let arena = ReaderArena::new(boxed_segments_ref, options);
        Reader { arena: arena, segments: boxed_segments, options: options }
    }

    fn get_root_internal(&self) -> Result<any_pointer::Reader> {
        unsafe {
            let segment : *const SegmentReader = &self.arena.segment0;

            let pointer_reader = try!(layout::PointerReader::get_root(
                segment, (*segment).get_start_ptr(), self.options.nesting_limit));

            Ok(any_pointer::Reader::new(pointer_reader))
        }
    }

    /// Gets the root of the message, interpreting it as the given type.
    pub fn get_root<'a, T : FromPointerReader<'a>>(&'a self) -> Result<T> {
        try!(self.get_root_internal()).get_as()
    }

    pub fn init_cap_table(&mut self, cap_table : Vec<Option<Box<ClientHook+Send>>>) {
        self.arena.init_cap_table(cap_table);
    }

    pub fn into_segments(self) -> S {
        *self.segments
    }
}

pub unsafe trait Allocator {
    // UNSAFETY ALERT: The callee is responsible for ensuring that the returned memory is valid
    // for the lifetime of the object.
    fn allocate_segment(&mut self, miniumum_size: u32) -> (*mut Word, u32);
    fn pre_drop(&mut self, _segment0_currently_allocated: u32) {}
}

/// A container used to build a message.
pub struct Builder<S> where S: Allocator {
    arena: Box<BuilderArena>,
    allocator: Box<S>,
}

unsafe impl <A> Send for Builder<A> where A: Send + Allocator {}

impl <S> Builder<S> where S: Allocator {
    pub fn new(allocator: S) -> Builder<S> {
        let mut boxed_allocator = Box::new(allocator);
        let boxed_allocator_ref = {
            let r: &mut S = &mut *boxed_allocator;
            unsafe { ::std::mem::transmute(r as &mut Allocator) }
        };
        let arena = BuilderArena::new(boxed_allocator_ref);
        Builder { arena: arena, allocator: boxed_allocator }
    }

    fn get_root_internal<'a>(&mut self) -> any_pointer::Builder<'a> {
        let root_segment: *mut SegmentBuilder = &mut self.arena.segment0;

        if self.arena.segment0.current_size() == 0 {
            match self.arena.segment0.allocate(WORDS_PER_POINTER as u32) {
                None => {panic!("could not allocate root pointer") }
                Some(location) => {
                    assert!(location == self.arena.segment0.get_ptr_unchecked(0),
                            "First allocated word of new segment was not at offset 0");

                    any_pointer::Builder::new(layout::PointerBuilder::get_root(root_segment, location))

                }
            }
        } else {
            any_pointer::Builder::new(
                layout::PointerBuilder::get_root(root_segment,
                                                 self.arena.segment0.get_ptr_unchecked(0)))
        }

    }

    /// Initializes the root as a value of the given type.
    pub fn init_root<'a, T : FromPointerBuilder<'a>>(&'a mut self) -> T {
        self.get_root_internal().init_as()
    }

    /// Gets the root, interpreting it as the given type.
    pub fn get_root<'a, T : FromPointerBuilder<'a>>(&'a mut self) -> Result<T> {
        self.get_root_internal().get_as()
    }

    /// Sets the root to a deep copy of the given value.
    pub fn set_root<To, From : SetPointerBuilder<To>>(&mut self, value : From) -> Result<()> {
        self.get_root_internal().set_as(value)
    }

    pub fn get_segments_for_output<'a>(&'a self) -> OutputSegments<'a> {
        self.arena.get_segments_for_output()
    }

    pub fn get_cap_table<'a>(&'a self) -> &'a [Option<Box<ClientHook+Send>>] {
        self.arena.get_cap_table()
    }
}

impl <S> Drop for Builder<S> where S: Allocator {
    fn drop(&mut self) {
        self.allocator.pre_drop(self.arena.segment0.current_size());
    }
}

pub struct HeapAllocator {
    owned_memory : Vec<Vec<Word>>,
    next_size: u32,
    allocation_strategy: AllocationStrategy,
}

#[derive(Clone, Copy)]
pub enum AllocationStrategy {
    FixedSize,
    GrowHeuristically
}

pub const SUGGESTED_FIRST_SEGMENT_WORDS : u32 = 1024;
pub const SUGGESTED_ALLOCATION_STRATEGY : AllocationStrategy = AllocationStrategy::GrowHeuristically;

impl HeapAllocator {
    pub fn new() -> HeapAllocator {
        HeapAllocator { owned_memory: Vec::new(),
                        next_size: SUGGESTED_FIRST_SEGMENT_WORDS,
                        allocation_strategy: SUGGESTED_ALLOCATION_STRATEGY }
    }

    pub fn first_segment_words(mut self, value: u32) -> HeapAllocator {
        self.next_size = value;
        return self;
    }

    pub fn allocation_strategy(mut self, value : AllocationStrategy) -> HeapAllocator {
        self.allocation_strategy = value;
        return self;
    }
}

unsafe impl Allocator for HeapAllocator {
    fn allocate_segment(&mut self, minimum_size: u32) -> (*mut Word, u32) {
        let size = ::std::cmp::max(minimum_size, self.next_size);
        let mut new_words = Word::allocate_zeroed_vec(size as usize);
        let ptr = new_words.as_mut_ptr();
        self.owned_memory.push(new_words);

        match self.allocation_strategy {
            AllocationStrategy::GrowHeuristically => { self.next_size += size; }
            _ => { }
        }
        (ptr, size as u32)
    }
}

impl Builder<HeapAllocator> {
    pub fn new_default() -> Builder<HeapAllocator> {
        Builder::new(HeapAllocator::new())
    }
}

pub struct ScratchSpace<'a> {
    slice: &'a mut [Word],
    in_use: bool,
}

impl <'a> ScratchSpace<'a> {
    pub fn new(slice: &'a mut [Word]) -> ScratchSpace<'a> {
        ScratchSpace { slice: slice, in_use: false }
    }
}

pub struct ScratchSpaceHeapAllocator<'a, 'b: 'a> {
    scratch_space: &'a mut ScratchSpace<'b>,
    allocator: HeapAllocator,
}

impl <'a, 'b: 'a> ScratchSpaceHeapAllocator<'a, 'b> {
    pub fn new(scratch_space: &'a mut ScratchSpace<'b>) -> ScratchSpaceHeapAllocator<'a, 'b> {
        ScratchSpaceHeapAllocator { scratch_space: scratch_space,
                                    allocator: HeapAllocator::new()}
    }

    pub fn second_segment_words(mut self, value: u32) -> ScratchSpaceHeapAllocator<'a, 'b> {
        ScratchSpaceHeapAllocator { scratch_space: self.scratch_space,
                                    allocator: self.allocator.first_segment_words(value) }

    }

    pub fn allocation_strategy(mut self, value: AllocationStrategy) -> ScratchSpaceHeapAllocator<'a, 'b> {
        ScratchSpaceHeapAllocator { scratch_space: self.scratch_space,
                                    allocator: self.allocator.allocation_strategy(value) }
    }

}

unsafe impl <'a, 'b: 'a> Allocator for ScratchSpaceHeapAllocator<'a, 'b> {
    fn allocate_segment(&mut self, minimum_size: u32) -> (*mut Word, u32) {
        if !self.scratch_space.in_use {
            self.scratch_space.in_use = true;
            (self.scratch_space.slice.as_mut_ptr(), self.scratch_space.slice.len() as u32)
        } else {
            self.allocator.allocate_segment(minimum_size)
        }
    }

    fn pre_drop(&mut self, segment0_currently_allocated: u32) {
        let ptr = self.scratch_space.slice.as_mut_ptr();
        unsafe {
            ::std::ptr::write_bytes(ptr, 0u8, segment0_currently_allocated as usize);
        }
        self.scratch_space.in_use = false;
    }
}
<commit_msg>Add some comments.<commit_after>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

//! Untyped root container for a Cap'n Proto value.

use any_pointer;
use private::capability::ClientHook;
use private::units::*;
use private::arena::{BuilderArena, ReaderArena, SegmentBuilder, SegmentReader};
use private::layout;
use traits::{FromPointerReader, FromPointerBuilder, SetPointerBuilder};
use {OutputSegments, Result, Word};

/// Options controlling how data is read.
#[derive(Clone, Copy)]
pub struct ReaderOptions {

    /// Limits how many total words of data are allowed to be traversed. Traversal is counted when
    /// a new struct or list builder is obtained, e.g. from a get() accessor. This means that calling
    /// the getter for the same sub-struct multiple times will cause it to be double-counted. Once
    /// the traversal limit is reached, an error will be reported.
    ///
    /// This limit exists for security reasons. It is possible for an attacker to construct a message
    /// in which multiple pointers point at the same location. This is technically invalid, but hard
    /// to detect. Using such a message, an attacker could cause a message which is small on the wire
    /// to appear much larger when actually traversed, possibly exhausting server resources leading to
    /// denial-of-service.
    ///
    /// It makes sense to set a traversal limit that is much larger than the underlying message.
    /// Together with sensible coding practices (e.g. trying to avoid calling sub-object getters
    /// multiple times, which is expensive anyway), this should provide adequate protection without
    /// inconvenience.
    pub traversal_limit_in_words : u64,

    /// Limits how deeply nested a message structure can be, e.g. structs containing other structs or
    /// lists of structs.
    ///
    /// Like the traversal limit, this limit exists for security reasons. Since it is common to use
    /// recursive code to traverse recursive data structures, an attacker could easily cause a stack
    /// overflow by sending a very-depply-nested (or even cyclic) message, without the message even
    /// being very large. The default limit of 64 is probably low enough to prevent any chance of
    /// stack overflow, yet high enough that it is never a problem in practice.
    pub nesting_limit : i32,
}

pub const DEFAULT_READER_OPTIONS : ReaderOptions =
    ReaderOptions { traversal_limit_in_words : 8 * 1024 * 1024, nesting_limit : 64 };

impl ReaderOptions {
    pub fn new() -> ReaderOptions { DEFAULT_READER_OPTIONS }

    pub fn nesting_limit<'a>(&'a mut self, value : i32) -> &'a mut ReaderOptions {
        self.nesting_limit = value;
        return self;
    }

    pub fn traversal_limit_in_words<'a>(&'a mut self, value : u64) -> &'a mut ReaderOptions {
        self.traversal_limit_in_words = value;
        return self;
    }
}

type SegmentId = u32;

pub unsafe trait ReaderSegments {
    // UNSAFETY ALERT: The callee is responsible for ensuring that the returned memory is valid
    // for the lifetime of the object. This could be much longer than `'a`.
    fn get_segment<'a>(&'a self, id: u32) -> Option<&'a [Word]>;
}

pub struct SegmentArray<'a> {
    segments: &'a [&'a [Word]],
}

impl <'a> SegmentArray<'a> {
    pub fn new(segments: &'a [&'a [Word]]) -> SegmentArray<'a> {
        SegmentArray { segments: segments }
    }
}

unsafe impl <'b> ReaderSegments for SegmentArray<'b> {
    fn get_segment<'a>(&'a self, id: u32) -> Option<&'a [Word]> {
        if id < self.segments.len() as u32 {
            Some(self.segments[id as usize])
        } else {
            None
        }
    }
}

/// An abstract container used to read a message.
///
/// The underlying implemention uses the `ReaderSegments` as a trait object. However, we
/// need to include `S` as concrete type parameter so that the typechecker can
/// correctly deduce appropriate bounds like `Send`.
pub struct Reader<S> where S: ReaderSegments {
    arena: Box<ReaderArena>,
    segments: Box<S>,
    options: ReaderOptions,
}

unsafe impl <S> Send for Reader<S> where S: Send + ReaderSegments {}

impl <S> Reader<S> where S: ReaderSegments {
    pub fn new(segments: S, options: ReaderOptions) -> Reader<S> {
        let boxed_segments = Box::new(segments);
        let boxed_segments_ref = {
            let r :&S = &*boxed_segments;
            unsafe { ::std::mem::transmute(r as &ReaderSegments) }
        };
        let arena = ReaderArena::new(boxed_segments_ref, options);
        Reader { arena: arena, segments: boxed_segments, options: options }
    }

    fn get_root_internal(&self) -> Result<any_pointer::Reader> {
        unsafe {
            let segment : *const SegmentReader = &self.arena.segment0;

            let pointer_reader = try!(layout::PointerReader::get_root(
                segment, (*segment).get_start_ptr(), self.options.nesting_limit));

            Ok(any_pointer::Reader::new(pointer_reader))
        }
    }

    /// Gets the root of the message, interpreting it as the given type.
    pub fn get_root<'a, T : FromPointerReader<'a>>(&'a self) -> Result<T> {
        try!(self.get_root_internal()).get_as()
    }

    pub fn init_cap_table(&mut self, cap_table : Vec<Option<Box<ClientHook+Send>>>) {
        self.arena.init_cap_table(cap_table);
    }

    pub fn into_segments(self) -> S {
        *self.segments
    }
}

pub unsafe trait Allocator {
    // UNSAFETY ALERT: The callee is responsible for ensuring that the returned memory is valid
    // for the lifetime of the object.
    fn allocate_segment(&mut self, miniumum_size: u32) -> (*mut Word, u32);
    fn pre_drop(&mut self, _segment0_currently_allocated: u32) {}
}

/// A container used to build a message.
///
/// The underlying implemention uses the `Allocator` as a trait object. However, we
/// need to include `A` as concrete type parameter so that the typechecker can
/// correctly deduce appropriate bounds like `Send`.
pub struct Builder<A> where A: Allocator {
    arena: Box<BuilderArena>,
    allocator: Box<A>,
}

unsafe impl <A> Send for Builder<A> where A: Send + Allocator {}

impl <A> Builder<A> where A: Allocator {
    pub fn new(allocator: A) -> Builder<A> {
        let mut boxed_allocator = Box::new(allocator);
        let boxed_allocator_ref = {
            let r: &mut A = &mut *boxed_allocator;
            unsafe { ::std::mem::transmute(r as &mut Allocator) }
        };
        let arena = BuilderArena::new(boxed_allocator_ref);
        Builder { arena: arena, allocator: boxed_allocator }
    }

    fn get_root_internal<'a>(&mut self) -> any_pointer::Builder<'a> {
        let root_segment: *mut SegmentBuilder = &mut self.arena.segment0;

        if self.arena.segment0.current_size() == 0 {
            match self.arena.segment0.allocate(WORDS_PER_POINTER as u32) {
                None => {panic!("could not allocate root pointer") }
                Some(location) => {
                    assert!(location == self.arena.segment0.get_ptr_unchecked(0),
                            "First allocated word of new segment was not at offset 0");

                    any_pointer::Builder::new(layout::PointerBuilder::get_root(root_segment, location))

                }
            }
        } else {
            any_pointer::Builder::new(
                layout::PointerBuilder::get_root(root_segment,
                                                 self.arena.segment0.get_ptr_unchecked(0)))
        }

    }

    /// Initializes the root as a value of the given type.
    pub fn init_root<'a, T : FromPointerBuilder<'a>>(&'a mut self) -> T {
        self.get_root_internal().init_as()
    }

    /// Gets the root, interpreting it as the given type.
    pub fn get_root<'a, T : FromPointerBuilder<'a>>(&'a mut self) -> Result<T> {
        self.get_root_internal().get_as()
    }

    /// Sets the root to a deep copy of the given value.
    pub fn set_root<To, From : SetPointerBuilder<To>>(&mut self, value : From) -> Result<()> {
        self.get_root_internal().set_as(value)
    }

    pub fn get_segments_for_output<'a>(&'a self) -> OutputSegments<'a> {
        self.arena.get_segments_for_output()
    }

    pub fn get_cap_table<'a>(&'a self) -> &'a [Option<Box<ClientHook+Send>>] {
        self.arena.get_cap_table()
    }
}

impl <A> Drop for Builder<A> where A: Allocator {
    fn drop(&mut self) {
        self.allocator.pre_drop(self.arena.segment0.current_size());
    }
}

pub struct HeapAllocator {
    owned_memory : Vec<Vec<Word>>,
    next_size: u32,
    allocation_strategy: AllocationStrategy,
}

#[derive(Clone, Copy)]
pub enum AllocationStrategy {
    FixedSize,
    GrowHeuristically
}

pub const SUGGESTED_FIRST_SEGMENT_WORDS : u32 = 1024;
pub const SUGGESTED_ALLOCATION_STRATEGY : AllocationStrategy = AllocationStrategy::GrowHeuristically;

impl HeapAllocator {
    pub fn new() -> HeapAllocator {
        HeapAllocator { owned_memory: Vec::new(),
                        next_size: SUGGESTED_FIRST_SEGMENT_WORDS,
                        allocation_strategy: SUGGESTED_ALLOCATION_STRATEGY }
    }

    pub fn first_segment_words(mut self, value: u32) -> HeapAllocator {
        self.next_size = value;
        return self;
    }

    pub fn allocation_strategy(mut self, value : AllocationStrategy) -> HeapAllocator {
        self.allocation_strategy = value;
        return self;
    }
}

unsafe impl Allocator for HeapAllocator {
    fn allocate_segment(&mut self, minimum_size: u32) -> (*mut Word, u32) {
        let size = ::std::cmp::max(minimum_size, self.next_size);
        let mut new_words = Word::allocate_zeroed_vec(size as usize);
        let ptr = new_words.as_mut_ptr();
        self.owned_memory.push(new_words);

        match self.allocation_strategy {
            AllocationStrategy::GrowHeuristically => { self.next_size += size; }
            _ => { }
        }
        (ptr, size as u32)
    }
}

impl Builder<HeapAllocator> {
    pub fn new_default() -> Builder<HeapAllocator> {
        Builder::new(HeapAllocator::new())
    }
}

pub struct ScratchSpace<'a> {
    slice: &'a mut [Word],
    in_use: bool,
}

impl <'a> ScratchSpace<'a> {
    pub fn new(slice: &'a mut [Word]) -> ScratchSpace<'a> {
        ScratchSpace { slice: slice, in_use: false }
    }
}

pub struct ScratchSpaceHeapAllocator<'a, 'b: 'a> {
    scratch_space: &'a mut ScratchSpace<'b>,
    allocator: HeapAllocator,
}

impl <'a, 'b: 'a> ScratchSpaceHeapAllocator<'a, 'b> {
    pub fn new(scratch_space: &'a mut ScratchSpace<'b>) -> ScratchSpaceHeapAllocator<'a, 'b> {
        ScratchSpaceHeapAllocator { scratch_space: scratch_space,
                                    allocator: HeapAllocator::new()}
    }

    pub fn second_segment_words(mut self, value: u32) -> ScratchSpaceHeapAllocator<'a, 'b> {
        ScratchSpaceHeapAllocator { scratch_space: self.scratch_space,
                                    allocator: self.allocator.first_segment_words(value) }

    }

    pub fn allocation_strategy(mut self, value: AllocationStrategy) -> ScratchSpaceHeapAllocator<'a, 'b> {
        ScratchSpaceHeapAllocator { scratch_space: self.scratch_space,
                                    allocator: self.allocator.allocation_strategy(value) }
    }

}

unsafe impl <'a, 'b: 'a> Allocator for ScratchSpaceHeapAllocator<'a, 'b> {
    fn allocate_segment(&mut self, minimum_size: u32) -> (*mut Word, u32) {
        if !self.scratch_space.in_use {
            self.scratch_space.in_use = true;
            (self.scratch_space.slice.as_mut_ptr(), self.scratch_space.slice.len() as u32)
        } else {
            self.allocator.allocate_segment(minimum_size)
        }
    }

    fn pre_drop(&mut self, segment0_currently_allocated: u32) {
        let ptr = self.scratch_space.slice.as_mut_ptr();
        unsafe {
            ::std::ptr::write_bytes(ptr, 0u8, segment0_currently_allocated as usize);
        }
        self.scratch_space.in_use = false;
    }
}
<|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(),
            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\"}");
    }
}
<commit_msg>bugfix: Add ip_addresses to child validator<commit_after>use failure::Error;
use std::str::FromStr;
use std::net::IpAddr;
use std::collections::BTreeMap;

use quire::validate::{Structure, Scalar, Numeric, Mapping, Sequence};
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"))
        .member("ip_addresses", Sequence::new(Scalar::new()))
    }
}

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>use decoders::*;
use decoders::tiff::*;
use decoders::basics::*;
use std::mem::transmute;
extern crate itertools;
use self::itertools::Itertools;
extern crate crossbeam;
use std::f32;
use std::cmp;

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

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

impl<'a> Decoder for ArwDecoder<'a> {
  fn identify(&self) -> Result<&Camera, String> {
    let make = fetch_tag!(self.tiff, Tag::Make, "ARW: Couldn't find Make").get_str();
    let model = fetch_tag!(self.tiff, Tag::Model, "ARW: Couldn't find Model").get_str();
    self.rawloader.check_supported(make, model)
  }

  fn image(&self) -> Result<Image,String> {
    let camera = try!(self.identify());
    let data = self.tiff.find_ifds_with_tag(Tag::StripOffsets);
    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, "ARW: Couldn't find width").get_u16(0) as u32;
    let height = fetch_tag!(raw, Tag::ImageLength, "ARW: Couldn't find height").get_u16(0) as u32;
    let offset = fetch_tag!(raw, Tag::StripOffsets, "ARW: Couldn't find offset").get_u32(0) as usize;
    let count = fetch_tag!(raw, Tag::StripByteCounts, "ARW: Couldn't find byte count").get_u32(0) as usize;
    let compression = fetch_tag!(raw, Tag::Compression, "ARW: Couldn't find Compression").get_u16(0);
    let bps = if camera.bps != 0 {
      camera.bps
    } else {
      fetch_tag!(raw, Tag::BitsPerSample, "ARW: Couldn't find bps").get_u16(0) as u32
    };
    let src = &self.buffer[offset .. self.buffer.len()];

    let image = match compression {
      1 => decode_16le(src, width as usize, height as usize),
      32767 => {
        if ((width*height*bps) as usize) != count*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 as usize, height as usize),
            _ => return Err(format!("ARW2: Don't know how to decode images with {} bps", bps)),
          }
        }
      },
      _ => return Err(format!("ARW: Don't know how to decode type {}", compression).to_string()),
    };

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

impl<'a> ArwDecoder<'a> {
  fn decode_arw1(buf: &[u8], width: u32, height: u32) -> Vec<u16> {
    let mut buffer: Vec<u16> = vec![0; (width*height) as usize];

    buffer[0] = buf[0] as u16; // Shut up the warnings for now

    buffer
  }

  fn decode_arw2(buf: &[u8], width: u32, height: u32, curve: &LookupTable) -> Vec<u16> {
    // Simple dual threaded implementation of the arw decode that could be
    // generalized into more threads easily. In practice doing much more than
    // two threads is probably suboptimal on most hardware. Could be worth a
    // shot though.
    let height1 = height/2;
    let height2 = height - height1;
    let split = (height1+1)*width;
    let src1 = &buf[0..(split as usize)];
    let src2 = &buf[(split as usize)..buf.len()];

    crossbeam::scope(|scope| {
      let handle1 = scope.spawn(move || {
        ArwDecoder::decode_arw2_slice(src1, width, height1, &curve)
      });
      let handle2 = scope.spawn(move || {
        ArwDecoder::decode_arw2_slice(src2, width, height2, &curve)
      });

      let mut out1 = handle1.join();
      let mut out2 = handle2.join();
      out1.append(&mut out2);
      out1
    })
  }

  fn decode_arw2_slice(buf: &[u8], width: u32, height: u32, curve: &LookupTable) -> Vec<u16> {
    println!("Decoding {}x{} with a {} bytes buf", width, height, buf.len());

    let mut buffer: Vec<u16> = vec![0; (width*height) as usize];
    let mut pump = BitPump::new(buf);

    for row in 0..height {
      // Process 16 pixels at a time in interleaved fashion
      let mut col = 0;
      let mut random = pump.peek_bits(24);
      while col < (width-30) {
        let max = pump.get_bits(11);
        let min = pump.get_bits(11);
        let imax = pump.get_bits(4);
        let imin = pump.get_bits(4);
        let mut sh = 0;
        while sh<4 && (0x80 << sh) <= (max - min) {sh = sh + 1;}
        for i in 0..16 {
          let val = if i == imax {
            max
          } else if i == imin {
            min
          } else {
            cmp::min(0x7ff,(pump.get_bits(7) << sh) + min)
          };
          buffer[(row*width+col+i*2) as usize] = curve.dither(val as u16, &mut random);
        }
        col += if (col & 1) != 0 {31} else {1};  // Skip to next 16 pixels
      }
    }

    buffer
  }

  fn get_wb(&self) -> Result<[f32;4], String> {
    let priv_offset = fetch_tag!(self.tiff, Tag::DNGPrivateArea, "ARW: Couldn't find private offset").get_u32(0);
    let priv_tiff = TiffIFD::new(self.buffer, priv_offset as usize, 0, 0, LITTLE_ENDIAN);
    let sony_offset = fetch_tag!(priv_tiff, Tag::SonyOffset, "ARW: Couldn't find sony offset").get_u32(0) as usize;
    let sony_length = fetch_tag!(priv_tiff, Tag::SonyLength, "ARW: Couldn't find sony length").get_u32(0) as usize;
    let sony_key = fetch_tag!(priv_tiff, Tag::SonyKey, "ARW: Couldn't find sony key").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, LITTLE_ENDIAN);
    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_u16(1) as f32, levels.get_u16(0) as f32, levels.get_u16(2) as f32, f32::NAN])
    } else if rggb_levels.is_some() {
      let levels = rggb_levels.unwrap();
      Ok([levels.get_u16(0) as f32, levels.get_u16(1) as f32, levels.get_u16(3) as f32, f32::NAN])
    } else {
      Err("ARW: Couldn't find GRGB or RGGB levels".to_string())
    }
  }

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

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

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

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

  fn sony_decrypt(buf: &[u8], offset: usize, length: usize, key: u32) -> Vec<u8>{
    let mut pad: [u32; 128] = [0 as u32; 128];
    let mut mkey = key;
    // Make sure we always have space for the final bytes even if the buffer
    // isn't a multiple of 4
    let mut out = vec![0;length+4];
    // 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]);
    }

    for i in (0..length).step(4) {
      let p = i/4 + 127;
      pad[p & 127] = pad[(p+1) & 127] ^ pad[(p+1+64) & 127];
      let output = LEu32(buf, offset+i) ^ pad[p & 127];
      let bytes: [u8; 4] = unsafe { transmute(output.to_le()) };
      out[i]   = bytes[0];
      out[i+1] = bytes[1];
      out[i+2] = bytes[2];
      out[i+3] = bytes[3];
    }
    out
  }
}
<commit_msg>ARW: Remove left over println!<commit_after>use decoders::*;
use decoders::tiff::*;
use decoders::basics::*;
use std::mem::transmute;
extern crate itertools;
use self::itertools::Itertools;
extern crate crossbeam;
use std::f32;
use std::cmp;

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

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

impl<'a> Decoder for ArwDecoder<'a> {
  fn identify(&self) -> Result<&Camera, String> {
    let make = fetch_tag!(self.tiff, Tag::Make, "ARW: Couldn't find Make").get_str();
    let model = fetch_tag!(self.tiff, Tag::Model, "ARW: Couldn't find Model").get_str();
    self.rawloader.check_supported(make, model)
  }

  fn image(&self) -> Result<Image,String> {
    let camera = try!(self.identify());
    let data = self.tiff.find_ifds_with_tag(Tag::StripOffsets);
    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, "ARW: Couldn't find width").get_u16(0) as u32;
    let height = fetch_tag!(raw, Tag::ImageLength, "ARW: Couldn't find height").get_u16(0) as u32;
    let offset = fetch_tag!(raw, Tag::StripOffsets, "ARW: Couldn't find offset").get_u32(0) as usize;
    let count = fetch_tag!(raw, Tag::StripByteCounts, "ARW: Couldn't find byte count").get_u32(0) as usize;
    let compression = fetch_tag!(raw, Tag::Compression, "ARW: Couldn't find Compression").get_u16(0);
    let bps = if camera.bps != 0 {
      camera.bps
    } else {
      fetch_tag!(raw, Tag::BitsPerSample, "ARW: Couldn't find bps").get_u16(0) as u32
    };
    let src = &self.buffer[offset .. self.buffer.len()];

    let image = match compression {
      1 => decode_16le(src, width as usize, height as usize),
      32767 => {
        if ((width*height*bps) as usize) != count*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 as usize, height as usize),
            _ => return Err(format!("ARW2: Don't know how to decode images with {} bps", bps)),
          }
        }
      },
      _ => return Err(format!("ARW: Don't know how to decode type {}", compression).to_string()),
    };

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

impl<'a> ArwDecoder<'a> {
  fn decode_arw1(buf: &[u8], width: u32, height: u32) -> Vec<u16> {
    let mut buffer: Vec<u16> = vec![0; (width*height) as usize];

    buffer[0] = buf[0] as u16; // Shut up the warnings for now

    buffer
  }

  fn decode_arw2(buf: &[u8], width: u32, height: u32, curve: &LookupTable) -> Vec<u16> {
    // Simple dual threaded implementation of the arw decode that could be
    // generalized into more threads easily. In practice doing much more than
    // two threads is probably suboptimal on most hardware. Could be worth a
    // shot though.
    let height1 = height/2;
    let height2 = height - height1;
    let split = (height1+1)*width;
    let src1 = &buf[0..(split as usize)];
    let src2 = &buf[(split as usize)..buf.len()];

    crossbeam::scope(|scope| {
      let handle1 = scope.spawn(move || {
        ArwDecoder::decode_arw2_slice(src1, width, height1, &curve)
      });
      let handle2 = scope.spawn(move || {
        ArwDecoder::decode_arw2_slice(src2, width, height2, &curve)
      });

      let mut out1 = handle1.join();
      let mut out2 = handle2.join();
      out1.append(&mut out2);
      out1
    })
  }

  fn decode_arw2_slice(buf: &[u8], width: u32, height: u32, curve: &LookupTable) -> Vec<u16> {
    let mut buffer: Vec<u16> = vec![0; (width*height) as usize];
    let mut pump = BitPump::new(buf);

    for row in 0..height {
      // Process 16 pixels at a time in interleaved fashion
      let mut col = 0;
      let mut random = pump.peek_bits(24);
      while col < (width-30) {
        let max = pump.get_bits(11);
        let min = pump.get_bits(11);
        let imax = pump.get_bits(4);
        let imin = pump.get_bits(4);
        let mut sh = 0;
        while sh<4 && (0x80 << sh) <= (max - min) {sh = sh + 1;}
        for i in 0..16 {
          let val = if i == imax {
            max
          } else if i == imin {
            min
          } else {
            cmp::min(0x7ff,(pump.get_bits(7) << sh) + min)
          };
          buffer[(row*width+col+i*2) as usize] = curve.dither(val as u16, &mut random);
        }
        col += if (col & 1) != 0 {31} else {1};  // Skip to next 16 pixels
      }
    }

    buffer
  }

  fn get_wb(&self) -> Result<[f32;4], String> {
    let priv_offset = fetch_tag!(self.tiff, Tag::DNGPrivateArea, "ARW: Couldn't find private offset").get_u32(0);
    let priv_tiff = TiffIFD::new(self.buffer, priv_offset as usize, 0, 0, LITTLE_ENDIAN);
    let sony_offset = fetch_tag!(priv_tiff, Tag::SonyOffset, "ARW: Couldn't find sony offset").get_u32(0) as usize;
    let sony_length = fetch_tag!(priv_tiff, Tag::SonyLength, "ARW: Couldn't find sony length").get_u32(0) as usize;
    let sony_key = fetch_tag!(priv_tiff, Tag::SonyKey, "ARW: Couldn't find sony key").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, LITTLE_ENDIAN);
    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_u16(1) as f32, levels.get_u16(0) as f32, levels.get_u16(2) as f32, f32::NAN])
    } else if rggb_levels.is_some() {
      let levels = rggb_levels.unwrap();
      Ok([levels.get_u16(0) as f32, levels.get_u16(1) as f32, levels.get_u16(3) as f32, f32::NAN])
    } else {
      Err("ARW: Couldn't find GRGB or RGGB levels".to_string())
    }
  }

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

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

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

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

  fn sony_decrypt(buf: &[u8], offset: usize, length: usize, key: u32) -> Vec<u8>{
    let mut pad: [u32; 128] = [0 as u32; 128];
    let mut mkey = key;
    // Make sure we always have space for the final bytes even if the buffer
    // isn't a multiple of 4
    let mut out = vec![0;length+4];
    // 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]);
    }

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


<commit_before>//! This module implements the root command of the CLI tool.

use std::io;
use std::io::Write;
use std::env;
use std::process;

use log;
use clap::{Arg, App, AppSettings};

use prelude::*;
use constants::VERSION;
use utils::{make_subcommand, Logger};
use config::{Config, Auth};

macro_rules! each_subcommand {
    ($mac:ident) => {
        $mac!(upload_dsym);
        $mac!(releases);
        $mac!(issues);
        $mac!(update);
        $mac!(uninstall);
        $mac!(info);
        $mac!(login);
        $mac!(send_event);
    }
}

macro_rules! import_subcommand {
    ($name:ident) => { pub mod $name; }
}
each_subcommand!(import_subcommand);

/// Given an argument vector and a `Config` this executes the
/// command line and returns the result.
pub fn execute(args: Vec<String>, config: &mut Config) -> Result<()> {
    let mut app = App::new("sentry-cli")
        .version(VERSION)
        .about("Command line utility for Sentry")
        .setting(AppSettings::VersionlessSubcommands)
        .setting(AppSettings::UnifiedHelpMessage)
        .setting(AppSettings::SubcommandRequiredElseHelp)
        .arg(Arg::with_name("url")
            .value_name("URL")
            .long("url")
            .help("The sentry API URL"))
        .arg(Arg::with_name("auth_token")
            .value_name("AUTH_TOKEN")
            .long("auth-token")
            .help("The sentry auth token to use"))
        .arg(Arg::with_name("api_key")
            .value_name("API_KEY")
            .long("api-key")
            .help("The sentry API key to use"))
        .arg(Arg::with_name("log_level")
            .value_name("LOG_LEVEL")
            .long("log-level")
            .help("The log level for the sentrycli"));

    macro_rules! add_subcommand {
        ($name:ident) => {{
            app = app.subcommand($name::make_app(
                make_subcommand(&stringify!($name).replace("_", "-"))));
        }}
    }
    each_subcommand!(add_subcommand);

    let matches = app.get_matches_from_safe(args)?;

    if let Some(url) = matches.value_of("url") {
        config.url = url.to_owned();
    }
    if let Some(api_key) = matches.value_of("api_key") {
        config.auth = Some(Auth::Key(api_key.to_owned()));
    }
    if let Some(auth_token) = matches.value_of("auth_token") {
        config.auth = Some(Auth::Token(auth_token.to_owned()));
    }
    if let Some(level_str) = matches.value_of("log_level") {
        match level_str.parse() {
            Ok(level) => {
                config.log_level = level;
            }
            Err(_) => {
                fail!("Unknown log level: {}", level_str);
            }
        }
    }

    log::set_logger(|max_log_level| {
            max_log_level.set(config.log_level);
            Box::new(Logger)
        })
        .ok();

    macro_rules! execute_subcommand {
        ($name:ident) => {{
            let cmd = stringify!($name).replace("_", "-");
            if let Some(sub_matches) = matches.subcommand_matches(cmd) {
                return Ok($name::execute(&sub_matches, &config)?);
            }
        }}
    }
    each_subcommand!(execute_subcommand);
    unreachable!();
}

fn run() -> Result<()> {
    execute(env::args().collect(), &mut Config::from_cli_config()?)
}

/// Executes the command line application and exists the process.
pub fn main() {
    match run() {
        Ok(()) => process::exit(0),
        Err(err) => {
            if let &ErrorKind::ClapError(ref clap_err) = err.kind() {
                clap_err.exit();
            }

            writeln!(&mut io::stderr(), "error: {}", err).ok();

            if env::var("RUST_BACKTRACE") == Ok("1".into()) {
                writeln!(&mut io::stderr(), "").ok();
                writeln!(&mut io::stderr(), "{:?}", err.backtrace()).ok();
            }

            process::exit(1);
        }
    }
}
<commit_msg>Added workaround for rust bug rust-lang/rust#37663<commit_after>//! This module implements the root command of the CLI tool.

use std::io;
use std::io::Write;
use std::env;
use std::process;

use log;
use clap::{Arg, App, AppSettings};

use prelude::*;
use constants::VERSION;
use utils::{make_subcommand, Logger};
use config::{Config, Auth};

macro_rules! each_subcommand {
    ($mac:ident) => {
        $mac!(upload_dsym);
        $mac!(releases);
        $mac!(issues);
        $mac!(update);
        $mac!(uninstall);
        $mac!(info);
        $mac!(login);
        $mac!(send_event);
    }
}

// it would be great if this could be a macro expansion as well
// but rust bug #37663 breaks location information then.
pub mod upload_dsym;
pub mod releases;
pub mod issues;
pub mod update;
pub mod uninstall;
pub mod info;
pub mod login;
pub mod send_event;

/// Given an argument vector and a `Config` this executes the
/// command line and returns the result.
pub fn execute(args: Vec<String>, config: &mut Config) -> Result<()> {
    let mut app = App::new("sentry-cli")
        .version(VERSION)
        .about("Command line utility for Sentry")
        .setting(AppSettings::VersionlessSubcommands)
        .setting(AppSettings::UnifiedHelpMessage)
        .setting(AppSettings::SubcommandRequiredElseHelp)
        .arg(Arg::with_name("url")
            .value_name("URL")
            .long("url")
            .help("The sentry API URL"))
        .arg(Arg::with_name("auth_token")
            .value_name("AUTH_TOKEN")
            .long("auth-token")
            .help("The sentry auth token to use"))
        .arg(Arg::with_name("api_key")
            .value_name("API_KEY")
            .long("api-key")
            .help("The sentry API key to use"))
        .arg(Arg::with_name("log_level")
            .value_name("LOG_LEVEL")
            .long("log-level")
            .help("The log level for the sentrycli"));

    macro_rules! add_subcommand {
        ($name:ident) => {{
            app = app.subcommand($name::make_app(
                make_subcommand(&stringify!($name).replace("_", "-"))));
        }}
    }
    each_subcommand!(add_subcommand);

    let matches = app.get_matches_from_safe(args)?;

    if let Some(url) = matches.value_of("url") {
        config.url = url.to_owned();
    }
    if let Some(api_key) = matches.value_of("api_key") {
        config.auth = Some(Auth::Key(api_key.to_owned()));
    }
    if let Some(auth_token) = matches.value_of("auth_token") {
        config.auth = Some(Auth::Token(auth_token.to_owned()));
    }
    if let Some(level_str) = matches.value_of("log_level") {
        match level_str.parse() {
            Ok(level) => {
                config.log_level = level;
            }
            Err(_) => {
                fail!("Unknown log level: {}", level_str);
            }
        }
    }

    log::set_logger(|max_log_level| {
            max_log_level.set(config.log_level);
            Box::new(Logger)
        })
        .ok();

    macro_rules! execute_subcommand {
        ($name:ident) => {{
            let cmd = stringify!($name).replace("_", "-");
            if let Some(sub_matches) = matches.subcommand_matches(cmd) {
                return Ok($name::execute(&sub_matches, &config)?);
            }
        }}
    }
    each_subcommand!(execute_subcommand);
    unreachable!();
}

fn run() -> Result<()> {
    execute(env::args().collect(), &mut Config::from_cli_config()?)
}

/// Executes the command line application and exists the process.
pub fn main() {
    match run() {
        Ok(()) => process::exit(0),
        Err(err) => {
            if let &ErrorKind::ClapError(ref clap_err) = err.kind() {
                clap_err.exit();
            }

            writeln!(&mut io::stderr(), "error: {}", err).ok();

            if env::var("RUST_BACKTRACE") == Ok("1".into()) {
                writeln!(&mut io::stderr(), "").ok();
                writeln!(&mut io::stderr(), "{:?}", err.backtrace()).ok();
            }

            process::exit(1);
        }
    }
}
<|endoftext|>


<commit_before>use screen::*;
use piston_window;
use conrod::{self, Widget, Frameable, Positionable, Sizeable, Labelable};
use conrod::color::Colorable;

pub struct MainMenuScreen {
    ui: conrod::Ui<piston_window::Glyphs>,
    new_game: bool,
    exit: bool,
}

impl MainMenuScreen {
    pub fn new(window: &piston_window::PistonWindow, context: &mut Context) -> MainMenuScreen {
        MainMenuScreen { ui: create_ui(&window, &context.assets_path), new_game: false, exit: false, }
    }
}

impl Screen for MainMenuScreen {
    fn on_input(&mut self,
                input: &piston_window::Input,
                window: &piston_window::PistonWindow,
                context: &mut Context)
                -> InputResults {
        use piston_window::{Input, Button, MouseButton};

        let mut result = Vec::new();

        match input {
            &Input::Press(Button::Mouse(MouseButton::Left)) => {
                result.push(InputResult::PopScreen);
                result.push(InputResult::PushScreen(Box::new(BattleScreen::new(&window, context))));
            }
            _ => {}
        }

        result
    }

    fn on_draw(&mut self, args: &piston_window::RenderArgs, window: &piston_window::PistonWindow) {
        self.ui.handle_event(window);
        window.draw_2d(|context, graphics| self.ui.draw_if_changed(context, graphics));
    }

    fn on_update(&mut self, args: &piston_window::UpdateArgs) {
        handle_on_update(&mut self.ui, &mut self.new_game, &mut self.exit);
    }
}

widget_ids! {
    CANVAS,
    NEW_GAME_BUTTON,
    EXIT_BUTTON,
}

fn create_ui(window: &piston_window::PistonWindow,
             assets_path: &::std::path::Path)
             -> conrod::Ui<piston_window::Glyphs> {
    let font_path = assets_path.join("NotoSans-Regular.ttf");
    let theme = conrod::Theme::default();
    let glyph_cache = piston_window::Glyphs::new(&font_path, window.factory.borrow().clone());
    conrod::Ui::new(glyph_cache.unwrap(), theme)
}

fn handle_on_update(ui: &mut conrod::Ui<piston_window::Glyphs>, new_game: &mut bool, exit: &mut bool) {
    ui.set_widgets(|ui| {
        conrod::Canvas::new()
            .pad(30.)
            .color(conrod::color::rgb(0.2, 0.35, 0.45))
            .scroll_kids()
            .set(CANVAS, ui);

        conrod::Button::new()
            .w_h(200.0, 50.0)
            .mid_left_of(CANVAS)
            .rgb(0.4, 0.75, 0.6)
            .label("New game")
            .react(|| *new_game = true)
            .set(NEW_GAME_BUTTON, ui);

        conrod::Button::new()
            .w_h(200.0, 50.0)
            .mid_left_of(CANVAS)
            .down_from(NEW_GAME_BUTTON, 45.0)
            .rgb(0.4, 0.75, 0.6)
            .label("Exit")
            .react(|| *exit = true)
            .set(EXIT_BUTTON, ui);
    });
}
<commit_msg>OPTIONS button id<commit_after>use screen::*;
use piston_window;
use conrod::{self, Widget, Frameable, Positionable, Sizeable, Labelable};
use conrod::color::Colorable;

pub struct MainMenuScreen {
    ui: conrod::Ui<piston_window::Glyphs>,
    new_game: bool,
    exit: bool,
}

impl MainMenuScreen {
    pub fn new(window: &piston_window::PistonWindow, context: &mut Context) -> MainMenuScreen {
        MainMenuScreen { ui: create_ui(&window, &context.assets_path), new_game: false, exit: false, }
    }
}

impl Screen for MainMenuScreen {
    fn on_input(&mut self,
                input: &piston_window::Input,
                window: &piston_window::PistonWindow,
                context: &mut Context)
                -> InputResults {
        use piston_window::{Input, Button, MouseButton};

        let mut result = Vec::new();

        match input {
            &Input::Press(Button::Mouse(MouseButton::Left)) => {
                result.push(InputResult::PopScreen);
                result.push(InputResult::PushScreen(Box::new(BattleScreen::new(&window, context))));
            }
            _ => {}
        }

        result
    }

    fn on_draw(&mut self, args: &piston_window::RenderArgs, window: &piston_window::PistonWindow) {
        self.ui.handle_event(window);
        window.draw_2d(|context, graphics| self.ui.draw_if_changed(context, graphics));
    }

    fn on_update(&mut self, args: &piston_window::UpdateArgs) {
        handle_on_update(&mut self.ui, &mut self.new_game, &mut self.exit);
    }
}

widget_ids! {
    CANVAS,
    NEW_GAME_BUTTON,
    OPTIONS_BUTTON,
    EXIT_BUTTON,
}

fn create_ui(window: &piston_window::PistonWindow,
             assets_path: &::std::path::Path)
             -> conrod::Ui<piston_window::Glyphs> {
    let font_path = assets_path.join("NotoSans-Regular.ttf");
    let theme = conrod::Theme::default();
    let glyph_cache = piston_window::Glyphs::new(&font_path, window.factory.borrow().clone());
    conrod::Ui::new(glyph_cache.unwrap(), theme)
}

fn handle_on_update(ui: &mut conrod::Ui<piston_window::Glyphs>, new_game: &mut bool, exit: &mut bool) {
    ui.set_widgets(|ui| {
        conrod::Canvas::new()
            .pad(30.)
            .color(conrod::color::rgb(0.2, 0.35, 0.45))
            .scroll_kids()
            .set(CANVAS, ui);

        conrod::Button::new()
            .w_h(200.0, 50.0)
            .mid_left_of(CANVAS)
            .rgb(0.4, 0.75, 0.6)
            .label("New game")
            .react(|| *new_game = true)
            .set(NEW_GAME_BUTTON, ui);

        conrod::Button::new()
            .w_h(200.0, 50.0)
            .mid_left_of(CANVAS)
            .down_from(NEW_GAME_BUTTON, 45.0)
            .rgb(0.4, 0.75, 0.6)
            .label("Exit")
            .react(|| *exit = true)
            .set(EXIT_BUTTON, ui);
    });
}
<|endoftext|>


<commit_before>extern crate mysql;
extern crate regex;

use mysql::conn::pool::MyPool;
use mysql::conn::pool::MyPooledConn;
use mysql::conn::Stmt;
use mysql::conn::QueryResult;
use mysql::error::MyError;

use std::process::{Command, Stdio};
use std::error::Error;
use std::io::prelude::*;
use std::io::BufReader;
use std::process::Child;
use std::ascii::AsciiExt;

macro_rules! regex(
    ($s:expr) => (regex::Regex::new($s).unwrap());
);

pub struct DownloadDB {
	pub url: String,
	pub quality: i16,
	pub playlist: bool,
	pub compress: bool,
	pub qid: i64,
	pub folder_format: String,
	pub pool: MyPool,
	pub download_limit: i16,
}

pub struct Downloader {
	ddb: DownloadDB,
	// pool: MyPool,
}

#[derive(Debug)]
pub enum DownloadError{
    ConsoleError(String),
    ReadError,
    DMCAError,
    InternalError,
    DBError(String),
}

impl From<MyError> for DownloadError {
	fn from(err: MyError) -> DownloadError {
		DownloadError::DBError(err.description().into())
	}
}

impl Downloader {
	pub fn new(ddb: DownloadDB) -> Downloader{
		Downloader {ddb: ddb}
	}

	///Regex: [download]  13.4% of 275.27MiB at 525.36KiB/s ETA 07:52
	///
	///Downloads a video, updates the DB
	///TODO: get the sql statements out of the class
	///TODO: wrap errors
	///Doesn't care about DMCAs, will emit errors on them
	pub fn download_video(&self, filename: &str) -> Result<bool,DownloadError> {
	    println!("{:?}", self.ddb.url);
	    let process = try!(self.run_download_process(filename));
	    let stdout = BufReader::new(process.stdout.unwrap());

	    let mut conn = self.ddb.pool.get_conn().unwrap();
	    let mut statement = self.prepare_progress_updater(&mut conn);
	    let mut conn = self.ddb.pool.get_conn().unwrap();
	    let re = regex!(r"\d+\.\d%");

	    try!(self.set_query_code(&mut conn, &1));

	    for line in stdout.lines(){
	        match line{
	            Err(why) => panic!("couldn't read cmd stdout: {}", Error::description(&why)),
	            Ok(text) => {
	                    println!("Out: {}",text);
	                    match re.find(&text) {
	                        Some(s) => { println!("Match at {}", s.0);
	                                    println!("{}", &text[s.0..s.1]); // ONLY with ASCII chars makeable!
	                                    self.update_progress(&mut statement, &text[s.0..s.1].to_string());
	                                },
	                        None => println!("Detected no % match."),
	                    }
	                    //if re.is_match(&text) {println!("Match: {:?}", text);}
	                },
	        }
	    }

	    self.update_progress(&mut statement, &"Finished".to_string());
	    try!(self.set_query_code(&mut conn, &3));

	    Ok(true)
	}

	pub fn get_file_name(&self) -> Result<String,DownloadError> {
		let process = try!(self.run_filename_process());
		let mut stdout_buffer = BufReader::new(process.stdout.unwrap());
		let mut stderr_buffer = BufReader::new(process.stderr.unwrap());

		let mut stdout: String = String::new();
		stdout_buffer.read_to_string(&mut stdout);
		let mut stderr: String = String::new();
		stderr_buffer.read_to_string(&mut stderr);
		println!("stderr: {:?}", stderr);
		println!("stdout: {:?}", stdout);
		stdout.trim();
		println!("get_file_name: {:?}", stdout);
		Ok(stdout)
		// "asd".to_string()
		// for line in stdout.lines(){
		// 	match line{
		// 		_ => { println!("Line: {:?}", line); }
		// 	}
		// }
	}

	pub fn url_encode(input: &str) -> String {
		// iterator over input, apply function to each element(function
		input.chars().map(|char| {
	        match char {
	            ' ' => '_',
	            '&' => '-',
	            c if c.is_ascii() => c,
	            _ => '_'
	        }
    	}).collect()
    	// match for each char, then do collect: loop through the iterator, collect all elements
    	// into container FromIterator
	}

	fn run_download_process(&self, filename: &str) -> Result<Child,DownloadError> {
		match Command::new("youtube-dl")
	                                .arg("--newline")
	                                .arg(format!("-r {}M",self.ddb.download_limit))
	                                .arg(format!("-o {}/{}",self.ddb.folder_format,filename))
	                                .arg(&self.ddb.url)
	                                .stdin(Stdio::null())
	                                .stdout(Stdio::piped())
	                                .spawn() {
	        Err(why) => Err(DownloadError::ConsoleError(Error::description(&why).into())),
	        Ok(process) => Ok(process),
	    }
	}

	fn run_filename_process(&self) -> Result<Child,DownloadError> {
		match Command::new("youtube-dl")
	                                .arg("--get-filename")
	                                .arg(&self.ddb.url)
	                                .stdin(Stdio::null())
	                                .stdout(Stdio::piped())
	                                .stderr(Stdio::piped())
	                                .spawn() {
	        Err(why) => Err(DownloadError::ConsoleError(Error::description(&why).into())),
	        Ok(process) => Ok(process),
	    }
	}

	// fn create_regex(expression: & str) -> regex::Regex {
	//     match regex::Regex::new(expression) {
	//         Ok(re) => re,
	//         Err(err) => panic!("regex {}", err),
	//     }
	// }

	// MyPooledConn does only live when MyOpts is alive -> lifetime needs to be declared
	fn prepare_progress_updater<'a>(&'a self,conn: &'a mut MyPooledConn) -> Stmt<'a> { // no livetime needed: struct livetime used
	    conn.prepare("UPDATE querydetails SET status = ? WHERE qid = ?").unwrap()
	}

	fn set_query_code(&self,conn: & mut MyPooledConn, code: &i8) -> Result<(), DownloadError> { // same here
	    let mut stmt = conn.prepare("UPDATE querydetails SET code = ? WHERE qid = ?").unwrap();
	    let result = stmt.execute(&[code,&self.ddb.qid]); // why is this var needed ?!
	    match result {
	    	Ok(_) => Ok(()),
	    	Err(why) => Err(DownloadError::DBError(why.description().into())),
	    }

	}

	///updater called from the stdout progress
	fn update_progress(&self,stmt: &mut Stmt, progress: &String){
	    stmt.execute(&[progress,&self.ddb.qid]);
	}
}<commit_msg>*internalerror more precise now +io Error wrapping +DMCA error scanning<commit_after>extern crate mysql;
extern crate regex;

use mysql::conn::pool::MyPool;
use mysql::conn::pool::MyPooledConn;
use mysql::conn::Stmt;
use mysql::error::MyError;

use std::process::{Command, Stdio, Child};
use std::error::Error;
use std::io::prelude::*;
use std::io::BufReader;
use std::io;
use std::ascii::AsciiExt;

macro_rules! regex(
    ($s:expr) => (regex::Regex::new($s).unwrap());
);

pub struct DownloadDB {
	pub url: String,
	pub quality: i16,
	pub playlist: bool,
	pub compress: bool,
	pub qid: i64,
	pub folder_format: String,
	pub pool: MyPool,
	pub download_limit: i16,
}

pub struct Downloader {
	ddb: DownloadDB,
	// pool: MyPool,
}

#[derive(Debug)]
pub enum DownloadError{
    ConsoleError(String),
    ReadError,
    DMCAError,
    InternalError(String),
    DBError(String),
}

impl From<MyError> for DownloadError {
	fn from(err: MyError) -> DownloadError {
		DownloadError::DBError(err.description().into())
	}
}

impl From<io::Error> for DownloadError {
	fn from(err: io::Error) -> DownloadError {
		DownloadError::InternalError(err.description().into())
	}
}

impl Downloader {
	pub fn new(ddb: DownloadDB) -> Downloader{
		Downloader {ddb: ddb}
	}

	///Regex: [download]  13.4% of 275.27MiB at 525.36KiB/s ETA 07:52
	///
	///Downloads a video, updates the DB
	///TODO: get the sql statements out of the class
	///TODO: wrap errors
	///Doesn't care about DMCAs, will emit errors on them
	pub fn download_video(&self, filename: &str) -> Result<bool,DownloadError> {
	    println!("{:?}", self.ddb.url);
	    let process = try!(self.run_download_process(filename));
	    let stdout = BufReader::new(process.stdout.unwrap());

	    let mut conn = self.ddb.pool.get_conn().unwrap();
	    let mut statement = self.prepare_progress_updater(&mut conn);
	    let mut conn = self.ddb.pool.get_conn().unwrap();
	    let re = regex!(r"\d+\.\d%");

	    try!(self.set_query_code(&mut conn, &1));

	    for line in stdout.lines(){
	        match line{
	            Err(why) => panic!("couldn't read cmd stdout: {}", Error::description(&why)),
	            Ok(text) => {
	                    println!("Out: {}",text);
	                    match re.find(&text) {
	                        Some(s) => { println!("Match at {}", s.0);
	                                    println!("{}", &text[s.0..s.1]); // ONLY with ASCII chars makeable!
	                                    self.update_progress(&mut statement, &text[s.0..s.1].to_string());
	                                },
	                        None => println!("Detected no % match."),
	                    }
	                    //if re.is_match(&text) {println!("Match: {:?}", text);}
	                },
	        }
	    }

	    self.update_progress(&mut statement, &"Finished".to_string());
	    try!(self.set_query_code(&mut conn, &3));

	    Ok(true)
	}

	pub fn get_file_name(&self) -> Result<String,DownloadError> {
		let process = try!(self.run_filename_process());
		let mut stdout_buffer = BufReader::new(process.stdout.unwrap());
		let mut stderr_buffer = BufReader::new(process.stderr.unwrap());

		let mut stdout: String = String::new();
		try!(stdout_buffer.read_to_string(&mut stdout));
		let mut stderr: String = String::new();
		try!(stderr_buffer.read_to_string(&mut stderr));
		println!("stderr: {:?}", stderr);
		println!("stdout: {:?}", stdout);
		if stderr.contains("not available in your country") {
			return Err(DownloadError::DMCAError);
		}
		stdout.trim();
		println!("get_file_name: {:?}", stdout);
		Ok(stdout)
		// "asd".to_string()
		// for line in stdout.lines(){
		// 	match line{
		// 		_ => { println!("Line: {:?}", line); }
		// 	}
		// }
	}

	pub fn url_encode(input: &str) -> String {
		// iterator over input, apply function to each element(function
		input.chars().map(|char| {
	        match char {
	            ' ' => '_',
	            '&' => '-',
	            c if c.is_ascii() => c,
	            _ => '_'
	        }
    	}).collect()
    	// match for each char, then do collect: loop through the iterator, collect all elements
    	// into container FromIterator
	}

	fn run_download_process(&self, filename: &str) -> Result<Child,DownloadError> {
		match Command::new("youtube-dl")
	                                .arg("--newline")
	                                .arg(format!("-r {}M",self.ddb.download_limit))
	                                .arg(format!("-o {}/{}",self.ddb.folder_format,filename))
	                                .arg(&self.ddb.url)
	                                .stdin(Stdio::null())
	                                .stdout(Stdio::piped())
	                                .spawn() {
	        Err(why) => Err(DownloadError::ConsoleError(Error::description(&why).into())),
	        Ok(process) => Ok(process),
	    }
	}

	fn run_filename_process(&self) -> Result<Child,DownloadError> {
		match Command::new("youtube-dl")
	                                .arg("--get-filename")
	                                .arg(&self.ddb.url)
	                                .stdin(Stdio::null())
	                                .stdout(Stdio::piped())
	                                .stderr(Stdio::piped())
	                                .spawn() {
	        Err(why) => Err(DownloadError::ConsoleError(Error::description(&why).into())),
	        Ok(process) => Ok(process),
	    }
	}

	// fn create_regex(expression: & str) -> regex::Regex {
	//     match regex::Regex::new(expression) {
	//         Ok(re) => re,
	//         Err(err) => panic!("regex {}", err),
	//     }
	// }

	// MyPooledConn does only live when MyOpts is alive -> lifetime needs to be declared
	fn prepare_progress_updater<'a>(&'a self,conn: &'a mut MyPooledConn) -> Stmt<'a> { // no livetime needed: struct livetime used
	    conn.prepare("UPDATE querydetails SET status = ? WHERE qid = ?").unwrap()
	}

	fn set_query_code(&self,conn: & mut MyPooledConn, code: &i8) -> Result<(), DownloadError> { // same here
	    let mut stmt = conn.prepare("UPDATE querydetails SET code = ? WHERE qid = ?").unwrap();
	    let result = stmt.execute(&[code,&self.ddb.qid]); // why is this var needed ?!
	    match result {
	    	Ok(_) => Ok(()),
	    	Err(why) => Err(DownloadError::DBError(why.description().into())),
	    }

	}

	///updater called from the stdout progress
	fn update_progress(&self,stmt: &mut Stmt, progress: &String){
	    stmt.execute(&[progress,&self.ddb.qid]);
	}
}<|endoftext|>