<commit_before>//! Connection pooling via r2d2.
//!
//! Note: This module requires enabling the `r2d2` feature
//!
//! # Example
//!
//! The below snippet is a contrived example emulating a web application,
//! where one would first initialize the pool in the `main()` function
//! (at the start of a long-running process). One would then pass this
//! pool struct around as shared state, which, here, we've emulated using
//! threads instead of routes.
//!
//! ```rust
//! # include!("doctest_setup.rs");
//! use diesel::prelude::*;
//! use diesel::r2d2::ConnectionManager;
//! # use diesel::r2d2::CustomizeConnection;
//! # use diesel::r2d2::Error as R2D2Error;
//! use diesel::r2d2::Pool;
//! use diesel::r2d2::PooledConnection;
//! use diesel::result::Error;
//! use std::env;
//! use std::thread;
//!
//! # #[derive(Copy, Clone, Debug)]
//! # pub struct SetupUserTableCustomizer;
//! #
//! # impl CustomizeConnection<DbConnection, R2D2Error> for SetupUserTableCustomizer
//! # {
//! #     fn on_acquire(&self, conn: &mut DbConnection) -> Result<(), R2D2Error> {
//! #         setup_database(conn);
//! #         Ok(())
//! #     }
//! # }
//!
//! pub fn get_connection_pool() -> Pool<ConnectionManager<DbConnection>> {
//!     let url = database_url_for_env();
//!     let manager = ConnectionManager::<DbConnection>::new(url);
//!     // Refer to the `r2d2` documentation for more methods to use
//!     // when building a connection pool
//!     Pool::builder()
//! #         .max_size(1)
//!         .test_on_check_out(true)
//! #         .connection_customizer(Box::new(SetupUserTableCustomizer))
//!         .build(manager)
//!         .expect("Could not build connection pool")
//! }
//!
//! pub fn create_user(conn: &mut DbConnection, user_name: &str) -> Result<usize, Error> {
//!     use schema::users::dsl::*;
//!
//!     diesel::insert_into(users)
//!         .values(name.eq(user_name))
//!         .execute(conn)
//! }
//!
//! fn main() {
//!     let pool = get_connection_pool();
//!     let mut threads = vec![];
//!     let max_users_to_create = 1;
//!
//!     for i in 0..max_users_to_create {
//!         let pool = pool.clone();
//!         threads.push(thread::spawn({
//!             move || {
//!                 let conn = &mut pool.get().unwrap();
//!                 let name = format!("Person {}", i);
//!                 create_user(conn, &name).unwrap();
//!             }
//!         }))
//!     }
//!
//!     for handle in threads {
//!         handle.join().unwrap();
//!     }
//! }
//! ```
//!
//! # A note on error handling
//!
//! When used inside a pool, if an individual connection becomes
//! broken (as determined by the [R2D2Connection::is_broken] method)
//! then `r2d2` will close and return the connection to the DB.
//!
//! `diesel` determines broken connections by whether or not the current
//! thread is panicking or if individual `Connection` structs are
//! broken (determined by the `is_broken()` method). Generically, these
//! are left to individual backends to implement themselves.
//!
//! For SQLite, PG, and MySQL backends, specifically, `is_broken()`
//! is determined by whether or not the `TransactionManagerStatus` (as a part
//! of the `AnsiTransactionManager` struct) is in an `InError` state.
//!

pub use r2d2::*;

/// A re-export of [`r2d2::Error`], which is only used by methods on [`r2d2::Pool`].
///
/// [`r2d2::Error`]: r2d2::Error
/// [`r2d2::Pool`]: r2d2::Pool
pub type PoolError = r2d2::Error;

use std::convert::Into;
use std::fmt;
use std::marker::PhantomData;

use crate::backend::Backend;
use crate::connection::commit_error_processor::{CommitErrorOutcome, CommitErrorProcessor};
use crate::connection::{
    ConnectionGatWorkaround, SimpleConnection, TransactionManager, TransactionManagerStatus,
};
use crate::expression::QueryMetadata;
use crate::prelude::*;
use crate::query_builder::{Query, QueryFragment, QueryId};

/// An r2d2 connection manager for use with Diesel.
///
/// See the [r2d2 documentation] for usage examples.
///
/// [r2d2 documentation]: r2d2
#[derive(Clone)]
pub struct ConnectionManager<T> {
    database_url: String,
    _marker: PhantomData<T>,
}

impl<T> fmt::Debug for ConnectionManager<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ConnectionManager<{}>", std::any::type_name::<T>())
    }
}

unsafe impl<T: Send + 'static> Sync for ConnectionManager<T> {}

impl<T> ConnectionManager<T> {
    /// Returns a new connection manager,
    /// which establishes connections to the given database URL.
    pub fn new<S: Into<String>>(database_url: S) -> Self {
        ConnectionManager {
            database_url: database_url.into(),
            _marker: PhantomData,
        }
    }
}

/// The error used when managing connections with `r2d2`.
#[derive(Debug)]
pub enum Error {
    /// An error occurred establishing the connection
    ConnectionError(ConnectionError),

    /// An error occurred pinging the database
    QueryError(crate::result::Error),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Error::ConnectionError(ref e) => e.fmt(f),
            Error::QueryError(ref e) => e.fmt(f),
        }
    }
}

impl ::std::error::Error for Error {}

/// A trait indicating a connection could be used inside a r2d2 pool
pub trait R2D2Connection: Connection {
    /// Check if a connection is still valid
    fn ping(&mut self) -> QueryResult<()>;

    /// Checks if the connection is broken and should not be reused
    ///
    /// This method should return only contain a fast non-blocking check
    /// if the connection is considered to be broken or not. See
    /// [ManageConnection::has_broken] for details.
    ///
    /// The default implementation does not consider any connection as broken
    fn is_broken(&mut self) -> bool {
        false
    }
}

impl<T> ManageConnection for ConnectionManager<T>
where
    T: R2D2Connection + Send + 'static,
{
    type Connection = T;
    type Error = Error;

    fn connect(&self) -> Result<T, Error> {
        T::establish(&self.database_url).map_err(Error::ConnectionError)
    }

    fn is_valid(&self, conn: &mut T) -> Result<(), Error> {
        conn.ping().map_err(Error::QueryError)
    }

    fn has_broken(&self, conn: &mut T) -> bool {
        std::thread::panicking() || conn.is_broken()
    }
}

impl<M> SimpleConnection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: R2D2Connection + Send + 'static,
{
    fn batch_execute(&mut self, query: &str) -> QueryResult<()> {
        (&mut **self).batch_execute(query)
    }
}

impl<'conn, 'query, DB, M> ConnectionGatWorkaround<'conn, 'query, DB> for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: Connection<Backend = DB>,
    DB: Backend,
{
    type Cursor = <M::Connection as ConnectionGatWorkaround<'conn, 'query, DB>>::Cursor;
    type Row = <M::Connection as ConnectionGatWorkaround<'conn, 'query, DB>>::Row;
}

impl<M> CommitErrorProcessor for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: R2D2Connection + CommitErrorProcessor + Send + 'static,
{
    fn process_commit_error(&self, error: crate::result::Error) -> CommitErrorOutcome {
        (&**self).process_commit_error(error)
    }
}

impl<M> Connection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: Connection + R2D2Connection + Send + 'static,
{
    type Backend = <M::Connection as Connection>::Backend;
    type TransactionManager =
        PoolTransactionManager<<M::Connection as Connection>::TransactionManager>;

    fn establish(_: &str) -> ConnectionResult<Self> {
        Err(ConnectionError::BadConnection(String::from(
            "Cannot directly establish a pooled connection",
        )))
    }

    fn load<'conn, 'query, T>(
        &'conn mut self,
        source: T,
    ) -> QueryResult<<Self as ConnectionGatWorkaround<'conn, 'query, Self::Backend>>::Cursor>
    where
        T: Query + QueryFragment<Self::Backend> + QueryId + 'query,
        Self::Backend: QueryMetadata<T::SqlType>,
    {
        (&mut **self).load(source)
    }

    fn execute_returning_count<T>(&mut self, source: &T) -> QueryResult<usize>
    where
        T: QueryFragment<Self::Backend> + QueryId,
    {
        (&mut **self).execute_returning_count(source)
    }

    fn transaction_state(
        &mut self,
    ) -> &mut <Self::TransactionManager as TransactionManager<Self>>::TransactionStateData {
        (&mut **self).transaction_state()
    }

    fn begin_test_transaction(&mut self) -> QueryResult<()> {
        (&mut **self).begin_test_transaction()
    }
}

#[doc(hidden)]
#[allow(missing_debug_implementations)]
pub struct PoolTransactionManager<T>(std::marker::PhantomData<T>);

impl<M, T> TransactionManager<PooledConnection<M>> for PoolTransactionManager<T>
where
    M: ManageConnection,
    M::Connection: Connection<TransactionManager = T> + R2D2Connection,
    T: TransactionManager<M::Connection>,
{
    type TransactionStateData = T::TransactionStateData;

    fn begin_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::begin_transaction(&mut **conn)
    }

    fn rollback_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::rollback_transaction(&mut **conn)
    }

    fn commit_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::commit_transaction(&mut **conn)
    }

    fn transaction_manager_status_mut(
        conn: &mut PooledConnection<M>,
    ) -> &mut TransactionManagerStatus {
        T::transaction_manager_status_mut(&mut **conn)
    }
}

impl<M> crate::migration::MigrationConnection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: crate::migration::MigrationConnection,
    Self: Connection,
{
    fn setup(&mut self) -> QueryResult<usize> {
        (&mut **self).setup()
    }
}

impl<Changes, Output, M> crate::query_dsl::UpdateAndFetchResults<Changes, Output>
    for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: crate::query_dsl::UpdateAndFetchResults<Changes, Output>,
    Self: Connection,
{
    fn update_and_fetch(&mut self, changeset: Changes) -> QueryResult<Output> {
        (&mut **self).update_and_fetch(changeset)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::mpsc;
    use std::sync::Arc;
    use std::thread;

    use crate::r2d2::*;
    use crate::test_helpers::*;

    #[test]
    fn establish_basic_connection() {
        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Arc::new(Pool::builder().max_size(2).build(manager).unwrap());

        let (s1, r1) = mpsc::channel();
        let (s2, r2) = mpsc::channel();

        let pool1 = Arc::clone(&pool);
        let t1 = thread::spawn(move || {
            let conn = pool1.get().unwrap();
            s1.send(()).unwrap();
            r2.recv().unwrap();
            drop(conn);
        });

        let pool2 = Arc::clone(&pool);
        let t2 = thread::spawn(move || {
            let conn = pool2.get().unwrap();
            s2.send(()).unwrap();
            r1.recv().unwrap();
            drop(conn);
        });

        t1.join().unwrap();
        t2.join().unwrap();

        pool.get().unwrap();
    }

    #[test]
    fn is_valid() {
        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Pool::builder()
            .max_size(1)
            .test_on_check_out(true)
            .build(manager)
            .unwrap();

        pool.get().unwrap();
    }

    #[test]
    fn pooled_connection_impls_connection() {
        use crate::select;
        use crate::sql_types::Text;

        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Pool::builder()
            .max_size(1)
            .test_on_check_out(true)
            .build(manager)
            .unwrap();
        let mut conn = pool.get().unwrap();

        let query = select("foo".into_sql::<Text>());
        assert_eq!("foo", query.get_result::<String>(&mut conn).unwrap());
    }
}

#[derive(QueryId)]
pub(crate) struct CheckConnectionQuery;

impl<DB> QueryFragment<DB> for CheckConnectionQuery
where
    DB: Backend,
{
    fn walk_ast<'b>(
        &'b self,
        mut pass: crate::query_builder::AstPass<'_, 'b, DB>,
    ) -> QueryResult<()> {
        pass.push_sql("SELECT 1");
        Ok(())
    }
}

