<commit_before>extern crate http;
extern crate hyper;

use std::fmt;
use std::marker::PhantomData;
use std::sync::{Arc, Mutex};

use self::{http::StatusCode, hyper::service::Service, hyper::Chunk, hyper::Server};

use futures::future::Either;
use futures::stream::Stream;
use futures::{future, Async, Future};

use base::types::{ArcType, Type};

use vm::{
    self,
    api::{Function, OpaqueValue, PushAsRef, Userdata, VmType, WithVM, IO},
    gc::{Gc, Traverseable},
    thread::{RootedThread, Thread},
    ExternModule,
};

mod std {
    pub use http;
}

// `Handler` is a type defined in http.glu but since we need to refer to it in the signature of
// listen we define a phantom type which we can use with `OpaqueValue` to store a `Handler` in Rust
struct Handler<T>(PhantomData<T>);

impl<T: VmType + 'static> VmType for Handler<T> {
    type Type = Self;
    fn make_type(vm: &Thread) -> ArcType {
        vm.find_type_info("std.http.types.Handler")
            .map(|typ| Type::app(typ.into_type(), collect![T::make_type(vm)]))
            .unwrap_or_else(|_| Type::hole())
    }
}

// Representation of a http body that is in the prograss of being read
pub struct Body(Arc<Mutex<Box<Stream<Item = PushAsRef<Chunk, [u8]>, Error = vm::Error> + Send>>>);

// By implementing `Userdata` on `Body` it can be automatically pushed and retrieved from gluon
// threads
impl Userdata for Body {}

// Types implementing `Userdata` requires a `std::fmt::Debug` implementation so it can be displayed
impl fmt::Debug for Body {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "hyper::Body")
    }
}

// `Traverseable` is required by `Userdata` so that the garbage collector knows how to scan the
// value for garbage collected references. Normally objects do not contain any references so this
// can be empty
impl Traverseable for Body {
    fn traverse(&self, _: &mut Gc) {}
}

// `VmType` is the last trait required for a type to implement `Userdata` and defines the type used
// in gluon for this Rust type. For opaque `Userdata` values this minimal implementation is enough
// as the default implementation of `make_type` will lookup `VmType::Type` from the virtual machine
// which should have been registered earlier with `Thread::register_type`
impl VmType for Body {
    type Type = Self;
}

// Since `Body` implements `Userdata` gluon will automatically marshal the gluon representation
// into `&Body` argument
fn read_chunk(
    body: &Body,
) -> impl Future<Item = IO<Option<PushAsRef<Chunk, [u8]>>>, Error = vm::Error> {
    use futures::future::poll_fn;

    let body = body.0.clone();
    poll_fn(move || {
        let mut stream = body.lock().unwrap();
        stream.poll().map(|async| async.map(IO::Value))
    })
}

// A http body that is being written
pub struct ResponseBody(Arc<Mutex<Option<hyper::body::Sender>>>);

impl fmt::Debug for ResponseBody {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "hyper::Response")
    }
}

impl Userdata for ResponseBody {}

impl Traverseable for ResponseBody {
    fn traverse(&self, _: &mut Gc) {}
}

impl VmType for ResponseBody {
    type Type = Self;
}

fn write_response(
    response: &ResponseBody,
    bytes: &[u8],
) -> impl Future<Item = IO<()>, Error = vm::Error> {
    use futures::future::poll_fn;

    // Turn `bytes´ into a `Chunk` which can be sent to the http body
    let mut unsent_chunk = Some(bytes.to_owned().into());
    let response = response.0.clone();
    poll_fn(move || {
        info!("Starting response send");
        let mut sender = response.lock().unwrap();
        let sender = sender
            .as_mut()
            .expect("Sender has been dropped while still in use");
        let chunk = unsent_chunk
            .take()
            .expect("Attempt to poll after chunk is sent");
        match sender.poll_ready() {
            Ok(Async::NotReady) => {
                unsent_chunk = Some(chunk);
                return Ok(Async::NotReady);
            }
            Ok(Async::Ready(_)) => (),
            Err(_) => {
                info!("Could not send http response");
                return Ok(Async::Ready(IO::Value(())));
            }
        }
        match sender.send_data(chunk) {
            Ok(()) => Ok(Async::Ready(IO::Value(()))),
            Err(chunk) => {
                info!("Could not send http response");
                unsent_chunk = Some(chunk);
                Ok(Async::NotReady)
            }
        }
    })
}

// Next we define some record types which are marshalled to and from gluon. These have equivalent
// definitions in http_types.glu
field_decl! { method, uri, status, body, request, response }

type Request = record_type!{
    method => String,
    uri => String,
    body => Body
};

type Response = record_type!{
    status => u16
};

type HttpState = record_type!{
    request => Request,
    response => ResponseBody
};

fn listen(
    port: i32,
    value: WithVM<OpaqueValue<RootedThread, Handler<Response>>>,
) -> impl Future<Item = IO<()>, Error = vm::Error> + Send + 'static {
    let WithVM {
        value: handler,
        vm: thread,
    } = value;

    let thread = match thread.new_thread() {
        Ok(thread) => thread,
        Err(err) => return Either::A(future::err(err)),
    };

    // Retrieve the `handle` function from the http module which we use to evaluate values of type
    // `Handler Response`
    type ListenFn = fn(OpaqueValue<RootedThread, Handler<Response>>, HttpState) -> IO<Response>;
    let handle: Function<RootedThread, ListenFn> = thread
        .get_global("std.http.http.handle")
        .unwrap_or_else(|err| panic!("{}", err));

    struct Listen {
        handle: Function<RootedThread, ListenFn>,
        handler: OpaqueValue<RootedThread, Handler<Response>>,
    }

    impl Service for Listen {
        type ReqBody = hyper::Body;
        type ResBody = hyper::Body;
        type Error = hyper::Error;
        type Future =
            Box<Future<Item = http::Response<hyper::Body>, Error = hyper::Error> + Send + 'static>;

        fn call(&mut self, request: http::Request<hyper::Body>) -> Self::Future {
            let gluon_request = record_no_decl! {
                method => request.method().as_str().to_owned(),
                uri => request.uri().to_string(),
                // Since `Body` implements `Userdata` it can be directly pushed to gluon
                body => Body(Arc::new(Mutex::new(Box::new(request.into_body()
                    .map_err(|err| vm::Error::Message(format!("{}", err)))
                    // `PushAsRef` makes the `body` parameter act as a `&[u8]` which means it is
                    // marshalled to `Array Byte` in gluon
                    .map(PushAsRef::<_, [u8]>::new)))))
            };
            let (response_sender, response_body) = hyper::Body::channel();
            let response_sender = Arc::new(Mutex::new(Some(response_sender)));
            let http_state = record_no_decl!{
                request => gluon_request,
                response => ResponseBody(response_sender.clone())
            };

            Box::new(
                self.handle
                    .clone()
                    .call_async(self.handler.clone(), http_state)
                    .then(move |result| match result {
                        Ok(value) => {
                            match value {
                                IO::Value(record_p!{ status }) => {
                                    // Drop the sender to so that it the receiver stops waiting for
                                    // more chunks
                                    *response_sender.lock().unwrap() = None;

                                    let status = StatusCode::from_u16(status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
                                    Ok(
                                        http::Response::builder()
                                            .status(status)
                                            .body(response_body).unwrap(),
                                    )
                                }
                                IO::Exception(err) => {
                                    error!("{}", err);
                                    Ok(
                                        http::Response::builder()
                                            .status(StatusCode::INTERNAL_SERVER_ERROR).body("".into()).unwrap()
                                    )
                                }
                            }
                        }
                        Err(err) => {
                            error!("{}", err);
                            Ok(http::Response::builder()
                                .status(StatusCode::INTERNAL_SERVER_ERROR)
                                .body("".into())
                                .unwrap())
                        }
                    }),
            )
        }
    }

    let addr = format!("0.0.0.0:{}", port).parse().unwrap();

    Either::B(
        Server::bind(&addr)
            .serve(move || -> Result<_, hyper::Error> {
                Ok(Listen {
                    handle: handle.clone(),
                    handler: handler.clone(),
                })
            }).map_err(|err| vm::Error::from(format!("Server error: {}", err)))
            .and_then(|_| Ok(IO::Value(()))),
    )
}

// To let the `http_types` module refer to `Body` and `ResponseBody` we register these types in a
// separate function which is called before loading `http_types`
pub fn load_types(vm: &Thread) -> vm::Result<ExternModule> {
    vm.register_type::<Body>("Body", &[])?;
    vm.register_type::<ResponseBody>("ResponseBody", &[])?;

    ExternModule::new(
        vm,
        record! {
            // Define the types so that they can be used from gluon
            type Body => Body,
            type ResponseBody => ResponseBody,
            type Method => String,
            type StatusCode => u16,
            type Request => Request,
            type Response => Response,
            type HttpState => HttpState
        },
    )
}

pub fn load(vm: &Thread) -> vm::Result<ExternModule> {
    ExternModule::new(
        vm,
        record! {
            listen => primitive!(2, async fn listen),
            read_chunk => primitive!(1, async fn read_chunk),
            write_response => primitive!(2, async fn write_response)
        },
    )
}
<commit_msg>fix: Handle concurrency by spawning a new gluon thread per request<commit_after>extern crate http;
extern crate hyper;

use std::fmt;
use std::marker::PhantomData;
use std::sync::{Arc, Mutex};

use self::{http::StatusCode, hyper::service::Service, hyper::Chunk, hyper::Server};

use futures::future::Either;
use futures::stream::Stream;
use futures::{future, Async, Future};

use base::types::{ArcType, Type};

use vm::{
    self,
    api::{Function, OpaqueValue, PushAsRef, Userdata, VmType, WithVM, IO},
    gc::{Gc, Traverseable},
    thread::{RootedThread, Thread},
    ExternModule,
};

macro_rules! try_future {
    ($e:expr) => {
        try_future!($e, Box::new)
    };
    ($e:expr, $f:expr) => {
        match $e {
            Ok(x) => x,
            Err(err) => return $f(::futures::future::err(err.into())),
        }
    };
}

mod std {
    pub use http;
}