impl Query for CheckConnectionQuery {
    type SqlType = crate::sql_types::Integer;
}

impl<C> RunQueryDsl<C> for CheckConnectionQuery {}
<commit_msg>Remove unneeded imports in r2d2 doctest<commit_after>//! Connection pooling via r2d2.
//!
//! Note: This module requires enabling the `r2d2` feature
//!
//! # Example
//!
//! The below snippet is a contrived example emulating a web application,
//! where one would first initialize the pool in the `main()` function
//! (at the start of a long-running process). One would then pass this
//! pool struct around as shared state, which, here, we've emulated using
//! threads instead of routes.
//!
//! ```rust
//! # include!("doctest_setup.rs");
//! use diesel::prelude::*;
//! use diesel::r2d2::ConnectionManager;
//! # use diesel::r2d2::CustomizeConnection;
//! # use diesel::r2d2::Error as R2D2Error;
//! use diesel::r2d2::Pool;
//! use diesel::result::Error;
//! use std::thread;
//!
//! # #[derive(Copy, Clone, Debug)]
//! # pub struct SetupUserTableCustomizer;
//! #
//! # impl CustomizeConnection<DbConnection, R2D2Error> for SetupUserTableCustomizer
//! # {
//! #     fn on_acquire(&self, conn: &mut DbConnection) -> Result<(), R2D2Error> {
//! #         setup_database(conn);
//! #         Ok(())
//! #     }
//! # }
//!
//! pub fn get_connection_pool() -> Pool<ConnectionManager<DbConnection>> {
//!     let url = database_url_for_env();
//!     let manager = ConnectionManager::<DbConnection>::new(url);
//!     // Refer to the `r2d2` documentation for more methods to use
//!     // when building a connection pool
//!     Pool::builder()
//! #         .max_size(1)
//!         .test_on_check_out(true)
//! #         .connection_customizer(Box::new(SetupUserTableCustomizer))
//!         .build(manager)
//!         .expect("Could not build connection pool")
//! }
//!
//! pub fn create_user(conn: &mut DbConnection, user_name: &str) -> Result<usize, Error> {
//!     use schema::users::dsl::*;
//!
//!     diesel::insert_into(users)
//!         .values(name.eq(user_name))
//!         .execute(conn)
//! }
//!
//! fn main() {
//!     let pool = get_connection_pool();
//!     let mut threads = vec![];
//!     let max_users_to_create = 1;
//!
//!     for i in 0..max_users_to_create {
//!         let pool = pool.clone();
//!         threads.push(thread::spawn({
//!             move || {
//!                 let conn = &mut pool.get().unwrap();
//!                 let name = format!("Person {}", i);
//!                 create_user(conn, &name).unwrap();
//!             }
//!         }))
//!     }
//!
//!     for handle in threads {
//!         handle.join().unwrap();
//!     }
//! }
//! ```
//!
//! # A note on error handling
//!
//! When used inside a pool, if an individual connection becomes
//! broken (as determined by the [R2D2Connection::is_broken] method)
//! then `r2d2` will close and return the connection to the DB.
//!
//! `diesel` determines broken connections by whether or not the current
//! thread is panicking or if individual `Connection` structs are
//! broken (determined by the `is_broken()` method). Generically, these
//! are left to individual backends to implement themselves.
//!
//! For SQLite, PG, and MySQL backends, specifically, `is_broken()`
//! is determined by whether or not the `TransactionManagerStatus` (as a part
//! of the `AnsiTransactionManager` struct) is in an `InError` state.
//!

pub use r2d2::*;

/// A re-export of [`r2d2::Error`], which is only used by methods on [`r2d2::Pool`].
///
/// [`r2d2::Error`]: r2d2::Error
/// [`r2d2::Pool`]: r2d2::Pool
pub type PoolError = r2d2::Error;

use std::convert::Into;
use std::fmt;
use std::marker::PhantomData;

use crate::backend::Backend;
use crate::connection::commit_error_processor::{CommitErrorOutcome, CommitErrorProcessor};
use crate::connection::{
    ConnectionGatWorkaround, SimpleConnection, TransactionManager, TransactionManagerStatus,
};
use crate::expression::QueryMetadata;
use crate::prelude::*;
use crate::query_builder::{Query, QueryFragment, QueryId};

/// An r2d2 connection manager for use with Diesel.
///
/// See the [r2d2 documentation] for usage examples.
///
/// [r2d2 documentation]: r2d2
#[derive(Clone)]
pub struct ConnectionManager<T> {
    database_url: String,
    _marker: PhantomData<T>,
}

impl<T> fmt::Debug for ConnectionManager<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ConnectionManager<{}>", std::any::type_name::<T>())
    }
}

unsafe impl<T: Send + 'static> Sync for ConnectionManager<T> {}

impl<T> ConnectionManager<T> {
    /// Returns a new connection manager,
    /// which establishes connections to the given database URL.
    pub fn new<S: Into<String>>(database_url: S) -> Self {
        ConnectionManager {
            database_url: database_url.into(),
            _marker: PhantomData,
        }
    }
}

/// The error used when managing connections with `r2d2`.
#[derive(Debug)]
pub enum Error {
    /// An error occurred establishing the connection
    ConnectionError(ConnectionError),

    /// An error occurred pinging the database
    QueryError(crate::result::Error),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Error::ConnectionError(ref e) => e.fmt(f),
            Error::QueryError(ref e) => e.fmt(f),
        }
    }
}

impl ::std::error::Error for Error {}

/// A trait indicating a connection could be used inside a r2d2 pool
pub trait R2D2Connection: Connection {
    /// Check if a connection is still valid
    fn ping(&mut self) -> QueryResult<()>;

    /// Checks if the connection is broken and should not be reused
    ///
    /// This method should return only contain a fast non-blocking check
    /// if the connection is considered to be broken or not. See
    /// [ManageConnection::has_broken] for details.
    ///
    /// The default implementation does not consider any connection as broken
    fn is_broken(&mut self) -> bool {
        false
    }
}

impl<T> ManageConnection for ConnectionManager<T>
where
    T: R2D2Connection + Send + 'static,
{
    type Connection = T;
    type Error = Error;

    fn connect(&self) -> Result<T, Error> {
        T::establish(&self.database_url).map_err(Error::ConnectionError)
    }

    fn is_valid(&self, conn: &mut T) -> Result<(), Error> {
        conn.ping().map_err(Error::QueryError)
    }

    fn has_broken(&self, conn: &mut T) -> bool {
        std::thread::panicking() || conn.is_broken()
    }
}

impl<M> SimpleConnection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: R2D2Connection + Send + 'static,
{
    fn batch_execute(&mut self, query: &str) -> QueryResult<()> {
        (&mut **self).batch_execute(query)
    }
}

impl<'conn, 'query, DB, M> ConnectionGatWorkaround<'conn, 'query, DB> for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: Connection<Backend = DB>,
    DB: Backend,
{
    type Cursor = <M::Connection as ConnectionGatWorkaround<'conn, 'query, DB>>::Cursor;
    type Row = <M::Connection as ConnectionGatWorkaround<'conn, 'query, DB>>::Row;
}

impl<M> CommitErrorProcessor for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: R2D2Connection + CommitErrorProcessor + Send + 'static,
{
    fn process_commit_error(&self, error: crate::result::Error) -> CommitErrorOutcome {
        (&**self).process_commit_error(error)
    }
}

impl<M> Connection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: Connection + R2D2Connection + Send + 'static,
{
    type Backend = <M::Connection as Connection>::Backend;
    type TransactionManager =
        PoolTransactionManager<<M::Connection as Connection>::TransactionManager>;

    fn establish(_: &str) -> ConnectionResult<Self> {
        Err(ConnectionError::BadConnection(String::from(
            "Cannot directly establish a pooled connection",
        )))
    }

    fn load<'conn, 'query, T>(
        &'conn mut self,
        source: T,
    ) -> QueryResult<<Self as ConnectionGatWorkaround<'conn, 'query, Self::Backend>>::Cursor>
    where
        T: Query + QueryFragment<Self::Backend> + QueryId + 'query,
        Self::Backend: QueryMetadata<T::SqlType>,
    {
        (&mut **self).load(source)
    }

    fn execute_returning_count<T>(&mut self, source: &T) -> QueryResult<usize>
    where
        T: QueryFragment<Self::Backend> + QueryId,
    {
        (&mut **self).execute_returning_count(source)
    }

    fn transaction_state(
        &mut self,
    ) -> &mut <Self::TransactionManager as TransactionManager<Self>>::TransactionStateData {
        (&mut **self).transaction_state()
    }

    fn begin_test_transaction(&mut self) -> QueryResult<()> {
        (&mut **self).begin_test_transaction()
    }
}

#[doc(hidden)]
#[allow(missing_debug_implementations)]
pub struct PoolTransactionManager<T>(std::marker::PhantomData<T>);

impl<M, T> TransactionManager<PooledConnection<M>> for PoolTransactionManager<T>
where
    M: ManageConnection,
    M::Connection: Connection<TransactionManager = T> + R2D2Connection,
    T: TransactionManager<M::Connection>,
{
    type TransactionStateData = T::TransactionStateData;

    fn begin_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::begin_transaction(&mut **conn)
    }

    fn rollback_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::rollback_transaction(&mut **conn)
    }

    fn commit_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::commit_transaction(&mut **conn)
    }

    fn transaction_manager_status_mut(
        conn: &mut PooledConnection<M>,
    ) -> &mut TransactionManagerStatus {
        T::transaction_manager_status_mut(&mut **conn)
    }
}

impl<M> crate::migration::MigrationConnection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: crate::migration::MigrationConnection,
    Self: Connection,
{
    fn setup(&mut self) -> QueryResult<usize> {
        (&mut **self).setup()
    }
}