// `Handler` is a type defined in http.glu but since we need to refer to it in the signature of
// listen we define a phantom type which we can use with `OpaqueValue` to store a `Handler` in Rust
struct Handler<T>(PhantomData<T>);

impl<T: VmType + 'static> VmType for Handler<T> {
    type Type = Self;
    fn make_type(vm: &Thread) -> ArcType {
        vm.find_type_info("std.http.types.Handler")
            .map(|typ| Type::app(typ.into_type(), collect![T::make_type(vm)]))
            .unwrap_or_else(|_| Type::hole())
    }
}

// Representation of a http body that is in the prograss of being read
pub struct Body(Arc<Mutex<Box<Stream<Item = PushAsRef<Chunk, [u8]>, Error = vm::Error> + Send>>>);

// By implementing `Userdata` on `Body` it can be automatically pushed and retrieved from gluon
// threads
impl Userdata for Body {}

// Types implementing `Userdata` requires a `std::fmt::Debug` implementation so it can be displayed
impl fmt::Debug for Body {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "hyper::Body")
    }
}

// `Traverseable` is required by `Userdata` so that the garbage collector knows how to scan the
// value for garbage collected references. Normally objects do not contain any references so this
// can be empty
impl Traverseable for Body {
    fn traverse(&self, _: &mut Gc) {}
}

// `VmType` is the last trait required for a type to implement `Userdata` and defines the type used
// in gluon for this Rust type. For opaque `Userdata` values this minimal implementation is enough
// as the default implementation of `make_type` will lookup `VmType::Type` from the virtual machine
// which should have been registered earlier with `Thread::register_type`
impl VmType for Body {
    type Type = Self;
}

// Since `Body` implements `Userdata` gluon will automatically marshal the gluon representation
// into `&Body` argument
fn read_chunk(
    body: &Body,
) -> impl Future<Item = IO<Option<PushAsRef<Chunk, [u8]>>>, Error = vm::Error> {
    use futures::future::poll_fn;

    let body = body.0.clone();
    poll_fn(move || {
        let mut stream = body.lock().unwrap();
        stream.poll().map(|async| async.map(IO::Value))
    })
}

// A http body that is being written
pub struct ResponseBody(Arc<Mutex<Option<hyper::body::Sender>>>);

impl fmt::Debug for ResponseBody {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "hyper::Response")
    }
}

impl Userdata for ResponseBody {}

impl Traverseable for ResponseBody {
    fn traverse(&self, _: &mut Gc) {}
}

impl VmType for ResponseBody {
    type Type = Self;
}

fn write_response(
    response: &ResponseBody,
    bytes: &[u8],
) -> impl Future<Item = IO<()>, Error = vm::Error> {
    use futures::future::poll_fn;

    // Turn `bytes´ into a `Chunk` which can be sent to the http body
    let mut unsent_chunk = Some(bytes.to_owned().into());
    let response = response.0.clone();
    poll_fn(move || {
        info!("Starting response send");
        let mut sender = response.lock().unwrap();
        let sender = sender
            .as_mut()
            .expect("Sender has been dropped while still in use");
        let chunk = unsent_chunk
            .take()
            .expect("Attempt to poll after chunk is sent");
        match sender.poll_ready() {
            Ok(Async::NotReady) => {
                unsent_chunk = Some(chunk);
                return Ok(Async::NotReady);
            }
            Ok(Async::Ready(_)) => (),
            Err(_) => {
                info!("Could not send http response");
                return Ok(Async::Ready(IO::Value(())));
            }
        }
        match sender.send_data(chunk) {
            Ok(()) => Ok(Async::Ready(IO::Value(()))),
            Err(chunk) => {
                info!("Could not send http response");
                unsent_chunk = Some(chunk);
                Ok(Async::NotReady)
            }
        }
    })
}

// Next we define some record types which are marshalled to and from gluon. These have equivalent
// definitions in http_types.glu
field_decl! { method, uri, status, body, request, response }

type Request = record_type!{
    method => String,
    uri => String,
    body => Body
};

type Response = record_type!{
    status => u16
};

type HttpState = record_type!{
    request => Request,
    response => ResponseBody
};

fn listen(
    port: i32,
    value: WithVM<OpaqueValue<RootedThread, Handler<Response>>>,
) -> impl Future<Item = IO<()>, Error = vm::Error> + Send + 'static {
    let WithVM {
        value: handler,
        vm: thread,
    } = value;

    let thread = match thread.new_thread() {
        Ok(thread) => thread,
        Err(err) => return Either::A(future::err(err)),
    };

    // Retrieve the `handle` function from the http module which we use to evaluate values of type
    // `Handler Response`
    type ListenFn = fn(OpaqueValue<RootedThread, Handler<Response>>, HttpState) -> IO<Response>;
    let handle: Function<RootedThread, ListenFn> = thread
        .get_global("std.http.http.handle")
        .unwrap_or_else(|err| panic!("{}", err));

    struct Listen {
        handle: Function<RootedThread, ListenFn>,
        handler: OpaqueValue<RootedThread, Handler<Response>>,
    }

    impl Service for Listen {
        type ReqBody = hyper::Body;
        type ResBody = hyper::Body;
        type Error = vm::Error;
        type Future =
            Box<Future<Item = http::Response<hyper::Body>, Error = Self::Error> + Send + 'static>;

        fn call(&mut self, request: http::Request<hyper::Body>) -> Self::Future {
            let gluon_request = record_no_decl! {
                method => request.method().as_str().to_owned(),
                uri => request.uri().to_string(),
                // Since `Body` implements `Userdata` it can be directly pushed to gluon
                body => Body(Arc::new(Mutex::new(Box::new(request.into_body()
                    .map_err(|err| vm::Error::Message(format!("{}", err)))
                    // `PushAsRef` makes the `body` parameter act as a `&[u8]` which means it is
                    // marshalled to `Array Byte` in gluon
                    .map(PushAsRef::<_, [u8]>::new)))))
            };
            let (response_sender, response_body) = hyper::Body::channel();
            let response_sender = Arc::new(Mutex::new(Some(response_sender)));
            let http_state = record_no_decl!{
                request => gluon_request,
                response => ResponseBody(response_sender.clone())
            };

            let child_thread = try_future!(self.handle.vm().new_thread());
            let mut handle = try_future!(self.handle.re_root(child_thread));

            Box::new(
                handle
                    .call_async(self.handler.clone(), http_state)
                    .then(move |result| match result {
                        Ok(value) => {
                            match value {
                                IO::Value(record_p!{ status }) => {
                                    // Drop the sender to so that it the receiver stops waiting for
                                    // more chunks
                                    *response_sender.lock().unwrap() = None;

                                    let status = StatusCode::from_u16(status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
                                    Ok(
                                        http::Response::builder()
                                            .status(status)
                                            .body(response_body).unwrap(),
                                    )
                                }
                                IO::Exception(err) => {
                                    error!("{}", err);
                                    Ok(
                                        http::Response::builder()
                                            .status(StatusCode::INTERNAL_SERVER_ERROR).body("".into()).unwrap()
                                    )
                                }
                            }
                        }
                        Err(err) => {
                            error!("{}", err);
                            Ok(http::Response::builder()
                                .status(StatusCode::INTERNAL_SERVER_ERROR)
                                .body("".into())
                                .unwrap())
                        }
                    }),
            )
        }
    }

    let addr = format!("0.0.0.0:{}", port).parse().unwrap();

    Either::B(
        Server::bind(&addr)
            .serve(move || -> Result<_, hyper::Error> {
                Ok(Listen {
                    handle: handle.clone(),
                    handler: handler.clone(),
                })
            }).map_err(|err| vm::Error::from(format!("Server error: {}", err)))
            .and_then(|_| Ok(IO::Value(()))),
    )
}

// To let the `http_types` module refer to `Body` and `ResponseBody` we register these types in a
// separate function which is called before loading `http_types`
pub fn load_types(vm: &Thread) -> vm::Result<ExternModule> {
    vm.register_type::<Body>("Body", &[])?;
    vm.register_type::<ResponseBody>("ResponseBody", &[])?;

    ExternModule::new(
        vm,
        record! {
            // Define the types so that they can be used from gluon
            type Body => Body,
            type ResponseBody => ResponseBody,
            type Method => String,
            type StatusCode => u16,
            type Request => Request,
            type Response => Response,
            type HttpState => HttpState
        },
    )
}

pub fn load(vm: &Thread) -> vm::Result<ExternModule> {
    ExternModule::new(
        vm,
        record! {
            listen => primitive!(2, async fn listen),
            read_chunk => primitive!(1, async fn read_chunk),
            write_response => primitive!(2, async fn write_response)
        },
    )
}
<|endoftext|>


<commit_before>use crate::progress::ProgressBar;

/// Wraps an iterator to display its progress.
pub trait ProgressIterator
where
    Self: Sized + Iterator,
{
    /// Wrap an iterator with default styling. Attempt to guess iterator
    /// length using `Iterator::size_hint`.
    fn progress(self) -> ProgressBarIter<Self> {
        let n = match self.size_hint() {
            (_, Some(n)) => n as u64,
            _ => 0,
        };
        self.progress_count(n)
    }

    /// Wrap an iterator with an explicit element count.
    fn progress_count(self, len: u64) -> ProgressBarIter<Self> {
        self.progress_with(ProgressBar::new(len))
    }

    /// Wrap an iterator with a custom progress bar.
    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self>;
}

/// Wraps an iterator to display its progress.
pub struct ProgressBarIter<T> {
    it: T,
    progress: ProgressBar,
}

impl<S, T: Iterator<Item = S>> Iterator for ProgressBarIter<T> {
    type Item = S;

    fn next(&mut self) -> Option<Self::Item> {
        let next = self.it.next();

        if next.is_some() {
            self.progress.inc(1);
        } else {
            self.progress.finish();
        }

        next
    }
}

impl<S, T: Iterator<Item = S>> ProgressIterator for T {
    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self> {
        ProgressBarIter { it: self, progress }
    }
}

#[cfg(feature = "rayon")]
pub mod rayon_support {
    use super::*;
    use rayon::iter::{
        plumbing::Consumer, plumbing::Folder, plumbing::UnindexedConsumer, ParallelIterator,
    };
    use std::sync::{Arc, Mutex};

    pub struct ParProgressBarIter<T> {
        it: T,
        progress: Arc<Mutex<ProgressBar>>,
    }

    /// Wraps a Rayon parallel iterator.
    ///
    /// See [`ProgressIterator`](trait.ProgressIterator.html) for method
    /// documentation.
    pub trait ParallelProgressIterator
    where
        Self: Sized,
    {
        /// Wrap an iterator with a custom progress bar.
        fn progress_with(self, progress: ProgressBar) -> ParProgressBarIter<Self>;

        /// Wrap an iterator with an explicit element count.
        fn progress_count(self, len: u64) -> ParProgressBarIter<Self> {
            self.progress_with(ProgressBar::new(len))
        }
    }

    impl<S: Send, T: ParallelIterator<Item = S>> ParallelProgressIterator for T {
        fn progress_with(self, progress: ProgressBar) -> ParProgressBarIter<Self> {
            ParProgressBarIter {
                it: self,
                progress: Arc::new(Mutex::new(progress)),
            }
        }
    }

    struct ProgressConsumer<C> {
        base: C,
        progress: Arc<Mutex<ProgressBar>>,
    }

    impl<C> ProgressConsumer<C> {
        fn new(base: C, progress: Arc<Mutex<ProgressBar>>) -> Self {
            ProgressConsumer { base, progress }
        }
    }

    impl<T, C: Consumer<T>> Consumer<T> for ProgressConsumer<C> {
        type Folder = ProgressFolder<C::Folder>;
        type Reducer = C::Reducer;
        type Result = C::Result;

        fn split_at(self, index: usize) -> (Self, Self, Self::Reducer) {
            let (left, right, reducer) = self.base.split_at(index);
            (
                ProgressConsumer::new(left, self.progress.clone()),
                ProgressConsumer::new(right, self.progress.clone()),
                reducer,
            )
        }

        fn into_folder(self) -> Self::Folder {
            ProgressFolder {
                base: self.base.into_folder(),
                progress: self.progress.clone(),
            }
        }

        fn full(&self) -> bool {
            self.base.full()
        }
    }

    impl<T, C: UnindexedConsumer<T>> UnindexedConsumer<T> for ProgressConsumer<C> {
        fn split_off_left(&self) -> Self {
            ProgressConsumer::new(self.base.split_off_left(), self.progress.clone())
        }

        fn to_reducer(&self) -> Self::Reducer {
            self.base.to_reducer()
        }
    }

    struct ProgressFolder<C> {
        base: C,
        progress: Arc<Mutex<ProgressBar>>,
    }

    impl<T, C: Folder<T>> Folder<T> for ProgressFolder<C> {
        type Result = C::Result;

        fn consume(self, item: T) -> Self {
            self.progress.lock().unwrap().inc(1);
            ProgressFolder {
                base: self.base.consume(item),
                progress: self.progress,
            }
        }

        fn complete(self) -> C::Result {
            self.base.complete()
        }

        fn full(&self) -> bool {
            self.base.full()
        }
    }

    impl<S: Send, T: ParallelIterator<Item = S>> ParallelIterator for ParProgressBarIter<T> {
        type Item = S;

        fn drive_unindexed<C: UnindexedConsumer<Self::Item>>(self, consumer: C) -> C::Result {
            let consumer1 = ProgressConsumer::new(consumer, self.progress.clone());
            self.it.drive_unindexed(consumer1)
        }
    }

    #[cfg(test)]
    mod test {
        use super::ParProgressBarIter;
        use crate::iter::rayon_support::ParallelProgressIterator;
        use crate::progress::ProgressBar;
        use rayon::iter::{IntoParallelRefIterator, ParallelIterator};

        #[test]
        fn it_can_wrap_a_parallel_iterator() {
            let v = vec![1, 2, 3];
            let wrap = |it: ParProgressBarIter<_>| {
                assert_eq!(it.map(|x| x * 2).collect::<Vec<_>>(), vec![2, 4, 6]);
            };

            wrap(v.par_iter().progress_count(3));
            wrap({
                let pb = ProgressBar::new(v.len() as u64);
                v.par_iter().progress_with(pb)
            });
        }
    }
}

#[cfg(test)]
mod test {
    use crate::iter::{ProgressBarIter, ProgressIterator};
    use crate::progress::ProgressBar;

    #[test]
    fn it_can_wrap_an_iterator() {
        let v = vec![1, 2, 3];
        let wrap = |it: ProgressBarIter<_>| {
            assert_eq!(it.map(|x| x * 2).collect::<Vec<_>>(), vec![2, 4, 6]);
        };

        wrap(v.iter().progress());
        wrap(v.iter().progress_count(3));
        wrap({
            let pb = ProgressBar::new(v.len() as u64);
            v.iter().progress_with(pb)
        });
    }
}
<commit_msg>Expose progress in ProgressBarIter<commit_after>use crate::progress::ProgressBar;

/// Wraps an iterator to display its progress.
pub trait ProgressIterator
where
    Self: Sized + Iterator,
{
    /// Wrap an iterator with default styling. Attempt to guess iterator
    /// length using `Iterator::size_hint`.
    fn progress(self) -> ProgressBarIter<Self> {
        let n = match self.size_hint() {
            (_, Some(n)) => n as u64,
            _ => 0,
        };
        self.progress_count(n)
    }

    /// Wrap an iterator with an explicit element count.
    fn progress_count(self, len: u64) -> ProgressBarIter<Self> {
        self.progress_with(ProgressBar::new(len))
    }

    /// Wrap an iterator with a custom progress bar.
    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self>;
}

/// Wraps an iterator to display its progress.
pub struct ProgressBarIter<T> {
    it: T,
    pub progress: ProgressBar,
}

impl<S, T: Iterator<Item = S>> Iterator for ProgressBarIter<T> {
    type Item = S;

    fn next(&mut self) -> Option<Self::Item> {
        let next = self.it.next();

        if next.is_some() {
            self.progress.inc(1);
        } else {
            self.progress.finish();
        }

        next
    }
}

impl<S, T: Iterator<Item = S>> ProgressIterator for T {
    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self> {
        ProgressBarIter { it: self, progress }
    }
}

#[cfg(feature = "rayon")]
pub mod rayon_support {
    use super::*;
    use rayon::iter::{
        plumbing::Consumer, plumbing::Folder, plumbing::UnindexedConsumer, ParallelIterator,
    };
    use std::sync::{Arc, Mutex};

    pub struct ParProgressBarIter<T> {
        it: T,
        progress: Arc<Mutex<ProgressBar>>,
    }

    /// Wraps a Rayon parallel iterator.
    ///
    /// See [`ProgressIterator`](trait.ProgressIterator.html) for method
    /// documentation.
    pub trait ParallelProgressIterator
    where
        Self: Sized,
    {
        /// Wrap an iterator with a custom progress bar.
        fn progress_with(self, progress: ProgressBar) -> ParProgressBarIter<Self>;

        /// Wrap an iterator with an explicit element count.
        fn progress_count(self, len: u64) -> ParProgressBarIter<Self> {
            self.progress_with(ProgressBar::new(len))
        }
    }

    impl<S: Send, T: ParallelIterator<Item = S>> ParallelProgressIterator for T {
        fn progress_with(self, progress: ProgressBar) -> ParProgressBarIter<Self> {
            ParProgressBarIter {
                it: self,
                progress: Arc::new(Mutex::new(progress)),
            }
        }
    }

    struct ProgressConsumer<C> {
        base: C,
        progress: Arc<Mutex<ProgressBar>>,
    }

    impl<C> ProgressConsumer<C> {
        fn new(base: C, progress: Arc<Mutex<ProgressBar>>) -> Self {
            ProgressConsumer { base, progress }
        }
    }

    impl<T, C: Consumer<T>> Consumer<T> for ProgressConsumer<C> {
        type Folder = ProgressFolder<C::Folder>;
        type Reducer = C::Reducer;
        type Result = C::Result;

        fn split_at(self, index: usize) -> (Self, Self, Self::Reducer) {
            let (left, right, reducer) = self.base.split_at(index);
            (
                ProgressConsumer::new(left, self.progress.clone()),
                ProgressConsumer::new(right, self.progress.clone()),
                reducer,
            )
        }

        fn into_folder(self) -> Self::Folder {
            ProgressFolder {
                base: self.base.into_folder(),
                progress: self.progress.clone(),
            }
        }

        fn full(&self) -> bool {
            self.base.full()
        }
    }

    impl<T, C: UnindexedConsumer<T>> UnindexedConsumer<T> for ProgressConsumer<C> {
        fn split_off_left(&self) -> Self {
            ProgressConsumer::new(self.base.split_off_left(), self.progress.clone())
        }

        fn to_reducer(&self) -> Self::Reducer {
            self.base.to_reducer()
        }
    }

    struct ProgressFolder<C> {
        base: C,
        progress: Arc<Mutex<ProgressBar>>,
    }

    impl<T, C: Folder<T>> Folder<T> for ProgressFolder<C> {
        type Result = C::Result;

        fn consume(self, item: T) -> Self {
            self.progress.lock().unwrap().inc(1);
            ProgressFolder {
                base: self.base.consume(item),
                progress: self.progress,
            }
        }

        fn complete(self) -> C::Result {
            self.base.complete()
        }

        fn full(&self) -> bool {
            self.base.full()
        }
    }

    impl<S: Send, T: ParallelIterator<Item = S>> ParallelIterator for ParProgressBarIter<T> {
        type Item = S;

        fn drive_unindexed<C: UnindexedConsumer<Self::Item>>(self, consumer: C) -> C::Result {
            let consumer1 = ProgressConsumer::new(consumer, self.progress.clone());
            self.it.drive_unindexed(consumer1)
        }
    }

    #[cfg(test)]
    mod test {
        use super::ParProgressBarIter;
        use crate::iter::rayon_support::ParallelProgressIterator;
        use crate::progress::ProgressBar;
        use rayon::iter::{IntoParallelRefIterator, ParallelIterator};

        #[test]
        fn it_can_wrap_a_parallel_iterator() {
            let v = vec![1, 2, 3];
            let wrap = |it: ParProgressBarIter<_>| {
                assert_eq!(it.map(|x| x * 2).collect::<Vec<_>>(), vec![2, 4, 6]);
            };

            wrap(v.par_iter().progress_count(3));
            wrap({
                let pb = ProgressBar::new(v.len() as u64);
                v.par_iter().progress_with(pb)
            });
        }
    }
}