impl<Changes, Output, M> crate::query_dsl::UpdateAndFetchResults<Changes, Output>
    for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: crate::query_dsl::UpdateAndFetchResults<Changes, Output>,
    Self: Connection,
{
    fn update_and_fetch(&mut self, changeset: Changes) -> QueryResult<Output> {
        (&mut **self).update_and_fetch(changeset)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::mpsc;
    use std::sync::Arc;
    use std::thread;

    use crate::r2d2::*;
    use crate::test_helpers::*;

    #[test]
    fn establish_basic_connection() {
        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Arc::new(Pool::builder().max_size(2).build(manager).unwrap());

        let (s1, r1) = mpsc::channel();
        let (s2, r2) = mpsc::channel();

        let pool1 = Arc::clone(&pool);
        let t1 = thread::spawn(move || {
            let conn = pool1.get().unwrap();
            s1.send(()).unwrap();
            r2.recv().unwrap();
            drop(conn);
        });

        let pool2 = Arc::clone(&pool);
        let t2 = thread::spawn(move || {
            let conn = pool2.get().unwrap();
            s2.send(()).unwrap();
            r1.recv().unwrap();
            drop(conn);
        });

        t1.join().unwrap();
        t2.join().unwrap();

        pool.get().unwrap();
    }

    #[test]
    fn is_valid() {
        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Pool::builder()
            .max_size(1)
            .test_on_check_out(true)
            .build(manager)
            .unwrap();

        pool.get().unwrap();
    }

    #[test]
    fn pooled_connection_impls_connection() {
        use crate::select;
        use crate::sql_types::Text;

        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Pool::builder()
            .max_size(1)
            .test_on_check_out(true)
            .build(manager)
            .unwrap();
        let mut conn = pool.get().unwrap();

        let query = select("foo".into_sql::<Text>());
        assert_eq!("foo", query.get_result::<String>(&mut conn).unwrap());
    }
}

#[derive(QueryId)]
pub(crate) struct CheckConnectionQuery;

impl<DB> QueryFragment<DB> for CheckConnectionQuery
where
    DB: Backend,
{
    fn walk_ast<'b>(
        &'b self,
        mut pass: crate::query_builder::AstPass<'_, 'b, DB>,
    ) -> QueryResult<()> {
        pass.push_sql("SELECT 1");
        Ok(())
    }
}

impl Query for CheckConnectionQuery {
    type SqlType = crate::sql_types::Integer;
}

impl<C> RunQueryDsl<C> for CheckConnectionQuery {}
<|endoftext|>


<commit_before>//! Zig-zag encoding of integral values

/// A trait for enabling zig-zag encoding of various values
pub trait ZigZag<T> {

    /// Encodes this ZigZag-enabled type into the type specified by implementation
    fn zigzag(&self) -> T;

}

impl ZigZag<u8> for i8 {

    /// Encodes this i8 as a zigzagged u8
    fn zigzag(&self) -> u8 {
        ((self << 1) ^ (self >> 7)) as u8
    }

}

impl ZigZag<i8> for u8 {

    /// Encodes this u8 as a zigzagged i8
    fn zigzag(&self) -> i8 {
        ((self >> 1) as i8) ^ (-((self & 1) as i8))
    }

}

impl ZigZag<u16> for i16 {

    /// Encodes this i16 as a zigzagged u16
    fn zigzag(&self) -> u16 {
        ((self << 1) ^ (self >> 15)) as u16
    }

}

impl ZigZag<i16> for u16 {

    /// Encodes this u16 as a zigzagged i16
    fn zigzag(&self) -> i16 {
        ((self >> 1) as i16) ^ (-((self & 1) as i16))
    }

}

impl ZigZag<u32> for i32 {

    /// Encodes this i32 as a zigzagged u32
    fn zigzag(&self) -> u32 {
        ((self << 1) ^ (self >> 31)) as u32
    }

}

impl ZigZag<i32> for u32 {

    /// Encodes this u32 as a zigzagged i32
    fn zigzag(&self) -> i32 {
        ((self >> 1) as i32) ^ (-((self & 1) as i32))
    }

}

impl ZigZag<u64> for i64 {

    /// Encodes this i64 as a zigzagged u64
    fn zigzag(&self) -> u64 {
        ((self << 1) ^ (self >> 63)) as u64
    }

}

impl ZigZag<i64> for u64 {

    /// Encodes this u64 as a zigzagged i64
    fn zigzag(&self) -> i64 {
        ((self >> 1) as i64) ^ (-((self & 1) as i64))
    }

}
<commit_msg>Basic tests for zigzag functionality<commit_after>//! Zig-zag encoding of integral values

/// A trait for enabling zig-zag encoding of various values
pub trait ZigZag<T> {

    /// Encodes this ZigZag-enabled type into the type specified by implementation
    fn zigzag(&self) -> T;

}

impl ZigZag<u8> for i8 {

    /// Encodes this i8 as a zigzagged u8
    fn zigzag(&self) -> u8 {
        ((self << 1) ^ (self >> 7)) as u8
    }

}

impl ZigZag<i8> for u8 {

    /// Encodes this u8 as a zigzagged i8
    fn zigzag(&self) -> i8 {
        ((self >> 1) as i8) ^ (-((self & 1) as i8))
    }

}

impl ZigZag<u16> for i16 {

    /// Encodes this i16 as a zigzagged u16
    fn zigzag(&self) -> u16 {
        ((self << 1) ^ (self >> 15)) as u16
    }

}

impl ZigZag<i16> for u16 {

    /// Encodes this u16 as a zigzagged i16
    fn zigzag(&self) -> i16 {
        ((self >> 1) as i16) ^ (-((self & 1) as i16))
    }

}

impl ZigZag<u32> for i32 {

    /// Encodes this i32 as a zigzagged u32
    fn zigzag(&self) -> u32 {
        ((self << 1) ^ (self >> 31)) as u32
    }

}

impl ZigZag<i32> for u32 {

    /// Encodes this u32 as a zigzagged i32
    fn zigzag(&self) -> i32 {
        ((self >> 1) as i32) ^ (-((self & 1) as i32))
    }

}

impl ZigZag<u64> for i64 {

    /// Encodes this i64 as a zigzagged u64
    fn zigzag(&self) -> u64 {
        ((self << 1) ^ (self >> 63)) as u64
    }

}

impl ZigZag<i64> for u64 {

    /// Encodes this u64 as a zigzagged i64
    fn zigzag(&self) -> i64 {
        ((self >> 1) as i64) ^ (-((self & 1) as i64))
    }

}

#[cfg(test)]
mod tests {

    use super::ZigZag;

    #[test]
    fn test_u8_i8_zigzag() {
        let mut unsigned: u8 = 0u8;

        assert_eq!(0i8, unsigned.zigzag());

        unsigned = 1;

        assert_eq!(-1i8, unsigned.zigzag());

        unsigned = 2;

        assert_eq!(1i8, unsigned.zigzag());

        let mut signed: i8 = 0i8;

        assert_eq!(0u8, signed.zigzag());

        signed = -1i8;

        assert_eq!(1u8, signed.zigzag());

        signed = 1i8;

        assert_eq!(2u8, signed.zigzag());

    }

    #[test]
    fn test_u16_i16_zigzag() {
        let mut unsigned: u16 = 0u16;

        assert_eq!(0i16, unsigned.zigzag());

        unsigned = 1;

        assert_eq!(-1i16, unsigned.zigzag());

        unsigned = 2;

        assert_eq!(1i16, unsigned.zigzag());

        let mut signed: i16 = 0i16;

        assert_eq!(0u16, signed.zigzag());

        signed = -1i16;

        assert_eq!(1u16, signed.zigzag());

        signed = 1i16;

        assert_eq!(2u16, signed.zigzag());

    }

    #[test]
    fn test_u32_i32_zigzag() {
        let mut unsigned: u32 = 0u32;

        assert_eq!(0i32, unsigned.zigzag());

        unsigned = 1;

        assert_eq!(-1i32, unsigned.zigzag());

        unsigned = 2;

        assert_eq!(1i32, unsigned.zigzag());

        let mut signed: i32 = 0i32;

        assert_eq!(0u32, signed.zigzag());

        signed = -1i32;

        assert_eq!(1u32, signed.zigzag());

        signed = 1i32;

        assert_eq!(2u32, signed.zigzag());

    }

    #[test]
    fn test_u64_i64_zigzag() {
        let mut unsigned: u64 = 0u64;

        assert_eq!(0i64, unsigned.zigzag());

        unsigned = 1;

        assert_eq!(-1i64, unsigned.zigzag());

        unsigned = 2;

        assert_eq!(1i64, unsigned.zigzag());

        let mut signed: i64 = 0i64;

        assert_eq!(0u64, signed.zigzag());

        signed = -1i64;

        assert_eq!(1u64, signed.zigzag());

        signed = 1i64;

        assert_eq!(2u64, signed.zigzag());

    }

}
<|endoftext|>


<commit_before>//! Simple module to store files in database.
//!
//! cratesfyi is generating more than 5 million files, they are small and mostly html files.
//! They are using so many inodes and it is better to store them in database instead of
//! filesystem. This module is adding files into database and retrieving them.

use std::path::Path;
use DocBuilderError;
use postgres::Connection;
use rustc_serialize::json::{Json, ToJson};
use std::fs::File;
use std::io::Read;

fn file_path(prefix: &str, name: &str) -> String {
    match prefix.is_empty() {
        true => name.to_owned(),
        false => format!("{}/{}", prefix, name),
    }
}

fn get_file_list_from_dir<P: AsRef<Path>>(path: P,
                                          prefix: &str,
                                          files: &mut Vec<String>)
                                          -> Result<(), DocBuilderError> {
    let path = path.as_ref();

    for file in try!(path.read_dir()) {
        let file = try!(file);

        if try!(file.file_type()).is_file() {
            file.file_name().to_str().map(|name| files.push(file_path(prefix, name)));
        } else if try!(file.file_type()).is_dir() {
            file.file_name()
                .to_str()
                .map(|name| get_file_list_from_dir(file.path(), &file_path(prefix, name), files));
        }
    }

    Ok(())
}

pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<String>, DocBuilderError> {
    let path = path.as_ref();
    let mut files: Vec<String> = Vec::new();

    if !path.exists() {
        return Err(DocBuilderError::FileNotFound);
    } else if path.is_file() {
        path.file_name()
            .and_then(|name| name.to_str())
            .map(|name| files.push(format!("{}", name)));
    } else if path.is_dir() {
        try!(get_file_list_from_dir(path, "", &mut files));
    }

    Ok(files)
}

/// Adds files into database and returns list of files with their mime type in Json
pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
                                              prefix: &str,
                                              path: P)
                                              -> Result<Json, DocBuilderError> {
    use magic::{Cookie, flags};
    let cookie = try!(Cookie::open(flags::MIME_TYPE));
    // FIXME: This is linux specific but idk any alternative
    try!(cookie.load(&vec!["/usr/share/misc/magic.mgc"]));

    let trans = try!(conn.transaction());

    let mut file_list_with_mimes: Vec<(String, String)> = Vec::new();

    for file_path_str in try!(get_file_list(&path)) {
        let (path, content, mime) = {
            let path = Path::new(path.as_ref()).join(&file_path_str);
            // Some files have insufficient permissions (like .lock file created by cargo in
            // documentation directory). We are skipping this files.
            let mut file = match File::open(path) {
                Ok(f) => f,
                Err(_) => continue,
            };
            let mut content: Vec<u8> = Vec::new();
            try!(file.read_to_end(&mut content));
            let mime = try!(cookie.buffer(&content));

            file_list_with_mimes.push((mime.clone(), file_path_str.clone()));

            (file_path(prefix, &file_path_str), content, mime)
        };

        // check if file already exists in database
        let rows = try!(conn.query("SELECT COUNT(*) FROM files WHERE path = $1", &[&path]));

        if rows.get(0).get::<usize, i64>(0) == 0 {
            try!(trans.query("INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                             &[&path, &mime, &content]));
        } else {
            try!(trans.query("UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                              WHERE path = $1",
                             &[&path, &mime, &content]));
        }
    }

    try!(trans.commit());

    file_list_to_json(file_list_with_mimes)
}

fn file_list_to_json(file_list: Vec<(String, String)>) -> Result<Json, DocBuilderError> {

    let mut file_list_json: Vec<Json> = Vec::new();

    for file in file_list {
        let mut v: Vec<String> = Vec::new();
        v.push(file.0.clone());
        v.push(file.1.clone());
        file_list_json.push(v.to_json());
    }

    Ok(file_list_json.to_json())
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use std::env;
    use super::{get_file_list, add_path_into_database};
    use super::super::connect_db;

    #[test]
    fn test_get_file_list() {
        let _ = env_logger::init();

        let files = get_file_list(env::current_dir().unwrap());
        assert!(files.is_ok());
        assert!(files.unwrap().len() > 0);

        let files = get_file_list(env::current_dir().unwrap().join("Cargo.toml")).unwrap();
        assert_eq!(files[0], "Cargo.toml");
    }

    #[test]
    #[ignore]
    fn test_add_path_into_database() {
        let _ = env_logger::init();

        let conn = connect_db().unwrap();
        let res = add_path_into_database(&conn, "example", env::current_dir().unwrap().join("src"));
        assert!(res.is_ok());
    }
}
<commit_msg>Use text/css mime type for css files<commit_after>//! Simple module to store files in database.
//!
//! cratesfyi is generating more than 5 million files, they are small and mostly html files.
//! They are using so many inodes and it is better to store them in database instead of
//! filesystem. This module is adding files into database and retrieving them.

use std::path::Path;
use DocBuilderError;
use postgres::Connection;
use rustc_serialize::json::{Json, ToJson};
use std::fs::File;
use std::io::Read;

fn file_path(prefix: &str, name: &str) -> String {
    match prefix.is_empty() {
        true => name.to_owned(),
        false => format!("{}/{}", prefix, name),
    }
}

fn get_file_list_from_dir<P: AsRef<Path>>(path: P,
                                          prefix: &str,
                                          files: &mut Vec<String>)
                                          -> Result<(), DocBuilderError> {
    let path = path.as_ref();

    for file in try!(path.read_dir()) {
        let file = try!(file);

        if try!(file.file_type()).is_file() {
            file.file_name().to_str().map(|name| files.push(file_path(prefix, name)));
        } else if try!(file.file_type()).is_dir() {
            file.file_name()
                .to_str()
                .map(|name| get_file_list_from_dir(file.path(), &file_path(prefix, name), files));
        }
    }

    Ok(())
}

pub fn get_file_list<P: AsRef<Path>>(path: P) -> Result<Vec<String>, DocBuilderError> {
    let path = path.as_ref();
    let mut files: Vec<String> = Vec::new();

    if !path.exists() {
        return Err(DocBuilderError::FileNotFound);
    } else if path.is_file() {
        path.file_name()
            .and_then(|name| name.to_str())
            .map(|name| files.push(format!("{}", name)));
    } else if path.is_dir() {
        try!(get_file_list_from_dir(path, "", &mut files));
    }

    Ok(files)
}

/// Adds files into database and returns list of files with their mime type in Json
pub fn add_path_into_database<P: AsRef<Path>>(conn: &Connection,
                                              prefix: &str,
                                              path: P)
                                              -> Result<Json, DocBuilderError> {
    use magic::{Cookie, flags};
    let cookie = try!(Cookie::open(flags::MIME_TYPE));
    // FIXME: This is linux specific but idk any alternative
    try!(cookie.load(&vec!["/usr/share/misc/magic.mgc"]));

    let trans = try!(conn.transaction());

    let mut file_list_with_mimes: Vec<(String, String)> = Vec::new();

    for file_path_str in try!(get_file_list(&path)) {
        let (path, content, mime) = {
            let path = Path::new(path.as_ref()).join(&file_path_str);
            // Some files have insufficient permissions (like .lock file created by cargo in
            // documentation directory). We are skipping this files.
            let mut file = match File::open(path) {
                Ok(f) => f,
                Err(_) => continue,
            };
            let mut content: Vec<u8> = Vec::new();
            try!(file.read_to_end(&mut content));
            let mime = {
                let mime = try!(cookie.buffer(&content));
                // css's are causing some problem in browsers
                // magic will return text/plain for css file types
                // convert them to text/css
                if mime == "text/plain" && file_path_str.ends_with(".css") {
                    "text/css".to_owned()
                } else {
                    mime.to_owned()
                }
            };

            file_list_with_mimes.push((mime.clone(), file_path_str.clone()));

            (file_path(prefix, &file_path_str), content, mime)
        };

        // check if file already exists in database
        let rows = try!(conn.query("SELECT COUNT(*) FROM files WHERE path = $1", &[&path]));

        if rows.get(0).get::<usize, i64>(0) == 0 {
            try!(trans.query("INSERT INTO files (path, mime, content) VALUES ($1, $2, $3)",
                             &[&path, &mime, &content]));
        } else {
            try!(trans.query("UPDATE files SET mime = $2, content = $3, date_updated = NOW() \
                              WHERE path = $1",
                             &[&path, &mime, &content]));
        }
    }

    try!(trans.commit());

    file_list_to_json(file_list_with_mimes)
}

fn file_list_to_json(file_list: Vec<(String, String)>) -> Result<Json, DocBuilderError> {

    let mut file_list_json: Vec<Json> = Vec::new();

    for file in file_list {
        let mut v: Vec<String> = Vec::new();
        v.push(file.0.clone());
        v.push(file.1.clone());
        file_list_json.push(v.to_json());
    }

    Ok(file_list_json.to_json())
}

#[cfg(test)]
mod test {
    extern crate env_logger;
    use std::env;
    use super::{get_file_list, add_path_into_database};
    use super::super::connect_db;

    #[test]
    fn test_get_file_list() {
        let _ = env_logger::init();

        let files = get_file_list(env::current_dir().unwrap());
        assert!(files.is_ok());
        assert!(files.unwrap().len() > 0);

        let files = get_file_list(env::current_dir().unwrap().join("Cargo.toml")).unwrap();
        assert_eq!(files[0], "Cargo.toml");
    }

    #[test]
    #[ignore]
    fn test_add_path_into_database() {
        let _ = env_logger::init();

        let conn = connect_db().unwrap();
        let res = add_path_into_database(&conn, "example", env::current_dir().unwrap().join("src"));
        assert!(res.is_ok());
    }
}
<|endoftext|>


<commit_before>
// The MIT License (MIT)
//
// Copyright (c) 2016 Marvin Böcker
//
// 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.

//! This module contains a wrapper around the libsodium implementation of ed25519.
//! It reduces the size of signatures to 64 byte.

use std::sync::Once;
use std::sync::ONCE_INIT;

static START: Once = ONCE_INIT;

use sodiumoxide;
use sodiumoxide::crypto::sign::ed25519;

/// This is the length of a ed25519 private key.
pub const PRIVATE_KEY_LEN: usize = 64;

/// This is the length of a ed25519 public key.
pub const PUBLIC_KEY_LEN: usize = 32;

/// This is the length of a ed25519 signature.
pub const SIGNATURE_LEN: usize = 64;

/// This method generates a random ed25519 keypair from a cryptographically secure source
/// (on unix this is /dev/urandom). Returns (`public_key`, `private_key`).
pub fn generate_keypair() -> ([u8; PUBLIC_KEY_LEN], [u8; PRIVATE_KEY_LEN]) {

    // Initialize the random number generator provided by libsodium.
    START.call_once(|| { sodiumoxide::init(); });

    let (pk, sk) = ed25519::gen_keypair();

    let public_key = pk.0;
    let private_key = sk.0;

    assert_eq!(PUBLIC_KEY_LEN, public_key.len());
    assert_eq!(PRIVATE_KEY_LEN, private_key.len());

    (public_key, private_key)
}

/// This method takes a data vector and a private key and computes the signature which can be
/// verified using the public key.
pub fn sign(data: &[u8], private_key: &[u8]) -> Vec<u8> {
    assert_eq!(private_key.len(), PRIVATE_KEY_LEN);

    let sk = ed25519::SecretKey::from_slice(private_key);
    let sk = sk.as_ref().unwrap();

    let s = ed25519::sign(data, sk);

    let mut v = Vec::new();
    v.extend_from_slice(&s[0..64]);
    v
}

/// This method takes a data vector, a signature and a public key and returns true, if the
/// signature has been created using the correct private key.
pub fn verify(data: &[u8], signature: &[u8], public_key: &[u8]) -> bool {

    let pk = ed25519::PublicKey::from_slice(public_key);
    let pk = pk.as_ref().unwrap();

    // allocate buffer on heap for signature and data.
    let mut vi = Vec::with_capacity(SIGNATURE_LEN + data.len());
    vi.extend_from_slice(signature);
    vi.extend_from_slice(data);

    let r = ed25519::verify(&vi, pk);

    match r {
        Err(_) => false,
        Ok(x) => x == data,
    }
}

#[test]
fn test_ed25519_simple() {
    let (pk, sk) = generate_keypair();

    let msg = &[0; 128][..];
    let mut sig = sign(msg, &sk);

    println!("signature: {:?}", sig);
    println!("");

    sig[0] = ((sig[0] as u16 + 1) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), false);

    sig[0] = ((sig[0] as u16 + 255) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), true);
}

#[test]
fn test_ed25519_shortmsg() {
    let (pk, sk) = generate_keypair();

    let msg = &[0; 32][..];
    let mut sig = sign(msg, &sk);

    println!("signature: {:?}", sig);
    println!("");

    sig[0] = ((sig[0] as u16 + 1) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), false);

    sig[0] = ((sig[0] as u16 + 255) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), true);
}