#[cfg(test)]
mod test {
    use crate::iter::{ProgressBarIter, ProgressIterator};
    use crate::progress::ProgressBar;

    #[test]
    fn it_can_wrap_an_iterator() {
        let v = vec![1, 2, 3];
        let wrap = |it: ProgressBarIter<_>| {
            assert_eq!(it.map(|x| x * 2).collect::<Vec<_>>(), vec![2, 4, 6]);
        };

        wrap(v.iter().progress());
        wrap(v.iter().progress_count(3));
        wrap({
            let pb = ProgressBar::new(v.len() as u64);
            v.iter().progress_with(pb)
        });
    }
}
<|endoftext|>


<commit_before>// An Iterator adapter that walks through all the elements in the Iterator,
// converts them to Strings and joins them to one big String, seperated by
// some seperator string slice.
trait StringJoiner {
  fn join(&mut self, sep: &str) -> String;
}

// Implement it for all Iterators with Elements convertable into a String
impl<A: ToString, T: Extend<A>> StringJoiner for T {
  fn join(&mut self, sep: &str) -> String {
    match self.next() {
      Some(elem) => {
        let mut output = String::from_str(elem.to_string().as_slice());
        for elem in *self {
          output.push_str(sep);
          output.push_str(elem.to_string().as_slice());
        }
        output
      }
      None => String::new()
    }
  }
}

#[test]
fn test_join() {
  let many = vec![1u,2,3];
  let one  = vec![1u];
  let none: Vec<uint> = vec![];

  assert_eq!(many.iter().join(", ").as_slice(), "1, 2, 3");
  assert_eq!(one .iter().join(", ").as_slice(), "1");
  assert_eq!(none.iter().join(", ").as_slice(), "");
}
<commit_msg>fix travis error and String::from_str() stupidity<commit_after>// An Iterator adapter that walks through all the elements in the Iterator,
// converts them to Strings and joins them to one big String, seperated by
// some seperator string slice.
trait StringJoiner {
  fn join(&mut self, sep: &str) -> String;
}

// Implement it for all Iterators with Elements convertable into a String
impl<A: ToString, T: Iterator<A>> StringJoiner for T {
  fn join(&mut self, sep: &str) -> String {
    match self.next() {
      Some(elem) => {
        let mut output = elem.to_string();
        for elem in *self {
          output.push_str(sep);
          output.push_str(elem.to_string().as_slice());
        }
        output
      }
      None => String::new()
    }
  }
}

#[test]
fn test_join() {
  let many = vec![1u,2,3];
  let one  = vec![1u];
  let none: Vec<uint> = vec![];

  assert_eq!(many.iter().join(", ").as_slice(), "1, 2, 3");
  assert_eq!(one .iter().join(", ").as_slice(), "1");
  assert_eq!(none.iter().join(", ").as_slice(), "");
}
<|endoftext|>


<commit_before>extern crate hex;
use hex::FromHex as VecFromHex;
use hex::ToHex as VecToHex;
use cryptography::crypto_box;

pub const PUBLIC_KEY_BYTES: usize = crypto_box::PUBLICKEYBYTES;
pub const SECRET_KEY_BYTES: usize = crypto_box::SECRETKEYBYTES;

/// Used to decode a CryptoAuth public key.
///
/// # Example
///
/// ```
/// use fcp_cryptoauth::cryptography::crypto_box::PublicKey;
/// use fcp_cryptoauth::keys::FromBase32;
/// let pk = PublicKey::from_base32(b"2wrpv8p4tjwm532sjxcbqzkp7kdwfwzzbg7g0n5l6g3s8df4kvv0.k");
/// assert!(pk.is_some());
/// ```
pub trait FromBase32: Sized {
    // decode a base32-encoded byte array into an other byte array
    // Returns Err(c) if a character c is not valid in a Base32 string.
    fn from_base32(characters: &[u8]) -> Option<Self>;
}

/// Used to encode a CryptoAuth public key.
///
/// # Example
///
/// ```
/// use fcp_cryptoauth::cryptography::crypto_box::PublicKey;
/// use fcp_cryptoauth::keys::FromBase32;
/// use fcp_cryptoauth::keys::ToBase32;
/// let representation = b"2wrpv8p4tjwm532sjxcbqzkp7kdwfwzzbg7g0n5l6g3s8df4kvv0.k";
/// let pk = PublicKey::from_base32(representation).unwrap();
/// assert_eq!(pk.to_base32(), String::from_utf8(representation.to_vec()).unwrap());
/// ```
pub trait ToBase32: Sized {
    // encode a byte array into an other base32-encoded byte array
    fn to_base32(&self) -> String;
}

/// Used to decode a CryptoAuth secret key.
///
/// # Example
///
/// ```
/// use fcp_cryptoauth::cryptography::crypto_box::SecretKey;
/// use fcp_cryptoauth::keys::FromHex;
/// let sk = SecretKey::from_hex(b"ac3e53b518e68449692b0b2f2926ef2fdc1eac5b9dbd10a48114263b8c8ed12e");
/// assert!(sk.is_some());
/// ```
pub trait FromHex: Sized {
    fn from_hex(characters: &[u8]) -> Option<Self>;
}

/// Used to encode a CryptoAuth secret key.
///
/// # Example
///
/// ```
/// use fcp_cryptoauth::cryptography::crypto_box::SecretKey;
/// use fcp_cryptoauth::keys::FromHex;
/// use fcp_cryptoauth::keys::ToHex;
/// let representation = b"ac3e53b518e68449692b0b2f2926ef2fdc1eac5b9dbd10a48114263b8c8ed12e";
/// let sk = SecretKey::from_hex(representation).unwrap();
/// assert_eq!(sk.to_hex(), String::from_utf8(representation.to_vec()).unwrap());
/// ```
pub trait ToHex: Sized {
    fn to_hex(&self) -> String;
}

// DNSCurve's Base32 decoding table
const BASE32_DECODING_TABLE: [u8; 128] = [
        99,99,99,99,99,99,99,99,99,99,99,99,99,99,99,99,
        99,99,99,99,99,99,99,99,99,99,99,99,99,99,99,99,
        99,99,99,99,99,99,99,99,99,99,99,99,99,99,99,99,
         0, 1, 2, 3, 4, 5, 6, 7, 8, 9,99,99,99,99,99,99,
        99,99,10,11,12,99,13,14,15,99,16,17,18,19,20,99,
        21,22,23,24,25,26,27,28,29,30,31,99,99,99,99,99,
        99,99,10,11,12,99,13,14,15,99,16,17,18,19,20,99,
        21,22,23,24,25,26,27,28,29,30,31,99,99,99,99,99
        ];
pub fn decode_base32(encoded: &[u8]) -> Result<Vec<u8>, u8> {
    let mut res = Vec::with_capacity(encoded.len()*5/8);
    let mut last_word_length = 0;
    let mut last_word = 0u16;
    for c in encoded.iter() {
        let c = c.clone();
        if c >= 128 {
            return Err(c);
        }
        let bits = BASE32_DECODING_TABLE[c as usize];
        if bits == 99 {
            return Err(c);
        }

        // Add the new bits at the left of the current word
        last_word += (bits as u16) << last_word_length;
        last_word_length += 5;

        // If the right byte of the current word is full, push the on the vector
        // and shift the current word to remove the right byte.
        if last_word_length >= 8 {
            last_word_length -= 8;
            res.push((last_word & 0b11111111) as u8);
            last_word >>= 8;
        }
    }
    if last_word > 0 {
        res.push(last_word as u8);
    }
    Ok(res)
}

#[test]
fn test_decode_base32() {
    assert_eq!(decode_base32(b"4321"), Ok(vec![0x64, 0x88]));
    decode_base32(b"2j1xz5k5y1xwz7kcczc4565jurhp8bbz1lqfu9kljw36p3nmb050").unwrap();
}

impl FromBase32 for crypto_box::PublicKey {
    /// Returns a Public key from its Base32 representation.
    ///
    /// TODO: better errors
    fn from_base32(characters: &[u8]) -> Option<crypto_box::PublicKey> {
        assert!(characters.len() >= 2);
        let length: usize = characters.len()-2;
        assert_eq!(characters[length..], b".k".to_owned());
        if let Ok(bytes) = decode_base32(&characters[..length]) {
            crypto_box::PublicKey::from_slice(&bytes)
        }
        else {
            None
        }
    }
}

const BASE32_ENCODING_TABLE: [u8; 32] = [
        '0' as u8, '1' as u8, '2' as u8, '3' as u8, '4' as u8,
        '5' as u8, '6' as u8, '7' as u8, '8' as u8, '9' as u8,
        'b' as u8, 'c' as u8, 'd' as u8, 'f' as u8, 'g' as u8,
        'h' as u8, 'j' as u8, 'k' as u8, 'l' as u8, 'm' as u8,
        'n' as u8, 'p' as u8, 'q' as u8, 'r' as u8, 's' as u8,
        't' as u8, 'u' as u8, 'v' as u8, 'w' as u8, 'x' as u8,
        'y' as u8, 'z' as u8];
pub fn encode_base32(bytes: &[u8]) -> Vec<u8> {
    let mut encoded: Vec<u8> = Vec::new();
    encoded.reserve(bytes.len()*8/5);
    let mut window = 0u16;
    let mut bits_in_window = 0;
    let mut bytes_offset = 0;
    assert!(BASE32_ENCODING_TABLE.len() == 32);
    while bytes_offset < bytes.len() || window > 0 {
        if bits_in_window < 5 {
            let byte = *bytes.get(bytes_offset).unwrap_or(&0);
            bytes_offset += 1;
            window += (byte as u16) << bits_in_window;
            bits_in_window += 8;
        }
        let char_index = (window & 0b11111) as usize;
        window >>= 5;
        bits_in_window -= 5;
        encoded.push(BASE32_ENCODING_TABLE[char_index]);
    }
    encoded
}

#[test]
fn test_encode_decode_base32() {
    let key = b"2j1xz5k5y1xwz7kcczc4565jurhp8bbz1lqfu9kljw36p3nmb050";
    let expected = "2j1xz5k5y1xwz7kcczc4565jurhp8bbz1lqfu9kljw36p3nmb05";
    assert_eq!(String::from_utf8(encode_base32(&decode_base32(key).unwrap())).unwrap(), expected)
}