#[test]
fn test_testvectors() {
    use rustc_serialize::json;
    use bytescontainer::BytesContainer;

    let decode = |t| {
        let t = json::decode::<BytesContainer>(&format!("\"{}\"", t)).unwrap();
        t.get().clone()
    };

    // test vectors from
    // https://github.com/cryptosphere/rbnacl/blob/master/lib/rbnacl/test_vectors.rb

    let prv = decode("b18e1d0045995ec3d010c387ccfeb984d783af8fbb0f40fa7db126d889f6dadd77f48b59cae\
                      da77751ed138b0ec667ff50f8768c25d48309a8f386a2bad187fb");

    let pbl = decode("77f48b59caeda77751ed138b0ec667ff50f8768c25d48309a8f386a2bad187fb");

    let msg =
        decode("916c7d1d268fc0e77c1bef238432573c39be577bbea0998936add2b50a653171ce18a542b0b7f96c1691a3be6031522894a8634183eda38798a0c5d5d79fbd01dd04a8646d71873b77b221998a81922d8105f892316369d5224c9983372d2313c6b1f4556ea26ba49d46e8b561e0fc76633ac9766e68e21fba7edca93c4c7460376d7f3ac22ff372c18f613f2ae2e856af40");

    let sig = decode("6bd710a368c1249923fc7a1610747403040f0cc30815a00f9ff548a896bbda0b4eb2ca19ebc\
                      f917f0f34200a9edbad3901b64ab09cc5ef7b9bcc3c40c0ff7509");

    let sig_my = sign(&msg, &prv);

    // both signatures must be equal
    assert_eq!(&sig_my, &sig);

    // the signature must be valid
    assert_eq!(true, verify(&msg, &sig_my, &pbl));
}
<commit_msg>removed one heap allocation<commit_after>
// The MIT License (MIT)
//
// Copyright (c) 2016 Marvin Böcker
//
// 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.

//! This module contains a wrapper around the libsodium implementation of ed25519.
//! It reduces the size of signatures to 64 byte.

use std::sync::Once;
use std::sync::ONCE_INIT;

static START: Once = ONCE_INIT;

use sodiumoxide;
use sodiumoxide::crypto::sign::ed25519;

/// This is the length of a ed25519 private key.
pub const PRIVATE_KEY_LEN: usize = 64;

/// This is the length of a ed25519 public key.
pub const PUBLIC_KEY_LEN: usize = 32;

/// This is the length of a ed25519 signature.
pub const SIGNATURE_LEN: usize = 64;

/// This method generates a random ed25519 keypair from a cryptographically secure source
/// (on unix this is /dev/urandom). Returns (`public_key`, `private_key`).
pub fn generate_keypair() -> ([u8; PUBLIC_KEY_LEN], [u8; PRIVATE_KEY_LEN]) {

    // Initialize the random number generator provided by libsodium.
    START.call_once(|| { sodiumoxide::init(); });

    let (pk, sk) = ed25519::gen_keypair();

    let public_key = pk.0;
    let private_key = sk.0;

    assert_eq!(PUBLIC_KEY_LEN, public_key.len());
    assert_eq!(PRIVATE_KEY_LEN, private_key.len());

    (public_key, private_key)
}

/// This method takes a data vector and a private key and computes the signature which can be
/// verified using the public key.
pub fn sign(data: &[u8], private_key: &[u8]) -> Vec<u8> {
    assert_eq!(private_key.len(), PRIVATE_KEY_LEN);

    use sodiumoxide::crypto::sign::ed25519::Signature;

    let sk = ed25519::SecretKey::from_slice(private_key);
    let sk = sk.as_ref().unwrap();

    let sig = match ed25519::sign_detached(data, sk) {
        Signature(signature) => signature
    };

    let mut vec = Vec::new();
    vec.extend_from_slice(&sig);
    vec
}

/// This method takes a data vector, a signature and a public key and returns true, if the
/// signature has been created using the correct private key.
pub fn verify(data: &[u8], signature: &[u8], public_key: &[u8]) -> bool {
    assert_eq!(signature.len(), SIGNATURE_LEN);
    assert_eq!(public_key.len(), PUBLIC_KEY_LEN);

    let pk = ed25519::PublicKey::from_slice(public_key);
    let pk = pk.as_ref().unwrap();

    let sig = ed25519::Signature::from_slice(signature);
    let sig = sig.as_ref().unwrap();

    ed25519::verify_detached(sig, data, pk)
}

#[test]
fn test_ed25519_simple() {
    let (pk, sk) = generate_keypair();

    let msg = &[0; 128][..];
    let mut sig = sign(msg, &sk);

    println!("signature: {:?}", sig);
    println!("");

    sig[0] = ((sig[0] as u16 + 1) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), false);

    sig[0] = ((sig[0] as u16 + 255) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), true);
}

#[test]
fn test_ed25519_shortmsg() {
    let (pk, sk) = generate_keypair();

    let msg = &[0; 32][..];
    let mut sig = sign(msg, &sk);

    println!("signature: {:?}", sig);
    println!("");

    sig[0] = ((sig[0] as u16 + 1) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), false);

    sig[0] = ((sig[0] as u16 + 255) % 256) as u8;

    assert_eq!(verify(msg, &sig, &pk), true);
}

#[test]
fn test_testvectors() {
    use rustc_serialize::json;
    use bytescontainer::BytesContainer;

    let decode = |t| {
        let t = json::decode::<BytesContainer>(&format!("\"{}\"", t)).unwrap();
        t.get().clone()
    };

    // test vectors from
    // https://github.com/cryptosphere/rbnacl/blob/master/lib/rbnacl/test_vectors.rb

    let prv = decode("b18e1d0045995ec3d010c387ccfeb984d783af8fbb0f40fa7db126d889f6dadd77f48b59cae\
                      da77751ed138b0ec667ff50f8768c25d48309a8f386a2bad187fb");

    let pbl = decode("77f48b59caeda77751ed138b0ec667ff50f8768c25d48309a8f386a2bad187fb");

    let msg =
        decode("916c7d1d268fc0e77c1bef238432573c39be577bbea0998936add2b50a653171ce18a542b0b7f96c1691a3be6031522894a8634183eda38798a0c5d5d79fbd01dd04a8646d71873b77b221998a81922d8105f892316369d5224c9983372d2313c6b1f4556ea26ba49d46e8b561e0fc76633ac9766e68e21fba7edca93c4c7460376d7f3ac22ff372c18f613f2ae2e856af40");

    let sig = decode("6bd710a368c1249923fc7a1610747403040f0cc30815a00f9ff548a896bbda0b4eb2ca19ebc\
                      f917f0f34200a9edbad3901b64ab09cc5ef7b9bcc3c40c0ff7509");

    let sig_my = sign(&msg, &prv);

    // both signatures must be equal
    assert_eq!(&sig_my, &sig);

    // the signature must be valid
    assert_eq!(true, verify(&msg, &sig_my, &pbl));
}
<|endoftext|>


<commit_before>use interpreter::*;

pub enum Program {
    IntLiteral(u8),
    Binary {op: u8, left: Box<Program>, right: Box<Program>}
}

pub fn emit(program: &Program) -> Vec<u8> {
    let mut code = vec![];
    emit_into(&mut code, program, 0);
    code.push(OP_RETURN);
    code.push(0);
    code
}

fn emit_into(code: &mut Vec<u8>, program: &Program, target: u8) {
    match *program {
        Program::IntLiteral(n) => {
            load_u8(code, n, target)
        },
        Program::Binary {op, ref left, ref right} => {
            emit_into(code, &*left, 1);
            emit_into(code, &*right, 2);
            code.push(op);
            code.push(1);
            code.push(2);
            code.push(target);
        }
    }
}

fn load_u8(code: &mut Vec<u8>, value: u8, dest: u8) {
    code.push(OP_LOAD_U8);
    code.push(value);
    code.push(dest);
}

fn int_literal(n: u8) -> Box<Program> {
    Box::new(Program::IntLiteral(n))
}

fn binary_add(a: u8, b: u8) -> Box<Program> {
    let program = Program::Binary{
        op: OP_ADD,
        left: int_literal(a),
        right: int_literal(b)
    };
    Box::new(program)
}

#[test]
fn test_emit_int_literal() {
    let code1 = emit(&Program::IntLiteral(37));
    assert_eq!(37, interpret(&code1));
}

#[test]
fn test_emit_simple_binary_op() {
    let program = Program::Binary {
        op: OP_ADD,
        left: int_literal(73),
        right: int_literal(68)
    };

    let code = emit(&program);
    assert_eq!(141, interpret(&code));
}

#[test]
fn test_emit_nested_binary_op() {
    // 73 + 68 + 50 = 191
    let program = Program::Binary {
        op: OP_ADD,
        left: binary_add(73, 68),
        right: int_literal(50)
    };

    let code = emit(&program);
    assert_eq!(191, interpret(&code));
}
<commit_msg>Fix issues with clobbering temporary registers<commit_after>use interpreter::*;

pub enum Program {
    IntLiteral(u8),
    Binary {op: u8, left: Box<Program>, right: Box<Program>}
}

pub fn emit(program: &Program) -> Vec<u8> {
    let mut code = vec![];
    emit_into(&mut code, program, 0);
    code.push(OP_RETURN);
    code.push(0);
    code
}

fn emit_into(code: &mut Vec<u8>, program: &Program, target: u8) {
    match *program {
        Program::IntLiteral(n) => {
            load_u8(code, n, target)
        },
        Program::Binary {op, ref left, ref right} => {
            emit_into(code, &*left, target);
            emit_into(code, &*right, target + 1);
            code.push(op);
            code.push(target);
            code.push(target + 1);
            code.push(target);
        }
    }
}

fn load_u8(code: &mut Vec<u8>, value: u8, dest: u8) {
    code.push(OP_LOAD_U8);
    code.push(value);
    code.push(dest);
}

fn int_literal(n: u8) -> Box<Program> {
    Box::new(Program::IntLiteral(n))
}

fn binary_add(a: u8, b: u8) -> Box<Program> {
    let program = Program::Binary{
        op: OP_ADD,
        left: int_literal(a),
        right: int_literal(b)
    };
    Box::new(program)
}

#[test]
fn test_emit_int_literal() {
    let code1 = emit(&Program::IntLiteral(37));
    assert_eq!(37, interpret(&code1));
}

#[test]
fn test_emit_simple_binary_op() {
    let program = Program::Binary {
        op: OP_ADD,
        left: int_literal(73),
        right: int_literal(68)
    };

    let code = emit(&program);
    assert_eq!(141, interpret(&code));
}