#[test]
fn test_encode_base32() {
    assert_eq!(String::from_utf8(encode_base32(&vec![0x64, 0x88])).unwrap(), "4321");
}

impl ToBase32 for crypto_box::PublicKey {
    /// Returns the Base32 representation of a Public key
    ///
    /// TODO: better errors
    fn to_base32(&self) -> String {
        let mut repr = encode_base32(&self.0);
        assert!(repr.len() <= 52);
        repr.resize(52, '0' as u8);
        repr.push('.' as u8);
        repr.push('k' as u8);
        String::from_utf8(repr).unwrap()
    }
}

impl FromHex for crypto_box::SecretKey {
    fn from_hex(characters: &[u8]) -> Option<crypto_box::SecretKey> {
        let res = Vec::from_hex(characters);
        match res {
            Ok(vec) => {
                if vec.len() == 32 {
                    let mut bytes = [0u8; 32];
                    bytes.copy_from_slice(&vec);
                    crypto_box::SecretKey::from_slice(&bytes)
                }
                else {
                    None
                }
            },
            Err(_) => None,
        }
    }
}

impl ToHex for crypto_box::SecretKey {
    fn to_hex(&self) -> String {
        self.0.to_hex()
    }
}

#[test]
fn test_hex() {
    // Required for the result of SecretKey::to_hex() to always be
    // 64 characters.
    assert_eq!(vec![0, 0].to_hex(), "0000");
}
<commit_msg>Add keys::publickey_to_ipv6addr.<commit_after>extern crate hex;

use std::net::Ipv6Addr;
use rust_sodium::crypto::hash::sha512;

use hex::FromHex as VecFromHex;
use hex::ToHex as VecToHex;
use cryptography::crypto_box;

pub const PUBLIC_KEY_BYTES: usize = crypto_box::PUBLICKEYBYTES;
pub const SECRET_KEY_BYTES: usize = crypto_box::SECRETKEYBYTES;

/// Used to decode a CryptoAuth public key.
///
/// # Example
///
/// ```
/// use fcp_cryptoauth::cryptography::crypto_box::PublicKey;
/// use fcp_cryptoauth::keys::FromBase32;
/// let pk = PublicKey::from_base32(b"2wrpv8p4tjwm532sjxcbqzkp7kdwfwzzbg7g0n5l6g3s8df4kvv0.k");
/// assert!(pk.is_some());
/// ```
pub trait FromBase32: Sized {
    // decode a base32-encoded byte array into an other byte array
    // Returns Err(c) if a character c is not valid in a Base32 string.
    fn from_base32(characters: &[u8]) -> Option<Self>;
}

/// Used to encode a CryptoAuth public key.
///
/// # Example
///
/// ```
/// use fcp_cryptoauth::cryptography::crypto_box::PublicKey;
/// use fcp_cryptoauth::keys::FromBase32;
/// use fcp_cryptoauth::keys::ToBase32;
/// let representation = b"2wrpv8p4tjwm532sjxcbqzkp7kdwfwzzbg7g0n5l6g3s8df4kvv0.k";
/// let pk = PublicKey::from_base32(representation).unwrap();
/// assert_eq!(pk.to_base32(), String::from_utf8(representation.to_vec()).unwrap());
/// ```
pub trait ToBase32: Sized {
    // encode a byte array into an other base32-encoded byte array
    fn to_base32(&self) -> String;
}

/// Used to decode a CryptoAuth secret key.
///
/// # Example
///
/// ```
/// use fcp_cryptoauth::cryptography::crypto_box::SecretKey;
/// use fcp_cryptoauth::keys::FromHex;
/// let sk = SecretKey::from_hex(b"ac3e53b518e68449692b0b2f2926ef2fdc1eac5b9dbd10a48114263b8c8ed12e");
/// assert!(sk.is_some());
/// ```
pub trait FromHex: Sized {
    fn from_hex(characters: &[u8]) -> Option<Self>;
}

/// Used to encode a CryptoAuth secret key.
///
/// # Example
///
/// ```
/// use fcp_cryptoauth::cryptography::crypto_box::SecretKey;
/// use fcp_cryptoauth::keys::FromHex;
/// use fcp_cryptoauth::keys::ToHex;
/// let representation = b"ac3e53b518e68449692b0b2f2926ef2fdc1eac5b9dbd10a48114263b8c8ed12e";
/// let sk = SecretKey::from_hex(representation).unwrap();
/// assert_eq!(sk.to_hex(), String::from_utf8(representation.to_vec()).unwrap());
/// ```
pub trait ToHex: Sized {
    fn to_hex(&self) -> String;
}

// DNSCurve's Base32 decoding table
const BASE32_DECODING_TABLE: [u8; 128] = [
        99,99,99,99,99,99,99,99,99,99,99,99,99,99,99,99,
        99,99,99,99,99,99,99,99,99,99,99,99,99,99,99,99,
        99,99,99,99,99,99,99,99,99,99,99,99,99,99,99,99,
         0, 1, 2, 3, 4, 5, 6, 7, 8, 9,99,99,99,99,99,99,
        99,99,10,11,12,99,13,14,15,99,16,17,18,19,20,99,
        21,22,23,24,25,26,27,28,29,30,31,99,99,99,99,99,
        99,99,10,11,12,99,13,14,15,99,16,17,18,19,20,99,
        21,22,23,24,25,26,27,28,29,30,31,99,99,99,99,99
        ];
pub fn decode_base32(encoded: &[u8]) -> Result<Vec<u8>, u8> {
    let mut res = Vec::with_capacity(encoded.len()*5/8);
    let mut last_word_length = 0;
    let mut last_word = 0u16;
    for c in encoded.iter() {
        let c = c.clone();
        if c >= 128 {
            return Err(c);
        }
        let bits = BASE32_DECODING_TABLE[c as usize];
        if bits == 99 {
            return Err(c);
        }

        // Add the new bits at the left of the current word
        last_word += (bits as u16) << last_word_length;
        last_word_length += 5;

        // If the right byte of the current word is full, push the on the vector
        // and shift the current word to remove the right byte.
        if last_word_length >= 8 {
            last_word_length -= 8;
            res.push((last_word & 0b11111111) as u8);
            last_word >>= 8;
        }
    }
    if last_word > 0 {
        res.push(last_word as u8);
    }
    Ok(res)
}

#[test]
fn test_decode_base32() {
    assert_eq!(decode_base32(b"4321"), Ok(vec![0x64, 0x88]));
    decode_base32(b"2j1xz5k5y1xwz7kcczc4565jurhp8bbz1lqfu9kljw36p3nmb050").unwrap();
}

impl FromBase32 for crypto_box::PublicKey {
    /// Returns a Public key from its Base32 representation.
    ///
    /// TODO: better errors
    fn from_base32(characters: &[u8]) -> Option<crypto_box::PublicKey> {
        assert!(characters.len() >= 2);
        let length: usize = characters.len()-2;
        assert_eq!(characters[length..], b".k".to_owned());
        if let Ok(bytes) = decode_base32(&characters[..length]) {
            crypto_box::PublicKey::from_slice(&bytes)
        }
        else {
            None
        }
    }
}

const BASE32_ENCODING_TABLE: [u8; 32] = [
        '0' as u8, '1' as u8, '2' as u8, '3' as u8, '4' as u8,
        '5' as u8, '6' as u8, '7' as u8, '8' as u8, '9' as u8,
        'b' as u8, 'c' as u8, 'd' as u8, 'f' as u8, 'g' as u8,
        'h' as u8, 'j' as u8, 'k' as u8, 'l' as u8, 'm' as u8,
        'n' as u8, 'p' as u8, 'q' as u8, 'r' as u8, 's' as u8,
        't' as u8, 'u' as u8, 'v' as u8, 'w' as u8, 'x' as u8,
        'y' as u8, 'z' as u8];
pub fn encode_base32(bytes: &[u8]) -> Vec<u8> {
    let mut encoded: Vec<u8> = Vec::new();
    encoded.reserve(bytes.len()*8/5);
    let mut window = 0u16;
    let mut bits_in_window = 0;
    let mut bytes_offset = 0;
    assert!(BASE32_ENCODING_TABLE.len() == 32);
    while bytes_offset < bytes.len() || window > 0 {
        if bits_in_window < 5 {
            let byte = *bytes.get(bytes_offset).unwrap_or(&0);
            bytes_offset += 1;
            window += (byte as u16) << bits_in_window;
            bits_in_window += 8;
        }
        let char_index = (window & 0b11111) as usize;
        window >>= 5;
        bits_in_window -= 5;
        encoded.push(BASE32_ENCODING_TABLE[char_index]);
    }
    encoded
}

#[test]
fn test_encode_decode_base32() {
    let key = b"2j1xz5k5y1xwz7kcczc4565jurhp8bbz1lqfu9kljw36p3nmb050";
    let expected = "2j1xz5k5y1xwz7kcczc4565jurhp8bbz1lqfu9kljw36p3nmb05";
    assert_eq!(String::from_utf8(encode_base32(&decode_base32(key).unwrap())).unwrap(), expected)
}

#[test]
fn test_encode_base32() {
    assert_eq!(String::from_utf8(encode_base32(&vec![0x64, 0x88])).unwrap(), "4321");
}

impl ToBase32 for crypto_box::PublicKey {
    /// Returns the Base32 representation of a Public key
    ///
    /// TODO: better errors
    fn to_base32(&self) -> String {
        let mut repr = encode_base32(&self.0);
        assert!(repr.len() <= 52);
        repr.resize(52, '0' as u8);
        repr.push('.' as u8);
        repr.push('k' as u8);
        String::from_utf8(repr).unwrap()
    }
}