#[test]
fn test_emit_nested_binary_op() {
    // (73 + 68) + 50 = 191
    let program0 = Program::Binary {
        op: OP_ADD,
        left: binary_add(73, 68),
        right: int_literal(50)
    };

    let code0 = emit(&program0);
    assert_eq!(191, interpret(&code0));

    // 1 + (2 + 3) = 6
    let program1 = Program::Binary {
        op: OP_ADD,
        left: int_literal(1),
        right: binary_add(2, 3)
    };

    let code1 = emit(&program1);
    assert_eq!(6, interpret(&code1));
}

<|endoftext|>


<commit_before>//! Integer generic traits and operations
//!
//! These traits and functions are for use in generic algorithms
//! that work indifferently to particular types, relying on
//! some common traits.

use super::{
    CheckedAdd, CheckedDiv, CheckedMul, CheckedSub, /*Zero, One, */ Integer, ToPrimitive,
};
use std::cmp::PartialOrd;
use Sign;

#[cfg(feature = "with-bigint")]
use super::{BigInt, BigUint, One, Signed, Zero};

use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Rem, RemAssign, Sub, SubAssign};

/// Methods common to all integer types that
/// could be used generically in abstract algorithms
pub trait GenericInteger:
    'static
    + Sized
    + Integer
    + ToPrimitive
    + CheckedAdd
    + CheckedDiv
    + CheckedMul
    + CheckedSub
    + Add
    + Div
    + Mul
    + Rem
    + Sub
    + AddAssign
    + DivAssign
    + MulAssign
    + RemAssign
    + SubAssign
    + for<'a> Add<&'a Self, Output = Self>
    + for<'a> Div<&'a Self, Output = Self>
    + for<'a> Mul<&'a Self, Output = Self>
    + for<'a> Rem<&'a Self, Output = Self>
    + for<'a> Sub<&'a Self, Output = Self>
    + for<'a> AddAssign<&'a Self>
    + for<'a> DivAssign<&'a Self>
    + for<'a> MulAssign<&'a Self>
    + for<'a> RemAssign<&'a Self>
    + for<'a> SubAssign<&'a Self>
{
    /// Returns value 0 of the type
    fn _0() -> Self;

    /// Returns value 1 of the type
    fn _1() -> Self;

    /// Returns value 10 of the type
    fn _10() -> Self;

    /// Returns Maybe<static reference> of 0
    fn _0r() -> Option<&'static Self>;

    /// Returns Maybe<static reference> of 1
    fn _1r() -> Option<&'static Self>;

    /// Returns Maybe<static reference> of 10
    fn _10r() -> Option<&'static Self>;

    /// Returns the sign and the value itself.
    /// Zero values must have [Sign::Plus]
    fn get_signed_value(self) -> (Sign, Self);
}

#[cfg(feature = "with-bigint")]
lazy_static! {
    static ref _0_BU: BigUint = { BigUint::zero() };
    static ref _1_BU: BigUint = { BigUint::one() };
    static ref _10_BU: BigUint = { BigUint::from(10u8) };
    static ref _0_BI: BigInt = { BigInt::zero() };
    static ref _1_BI: BigInt = { BigInt::one() };
    static ref _10_BI: BigInt = { BigInt::from(10i8) };
}

#[cfg(feature = "with-bigint")]
impl GenericInteger for BigUint {
    #[inline]
    fn _0() -> Self {
        BigUint::zero()
    }
    #[inline]
    fn _1() -> Self {
        BigUint::one()
    }
    #[inline]
    fn _10() -> Self {
        _10_BU.clone()
    }
    #[inline]
    fn _0r() -> Option<&'static Self> {
        Some(&_0_BU)
    }
    #[inline]
    fn _1r() -> Option<&'static Self> {
        Some(&_1_BU)
    }
    #[inline]
    fn _10r() -> Option<&'static Self> {
        Some(&_10_BU)
    }

    #[inline]
    fn get_signed_value(self) -> (Sign, Self) {
        (Sign::Plus, self)
    }
}

#[cfg(feature = "with-bigint")]
impl GenericInteger for BigInt {
    #[inline]
    fn _0() -> Self {
        BigInt::zero()
    }
    #[inline]
    fn _1() -> Self {
        BigInt::one()
    }
    #[inline]
    fn _10() -> Self {
        _10_BI.clone()
    }
    #[inline]
    fn _0r() -> Option<&'static Self> {
        Some(&_0_BI)
    }
    #[inline]
    fn _1r() -> Option<&'static Self> {
        Some(&_1_BI)
    }
    #[inline]
    fn _10r() -> Option<&'static Self> {
        Some(&_10_BI)
    }

    #[inline]
    fn get_signed_value(self) -> (Sign, Self) {
        (
            if self.is_negative() {
                Sign::Minus
            } else {
                Sign::Plus
            },
            self,
        )
    }
}

macro_rules! generic_integer_for_uint {
    ($($t:ty),*) => {
        $(
            impl GenericInteger for $t {
                #[inline]
                fn _0() -> Self { 0 }
                #[inline]
                fn _1() -> Self { 1 }
                #[inline]
                fn _10() -> Self { 10 }
                #[inline]
                fn _0r() -> Option<&'static Self> { None }
                #[inline]
                fn _1r() -> Option<&'static Self> { None }
                #[inline]
                fn _10r() -> Option<&'static Self> { None }

                #[inline]
                fn get_signed_value(self) -> (Sign, Self) { (Sign::Plus, self) }
            }
        )*
    };
}

generic_integer_for_uint!(u8, u16, u32, u64, u128, usize);

macro_rules! generic_integer_for_int {
    ($($t:ty),*) => {
        $(
            impl GenericInteger for $t {
                #[inline]
                fn _0() -> Self { 0 }
                #[inline]
                fn _1() -> Self { 1 }
                #[inline]
                fn _10() -> Self { 10 }
                #[inline]
                fn _0r() -> Option<&'static Self> { None }
                #[inline]
                fn _1r() -> Option<&'static Self> { None }
                #[inline]
                fn _10r() -> Option<&'static Self> { None }

                #[inline]
                fn get_signed_value(self) -> (Sign, Self) {
                    (if self.is_negative() { Sign::Minus } else { Sign::Plus }, self)
                }
            }
        )*
    };
}

generic_integer_for_int!(i8, i16, i32, i64, i128, isize);

/// Builds integer of type T from another integer of type F in a generic way
///
/// Allows safe runtime conversions between integer types when it's not possible
/// statically. E.g: `i8 -> u8`, `u8 -> i8`, `usize -> u8` or even `BigUint -> u8` and so on.
///
/// # Examples
///
/// ```
/// use fraction::{Sign, generic::read_generic_integer};
///
/// assert_eq!((Sign::Plus, 127i8), read_generic_integer(127u8).unwrap());
/// assert_eq!((Sign::Minus, 128u8), read_generic_integer(-128i8).unwrap());
/// assert_eq!((Sign::Minus, 255u8), read_generic_integer(-255isize).unwrap());
/// ```
pub fn read_generic_integer<T, F>(val: F) -> Option<(Sign, T)>
where
    F: GenericInteger + PartialOrd,
    T: GenericInteger,
{
    let (sign, mut val) = val.get_signed_value();

    let mut vptr: F = F::_1();
    let mut rptr: T = T::_1();
    let mut result: T = T::_0();

    loop {
        vptr = match F::_10r().map_or_else(
            || vptr.checked_mul(&GenericInteger::_10()),
            |_10| vptr.checked_mul(_10),
        ) {
            Some(v) => v,
            None => break,
        };

        let vdelta: F = val.checked_sub(&val.checked_div(&vptr)?.checked_mul(&vptr)?)?;

        let mut rdelta: T = T::_0();

        let mut vldelta: F = vdelta.checked_div(&F::_10r().map_or_else(
            || vptr.checked_div(&GenericInteger::_10()),
            |_10| vptr.checked_div(_10),
        )?)?;

        loop {
            if F::_0r().map_or_else(|| vldelta == GenericInteger::_0(), |v| vldelta == *v) {
                break;
            }

            rdelta = T::_1r()
                .map_or_else(|| rdelta.checked_add(&T::_1()), |_1| rdelta.checked_add(_1))?;

            vldelta = F::_1r().map_or_else(
                || {
                    if sign == Sign::Plus {
                        vldelta.checked_sub(&GenericInteger::_1())
                    } else {
                        vldelta.checked_add(&GenericInteger::_1())
                    }
                },
                |_1| {
                    if sign == Sign::Plus {
                        vldelta.checked_sub(_1)
                    } else {
                        vldelta.checked_add(_1)
                    }
                },
            )?;
        }

        result = result.checked_add(&rdelta.checked_mul(&rptr)?)?;
        val = val.checked_sub(&vdelta)?;

        if F::_0r().map_or_else(|| val == GenericInteger::_0(), |_0| val == *_0) {
            break;
        }

        rptr = T::_10r().map_or_else(
            || rptr.checked_mul(&GenericInteger::_10()),
            |_10| rptr.checked_mul(_10),
        )?;
    }

    if F::_0r().map_or_else(|| val != GenericInteger::_0(), |_0| val != *_0) {
        let mut vldelta: F = val.checked_div(&vptr)?;

        let mut rdelta: T = T::_0();

        loop {
            if F::_0r().map_or_else(|| vldelta == GenericInteger::_0(), |_0| vldelta == *_0) {
                break;
            }

            rdelta = T::_1r().map_or_else(
                || rdelta.checked_add(&GenericInteger::_1()),
                |_1| rdelta.checked_add(_1),
            )?;

            vldelta = F::_1r().map_or_else(
                || {
                    if sign == Sign::Plus {
                        vldelta.checked_sub(&GenericInteger::_1())
                    } else {
                        vldelta.checked_add(&GenericInteger::_1())
                    }
                },
                |_1| {
                    if sign == Sign::Plus {
                        vldelta.checked_sub(_1)
                    } else {
                        vldelta.checked_add(_1)
                    }
                },
            )?;
        }
        result = result.checked_add(&rdelta.checked_mul(&rptr)?)?;
    }

    Some((sign, result))
}

#[cfg(test)]
mod tests {
    // TODO: tests

    use super::*;

    #[test]
    fn max_to_max() {
        let (s, v) = read_generic_integer::<u32, u32>(u32::max_value()).unwrap();
        assert_eq!(s, Sign::Plus);
        assert_eq!(v, u32::max_value());
    }

    #[test]
    fn sign() {
        let (s, _) = read_generic_integer::<i8, i8>(0i8).unwrap();
        assert_eq!(s, Sign::Plus);

        let (s, _) = read_generic_integer::<i8, i8>(1i8).unwrap();
        assert_eq!(s, Sign::Plus);

        let (s, _) = read_generic_integer::<i8, i8>(-1i8).unwrap();
        assert_eq!(s, Sign::Minus);
    }
}
<commit_msg>Removed unnecessary braces.<commit_after>//! Integer generic traits and operations
//!
//! These traits and functions are for use in generic algorithms
//! that work indifferently to particular types, relying on
//! some common traits.

use super::{
    CheckedAdd, CheckedDiv, CheckedMul, CheckedSub, /*Zero, One, */ Integer, ToPrimitive,
};
use std::cmp::PartialOrd;
use Sign;

#[cfg(feature = "with-bigint")]
use super::{BigInt, BigUint, One, Signed, Zero};

use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Rem, RemAssign, Sub, SubAssign};

/// Methods common to all integer types that
/// could be used generically in abstract algorithms
pub trait GenericInteger:
    'static
    + Sized
    + Integer
    + ToPrimitive
    + CheckedAdd
    + CheckedDiv
    + CheckedMul
    + CheckedSub
    + Add
    + Div
    + Mul
    + Rem
    + Sub
    + AddAssign
    + DivAssign
    + MulAssign
    + RemAssign
    + SubAssign
    + for<'a> Add<&'a Self, Output = Self>
    + for<'a> Div<&'a Self, Output = Self>
    + for<'a> Mul<&'a Self, Output = Self>
    + for<'a> Rem<&'a Self, Output = Self>
    + for<'a> Sub<&'a Self, Output = Self>
    + for<'a> AddAssign<&'a Self>
    + for<'a> DivAssign<&'a Self>
    + for<'a> MulAssign<&'a Self>
    + for<'a> RemAssign<&'a Self>
    + for<'a> SubAssign<&'a Self>
{
    /// Returns value 0 of the type
    fn _0() -> Self;

    /// Returns value 1 of the type
    fn _1() -> Self;

    /// Returns value 10 of the type
    fn _10() -> Self;

    /// Returns Maybe<static reference> of 0
    fn _0r() -> Option<&'static Self>;

    /// Returns Maybe<static reference> of 1
    fn _1r() -> Option<&'static Self>;

    /// Returns Maybe<static reference> of 10
    fn _10r() -> Option<&'static Self>;

    /// Returns the sign and the value itself.
    /// Zero values must have [Sign::Plus]
    fn get_signed_value(self) -> (Sign, Self);
}

#[cfg(feature = "with-bigint")]
lazy_static! {
    static ref _0_BU: BigUint =  BigUint::zero() ;
    static ref _1_BU: BigUint =  BigUint::one() ;
    static ref _10_BU: BigUint =  BigUint::from(10u8) ;
    static ref _0_BI: BigInt =  BigInt::zero() ;
    static ref _1_BI: BigInt =  BigInt::one() ;
    static ref _10_BI: BigInt =  BigInt::from(10i8) ;
}

#[cfg(feature = "with-bigint")]
impl GenericInteger for BigUint {
    #[inline]
    fn _0() -> Self {
        BigUint::zero()
    }
    #[inline]
    fn _1() -> Self {
        BigUint::one()
    }
    #[inline]
    fn _10() -> Self {
        _10_BU.clone()
    }
    #[inline]
    fn _0r() -> Option<&'static Self> {
        Some(&_0_BU)
    }
    #[inline]
    fn _1r() -> Option<&'static Self> {
        Some(&_1_BU)
    }
    #[inline]
    fn _10r() -> Option<&'static Self> {
        Some(&_10_BU)
    }

    #[inline]
    fn get_signed_value(self) -> (Sign, Self) {
        (Sign::Plus, self)
    }
}

#[cfg(feature = "with-bigint")]
impl GenericInteger for BigInt {
    #[inline]
    fn _0() -> Self {
        BigInt::zero()
    }
    #[inline]
    fn _1() -> Self {
        BigInt::one()
    }
    #[inline]
    fn _10() -> Self {
        _10_BI.clone()
    }
    #[inline]
    fn _0r() -> Option<&'static Self> {
        Some(&_0_BI)
    }
    #[inline]
    fn _1r() -> Option<&'static Self> {
        Some(&_1_BI)
    }
    #[inline]
    fn _10r() -> Option<&'static Self> {
        Some(&_10_BI)
    }

    #[inline]
    fn get_signed_value(self) -> (Sign, Self) {
        (
            if self.is_negative() {
                Sign::Minus
            } else {
                Sign::Plus
            },
            self,
        )
    }
}

macro_rules! generic_integer_for_uint {
    ($($t:ty),*) => {
        $(
            impl GenericInteger for $t {
                #[inline]
                fn _0() -> Self { 0 }
                #[inline]
                fn _1() -> Self { 1 }
                #[inline]
                fn _10() -> Self { 10 }
                #[inline]
                fn _0r() -> Option<&'static Self> { None }
                #[inline]
                fn _1r() -> Option<&'static Self> { None }
                #[inline]
                fn _10r() -> Option<&'static Self> { None }

                #[inline]
                fn get_signed_value(self) -> (Sign, Self) { (Sign::Plus, self) }
            }
        )*
    };
}

generic_integer_for_uint!(u8, u16, u32, u64, u128, usize);

macro_rules! generic_integer_for_int {
    ($($t:ty),*) => {
        $(
            impl GenericInteger for $t {
                #[inline]
                fn _0() -> Self { 0 }
                #[inline]
                fn _1() -> Self { 1 }
                #[inline]
                fn _10() -> Self { 10 }
                #[inline]
                fn _0r() -> Option<&'static Self> { None }
                #[inline]
                fn _1r() -> Option<&'static Self> { None }
                #[inline]
                fn _10r() -> Option<&'static Self> { None }

                #[inline]
                fn get_signed_value(self) -> (Sign, Self) {
                    (if self.is_negative() { Sign::Minus } else { Sign::Plus }, self)
                }
            }
        )*
    };
}

generic_integer_for_int!(i8, i16, i32, i64, i128, isize);

/// Builds integer of type T from another integer of type F in a generic way
///
/// Allows safe runtime conversions between integer types when it's not possible
/// statically. E.g: `i8 -> u8`, `u8 -> i8`, `usize -> u8` or even `BigUint -> u8` and so on.
///
/// # Examples
///
/// ```
/// use fraction::{Sign, generic::read_generic_integer};
///
/// assert_eq!((Sign::Plus, 127i8), read_generic_integer(127u8).unwrap());
/// assert_eq!((Sign::Minus, 128u8), read_generic_integer(-128i8).unwrap());
/// assert_eq!((Sign::Minus, 255u8), read_generic_integer(-255isize).unwrap());
/// ```
pub fn read_generic_integer<T, F>(val: F) -> Option<(Sign, T)>
where
    F: GenericInteger + PartialOrd,
    T: GenericInteger,
{
    let (sign, mut val) = val.get_signed_value();

    let mut vptr: F = F::_1();
    let mut rptr: T = T::_1();
    let mut result: T = T::_0();

    loop {
        vptr = match F::_10r().map_or_else(
            || vptr.checked_mul(&GenericInteger::_10()),
            |_10| vptr.checked_mul(_10),
        ) {
            Some(v) => v,
            None => break,
        };

        let vdelta: F = val.checked_sub(&val.checked_div(&vptr)?.checked_mul(&vptr)?)?;

        let mut rdelta: T = T::_0();

        let mut vldelta: F = vdelta.checked_div(&F::_10r().map_or_else(
            || vptr.checked_div(&GenericInteger::_10()),
            |_10| vptr.checked_div(_10),
        )?)?;

        loop {
            if F::_0r().map_or_else(|| vldelta == GenericInteger::_0(), |v| vldelta == *v) {
                break;
            }

            rdelta = T::_1r()
                .map_or_else(|| rdelta.checked_add(&T::_1()), |_1| rdelta.checked_add(_1))?;

            vldelta = F::_1r().map_or_else(
                || {
                    if sign == Sign::Plus {
                        vldelta.checked_sub(&GenericInteger::_1())
                    } else {
                        vldelta.checked_add(&GenericInteger::_1())
                    }
                },
                |_1| {
                    if sign == Sign::Plus {
                        vldelta.checked_sub(_1)
                    } else {
                        vldelta.checked_add(_1)
                    }
                },
            )?;
        }

        result = result.checked_add(&rdelta.checked_mul(&rptr)?)?;
        val = val.checked_sub(&vdelta)?;

        if F::_0r().map_or_else(|| val == GenericInteger::_0(), |_0| val == *_0) {
            break;
        }

        rptr = T::_10r().map_or_else(
            || rptr.checked_mul(&GenericInteger::_10()),
            |_10| rptr.checked_mul(_10),
        )?;
    }

    if F::_0r().map_or_else(|| val != GenericInteger::_0(), |_0| val != *_0) {
        let mut vldelta: F = val.checked_div(&vptr)?;

        let mut rdelta: T = T::_0();

        loop {
            if F::_0r().map_or_else(|| vldelta == GenericInteger::_0(), |_0| vldelta == *_0) {
                break;
            }

            rdelta = T::_1r().map_or_else(
                || rdelta.checked_add(&GenericInteger::_1()),
                |_1| rdelta.checked_add(_1),
            )?;

            vldelta = F::_1r().map_or_else(
                || {
                    if sign == Sign::Plus {
                        vldelta.checked_sub(&GenericInteger::_1())
                    } else {
                        vldelta.checked_add(&GenericInteger::_1())
                    }
                },
                |_1| {
                    if sign == Sign::Plus {
                        vldelta.checked_sub(_1)
                    } else {
                        vldelta.checked_add(_1)
                    }
                },
            )?;
        }
        result = result.checked_add(&rdelta.checked_mul(&rptr)?)?;
    }

    Some((sign, result))
}

#[cfg(test)]
mod tests {
    // TODO: tests

    use super::*;

    #[test]
    fn max_to_max() {
        let (s, v) = read_generic_integer::<u32, u32>(u32::max_value()).unwrap();
        assert_eq!(s, Sign::Plus);
        assert_eq!(v, u32::max_value());
    }

    #[test]
    fn sign() {
        let (s, _) = read_generic_integer::<i8, i8>(0i8).unwrap();
        assert_eq!(s, Sign::Plus);

        let (s, _) = read_generic_integer::<i8, i8>(1i8).unwrap();
        assert_eq!(s, Sign::Plus);

        let (s, _) = read_generic_integer::<i8, i8>(-1i8).unwrap();
        assert_eq!(s, Sign::Minus);
    }
}
<|endoftext|>


<commit_before>extern crate ws;

use message::PeerMessage;

use std::cmp;
use std::collections::hash_map::HashMap;
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::thread;

use bincode::rustc_serialize::decode;

static LIB_NAME: &'static str = "ws-p2p";

/// Connection handler for incoming messages.
pub struct MessageHandler {
    /// Sender that is used to communicate for the handler.
    pub ws: ws::Sender,
    /// An arc clone of my local vector clocks.
    pub clocks: Arc<Mutex<HashMap<String, u32>>>,
    /// My address/name.
    pub me: String,
    /// Buffer of messages that caused conflicts when initially received.
    pub buffer: Arc<Mutex<VecDeque<PeerMessage>>>,
}