/// Converts a public key to an Ipv6Addr, according to the first paragraph of
/// https://github.com/cjdelisle/cjdns/blob/cjdns-v19.1/doc/Whitepaper.md#pulling-it-all-together
///
/// # Example
///
/// ```
/// use std::net::Ipv6Addr;
/// use std::str::FromStr;
/// use fcp_cryptoauth::cryptography::crypto_box::PublicKey;
/// use fcp_cryptoauth::keys::{FromBase32, publickey_to_ipv6addr};
/// let pk = PublicKey::from_base32(b"2wrpv8p4tjwm532sjxcbqzkp7kdwfwzzbg7g0n5l6g3s8df4kvv0.k").unwrap();
/// let ip6 = Ipv6Addr::from_str("fc8f:a188:1b5:4de9:b0cb:5729:23a1:60f9").unwrap();
/// assert_eq!(publickey_to_ipv6addr(&pk), ip6);
/// ```
pub fn publickey_to_ipv6addr(pk: &crypto_box::PublicKey) -> Ipv6Addr {
    let digest1 = sha512::hash(&pk.0);
    let digest2 = sha512::hash(&digest1.0);
    let mut octets = [0u8; 16];
    octets.copy_from_slice(&digest2.0[0..16]);
    Ipv6Addr::from(octets)
}

impl FromHex for crypto_box::SecretKey {
    fn from_hex(characters: &[u8]) -> Option<crypto_box::SecretKey> {
        let res = Vec::from_hex(characters);
        match res {
            Ok(vec) => {
                if vec.len() == 32 {
                    let mut bytes = [0u8; 32];
                    bytes.copy_from_slice(&vec);
                    crypto_box::SecretKey::from_slice(&bytes)
                }
                else {
                    None
                }
            },
            Err(_) => None,
        }
    }
}

impl ToHex for crypto_box::SecretKey {
    fn to_hex(&self) -> String {
        self.0.to_hex()
    }
}

#[test]
fn test_hex() {
    // Required for the result of SecretKey::to_hex() to always be
    // 64 characters.
    assert_eq!(vec![0, 0].to_hex(), "0000");
}
<|endoftext|>


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

use error::{Error,Result};

use std::collections::HashMap;

use rustc_serialize::json;

/// Struct used to store localization information for a language.
pub struct Lang {
    /// The lang code
    pub lang: String,
    /// The content of localization
    pub content: HashMap<String, String>,
}

impl Lang {
    /// Create a new empty Lang with no content
    pub fn new<S>(lang: S) -> Lang
        where S: Into<String> {
        Lang {
            lang: lang.into(),
            content: HashMap::new(),
        }
    }

    /// Create a new Lang from a string
    ///
    /// This string should vaguely follow .po/.mo files: it can contain
    /// comments starting by a `#`, and an entry should be of the form:
    ///
    /// ```text, no_run
    /// msgid "Initial string"
    /// msgstr "Translated string"
    /// ```
    pub fn new_from_str<S>(lang: S, s: &str) -> Result<Lang>
        where S: Into<String> {
        let mut lang = Self::new(lang);
        let lines:Vec<_> = s.lines()
            .map(|s| s.trim())
            .collect();
        let mut i = 0;
        while i < lines.len() {
            if lines[i].is_empty() || lines[i].starts_with("#") {
                // empty line or comment, ignore
                i += 1;
                continue;
            }
            if let Some(begin) = lines[i].find("msgid") {
                let end = begin + "msgid".len();
                let s = &lines[i][end..];
                let key: String = try!(json::decode(s).map_err(|e| {
                    Error::parse(format!("initializing lang '{}' at line {}, could not parse {} as a String: {}",
                                         &lang.lang,
                                         i,
                                         s,
                                         e))
                }));
                if i >= lines.len() - 1 || ! lines[i+1].starts_with("msgstr") {
                    return Err(Error::parse(format!("initializing lang '{}' at line {}, found 'msgid' without matching 'msgstr on next line",
                                                &lang.lang,
                                                i)));
                }
                i += 1;
                if let Some(begin) = lines[i].find("msgstr") {
                    let end = begin + "msgstr".len();
                    let s = &lines[i][end..];
                    let value: String = try!(json::decode(s).map_err(|e| {
                        Error::parse(format!("initializing lang '{}' at line {}, could not parse {} as a String: {}",
                                             &lang.lang,
                                             i,
                                             s,
                                             e))
                    }));
                    lang.insert(key, value);
                } else {
                    unreachable!()
                }
                i += 1;
            } else {
                return Err(Error::parse(format!("initializing lang '{}' at line {}, unexected input: '{}'",
                                                &lang.lang, i, lines[i])));
            }
        }
        Ok(lang)
    }

    /// Insert a (key, value) pair in the HashMap containing localization strings
    ///
    /// # Arguments:
    /// * `key`: the string in default language
    /// * `value`: the translation in this language
    pub fn insert<S1, S2>(&mut self, key: S1, value: S2)
        where S1: Into<String>,
              S2: Into<String> {
        self.content.insert(key.into(), value.into());
    }
}

#[test]
fn lang_new_valid_1() {
    let s = r#"
# Some comment
msgid "Some string"
msgstr "Une chaîne"

# Other comment
msgid "Other string"
msgstr "Autre chaîne"
"#;
    Lang::new_from_str("fr", s).unwrap();
}

#[test]
fn lang_new_invalid_1() {
    let s = r#"
msgstr "Msgstr first"
msgid "Some string"
msgstr "Une chaîne"

# Other comment
msgid "Other string"
msgstr "Autre chaîne"
"#;
    let lang = Lang::new_from_str("fr", s);
    assert!(lang.is_err());
}

#[test]
fn lang_new_invalid_2() {
    let s = r#"
msgid "Some string"
msgid "Two consecutive msgid without msgstr"

# Other comment
msgid "Other string"
msgstr "Autre chaîne"
"#;
    let lang = Lang::new_from_str("fr", s);
    assert!(lang.is_err());
}
<commit_msg>Only add msgstr to translation if not empty<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with
// this file, You can obtain one at https://mozilla.org/MPL/2.0/.

use error::{Error,Result};

use std::collections::HashMap;

use rustc_serialize::json;

/// Struct used to store localization information for a language.
pub struct Lang {
    /// The lang code
    pub lang: String,
    /// The content of localization
    pub content: HashMap<String, String>,
}

impl Lang {
    /// Create a new empty Lang with no content
    pub fn new<S>(lang: S) -> Lang
        where S: Into<String> {
        Lang {
            lang: lang.into(),
            content: HashMap::new(),
        }
    }

    /// Create a new Lang from a string
    ///
    /// This string should vaguely follow .po/.mo files: it can contain
    /// comments starting by a `#`, and an entry should be of the form:
    ///
    /// ```text, no_run
    /// msgid "Initial string"
    /// msgstr "Translated string"
    /// ```
    pub fn new_from_str<S>(lang: S, s: &str) -> Result<Lang>
        where S: Into<String> {
        let mut lang = Self::new(lang);
        let lines:Vec<_> = s.lines()
            .map(|s| s.trim())
            .collect();
        let mut i = 0;
        while i < lines.len() {
            if lines[i].is_empty() || lines[i].starts_with("#") {
                // empty line or comment, ignore
                i += 1;
                continue;
            }
            if let Some(begin) = lines[i].find("msgid") {
                let end = begin + "msgid".len();
                let s = &lines[i][end..];
                let key: String = try!(json::decode(s).map_err(|e| {
                    Error::parse(format!("initializing lang '{}' at line {}, could not parse {} as a String: {}",
                                         &lang.lang,
                                         i,
                                         s,
                                         e))
                }));
                if i >= lines.len() - 1 || ! lines[i+1].starts_with("msgstr") {
                    return Err(Error::parse(format!("initializing lang '{}' at line {}, found 'msgid' without matching 'msgstr on next line",
                                                &lang.lang,
                                                i)));
                }
                i += 1;
                if let Some(begin) = lines[i].find("msgstr") {
                    let end = begin + "msgstr".len();
                    let s = &lines[i][end..];
                    let value: String = try!(json::decode(s).map_err(|e| {
                        Error::parse(format!("initializing lang '{}' at line {}, could not parse {} as a String: {}",
                                             &lang.lang,
                                             i,
                                             s,
                                             e))
                    }));
                    if !value.is_empty() {
                        lang.insert(key, value);
                    }
                } else {
                    unreachable!()
                }
                i += 1;
            } else {
                return Err(Error::parse(format!("initializing lang '{}' at line {}, unexected input: '{}'",
                                                &lang.lang, i, lines[i])));
            }
        }
        Ok(lang)
    }

    /// Insert a (key, value) pair in the HashMap containing localization strings
    ///
    /// # Arguments:
    /// * `key`: the string in default language
    /// * `value`: the translation in this language
    pub fn insert<S1, S2>(&mut self, key: S1, value: S2)
        where S1: Into<String>,
              S2: Into<String> {
        self.content.insert(key.into(), value.into());
    }
}

#[test]
fn lang_new_valid_1() {
    let s = r#"
# Some comment
msgid "Some string"
msgstr "Une chaîne"

# Other comment
msgid "Other string"
msgstr "Autre chaîne"
"#;
    Lang::new_from_str("fr", s).unwrap();
}

#[test]
fn lang_new_invalid_1() {
    let s = r#"
msgstr "Msgstr first"
msgid "Some string"
msgstr "Une chaîne"

# Other comment
msgid "Other string"
msgstr "Autre chaîne"
"#;
    let lang = Lang::new_from_str("fr", s);
    assert!(lang.is_err());
}

#[test]
fn lang_new_invalid_2() {
    let s = r#"
msgid "Some string"
msgid "Two consecutive msgid without msgstr"

# Other comment
msgid "Other string"
msgstr "Autre chaîne"
"#;
    let lang = Lang::new_from_str("fr", s);
    assert!(lang.is_err());
}
<|endoftext|>


<commit_before>use std::{
    io::{self, BufReader},
    net::TcpStream,
    thread,
};

use crossbeam_channel::{bounded, Receiver, Sender};

use crate::{
    stdio::{make_io_threads, IoThreads},
    Message,
};

pub(crate) fn socket_transport(
    stream: TcpStream,
) -> (Sender<Message>, Receiver<Message>, IoThreads) {
    let (reader_receiver, reader) = make_reader(stream.try_clone().unwrap());
    let (writer_sender, writer) = make_write(stream.try_clone().unwrap());
    let io_threads = make_io_threads(reader, writer);
    (writer_sender, reader_receiver, io_threads)
}

fn make_reader(stream: TcpStream) -> (Receiver<Message>, thread::JoinHandle<io::Result<()>>) {
    let (reader_sender, reader_receiver) = bounded::<Message>(0);
    let reader = thread::spawn(move || {
        let mut buf_read = BufReader::new(stream);
        while let Some(msg) = Message::read(&mut buf_read).unwrap() {
            let is_exit = matches!(&msg, Message::Notification(n) if n.is_exit());
            reader_sender.send(msg).unwrap();
            if is_exit {
                break;
            }
        }
        Ok(())
    });
    (reader_receiver, reader)
}

fn make_write(mut stream: TcpStream) -> (Sender<Message>, thread::JoinHandle<io::Result<()>>) {
    let (writer_sender, writer_receiver) = bounded::<Message>(0);
    let writer = thread::spawn(move || {
        writer_receiver.into_iter().try_for_each(|it| it.write(&mut stream)).unwrap();
        Ok(())
    });
    (writer_sender, writer)
}
<commit_msg>Auto merge of #13078 - pymongo:master, r=Veykril<commit_after>use std::{
    io::{self, BufReader},
    net::TcpStream,
    thread,
};

use crossbeam_channel::{bounded, Receiver, Sender};

use crate::{
    stdio::{make_io_threads, IoThreads},
    Message,
};

pub(crate) fn socket_transport(
    stream: TcpStream,
) -> (Sender<Message>, Receiver<Message>, IoThreads) {
    let (reader_receiver, reader) = make_reader(stream.try_clone().unwrap());
    let (writer_sender, writer) = make_write(stream);
    let io_threads = make_io_threads(reader, writer);
    (writer_sender, reader_receiver, io_threads)
}

fn make_reader(stream: TcpStream) -> (Receiver<Message>, thread::JoinHandle<io::Result<()>>) {
    let (reader_sender, reader_receiver) = bounded::<Message>(0);
    let reader = thread::spawn(move || {
        let mut buf_read = BufReader::new(stream);
        while let Some(msg) = Message::read(&mut buf_read).unwrap() {
            let is_exit = matches!(&msg, Message::Notification(n) if n.is_exit());
            reader_sender.send(msg).unwrap();
            if is_exit {
                break;
            }
        }
        Ok(())
    });
    (reader_receiver, reader)
}

fn make_write(mut stream: TcpStream) -> (Sender<Message>, thread::JoinHandle<io::Result<()>>) {
    let (writer_sender, writer_receiver) = bounded::<Message>(0);
    let writer = thread::spawn(move || {
        writer_receiver.into_iter().try_for_each(|it| it.write(&mut stream)).unwrap();
        Ok(())
    });
    (writer_sender, writer)
}
<|endoftext|>


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

// Code to handle a collection of block devices.

use std::io;
use std::collections::{HashMap, HashSet};
use std::fs::{File, OpenOptions};
use std::path::{Path, PathBuf};
use std::str::FromStr;

use time::Timespec;
use uuid::Uuid;

use devicemapper::{Bytes, Device, Sectors, Segment};

use super::super::consts::IEC;
use super::super::errors::{EngineError, EngineResult, ErrorEnum};
use super::super::types::{DevUuid, PoolUuid};

use super::cleanup::wipe_blockdevs;
use super::blockdev::BlockDev;
use super::device::blkdev_size;
use super::engine::DevOwnership;
use super::metadata::{BDA, MIN_MDA_SECTORS, StaticHeader, validate_mda_size};
use super::range_alloc::RangeAllocator;
use super::serde_structs::{BlockDevSave, Recordable};

const MIN_DEV_SIZE: Bytes = Bytes(IEC::Gi);

/// Resolve a list of Paths of some sort to a set of unique Devices.
/// Return an IOError if there was a problem resolving any particular device.
pub fn resolve_devices(paths: &[&Path]) -> io::Result<HashSet<Device>> {
    paths
        .iter()
        .map(|p| Device::from_str(&p.to_string_lossy()))
        .collect()
}

#[derive(Debug)]
pub struct BlockDevMgr {
    block_devs: Vec<BlockDev>,
}

impl BlockDevMgr {
    pub fn new(block_devs: Vec<BlockDev>) -> BlockDevMgr {
        BlockDevMgr { block_devs: block_devs }
    }

    /// Initialize a new BlockDevMgr with specified pool and devices.
    pub fn initialize(pool_uuid: &PoolUuid,
                      paths: &[&Path],
                      mda_size: Sectors,
                      force: bool)
                      -> EngineResult<BlockDevMgr> {
        let devices = try!(resolve_devices(paths));
        Ok(BlockDevMgr::new(try!(initialize(pool_uuid, devices, mda_size, force))))
    }

    /// Obtain a BlockDev by its Device.
    pub fn get_by_device(&self, device: Device) -> Option<&BlockDev> {
        self.block_devs.iter().find(|d| d.device() == &device)
    }

    // Obtain a BlockDev by its UUID.
    pub fn get_by_uuid(&self, uuid: &DevUuid) -> Option<&BlockDev> {
        self.block_devs.iter().find(|d| d.uuid() == uuid)
    }

    pub fn add(&mut self,
               pool_uuid: &PoolUuid,
               paths: &[&Path],
               force: bool)
               -> EngineResult<Vec<PathBuf>> {
        let devices = try!(resolve_devices(paths));
        let bds = try!(initialize(pool_uuid, devices, MIN_MDA_SECTORS, force));
        let bdev_paths = bds.iter().map(|p| p.devnode.clone()).collect();
        for bd in bds {
            self.block_devs.push(bd);
        }
        Ok(bdev_paths)
    }

    pub fn destroy_all(self) -> EngineResult<()> {
        Ok(try!(wipe_blockdevs(self.block_devs)))
    }

    // Unused space left on blockdevs
    pub fn avail_space(&self) -> Sectors {
        self.block_devs.iter().map(|bd| bd.available()).sum()
    }

    /// If available space is less than size, return None, else return
    /// the segments allocated.
    pub fn alloc_space(&mut self, size: Sectors) -> Option<Vec<Segment>> {
        let mut needed: Sectors = size;
        let mut segs = Vec::new();

        if self.avail_space() < size {
            return None;
        }

        for mut bd in self.block_devs.iter_mut() {
            if needed == Sectors(0) {
                break;
            }

            let (gotten, r_segs) = bd.request_space(needed);
            segs.extend(r_segs);
            needed = needed - gotten;
        }

        assert_eq!(needed, Sectors(0));

        Some(segs)
    }

    pub fn devnodes(&self) -> Vec<PathBuf> {
        self.block_devs
            .iter()
            .map(|d| d.devnode.clone())
            .collect()
    }

    /// Write the given data to all blockdevs marking with specified time.
    // TODO: Cap # of blockdevs written to, as described in SWDD
    pub fn save_state(&mut self, time: &Timespec, metadata: &[u8]) -> EngineResult<()> {
        // TODO: Do something better than panic when saving to blockdev fails.
        // Panic can occur for a the usual IO reasons, but also:
        // 1. If the timestamp is older than a previously written timestamp.
        // 2. If the variable length metadata is too large.
        for mut bd in self.block_devs.iter_mut() {
            bd.save_state(time, metadata).unwrap();
        }
        Ok(())
    }
}

impl Recordable<HashMap<Uuid, BlockDevSave>> for BlockDevMgr {
    fn record(&self) -> EngineResult<HashMap<Uuid, BlockDevSave>> {
        self.block_devs
            .iter()
            .map(|bd| bd.record().and_then(|bdsave| Ok((*bd.uuid(), bdsave))))
            .collect()
    }
}

/// Initialize multiple blockdevs at once. This allows all of them
/// to be checked for usability before writing to any of them.
pub fn initialize(pool_uuid: &PoolUuid,
                  devices: HashSet<Device>,
                  mda_size: Sectors,
                  force: bool)
                  -> EngineResult<Vec<BlockDev>> {

    /// Get device information, returns an error if problem with obtaining
    /// that information.
    /// Returns a tuple with the device's path, its size in bytes,
    /// its ownership as determined by calling determine_ownership(),
    /// and an open File handle, all of which are needed later.
    pub fn dev_info(dev: &Device) -> EngineResult<(PathBuf, Bytes, DevOwnership, File)> {
        let devnode = try!(dev.devnode().ok_or_else(|| {
            EngineError::Engine(ErrorEnum::NotFound,
                                format!("could not get device node from dev {}", dev.dstr()))
        }));

        let mut f = try!(OpenOptions::new().read(true).write(true).open(&devnode));
        let dev_size = try!(blkdev_size(&f));
        let ownership = try!(StaticHeader::determine_ownership(&mut f));

        Ok((devnode, dev_size, ownership, f))
    }

    /// Filter devices for admission to pool based on dev_infos.
    /// If there is an error finding out the info, return that error.
    /// Also, return an error if a device is not appropriate for this pool.
    #[allow(type_complexity)]
    fn filter_devs<I>(dev_infos: I,
                      pool_uuid: &PoolUuid,
                      force: bool)
                      -> EngineResult<Vec<(Device, (PathBuf, Bytes, File))>>
        where I: Iterator<Item = (Device, EngineResult<(PathBuf, Bytes, DevOwnership, File)>)>
    {
        let mut add_devs = Vec::new();
        for (dev, dev_result) in dev_infos {
            let (devnode, dev_size, ownership, f) = try!(dev_result);
            if dev_size < MIN_DEV_SIZE {
                let error_message = format!("{} too small, minimum {} bytes",
                                            devnode.display(),
                                            MIN_DEV_SIZE);
                return Err(EngineError::Engine(ErrorEnum::Invalid, error_message));
            };
            match ownership {
                DevOwnership::Unowned => add_devs.push((dev, (devnode, dev_size, f))),
                DevOwnership::Theirs => {
                    if !force {
                        let err_str = format!("Device {} appears to belong to another application",
                                              devnode.display());
                        return Err(EngineError::Engine(ErrorEnum::Invalid, err_str));
                    } else {
                        add_devs.push((dev, (devnode, dev_size, f)))
                    }
                }
                DevOwnership::Ours(uuid) => {
                    if *pool_uuid != uuid {
                        let error_str = format!("Device {} already belongs to Stratis pool {}",
                                                devnode.display(),
                                                uuid);
                        return Err(EngineError::Engine(ErrorEnum::Invalid, error_str));
                    } else {
                        // Already in this pool (according to its header)
                        // TODO: Check we already know about it
                        // if yes, ignore. If no, add it w/o initializing?
                    }
                }
            }
        }
        Ok(add_devs)
    }

    try!(validate_mda_size(mda_size));

    let dev_infos = devices.into_iter().map(|d: Device| (d, dev_info(&d)));

    let add_devs = try!(filter_devs(dev_infos, pool_uuid, force));

    let mut bds: Vec<BlockDev> = Vec::new();
    for (dev, (devnode, dev_size, mut f)) in add_devs {

        let bda = BDA::initialize(&mut f,
                                  pool_uuid,
                                  &Uuid::new_v4(),
                                  mda_size,
                                  dev_size.sectors());
        if let Ok(bda) = bda {
            let allocator = RangeAllocator::new(bda.dev_size(), &[(Sectors(0), bda.size())])
                .expect("bda.size() < bda.dev_size() and single range");

            bds.push(BlockDev::new(dev, devnode, bda, allocator));
        } else {
            // TODO: check the return values and update state machine on failure
            let _ = BDA::wipe(&mut f);
            let _ = wipe_blockdevs(bds);

            return Err(bda.unwrap_err());
        }
    }
    Ok(bds)
}
<commit_msg>Improve BlockDevMgr::save_state()<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