impl ws::Handler for MessageHandler {
    fn on_message(&mut self, msg: ws::Message) -> ws::Result<()> {
        match msg {
            ws::Message::Binary(vector) => {
                let encoded_msg = &*vector.into_boxed_slice();
                let message: PeerMessage = decode(encoded_msg).unwrap();

                info!(target: LIB_NAME, "Peer {} with clocks: {:?} got message: {}",
                        message.sender, message.clocks, message.message);
                if message.sender.as_str() == "127.0.0.1:3013" {
                    info!(target: LIB_NAME, "Faking delay!");
                    thread::sleep(Duration::from_millis(4000))
                }
                //message_checking(&mut clocks, message.clone());
				message_handler(message, self.clocks.clone());
                Ok(())
            },
            ws::Message::Text(string) => {
                info!(target: LIB_NAME, "Received peer's name. Adding {} to the client list", string);
                let mut clocks = self.clocks.lock().unwrap();
                clocks.insert(string, 0u32);
                Ok(())
            }
        }
    }
    fn on_close(&mut self, code: ws::CloseCode, reason: &str) {
        //TODO: Remove disconnected client's clock from our vclock copy.
        if reason.is_empty() {
            //This works: CloseCode::Abnormal
            info!(target: LIB_NAME, "Client disconnected with code: {:?}", code);
        } else {
            info!(target: LIB_NAME, "Client disconnected with code: {:?} and reason: {}", code, reason);
        }
    }
    fn on_error(&mut self, err: ws::Error) {
        warn!(target: LIB_NAME, "Error family robinson! {:?}", err.kind);
    }
}

pub struct MessageFactory {
    vclocks: Arc<Mutex<HashMap<String, u32>>>,
    me: String,
    buffer: Arc<Mutex<VecDeque<PeerMessage>>>,
}

impl MessageFactory {
    pub fn build(vclocks: Arc<Mutex<HashMap<String, u32>>>) -> MessageFactory {
        MessageFactory {
            vclocks: vclocks,
            me: String::from("undefined"),
            buffer: Arc::new(Mutex::new(VecDeque::new())),
        }
    }
    pub fn me(&self, me: &str) -> MessageFactory {
        MessageFactory {
            vclocks: self.vclocks.clone(),
            me: String::from(me),
            buffer: self.buffer.clone(),
        }
    }
}

impl ws::Factory for MessageFactory {
    type Handler = MessageHandler;

    fn connection_made(&mut self, ws: ws::Sender) -> Self::Handler {
        MessageHandler {
            ws: ws,
            clocks: self.vclocks.clone(),
            me: self.me.clone(),
            buffer: self.buffer.clone(),
        }
    }

    fn server_connected(&mut self, ws: ws::Sender) -> Self::Handler {
        ws.send(self.me.clone()).unwrap();
        MessageHandler {
            ws: ws,
            clocks: self.vclocks.clone(),
            me: self.me.clone(),
            buffer: self.buffer.clone(),
        }
    }

    fn client_connected(&mut self, ws: ws::Sender) -> Self::Handler {
        ws.send(self.me.clone()).unwrap();
        MessageHandler {
            ws: ws,
            clocks: self.vclocks.clone(),
            me: self.me.clone(),
            buffer: self.buffer.clone(),
        }
    }
}

fn message_handler(message: PeerMessage, vclocks: Arc<Mutex<HashMap<String, u32>>>) {
    let mut vclocks = vclocks.lock().unwrap();
    for (key, val) in message.clocks.iter() {
        if let Some(lval) = vclocks.get(&key.clone()) {
            debug!(target: LIB_NAME, "key: {} incoming val: {} and lval: {}", key, val, lval);
            if lval <= val {
                debug!(target: LIB_NAME, "val <= lval");
                continue;
            } else { // There's a clock that is greater than what we have
                // Push to local buffer
                // return
                debug!(target: LIB_NAME, "clock discrepency");
                ()
            }
        }
    }

    // Update clocks
    for (key, val) in message.clocks.iter() {
        let lval = vclocks.get(&key.clone()).unwrap().clone();
        let max_val = cmp::max(val, &lval);
        vclocks.insert(key.clone(), *max_val);
    }
//	macro_rules! hashmap {
//		($( $key: expr => $val: expr ),*) => {{
//			 let mut map = ::std::collections::HashMap::new();
//			 $( map.insert($key, $val); )*
//			 map
//		}}
//	}
    //let map = hashmap!['a' => 0, 'b' => 1];
}

<commit_msg>Attempt at buffer_check + move msg_handler to self<commit_after>extern crate ws;

use message::PeerMessage;

use std::cmp;
use std::collections::hash_map::HashMap;
use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::thread;

use bincode::rustc_serialize::decode;

static LIB_NAME: &'static str = "ws-p2p";

/// Connection handler for incoming messages.
pub struct MessageHandler {
    /// Sender that is used to communicate for the handler.
    pub ws: ws::Sender,
    /// An arc clone of my local vector clocks.
    pub clocks: Arc<Mutex<HashMap<String, u32>>>,
    /// My address/name.
    pub me: String,
    /// Buffer of messages that caused conflicts when initially received.
    pub buffer: Arc<Mutex<VecDeque<PeerMessage>>>,
}

impl ws::Handler for MessageHandler {
    fn on_message(&mut self, msg: ws::Message) -> ws::Result<()> {
        match msg {
            ws::Message::Binary(vector) => {
                let encoded_msg = &*vector.into_boxed_slice();
                let message: PeerMessage = decode(encoded_msg).unwrap();

                info!(target: LIB_NAME, "Peer {} with clocks: {:?} got message: {}",
                        message.sender, message.clocks, message.message);
                if message.sender.as_str() == "127.0.0.1:3013" {
                    info!(target: LIB_NAME, "Faking delay!");
                    thread::sleep(Duration::from_millis(4000))
                }
                //message_checking(&mut clocks, message.clone());
				self.message_handler(message);
                self.buffer_check();
                Ok(())
            },
            ws::Message::Text(string) => {
                info!(target: LIB_NAME, "Received peer's name. Adding {} to the client list", string);
                let mut clocks = self.clocks.lock().unwrap();
                clocks.insert(string, 0u32);
                Ok(())
            }
        }
    }
    fn on_close(&mut self, code: ws::CloseCode, reason: &str) {
        //TODO: Remove disconnected client's clock from our vclock copy.
        if reason.is_empty() {
            //This works: CloseCode::Abnormal
            info!(target: LIB_NAME, "Client disconnected with code: {:?}", code);
        } else {
            info!(target: LIB_NAME, "Client disconnected with code: {:?} and reason: {}", code, reason);
        }
    }
    fn on_error(&mut self, err: ws::Error) {
        warn!(target: LIB_NAME, "Error family robinson! {:?}", err.kind);
    }
}

impl MessageHandler {
    fn buffer_check(&self) {
        debug!("trying to get lock");
        let mut buffer = self.buffer.lock().unwrap();
        let buffer_clone = buffer.clone();
        for message in buffer_clone.iter() {
            debug!("Are we getting here?");
            let mut vclocks = self.clocks.lock().unwrap();
            for (key, val) in message.clocks.iter() {
                if let Some(lval) = vclocks.get(&key.clone()) {
                    debug!(target: LIB_NAME, "key: {} incoming val: {} and lval: {}", key, val, lval);
                    if lval <= val {
                        debug!(target: LIB_NAME, "val <= lval");
                        continue;
                    } else { // There's a clock that is greater than what we have
                        debug!(target: LIB_NAME, "discrepency still exists");
                        ()
                    }
                }
            }

            // Update clocks
            for (key, val) in message.clocks.iter() {
                let lval = vclocks.get(&key.clone()).unwrap().clone();
                let max_val = cmp::max(val, &lval);
                vclocks.insert(key.clone(), *max_val);
                info!(target: LIB_NAME, "Peer {} with clocks: {:?} got message: {}",
                        message.sender, message.clocks, message.message);
            }
            buffer.pop_front();
        }
    }
    fn message_handler(&self, message: PeerMessage) {
        let mut vclocks = self.clocks.lock().unwrap();
        for (key, val) in message.clocks.iter() {
            if let Some(lval) = vclocks.get(&key.clone()) {
                debug!(target: LIB_NAME, "key: {} incoming val: {} and lval: {}", key, val, lval);
                if lval <= val {
                    debug!(target: LIB_NAME, "val <= lval");
                    continue;
                } else { // There's a clock that is greater than what we have
                    // Push to local buffer
                    let mut buffer = self.buffer.lock().unwrap();
                    // return
                    debug!(target: LIB_NAME, "clock discrepency");
                    //self.buffer_check();
                    ()
                }
            }
        }

        // Update clocks
        for (key, val) in message.clocks.iter() {
            let lval = vclocks.get(&key.clone()).unwrap().clone();
            let max_val = cmp::max(val, &lval);
            vclocks.insert(key.clone(), *max_val);
        }
    //	macro_rules! hashmap {
    //		($( $key: expr => $val: expr ),*) => {{
    //			 let mut map = ::std::collections::HashMap::new();
    //			 $( map.insert($key, $val); )*
    //			 map
    //		}}
    //	}
        //let map = hashmap!['a' => 0, 'b' => 1];
    }
}

pub struct MessageFactory {
    vclocks: Arc<Mutex<HashMap<String, u32>>>,
    me: String,
    buffer: Arc<Mutex<VecDeque<PeerMessage>>>,
}

impl MessageFactory {
    pub fn build(vclocks: Arc<Mutex<HashMap<String, u32>>>) -> MessageFactory {
        MessageFactory {
            vclocks: vclocks,
            me: String::from("undefined"),
            buffer: Arc::new(Mutex::new(VecDeque::new())),
        }
    }
    pub fn me(&self, me: &str) -> MessageFactory {
        MessageFactory {
            vclocks: self.vclocks.clone(),
            me: String::from(me),
            buffer: self.buffer.clone(),
        }
    }
}

impl ws::Factory for MessageFactory {
    type Handler = MessageHandler;

    fn connection_made(&mut self, ws: ws::Sender) -> Self::Handler {
        MessageHandler {
            ws: ws,
            clocks: self.vclocks.clone(),
            me: self.me.clone(),
            buffer: self.buffer.clone(),
        }
    }

    fn server_connected(&mut self, ws: ws::Sender) -> Self::Handler {
        ws.send(self.me.clone()).unwrap();
        MessageHandler {
            ws: ws,
            clocks: self.vclocks.clone(),
            me: self.me.clone(),
            buffer: self.buffer.clone(),
        }
    }

    fn client_connected(&mut self, ws: ws::Sender) -> Self::Handler {
        ws.send(self.me.clone()).unwrap();
        MessageHandler {
            ws: ws,
            clocks: self.vclocks.clone(),
            me: self.me.clone(),
            buffer: self.buffer.clone(),
        }
    }
}

<|endoftext|>