// Code to handle a collection of block devices.

use std::io;
use std::collections::{HashMap, HashSet};
use std::fs::{File, OpenOptions};
use std::path::{Path, PathBuf};
use std::str::FromStr;

use time::Timespec;
use uuid::Uuid;

use devicemapper::{Bytes, Device, Sectors, Segment};

use super::super::consts::IEC;
use super::super::errors::{EngineError, EngineResult, ErrorEnum};
use super::super::types::{DevUuid, PoolUuid};

use super::cleanup::wipe_blockdevs;
use super::blockdev::BlockDev;
use super::device::blkdev_size;
use super::engine::DevOwnership;
use super::metadata::{BDA, MIN_MDA_SECTORS, StaticHeader, validate_mda_size};
use super::range_alloc::RangeAllocator;
use super::serde_structs::{BlockDevSave, Recordable};

const MIN_DEV_SIZE: Bytes = Bytes(IEC::Gi);

/// Resolve a list of Paths of some sort to a set of unique Devices.
/// Return an IOError if there was a problem resolving any particular device.
pub fn resolve_devices(paths: &[&Path]) -> io::Result<HashSet<Device>> {
    paths
        .iter()
        .map(|p| Device::from_str(&p.to_string_lossy()))
        .collect()
}

#[derive(Debug)]
pub struct BlockDevMgr {
    block_devs: Vec<BlockDev>,
}

impl BlockDevMgr {
    pub fn new(block_devs: Vec<BlockDev>) -> BlockDevMgr {
        BlockDevMgr { block_devs: block_devs }
    }

    /// Initialize a new BlockDevMgr with specified pool and devices.
    pub fn initialize(pool_uuid: &PoolUuid,
                      paths: &[&Path],
                      mda_size: Sectors,
                      force: bool)
                      -> EngineResult<BlockDevMgr> {
        let devices = try!(resolve_devices(paths));
        Ok(BlockDevMgr::new(try!(initialize(pool_uuid, devices, mda_size, force))))
    }

    /// Obtain a BlockDev by its Device.
    pub fn get_by_device(&self, device: Device) -> Option<&BlockDev> {
        self.block_devs.iter().find(|d| d.device() == &device)
    }

    // Obtain a BlockDev by its UUID.
    pub fn get_by_uuid(&self, uuid: &DevUuid) -> Option<&BlockDev> {
        self.block_devs.iter().find(|d| d.uuid() == uuid)
    }

    pub fn add(&mut self,
               pool_uuid: &PoolUuid,
               paths: &[&Path],
               force: bool)
               -> EngineResult<Vec<PathBuf>> {
        let devices = try!(resolve_devices(paths));
        let bds = try!(initialize(pool_uuid, devices, MIN_MDA_SECTORS, force));
        let bdev_paths = bds.iter().map(|p| p.devnode.clone()).collect();
        for bd in bds {
            self.block_devs.push(bd);
        }
        Ok(bdev_paths)
    }

    pub fn destroy_all(self) -> EngineResult<()> {
        Ok(try!(wipe_blockdevs(self.block_devs)))
    }

    // Unused space left on blockdevs
    pub fn avail_space(&self) -> Sectors {
        self.block_devs.iter().map(|bd| bd.available()).sum()
    }

    /// If available space is less than size, return None, else return
    /// the segments allocated.
    pub fn alloc_space(&mut self, size: Sectors) -> Option<Vec<Segment>> {
        let mut needed: Sectors = size;
        let mut segs = Vec::new();

        if self.avail_space() < size {
            return None;
        }

        for mut bd in self.block_devs.iter_mut() {
            if needed == Sectors(0) {
                break;
            }

            let (gotten, r_segs) = bd.request_space(needed);
            segs.extend(r_segs);
            needed = needed - gotten;
        }

        assert_eq!(needed, Sectors(0));

        Some(segs)
    }

    pub fn devnodes(&self) -> Vec<PathBuf> {
        self.block_devs
            .iter()
            .map(|d| d.devnode.clone())
            .collect()
    }

    /// Write the given data to all blockdevs marking with specified time.
    /// Return an error if data was not written to any blockdev.
    // TODO: Cap # of blockdevs written to, as described in SWDD
    pub fn save_state(&mut self, time: &Timespec, metadata: &[u8]) -> EngineResult<()> {
        let mut saved = false;
        for mut bd in &mut self.block_devs {
            saved |= bd.save_state(time, metadata).is_ok();
        }
        if saved {
            Ok(())
        } else {
            let err_msg = "Failed to save metadata to even one device in pool";
            Err(EngineError::Engine(ErrorEnum::Error, err_msg.into()))
        }
    }
}

impl Recordable<HashMap<Uuid, BlockDevSave>> for BlockDevMgr {
    fn record(&self) -> EngineResult<HashMap<Uuid, BlockDevSave>> {
        self.block_devs
            .iter()
            .map(|bd| bd.record().and_then(|bdsave| Ok((*bd.uuid(), bdsave))))
            .collect()
    }
}

/// Initialize multiple blockdevs at once. This allows all of them
/// to be checked for usability before writing to any of them.
pub fn initialize(pool_uuid: &PoolUuid,
                  devices: HashSet<Device>,
                  mda_size: Sectors,
                  force: bool)
                  -> EngineResult<Vec<BlockDev>> {

    /// Get device information, returns an error if problem with obtaining
    /// that information.
    /// Returns a tuple with the device's path, its size in bytes,
    /// its ownership as determined by calling determine_ownership(),
    /// and an open File handle, all of which are needed later.
    pub fn dev_info(dev: &Device) -> EngineResult<(PathBuf, Bytes, DevOwnership, File)> {
        let devnode = try!(dev.devnode().ok_or_else(|| {
            EngineError::Engine(ErrorEnum::NotFound,
                                format!("could not get device node from dev {}", dev.dstr()))
        }));

        let mut f = try!(OpenOptions::new().read(true).write(true).open(&devnode));
        let dev_size = try!(blkdev_size(&f));
        let ownership = try!(StaticHeader::determine_ownership(&mut f));

        Ok((devnode, dev_size, ownership, f))
    }

    /// Filter devices for admission to pool based on dev_infos.
    /// If there is an error finding out the info, return that error.
    /// Also, return an error if a device is not appropriate for this pool.
    #[allow(type_complexity)]
    fn filter_devs<I>(dev_infos: I,
                      pool_uuid: &PoolUuid,
                      force: bool)
                      -> EngineResult<Vec<(Device, (PathBuf, Bytes, File))>>
        where I: Iterator<Item = (Device, EngineResult<(PathBuf, Bytes, DevOwnership, File)>)>
    {
        let mut add_devs = Vec::new();
        for (dev, dev_result) in dev_infos {
            let (devnode, dev_size, ownership, f) = try!(dev_result);
            if dev_size < MIN_DEV_SIZE {
                let error_message = format!("{} too small, minimum {} bytes",
                                            devnode.display(),
                                            MIN_DEV_SIZE);
                return Err(EngineError::Engine(ErrorEnum::Invalid, error_message));
            };
            match ownership {
                DevOwnership::Unowned => add_devs.push((dev, (devnode, dev_size, f))),
                DevOwnership::Theirs => {
                    if !force {
                        let err_str = format!("Device {} appears to belong to another application",
                                              devnode.display());
                        return Err(EngineError::Engine(ErrorEnum::Invalid, err_str));
                    } else {
                        add_devs.push((dev, (devnode, dev_size, f)))
                    }
                }
                DevOwnership::Ours(uuid) => {
                    if *pool_uuid != uuid {
                        let error_str = format!("Device {} already belongs to Stratis pool {}",
                                                devnode.display(),
                                                uuid);
                        return Err(EngineError::Engine(ErrorEnum::Invalid, error_str));
                    } else {
                        // Already in this pool (according to its header)
                        // TODO: Check we already know about it
                        // if yes, ignore. If no, add it w/o initializing?
                    }
                }
            }
        }
        Ok(add_devs)
    }

    try!(validate_mda_size(mda_size));

    let dev_infos = devices.into_iter().map(|d: Device| (d, dev_info(&d)));

    let add_devs = try!(filter_devs(dev_infos, pool_uuid, force));

    let mut bds: Vec<BlockDev> = Vec::new();
    for (dev, (devnode, dev_size, mut f)) in add_devs {

        let bda = BDA::initialize(&mut f,
                                  pool_uuid,
                                  &Uuid::new_v4(),
                                  mda_size,
                                  dev_size.sectors());
        if let Ok(bda) = bda {
            let allocator = RangeAllocator::new(bda.dev_size(), &[(Sectors(0), bda.size())])
                .expect("bda.size() < bda.dev_size() and single range");

            bds.push(BlockDev::new(dev, devnode, bda, allocator));
        } else {
            // TODO: check the return values and update state machine on failure
            let _ = BDA::wipe(&mut f);
            let _ = wipe_blockdevs(bds);

            return Err(bda.unwrap_err());
        }
    }
    Ok(bds)
}
<|endoftext|>