<commit_before>// Copyright (c) 2016 Anatoly Ikorsky
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use ::futures::{
    task::{self, Task},
    Async::{self, NotReady, Ready},
    Future,
};

use std::{
    fmt,
    str::FromStr,
    sync::{Arc, Mutex, MutexGuard},
};

use crate::{
    conn::{pool::futures::*, Conn},
    error::*,
    opts::{Opts, PoolConstraints},
    queryable::{
        transaction::{Transaction, TransactionOptions},
        Queryable,
    },
    BoxFuture, MyFuture,
};

pub mod futures;

pub struct Inner {
    closed: bool,
    new: Vec<BoxFuture<Conn>>,
    queue: Vec<BoxFuture<Conn>>,
    idle: Vec<Conn>,
    ongoing: usize,
    tasks: Vec<Task>,
}

impl Inner {
    fn conn_count(&self) -> usize {
        self.new.len() + self.idle.len() + self.queue.len() + self.ongoing
    }
}

#[derive(Clone)]
/// Asynchronous pool of MySql connections.
pub struct Pool {
    opts: Opts,
    inner: Arc<Mutex<Inner>>,
    pool_constraints: PoolConstraints,
}

impl fmt::Debug for Pool {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let (new_len, idle_len, queue_len, ongoing, tasks_len) = self.with_inner(|inner| {
            (
                inner.new.len(),
                inner.idle.len(),
                inner.queue.len(),
                inner.ongoing,
                inner.tasks.len(),
            )
        });
        f.debug_struct("Pool")
            .field("pool_constraints", &self.pool_constraints)
            .field("new connections count", &new_len)
            .field("idle connections count", &idle_len)
            .field("queue length", &queue_len)
            .field("ongoing connections count", &ongoing)
            .field("tasks count", &tasks_len)
            .finish()
    }
}

impl Pool {
    /// Creates new pool of connections.
    pub fn new<O: Into<Opts>>(opts: O) -> Pool {
        let opts = opts.into();
        let pool_constraints = opts.get_pool_constraints().clone();
        let pool = Pool {
            opts,
            inner: Arc::new(Mutex::new(Inner {
                closed: false,
                new: Vec::with_capacity(pool_constraints.min()),
                idle: Vec::new(),
                queue: Vec::new(),
                ongoing: 0,
                tasks: Vec::new(),
            })),
            pool_constraints,
        };

        pool
    }

    /// Creates new pool of connections.
    pub fn from_url<T: AsRef<str>>(url: T) -> Result<Pool> {
        let opts = Opts::from_str(url.as_ref())?;
        Ok(Pool::new(opts))
    }

    /// Returns future that resolves to `Conn`.
    pub fn get_conn(&self) -> GetConn {
        new_get_conn(self)
    }

    /// Shortcut for `get_conn` followed by `start_transaction`.
    pub fn start_transaction(
        &self,
        options: TransactionOptions,
    ) -> impl MyFuture<Transaction<Conn>> {
        self.get_conn()
            .and_then(|conn| Queryable::start_transaction(conn, options))
    }

    /// Returns future that disconnects this pool from server and resolves to `()`.
    ///
    /// Active connections taken from this pool should be disconnected manually.
    /// Also all pending and new `GetConn`'s will resolve to error.
    pub fn disconnect(mut self) -> DisconnectPool {
        let become_closed = self.with_inner(|mut inner| {
            if !inner.closed {
                inner.closed = true;
                return true;
            } else {
                return false;
            }
        });
        if become_closed {
            while let Some(conn) = self.take_conn() {
                crate::conn::disconnect(conn);
            }
        }
        new_disconnect_pool(self)
    }

    /// Returns true if futures is in queue.
    fn in_queue(&self) -> bool {
        self.with_inner(|inner| {
            let count = inner.new.len() + inner.queue.len();
            count > 0
        })
    }

    /// A way to take connection from a pool.
    fn take_conn(&mut self) -> Option<Conn> {
        self.with_inner(|mut inner| {
            while let Some(mut conn) = inner.idle.pop() {
                if conn.expired() {
                    crate::conn::disconnect(conn);
                } else {
                    conn.inner.pool = Some(self.clone());
                    inner.ongoing += 1;
                    return Some(conn);
                }
            }
            None
        })
    }

    /// A way to return connection taken from a pool.
    fn return_conn(&mut self, conn: Conn) {
        let min = self.pool_constraints.min();

        self.with_inner(|mut inner| {
            inner.ongoing -= 1;

            if inner.closed {
                return;
            }

            if conn.inner.stream.is_none() {
                // drop incomplete connection
                return;
            }

            if conn.inner.in_transaction || conn.inner.has_result.is_some() {
                inner.queue.push(conn.cleanup());
            } else {
                if inner.idle.len() >= min {
                    crate::conn::disconnect(conn);
                } else {
                    inner.idle.push(conn);
                }
            }

            while let Some(task) = inner.tasks.pop() {
                task.notify()
            }
        });
    }

    fn with_inner<F, T>(&self, fun: F) -> T
    where
        F: FnOnce(MutexGuard<'_, Inner>) -> T,
        T: 'static,
    {
        fun(self.inner.lock().unwrap())
    }

    /// Will manage lifetime of futures stored in a pool.
    fn handle_futures(&mut self) -> Result<()> {
        if !self.in_queue() {
            // There is no futures in queue
            return Ok(());
        }

        let mut handled = false;

        let mut returned_conns: Vec<Conn> = vec![];

        self.with_inner(|mut inner| {
            macro_rules! handle {
                ($vec:ident { $($p:pat => $b:block,)+ }) => ({
                    let len = inner.$vec.len();
                    let mut done_fut_idxs = Vec::new();
                    for i in 0..len {
                        let result = inner.$vec.get_mut(i).unwrap().poll();
                        match result {
                            Ok(Ready(_)) | Err(_) => done_fut_idxs.push(i),
                            _ => (),
                        }

                        let out: Result<()> = match result {
                            $($p => $b),+
                            _ => {
                                Ok(())
                            }
                        };

                        match out {
                            Err(err) => {
                                // early return in case of error
                                while let Some(i) = done_fut_idxs.pop() {
                                    inner.$vec.swap_remove(i);
                                }
                                return Err(err)
                            }
                            _ => (),
                        }
                    }

                    while let Some(i) = done_fut_idxs.pop() {
                        inner.$vec.swap_remove(i);
                    }
                });
            }

            // Handle dirty connections.
            handle!(queue {
                Ok(Ready(conn)) => {
                    if inner.closed {
                crate::conn::disconnect(conn);
                    } else {
                        inner.ongoing += 1;
                        returned_conns.push(conn);
                    }
                    handled = true;
                    Ok(())
                },
                Err(_) => { Ok(()) },
            });

            // Handle connecting connections.
            handle!(new {
                Ok(Ready(conn)) => {
                    if inner.closed {
                        crate::conn::disconnect(conn);
                    } else {
                        inner.ongoing += 1;
                        returned_conns.push(conn);
                    }
                    handled = true;
                    Ok(())
                },
                Err(err) => {
                    if !inner.closed {
                        Err(err)
                    } else {
                        Ok(())
                    }
                },
            });

            Ok(())
        })?;

        for conn in returned_conns {
            self.return_conn(conn);
        }

        if handled {
            self.handle_futures()
        } else {
            Ok(())
        }
    }

    /// Will poll pool for connection.
    fn poll(&mut self) -> Result<Async<Conn>> {
        if self.with_inner(|inner| inner.closed) {
            return Err(DriverError::PoolDisconnected.into());
        }

        self.handle_futures()?;

        match self.take_conn() {
            Some(conn) => Ok(Ready(conn)),
            None => {
                let new_conn_created = self.with_inner(|mut inner| {
                    if inner.new.len() == 0 && inner.conn_count() < self.pool_constraints.max() {
                        let new_conn = Conn::new(self.opts.clone());
                        inner.new.push(Box::new(new_conn));
                        true
                    } else {
                        inner.tasks.push(task::current());
                        false
                    }
                });
                if new_conn_created {
                    self.poll()
                } else {
                    Ok(NotReady)
                }
            }
        }
    }
}

impl Drop for Conn {
    fn drop(&mut self) {
        if let Some(mut pool) = self.inner.pool.take() {
            pool.return_conn(self.take());
        } else if self.inner.stream.is_some() && !self.inner.disconnected {
            crate::conn::disconnect(self.take());
        }
    }
}

#[cfg(test)]
mod test {
    use futures::Future;

    use crate::{
        conn::pool::Pool, queryable::Queryable, test_misc::DATABASE_URL, TransactionOptions,
    };

    /// Same as `tokio::run`, but will panic if future panics and will return the result
    /// of future execution.
    fn run<F, T, U>(future: F) -> Result<T, U>
    where
        F: Future<Item = T, Error = U> + Send + 'static,
        T: Send + 'static,
        U: Send + 'static,
    {
        let mut runtime = tokio::runtime::Runtime::new().unwrap();
        let result = runtime.block_on(future);
        runtime.shutdown_on_idle().wait().unwrap();
        result
    }

    #[test]
    fn should_connect() {
        let pool = Pool::new(&**DATABASE_URL);
        let fut = pool
            .get_conn()
            .and_then(|conn| conn.ping().map(|_| ()))
            .and_then(|_| pool.disconnect());

        run(fut).unwrap();
    }

    #[test]
    fn should_start_transaction() {
        let pool = Pool::new(format!("{}?pool_min=1&pool_max=1", &**DATABASE_URL));
        let fut = pool
            .get_conn()
            .and_then(|conn| conn.drop_query("CREATE TABLE IF NOT EXISTS tmp(id int)"))
            .and_then({
                let pool = pool.clone();
                move |_| pool.start_transaction(TransactionOptions::default())
            })
            .and_then(|transaction| {
                transaction.batch_exec("INSERT INTO tmp (id) VALUES (?)", vec![(1,), (2,)])
            })
            .and_then(|transaction| transaction.prep_exec("SELECT * FROM tmp", ()))
            .map(|_| ())
            .and_then({
                let pool = pool.clone();
                move |_| pool.get_conn()
            })
            .and_then(|conn| conn.first("SELECT COUNT(*) FROM tmp"))
            .and_then(|(_, row_opt)| {
                assert_eq!(row_opt, Some((0u8,)));
                pool.get_conn()
                    .and_then(|conn| conn.drop_query("DROP TABLE tmp"))
                    .and_then(move |_| pool.disconnect())
            });

        run(fut).unwrap();
    }

    #[test]
    fn should_hold_bounds2() {
        use std::cmp::max;

        const POOL_MIN: usize = 5;
        const POOL_MAX: usize = 10;

        let url = format!("{}?pool_min={}&pool_max={}", &**DATABASE_URL, POOL_MIN, POOL_MAX);

        // Clean
        let pool = Pool::new(url.clone());
        let pool_clone = pool.clone();
        let conns = (0..POOL_MAX).map(|_| pool.get_conn()).collect::<Vec<_>>();

        let fut = ::futures::future::join_all(conns)
            .map(move |mut conns| {
                let mut popped = 0;
                assert_eq!(pool_clone.inner.lock().unwrap().conn_count(), POOL_MAX);

                while let Some(_) = conns.pop().map(drop) {
                    popped += 1;
                    assert_eq!(pool_clone.inner.lock().unwrap().conn_count(), POOL_MAX + POOL_MIN - max(popped, POOL_MIN));
                }

                pool_clone
            })
            .and_then(|pool| pool.disconnect());

        run(fut).unwrap();

        // Dirty
        let pool = Pool::new(url.clone());
        let pool_clone = pool.clone();
        let conns = (0..POOL_MAX)
            .map(|_| pool
                .get_conn()
                .and_then(|conn| conn.start_transaction(TransactionOptions::new())))
            .collect::<Vec<_>>();

        let fut = ::futures::future::join_all(conns)
            .map(move |mut conns| {
                assert_eq!(pool_clone.inner.lock().unwrap().conn_count(), POOL_MAX);

                while let Some(_) = conns.pop().map(drop) {
                    assert_eq!(pool_clone.inner.lock().unwrap().conn_count(), POOL_MAX);
                }
            });

        run(fut).unwrap();
    }

    #[test]
    fn should_hold_bounds() {
        let pool = Pool::new(format!("{}?pool_min=1&pool_max=2", &**DATABASE_URL));
        let pool_clone = pool.clone();
        let fut = pool
            .get_conn()
            .join(pool.get_conn())
            .and_then(move |(mut conn1, conn2)| {
                let new_conn = pool_clone.get_conn();
                conn1.inner.pool.as_mut().unwrap().handle_futures().unwrap();
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.new.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.idle.len()),
                    0
                );
                assert_eq!(
                    conn2
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.queue.len()),
                    0
                );
                new_conn
            })
            .and_then(|conn1| {
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.new.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.idle.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.queue.len()),
                    0
                );
                Ok(())
            })
            .and_then(|_| {
                assert_eq!(pool.with_inner(|inner| inner.new.len()), 0);
                assert_eq!(pool.with_inner(|inner| inner.idle.len()), 1);
                assert_eq!(pool.with_inner(|inner| inner.queue.len()), 0);
                pool.disconnect()
            });

        run(fut).unwrap();
    }

    #[cfg(feature = "nightly")]
    mod bench {
        use futures::Future;
        use tokio::runtime::Runtime;

        use crate::{conn::pool::Pool, queryable::Queryable, test_misc::DATABASE_URL};

        #[bench]
        fn connect(bencher: &mut test::Bencher) {
            let mut runtime = Runtime::new().expect("3");
            let pool = Pool::new(&**DATABASE_URL);

            bencher.iter(|| {
                let fut = pool.get_conn().and_then(|conn| conn.ping());
                runtime.block_on(fut).expect("1");
            });

            runtime.block_on(pool.disconnect()).unwrap();
        }
    }
}
<commit_msg>Remove duplicated code in Pool::handle_futures<commit_after>// Copyright (c) 2016 Anatoly Ikorsky
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use ::futures::{
    task::{self, Task},
    Async::{self, NotReady, Ready},
    Future,
};

use std::{
    fmt,
    str::FromStr,
    sync::{Arc, Mutex, MutexGuard},
};

use crate::{
    conn::{pool::futures::*, Conn},
    error::*,
    opts::{Opts, PoolConstraints},
    queryable::{
        transaction::{Transaction, TransactionOptions},
        Queryable,
    },
    BoxFuture, MyFuture,
};

pub mod futures;

pub struct Inner {
    closed: bool,
    new: Vec<BoxFuture<Conn>>,
    queue: Vec<BoxFuture<Conn>>,
    idle: Vec<Conn>,
    ongoing: usize,
    tasks: Vec<Task>,
}

impl Inner {
    fn conn_count(&self) -> usize {
        self.new.len() + self.idle.len() + self.queue.len() + self.ongoing
    }
}

#[derive(Clone)]
/// Asynchronous pool of MySql connections.
pub struct Pool {
    opts: Opts,
    inner: Arc<Mutex<Inner>>,
    pool_constraints: PoolConstraints,
}

impl fmt::Debug for Pool {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let (new_len, idle_len, queue_len, ongoing, tasks_len) = self.with_inner(|inner| {
            (
                inner.new.len(),
                inner.idle.len(),
                inner.queue.len(),
                inner.ongoing,
                inner.tasks.len(),
            )
        });
        f.debug_struct("Pool")
            .field("pool_constraints", &self.pool_constraints)
            .field("new connections count", &new_len)
            .field("idle connections count", &idle_len)
            .field("queue length", &queue_len)
            .field("ongoing connections count", &ongoing)
            .field("tasks count", &tasks_len)
            .finish()
    }
}

impl Pool {
    /// Creates new pool of connections.
    pub fn new<O: Into<Opts>>(opts: O) -> Pool {
        let opts = opts.into();
        let pool_constraints = opts.get_pool_constraints().clone();
        let pool = Pool {
            opts,
            inner: Arc::new(Mutex::new(Inner {
                closed: false,
                new: Vec::with_capacity(pool_constraints.min()),
                idle: Vec::new(),
                queue: Vec::new(),
                ongoing: 0,
                tasks: Vec::new(),
            })),
            pool_constraints,
        };

        pool
    }

    /// Creates new pool of connections.
    pub fn from_url<T: AsRef<str>>(url: T) -> Result<Pool> {
        let opts = Opts::from_str(url.as_ref())?;
        Ok(Pool::new(opts))
    }

    /// Returns future that resolves to `Conn`.
    pub fn get_conn(&self) -> GetConn {
        new_get_conn(self)
    }

    /// Shortcut for `get_conn` followed by `start_transaction`.
    pub fn start_transaction(
        &self,
        options: TransactionOptions,
    ) -> impl MyFuture<Transaction<Conn>> {
        self.get_conn()
            .and_then(|conn| Queryable::start_transaction(conn, options))
    }

    /// Returns future that disconnects this pool from server and resolves to `()`.
    ///
    /// Active connections taken from this pool should be disconnected manually.
    /// Also all pending and new `GetConn`'s will resolve to error.
    pub fn disconnect(mut self) -> DisconnectPool {
        let become_closed = self.with_inner(|mut inner| {
            if !inner.closed {
                inner.closed = true;
                return true;
            } else {
                return false;
            }
        });
        if become_closed {
            while let Some(conn) = self.take_conn() {
                crate::conn::disconnect(conn);
            }
        }
        new_disconnect_pool(self)
    }

    /// Returns true if futures is in queue.
    fn in_queue(&self) -> bool {
        self.with_inner(|inner| {
            let count = inner.new.len() + inner.queue.len();
            count > 0
        })
    }

    /// A way to take connection from a pool.
    fn take_conn(&mut self) -> Option<Conn> {
        self.with_inner(|mut inner| {
            while let Some(mut conn) = inner.idle.pop() {
                if conn.expired() {
                    crate::conn::disconnect(conn);
                } else {
                    conn.inner.pool = Some(self.clone());
                    inner.ongoing += 1;
                    return Some(conn);
                }
            }
            None
        })
    }

    /// A way to return connection taken from a pool.
    fn return_conn(&mut self, conn: Conn) {
        let min = self.pool_constraints.min();

        self.with_inner(|mut inner| {
            inner.ongoing -= 1;

            if inner.closed {
                return;
            }

            if conn.inner.stream.is_none() {
                // drop incomplete connection
                return;
            }

            if conn.inner.in_transaction || conn.inner.has_result.is_some() {
                inner.queue.push(conn.cleanup());
            } else {
                if inner.idle.len() >= min {
                    crate::conn::disconnect(conn);
                } else {
                    inner.idle.push(conn);
                }
            }

            while let Some(task) = inner.tasks.pop() {
                task.notify()
            }
        });
    }

    fn with_inner<F, T>(&self, fun: F) -> T
    where
        F: FnOnce(MutexGuard<'_, Inner>) -> T,
        T: 'static,
    {
        fun(self.inner.lock().unwrap())
    }

    /// Will manage lifetime of futures stored in a pool.
    fn handle_futures(&mut self) -> Result<()> {
        if !self.in_queue() {
            // There is no futures in queue
            return Ok(());
        }

        let mut handled = false;

        let mut returned_conns: Vec<Conn> = vec![];

        self.with_inner(|mut inner| {
            macro_rules! handle {
                ($vec:ident { $($p:pat => $b:block,)+ }) => ({
                    let len = inner.$vec.len();
                    let mut done_fut_idxs = Vec::new();

                    for i in 0..len {
                        let result = inner.$vec.get_mut(i).unwrap().poll();
                        match result {
                            Ok(Ready(_)) | Err(_) => done_fut_idxs.push(i),
                            _ => (),
                        }

                        let out: Result<()> = match result {
                            Ok(Ready(conn)) => {
                                if inner.closed {
                                    crate::conn::disconnect(conn);
                                } else {
                                    inner.ongoing += 1;
                                    returned_conns.push(conn);
                                }
                                handled = true;
                                Ok(())
                            }
                            $($p => $b),+
                            _ => {
                                Ok(())
                            }
                        };

                        match out {
                            Err(err) => {
                                // early return in case of error
                                while let Some(i) = done_fut_idxs.pop() {
                                    inner.$vec.swap_remove(i);
                                }
                                return Err(err)
                            }
                            _ => (),
                        }
                    }

                    while let Some(i) = done_fut_idxs.pop() {
                        inner.$vec.swap_remove(i);
                    }
                });
            }

            // Handle dirty connections.
            handle!(queue {
                // Drop it in case of error.
                Err(_) => { Ok(()) },
            });

            // Handle connecting connections.
            handle!(new {
                Err(err) => {
                    if !inner.closed {
                        Err(err)
                    } else {
                        Ok(())
                    }
                },
            });

            Ok(())
        })?;

        for conn in returned_conns {
            self.return_conn(conn);
        }

        if handled {
            self.handle_futures()
        } else {
            Ok(())
        }
    }

    /// Will poll pool for connection.
    fn poll(&mut self) -> Result<Async<Conn>> {
        if self.with_inner(|inner| inner.closed) {
            return Err(DriverError::PoolDisconnected.into());
        }

        self.handle_futures()?;

        match self.take_conn() {
            Some(conn) => Ok(Ready(conn)),
            None => {
                let new_conn_created = self.with_inner(|mut inner| {
                    if inner.new.len() == 0 && inner.conn_count() < self.pool_constraints.max() {
                        let new_conn = Conn::new(self.opts.clone());
                        inner.new.push(Box::new(new_conn));
                        true
                    } else {
                        inner.tasks.push(task::current());
                        false
                    }
                });
                if new_conn_created {
                    self.poll()
                } else {
                    Ok(NotReady)
                }
            }
        }
    }
}

impl Drop for Conn {
    fn drop(&mut self) {
        if let Some(mut pool) = self.inner.pool.take() {
            pool.return_conn(self.take());
        } else if self.inner.stream.is_some() && !self.inner.disconnected {
            crate::conn::disconnect(self.take());
        }
    }
}

#[cfg(test)]
mod test {
    use futures::Future;

    use crate::{
        conn::pool::Pool, queryable::Queryable, test_misc::DATABASE_URL, TransactionOptions,
    };

    /// Same as `tokio::run`, but will panic if future panics and will return the result
    /// of future execution.
    fn run<F, T, U>(future: F) -> Result<T, U>
    where
        F: Future<Item = T, Error = U> + Send + 'static,
        T: Send + 'static,
        U: Send + 'static,
    {
        let mut runtime = tokio::runtime::Runtime::new().unwrap();
        let result = runtime.block_on(future);
        runtime.shutdown_on_idle().wait().unwrap();
        result
    }

    #[test]
    fn should_connect() {
        let pool = Pool::new(&**DATABASE_URL);
        let fut = pool
            .get_conn()
            .and_then(|conn| conn.ping().map(|_| ()))
            .and_then(|_| pool.disconnect());

        run(fut).unwrap();
    }

    #[test]
    fn should_start_transaction() {
        let pool = Pool::new(format!("{}?pool_min=1&pool_max=1", &**DATABASE_URL));
        let fut = pool
            .get_conn()
            .and_then(|conn| conn.drop_query("CREATE TABLE IF NOT EXISTS tmp(id int)"))
            .and_then({
                let pool = pool.clone();
                move |_| pool.start_transaction(TransactionOptions::default())
            })
            .and_then(|transaction| {
                transaction.batch_exec("INSERT INTO tmp (id) VALUES (?)", vec![(1,), (2,)])
            })
            .and_then(|transaction| transaction.prep_exec("SELECT * FROM tmp", ()))
            .map(|_| ())
            .and_then({
                let pool = pool.clone();
                move |_| pool.get_conn()
            })
            .and_then(|conn| conn.first("SELECT COUNT(*) FROM tmp"))
            .and_then(|(_, row_opt)| {
                assert_eq!(row_opt, Some((0u8,)));
                pool.get_conn()
                    .and_then(|conn| conn.drop_query("DROP TABLE tmp"))
                    .and_then(move |_| pool.disconnect())
            });

        run(fut).unwrap();
    }

    #[test]
    fn should_hold_bounds2() {
        use std::cmp::max;

        const POOL_MIN: usize = 5;
        const POOL_MAX: usize = 10;

        let url = format!("{}?pool_min={}&pool_max={}", &**DATABASE_URL, POOL_MIN, POOL_MAX);

        // Clean
        let pool = Pool::new(url.clone());
        let pool_clone = pool.clone();
        let conns = (0..POOL_MAX).map(|_| pool.get_conn()).collect::<Vec<_>>();

        let fut = ::futures::future::join_all(conns)
            .map(move |mut conns| {
                let mut popped = 0;
                assert_eq!(pool_clone.inner.lock().unwrap().conn_count(), POOL_MAX);

                while let Some(_) = conns.pop().map(drop) {
                    popped += 1;
                    assert_eq!(pool_clone.inner.lock().unwrap().conn_count(), POOL_MAX + POOL_MIN - max(popped, POOL_MIN));
                }

                pool_clone
            })
            .and_then(|pool| pool.disconnect());

        run(fut).unwrap();

        // Dirty
        let pool = Pool::new(url.clone());
        let pool_clone = pool.clone();
        let conns = (0..POOL_MAX)
            .map(|_| pool
                .get_conn()
                .and_then(|conn| conn.start_transaction(TransactionOptions::new())))
            .collect::<Vec<_>>();

        let fut = ::futures::future::join_all(conns)
            .map(move |mut conns| {
                assert_eq!(pool_clone.inner.lock().unwrap().conn_count(), POOL_MAX);

                while let Some(_) = conns.pop().map(drop) {
                    assert_eq!(pool_clone.inner.lock().unwrap().conn_count(), POOL_MAX);
                }
            });

        run(fut).unwrap();
    }

    #[test]
    fn should_hold_bounds() {
        let pool = Pool::new(format!("{}?pool_min=1&pool_max=2", &**DATABASE_URL));
        let pool_clone = pool.clone();
        let fut = pool
            .get_conn()
            .join(pool.get_conn())
            .and_then(move |(mut conn1, conn2)| {
                let new_conn = pool_clone.get_conn();
                conn1.inner.pool.as_mut().unwrap().handle_futures().unwrap();
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.new.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.idle.len()),
                    0
                );
                assert_eq!(
                    conn2
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.queue.len()),
                    0
                );
                new_conn
            })
            .and_then(|conn1| {
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.new.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.idle.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.queue.len()),
                    0
                );
                Ok(())
            })
            .and_then(|_| {
                assert_eq!(pool.with_inner(|inner| inner.new.len()), 0);
                assert_eq!(pool.with_inner(|inner| inner.idle.len()), 1);
                assert_eq!(pool.with_inner(|inner| inner.queue.len()), 0);
                pool.disconnect()
            });

        run(fut).unwrap();
    }

    #[cfg(feature = "nightly")]
    mod bench {
        use futures::Future;
        use tokio::runtime::Runtime;

        use crate::{conn::pool::Pool, queryable::Queryable, test_misc::DATABASE_URL};

        #[bench]
        fn connect(bencher: &mut test::Bencher) {
            let mut runtime = Runtime::new().expect("3");
            let pool = Pool::new(&**DATABASE_URL);

            bencher.iter(|| {
                let fut = pool.get_conn().and_then(|conn| conn.ping());
                runtime.block_on(fut).expect("1");
            });

            runtime.block_on(pool.disconnect()).unwrap();
        }
    }
}
<|endoftext|>


<commit_before>extern crate hoedown;
use hoedown::renderer::Render;

use std::env::args as cmd_args;
use std::io::{Read, Write, BufReader};

const USAGE: &'static str  = "Usage: md_render <files>... [options]

Options:
--css=<file_name> - specify name of css file to link with.
--inline-css=<file_name> - specify the name of css file which style will be inlined.
";

const EXT_TABLES: u32            = 1 << 0;
const EXT_FENCED_CODE: u32       = 1 << 1;
const EXT_FOOTNOTES: u32         = 1 << 2;
const EXT_AUTOLINK: u32          = 1 << 3;
const EXT_STRIKETHROUGH: u32     = 1 << 4;
const EXT_UNDERLINE: u32         = 1 << 5;
const EXT_HIGHLIGHT: u32         = 1 << 6;
const EXT_QUOTE: u32             = 1 << 7;
const EXT_NO_INTRA_EMPHASIS: u32 = 1 << 11;
const EXT_SPACE_HEADERS: u32     = 1 << 12;

macro_rules! is_file {
    ($path:expr) => { std::fs::metadata($path).unwrap().is_file() }
}

#[inline]
fn usage() {
    println!("{}", USAGE);
}

fn main() {
    if cmd_args().len() < 2 { return usage(); }

    let mut css = String::new();
    let mut md_files: Vec<String> = vec![];
    for arg in cmd_args().skip(1) {
        if arg.starts_with("--") {
            let arg = &arg[2..];
            let eq_pos = arg.find('=');

            if arg.starts_with("css") {
                if let Some(eq_pos) = eq_pos {
                    std::fmt::Write::write_fmt(&mut css,
                                               format_args!("<link rel=\"stylesheet\" href=\"{}\">\n", &arg[eq_pos+1..])).unwrap();
                }
                else {
                    println!(">>>Invalid use of options. CSS File is not specified!");
                }

            }
            else if arg.starts_with("inline-css") {
                if let Some(eq_pos) = eq_pos {
                    if let Ok(mut css_file) = std::fs::File::open(&arg[eq_pos+1..]) {
                        css.push_str("\n");
                    }
                    else {
                        println!(">>>{}: no such CSS", &arg[eq_pos+1..]);
                    }
                }
                else {
                    println!(">>>Invalid use of options. CSS File is not specified!");
                }
            }
        }
        else if is_file!(&arg) {
            md_files.push(arg);
        }
        else {
            println!(">>>{}: No such file!\n", arg);
        }
    };

    if md_files.len() == 0 {
        println!(">>>No Markdown files are  specified!\n");
        usage();
        return;
    }

    let md_ext = hoedown::Extension::from_bits_truncate(EXT_TABLES |
                                                        EXT_FENCED_CODE |
                                                        EXT_FOOTNOTES |
                                                        EXT_AUTOLINK |
                                                        EXT_STRIKETHROUGH |
                                                        EXT_UNDERLINE |
                                                        EXT_HIGHLIGHT |
                                                        EXT_QUOTE |
                                                        EXT_NO_INTRA_EMPHASIS |
                                                        EXT_SPACE_HEADERS);
    let flags = hoedown::renderer::html::Flags::empty();
    let mut html_render = hoedown::renderer::html::Html::new(flags, 0);
    for arg in md_files {
        println!("Render - {}", &arg);
        let html_path = std::path::Path::new(&arg).with_extension("html");

        if let Ok(file) = std::fs::File::open(&arg) {
            let mut html_file = std::fs::File::create(&html_path).unwrap();

            let file = BufReader::new(file);
            let file = hoedown::Markdown::read_from(file).extensions(md_ext);
            let render_result = html_render.render(&file);

            html_file.write_all(css.as_bytes()).unwrap();
            html_file.write_all(&render_result).unwrap();
            println!(">>>Result: {}", &html_path.to_str().unwrap_or(""));
        }
        else {
            println!(">>>{}: failed to open", &arg);
        }
        println!("");
    }
}
<commit_msg>Rust: md_render add ToC<commit_after>extern crate hoedown;
use hoedown::renderer::Render;

use std::env::args as cmd_args;
use std::io::{Read, Write, BufReader};

const USAGE: &'static str  = "Usage: md_render <files>... [options]

Options:
--css=<file_name> - specify name of css file to link with.
--inline-css=<file_name> - specify the name of css file which style will be inlined.
--toc=<num> - specify the level of deepness for table of contents. Default 0.
";

const EXT_TABLES: u32            = 1 << 0;
const EXT_FENCED_CODE: u32       = 1 << 1;
const EXT_FOOTNOTES: u32         = 1 << 2;
const EXT_AUTOLINK: u32          = 1 << 3;
const EXT_STRIKETHROUGH: u32     = 1 << 4;
const EXT_UNDERLINE: u32         = 1 << 5;
const EXT_HIGHLIGHT: u32         = 1 << 6;
const EXT_QUOTE: u32             = 1 << 7;
const EXT_NO_INTRA_EMPHASIS: u32 = 1 << 11;
const EXT_SPACE_HEADERS: u32     = 1 << 12;

macro_rules! is_file {
    ($path:expr) => { std::fs::metadata($path).unwrap().is_file() }
}

#[inline]
fn usage() {
    println!("{}", USAGE);
}

fn main() {
    if cmd_args().len() < 2 { return usage(); }

    let mut toc: i32 = 0;
    let mut css = String::new();
    let mut md_files: Vec<String> = vec![];
    for arg in cmd_args().skip(1) {
        if arg.starts_with("--") {
            let arg = &arg[2..];
            let eq_pos = arg.find('=');

            if arg.starts_with("css") {
                if let Some(eq_pos) = eq_pos {
                    std::fmt::Write::write_fmt(&mut css,
                                               format_args!("<link rel=\"stylesheet\" href=\"{}\">\n", &arg[eq_pos+1..])).unwrap();
                }
                else {
                    println!(">>>Invalid use of options. CSS File is not specified!");
                }

            }
            else if arg.starts_with("toc") {
                if let Some(eq_pos) = eq_pos {
                    if let Ok(eq_pos) = arg[eq_pos+1..].parse::<i32>() {
                        toc = eq_pos;
                    }
                    else {
                        println!(">>>Invalid use of options. ToC number is invalid!");
                    }
                }
                else {
                    println!(">>>Invalid use of options. ToC number is not specified!");
                }
            }
            else if arg.starts_with("inline-css") {
                if let Some(eq_pos) = eq_pos {
                    if let Ok(mut css_file) = std::fs::File::open(&arg[eq_pos+1..]) {
                        css.push_str("\n");
                    }
                    else {
                        println!(">>>{}: no such CSS", &arg[eq_pos+1..]);
                    }
                }
                else {
                    println!(">>>Invalid use of options. CSS File is not specified!");
                }
            }
        }
        else if is_file!(&arg) {
            md_files.push(arg);
        }
        else {
            println!(">>>{}: No such file!\n", arg);
        }
    };

    if md_files.len() == 0 {
        println!(">>>No Markdown files are  specified!\n");
        usage();
        return;
    }

    let md_ext = hoedown::Extension::from_bits_truncate(EXT_TABLES |
                                                        EXT_FENCED_CODE |
                                                        EXT_FOOTNOTES |
                                                        EXT_AUTOLINK |
                                                        EXT_STRIKETHROUGH |
                                                        EXT_UNDERLINE |
                                                        EXT_HIGHLIGHT |
                                                        EXT_QUOTE |
                                                        EXT_NO_INTRA_EMPHASIS |
                                                        EXT_SPACE_HEADERS);
    let flags = hoedown::renderer::html::Flags::empty();
    let mut html_render = hoedown::renderer::html::Html::new(flags, toc);
    let mut html_toc = hoedown::renderer::html::Html::toc(toc);
    for arg in md_files {
        println!("Render - {}", &arg);
        let html_path = std::path::Path::new(&arg).with_extension("html");

        if let Ok(file) = std::fs::File::open(&arg) {
            let mut html_file = std::fs::File::create(&html_path).unwrap();

            let file = BufReader::new(file);
            let file = hoedown::Markdown::read_from(file).extensions(md_ext);
            let render_result = html_render.render(&file);

            html_file.write_all(css.as_bytes()).unwrap();
            if toc > 0 {
                let render_toc = html_toc.render(&file);
                html_file.write_all(b"
Table of Content
").unwrap();
                html_file.write_all(&render_toc).unwrap();
            }
            html_file.write_all(&render_result).unwrap();
            println!(">>>Result: {}", &html_path.to_str().unwrap_or(""));
        }
        else {
            println!(">>>{}: failed to open", &arg);
        }
        println!("");
    }
}
<|endoftext|>


<commit_before>//! Evolve a Population of Individuals in an Environment
//!
//! Evolution is the culmination of individuals in a population competing in a
//! harsh environment for limited resources. Here the environment is determined
//! by the fitness function, the individuals are programs.

use super::super::grammar::Program;
use super::super::population::Population;
use super::super::environment::Environment;
use super::super::fitness::{Config, ScoreCard, Score};

/// Return the next generation of the population
pub fn succesion(environment: Environment, population: Population) -> Population {
    let config: Config = Config {
        program_length_weight: -1,
        food_distance_weight: -1,
    };
    let Population(last_generation) = population;
    let mut scored_last_generation: Vec<(Program, Score)> = last_generation
        .iter()
        .map(|program| (program.clone(), config.score(program.clone(), environment.clone())))
        .collect();
    scored_last_generation.sort_by(|a,b| a.1.cmp(&b.1));
    let next_generation = scored_last_generation
        .iter()
        .map(|a| a.0.clone())
        .collect();

    Population(next_generation)
}
<commit_msg>Sort so the fittest program comes first<commit_after>//! Evolve a Population of Individuals in an Environment
//!
//! Evolution is the culmination of individuals in a population competing in a
//! harsh environment for limited resources. Here the environment is determined
//! by the fitness function, the individuals are programs.

use super::super::grammar::Program;
use super::super::population::Population;
use super::super::environment::Environment;
use super::super::fitness::{Config, ScoreCard, Score};

/// Return the next generation of the population
pub fn succesion(environment: Environment, population: Population) -> Population {
    let config: Config = Config {
        program_length_weight: -1,
        food_distance_weight: -1,
    };
    let Population(last_generation) = population;
    let mut scored_last_generation: Vec<(Program, Score)> = last_generation
        .iter()
        .map(|program| (program.clone(), config.score(program.clone(), environment.clone())))
        .collect();
    scored_last_generation.sort_by(|a,b| b.1.cmp(&a.1));
    let next_generation = scored_last_generation
        .iter()
        .map(|a| a.0.clone())
        .collect();

    Population(next_generation)
}
<|endoftext|>


<commit_before>// Copyright 2016 Adam Perry. Dual-licensed MIT and Apache 2.0 (see LICENSE files for details).

use std::collections::BTreeMap;
use std::io::Read;
use std::thread::sleep;
use std::time::Duration;
use std::u32;

use chrono::{DateTime, UTC};
use hyper;
use hyper::client::{RedirectPolicy, RequestBuilder, Response};
use hyper::header::Headers;
use hyper::net::HttpsConnector;
use hyper::status::StatusCode;
use hyper_native_tls::NativeTlsClient;
use serde::Deserialize;
use serde_json;

use config::CONFIG;
use error::{DashError, DashResult};
use github::models::{CommentFromJson, IssueFromJson, PullRequestFromJson, PullRequestUrls};

pub const BASE_URL: &'static str = "https://api.github.com";

pub const DELAY: u64 = 300;

type ParameterMap = BTreeMap<&'static str, String>;

header! { (Auth, "Authorization") => [String] }
header! { (UA, "User-Agent") => [String] }
header! { (TZ, "Time-Zone") => [String] }
header! { (Accept, "Accept") => [String] }
header! { (RateLimitRemaining, "X-RateLimit-Remaining") => [u32] }
header! { (RateLimitReset, "X-RateLimit-Reset") => [i64] }
header! { (Link, "Link") => [String] }

const PER_PAGE: u32 = 100;

#[derive(Debug)]
pub struct Client {
    token: String,
    ua: String,
    client: hyper::Client,
    rate_limit: u32,
    rate_limit_timeout: DateTime<UTC>,
}

impl Client {
    pub fn new() -> Self {
        let tls_connector = HttpsConnector::new(NativeTlsClient::new().unwrap());
        let mut client = hyper::Client::with_connector(tls_connector);
        client.set_redirect_policy(RedirectPolicy::FollowAll);

        Client {
            token: CONFIG.github_access_token.clone(),
            ua: CONFIG.github_user_agent.clone(),
            client: client,
            rate_limit: u32::MAX,
            rate_limit_timeout: UTC::now(),
        }
    }

    pub fn org_repos(&self, org: &str) -> DashResult<Vec<String>> {
        let url = format!("{}/orgs/{}/repos", BASE_URL, org);
        let vals: Vec<serde_json::Value> = try!(self.get_models(&url, None));

        let mut repos = Vec::new();
        for v in vals {

            let v = match v.as_object() {
                Some(v) => v,
                None => return Err(DashError::Misc(None)),
            };

            let repo = match v.get("name") {
                Some(n) => {
                    match n.as_str() {
                        Some(s) => format!("{}/{}", org, s),
                        None => return Err(DashError::Misc(None)),
                    }
                }
                None => return Err(DashError::Misc(None)),
            };

            repos.push(repo);

        }
        Ok(repos)
    }

    pub fn issues_since(&self, repo: &str, start: DateTime<UTC>) -> DashResult<Vec<IssueFromJson>> {

        let url = format!("{}/repos/{}/issues", BASE_URL, repo);
        let mut params = ParameterMap::new();

        params.insert("state", "all".to_string());
        params.insert("since", format!("{:?}", start));
        params.insert("state", "all".to_string());
        params.insert("per_page", format!("{}", PER_PAGE));
        params.insert("direction", "asc".to_string());

        self.get_models(&url, Some(&params))
    }

    pub fn comments_since(&self,
                          repo: &str,
                          start: DateTime<UTC>)
                          -> DashResult<Vec<CommentFromJson>> {
        let url = format!("{}/repos/{}/issues/comments", BASE_URL, repo);
        let mut params = ParameterMap::new();

        params.insert("sort", "created".to_string());
        params.insert("direction", "asc".to_string());
        params.insert("since", format!("{:?}", start));
        params.insert("per_page", format!("{}", PER_PAGE));

        self.get_models(&url, Some(&params))
    }

    fn get_models<M: Deserialize>(&self,
                                  start_url: &str,
                                  params: Option<&ParameterMap>)
                                  -> DashResult<Vec<M>> {

        let mut res = try!(self.get(start_url, params));
        let mut models = self.deserialize::<Vec<M>>(&mut res)?;
        while let Some(url) = Self::next_page(&res.headers) {
            sleep(Duration::from_millis(DELAY));
            res = self.get(&url, None)?;
            models.extend(self.deserialize::<Vec<M>>(&mut res)?);
        }
        Ok(models)
    }

    pub fn fetch_issue(&self, repo: &str, number: i32) -> DashResult<IssueFromJson> {
        let url = format!("{}/repos/{}/issues/{}", BASE_URL, repo, number);
        self.deserialize(&mut self.get(&url, None)?)
    }

    pub fn fetch_comments(&self, repo: &str, number: i32) -> DashResult<Vec<CommentFromJson>> {
        let url = format!("{}/repos/{}/issues/{}/comments", BASE_URL, repo, number);
        self.get_models(&url, None)
    }

    pub fn fetch_pr(&self, repo: &str, number: i32) -> DashResult<PullRequestFromJson> {
        let url = format!("{}/repos/{}/pulls/{}", BASE_URL, repo, number);
        self.deserialize(&mut self.get(&url, None)?)
    }

    pub fn fetch_pull_request(&self, pr_info: &PullRequestUrls) -> DashResult<PullRequestFromJson> {
        let url = pr_info.get("url");

        if let Some(url) = url {
            let mut res = try!(self.get(url, None));
            self.deserialize(&mut res)
        } else {
            Err(DashError::Misc(None))
        }
    }

    fn next_page(h: &Headers) -> Option<String> {
        if let Some(lh) = h.get::<Link>() {
            for link in (**lh).split(',').map(|s| s.trim()) {

                let tokens = link.split(';').map(|s| s.trim()).collect::<Vec<_>>();

                if tokens.len() != 2 {
                    continue;
                }

                if tokens[1] == "rel=\"next\"" {
                    let url = tokens[0].trim_left_matches('<').trim_right_matches('>').to_string();
                    return Some(url);
                }
            }
        }

        None
    }

    pub fn add_label(&self, repo: &str, issue_num: i32, label: &str) -> DashResult<()> {
        let url = format!("{}/repos/{}/issues/{}/labels",
                          BASE_URL, repo, issue_num);
        let payload = serde_json::to_string(&[label])?;

        let mut res = self.post(&url, &payload)?;

        match res.status {
            StatusCode::Ok => Ok(()),
            _ => {
                let mut body = String::new();
                res.read_to_string(&mut body)?;
                Err(DashError::Misc(Some(body)))
            }
        }
    }

    pub fn new_comment(&self,
                       repo: &str,
                       issue_num: i32,
                       text: &str)
                       -> DashResult<CommentFromJson> {
        let url = format!("{}/repos/{}/issues/{}/comments", BASE_URL, repo, issue_num);

        let mut obj = BTreeMap::new();
        obj.insert("body", text);

        let payload = serde_json::to_string(&obj)?;

        // FIXME propagate an error if it's a 404 or other error
        self.deserialize(&mut self.post(&url, &payload)?)
    }

    pub fn edit_comment(&self,
                        repo: &str,
                        comment_num: i32,
                        text: &str)
                        -> DashResult<CommentFromJson> {
        let url = format!("{}/repos/{}/issues/comments/{}", BASE_URL, repo, comment_num);

        let mut obj = BTreeMap::new();
        obj.insert("body", text);

        let payload = serde_json::to_string(&obj)?;

        // FIXME propagate an error if it's a 404 or other error
        self.deserialize(&mut self.patch(&url, &payload)?)
    }

    fn patch(&self, url: &str, payload: &str) -> Result<Response, hyper::error::Error> {
        self.set_headers(self.client.patch(url).body(payload)).send()
    }

    fn post(&self, url: &str, payload: &str) -> Result<Response, hyper::error::Error> {
        self.set_headers(self.client.post(url).body(payload)).send()
    }

    fn get(&self, url: &str, params: Option<&ParameterMap>)
           -> Result<Response, hyper::error::Error> {

        let qp_string = match params {
            Some(p) => {
                let mut qp = String::from("?");
                for (k, v) in p {
                    if qp.len() > 1 {
                        qp.push('&');
                    }
                    qp.push_str(&format!("{}={}", k, v));
                }
                qp
            }
            None => "".to_string(),
        };

        let url = format!("{}{}", url, qp_string);

        debug!("GETing: {}", &url);

        self.set_headers(self.client.get(&url))
            .send()
    }

    fn deserialize<M: Deserialize>(&self, res: &mut Response) -> DashResult<M> {
        let mut buf = String::new();
        res.read_to_string(&mut buf)?;

        match serde_json::from_str(&buf) {
            Ok(m) => Ok(m),
            Err(why) => {
                error!("Unable to parse from JSON ({:?}): {}", why, buf);
                Err(why.into())
            }
        }
    }

    fn set_headers<'a>(&self, req: RequestBuilder<'a>) -> RequestBuilder<'a> {
        req.header(Auth(format!("token {}", &self.token)))
            .header(UA(self.ua.clone()))
            .header(TZ("UTC".to_string()))
            .header(Accept("application/vnd.github.v3".to_string()))
            .header(hyper::header::Connection::close())
    }
}
<commit_msg>Simplify github::client::Client::org_repos<commit_after>// Copyright 2016 Adam Perry. Dual-licensed MIT and Apache 2.0 (see LICENSE files for details).

use std::collections::BTreeMap;
use std::io::Read;
use std::thread::sleep;
use std::time::Duration;
use std::u32;

use chrono::{DateTime, UTC};
use hyper;
use hyper::client::{RedirectPolicy, RequestBuilder, Response};
use hyper::header::Headers;
use hyper::net::HttpsConnector;
use hyper::status::StatusCode;
use hyper_native_tls::NativeTlsClient;
use serde::Deserialize;
use serde_json;

use config::CONFIG;
use error::{DashError, DashResult};
use github::models::{CommentFromJson, IssueFromJson, PullRequestFromJson, PullRequestUrls};

pub const BASE_URL: &'static str = "https://api.github.com";

pub const DELAY: u64 = 300;

type ParameterMap = BTreeMap<&'static str, String>;

header! { (Auth, "Authorization") => [String] }
header! { (UA, "User-Agent") => [String] }
header! { (TZ, "Time-Zone") => [String] }
header! { (Accept, "Accept") => [String] }
header! { (RateLimitRemaining, "X-RateLimit-Remaining") => [u32] }
header! { (RateLimitReset, "X-RateLimit-Reset") => [i64] }
header! { (Link, "Link") => [String] }

const PER_PAGE: u32 = 100;

#[derive(Debug)]
pub struct Client {
    token: String,
    ua: String,
    client: hyper::Client,
    rate_limit: u32,
    rate_limit_timeout: DateTime<UTC>,
}

impl Client {
    pub fn new() -> Self {
        let tls_connector = HttpsConnector::new(NativeTlsClient::new().unwrap());
        let mut client = hyper::Client::with_connector(tls_connector);
        client.set_redirect_policy(RedirectPolicy::FollowAll);

        Client {
            token: CONFIG.github_access_token.clone(),
            ua: CONFIG.github_user_agent.clone(),
            client: client,
            rate_limit: u32::MAX,
            rate_limit_timeout: UTC::now(),
        }
    }

    pub fn org_repos(&self, org: &str) -> DashResult<Vec<String>> {
        let url = format!("{}/orgs/{}/repos", BASE_URL, org);
        let vals: Vec<serde_json::Value> = try!(self.get_models(&url, None));

        let mut repos = Vec::new();
        for v in vals {
            if let Some(v) = v.as_object() {
                if let Some(n) = v.get("name") {
                    if let Some(s) = n.as_str() {
                        repos.push(format!("{}/{}", org, s));
                        continue;
                    }
                }
            }
            return Err(DashError::Misc(None));

        }
        Ok(repos)
    }

    pub fn issues_since(&self, repo: &str, start: DateTime<UTC>) -> DashResult<Vec<IssueFromJson>> {

        let url = format!("{}/repos/{}/issues", BASE_URL, repo);
        let mut params = ParameterMap::new();

        params.insert("state", "all".to_string());
        params.insert("since", format!("{:?}", start));
        params.insert("state", "all".to_string());
        params.insert("per_page", format!("{}", PER_PAGE));
        params.insert("direction", "asc".to_string());

        self.get_models(&url, Some(&params))
    }

    pub fn comments_since(&self,
                          repo: &str,
                          start: DateTime<UTC>)
                          -> DashResult<Vec<CommentFromJson>> {
        let url = format!("{}/repos/{}/issues/comments", BASE_URL, repo);
        let mut params = ParameterMap::new();

        params.insert("sort", "created".to_string());
        params.insert("direction", "asc".to_string());
        params.insert("since", format!("{:?}", start));
        params.insert("per_page", format!("{}", PER_PAGE));

        self.get_models(&url, Some(&params))
    }

    fn get_models<M: Deserialize>(&self,
                                  start_url: &str,
                                  params: Option<&ParameterMap>)
                                  -> DashResult<Vec<M>> {

        let mut res = try!(self.get(start_url, params));
        let mut models = self.deserialize::<Vec<M>>(&mut res)?;
        while let Some(url) = Self::next_page(&res.headers) {
            sleep(Duration::from_millis(DELAY));
            res = self.get(&url, None)?;
            models.extend(self.deserialize::<Vec<M>>(&mut res)?);
        }
        Ok(models)
    }

    pub fn fetch_issue(&self, repo: &str, number: i32) -> DashResult<IssueFromJson> {
        let url = format!("{}/repos/{}/issues/{}", BASE_URL, repo, number);
        self.deserialize(&mut self.get(&url, None)?)
    }

    pub fn fetch_comments(&self, repo: &str, number: i32) -> DashResult<Vec<CommentFromJson>> {
        let url = format!("{}/repos/{}/issues/{}/comments", BASE_URL, repo, number);
        self.get_models(&url, None)
    }

    pub fn fetch_pr(&self, repo: &str, number: i32) -> DashResult<PullRequestFromJson> {
        let url = format!("{}/repos/{}/pulls/{}", BASE_URL, repo, number);
        self.deserialize(&mut self.get(&url, None)?)
    }

    pub fn fetch_pull_request(&self, pr_info: &PullRequestUrls) -> DashResult<PullRequestFromJson> {
        let url = pr_info.get("url");

        if let Some(url) = url {
            let mut res = try!(self.get(url, None));
            self.deserialize(&mut res)
        } else {
            Err(DashError::Misc(None))
        }
    }

    fn next_page(h: &Headers) -> Option<String> {
        if let Some(lh) = h.get::<Link>() {
            for link in (**lh).split(',').map(|s| s.trim()) {

                let tokens = link.split(';').map(|s| s.trim()).collect::<Vec<_>>();

                if tokens.len() != 2 {
                    continue;
                }

                if tokens[1] == "rel=\"next\"" {
                    let url = tokens[0].trim_left_matches('<').trim_right_matches('>').to_string();
                    return Some(url);
                }
            }
        }

        None
    }

    pub fn add_label(&self, repo: &str, issue_num: i32, label: &str) -> DashResult<()> {
        let url = format!("{}/repos/{}/issues/{}/labels",
                          BASE_URL, repo, issue_num);
        let payload = serde_json::to_string(&[label])?;

        let mut res = self.post(&url, &payload)?;

        match res.status {
            StatusCode::Ok => Ok(()),
            _ => {
                let mut body = String::new();
                res.read_to_string(&mut body)?;
                Err(DashError::Misc(Some(body)))
            }
        }
    }

    pub fn new_comment(&self,
                       repo: &str,
                       issue_num: i32,
                       text: &str)
                       -> DashResult<CommentFromJson> {
        let url = format!("{}/repos/{}/issues/{}/comments", BASE_URL, repo, issue_num);

        let mut obj = BTreeMap::new();
        obj.insert("body", text);

        let payload = serde_json::to_string(&obj)?;

        // FIXME propagate an error if it's a 404 or other error
        self.deserialize(&mut self.post(&url, &payload)?)
    }

    pub fn edit_comment(&self,
                        repo: &str,
                        comment_num: i32,
                        text: &str)
                        -> DashResult<CommentFromJson> {
        let url = format!("{}/repos/{}/issues/comments/{}", BASE_URL, repo, comment_num);

        let mut obj = BTreeMap::new();
        obj.insert("body", text);

        let payload = serde_json::to_string(&obj)?;

        // FIXME propagate an error if it's a 404 or other error
        self.deserialize(&mut self.patch(&url, &payload)?)
    }

    fn patch(&self, url: &str, payload: &str) -> Result<Response, hyper::error::Error> {
        self.set_headers(self.client.patch(url).body(payload)).send()
    }

    fn post(&self, url: &str, payload: &str) -> Result<Response, hyper::error::Error> {
        self.set_headers(self.client.post(url).body(payload)).send()
    }

    fn get(&self, url: &str, params: Option<&ParameterMap>)
           -> Result<Response, hyper::error::Error> {

        let qp_string = match params {
            Some(p) => {
                let mut qp = String::from("?");
                for (k, v) in p {
                    if qp.len() > 1 {
                        qp.push('&');
                    }
                    qp.push_str(&format!("{}={}", k, v));
                }
                qp
            }
            None => "".to_string(),
        };

        let url = format!("{}{}", url, qp_string);

        debug!("GETing: {}", &url);

        self.set_headers(self.client.get(&url))
            .send()
    }

    fn deserialize<M: Deserialize>(&self, res: &mut Response) -> DashResult<M> {
        let mut buf = String::new();
        res.read_to_string(&mut buf)?;

        match serde_json::from_str(&buf) {
            Ok(m) => Ok(m),
            Err(why) => {
                error!("Unable to parse from JSON ({:?}): {}", why, buf);
                Err(why.into())
            }
        }
    }

    fn set_headers<'a>(&self, req: RequestBuilder<'a>) -> RequestBuilder<'a> {
        req.header(Auth(format!("token {}", &self.token)))
            .header(UA(self.ua.clone()))
            .header(TZ("UTC".to_string()))
            .header(Accept("application/vnd.github.v3".to_string()))
            .header(hyper::header::Connection::close())
    }
}
<|endoftext|>


<commit_before>use std::io::{Result, Write};

use super::super::general::write_vec;

pub fn begin<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"#![allow(non_camel_case_types)]",
"",
"extern crate libc;",
"#[macro_use] extern crate bitflags;",
    ];

    write_vec(w, &v)
}

pub fn after_extern_crates<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"#[allow(unused_imports)]",
"use libc::{c_int, c_char, c_float, c_uint, c_double,",
"    c_short, c_ushort, c_long, c_ulong,",
"    c_void, size_t, ssize_t, time_t};",
    ];

    write_vec(w, &v)
}

pub fn use_glib_ffi<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"#[allow(unused_imports)]",
"use glib_ffi::{gboolean, gconstpointer, gpointer, GType};",
    ];

    write_vec(w, &v)
}

pub fn only_for_glib<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"pub type gboolean = c_int;",
"pub const GFALSE:  c_int = 0;",
"pub const GTRUE:   c_int = 1;",
"",
"pub type gconstpointer = *const c_void;",
"pub type gpointer = *mut c_void;",
    ];

    write_vec(w, &v)
}

pub fn only_for_gtk<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"pub const GTK_ENTRY_BUFFER_MAX_SIZE: u16 = ::std::u16::MAX;",
"",
"//pub type GtkTreeModelForeachFunc = fn(model: *mut GtkTreeModel, path: *mut GtkTreePath, iter: *mut GtkTreeIter, data: gpointer) -> gboolean;",
"",
"pub const GTK_STYLE_PROVIDER_PRIORITY_FALLBACK: u32 = 1;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_THEME: u32 = 200;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_SETTINGS: u32 = 400;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_APPLICATION: u32 = 600;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_USER: u32 = 800;",
    ];

    write_vec(w, &v)
}
<commit_msg>Use libc::FILE<commit_after>use std::io::{Result, Write};

use super::super::general::write_vec;

pub fn begin<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"#![allow(non_camel_case_types)]",
"",
"extern crate libc;",
"#[macro_use] extern crate bitflags;",
    ];

    write_vec(w, &v)
}

pub fn after_extern_crates<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"#[allow(unused_imports)]",
"use libc::{c_int, c_char, c_float, c_uint, c_double,",
"    c_short, c_ushort, c_long, c_ulong,",
"    c_void, size_t, ssize_t, time_t, FILE};",
    ];

    write_vec(w, &v)
}

pub fn use_glib_ffi<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"#[allow(unused_imports)]",
"use glib_ffi::{gboolean, gconstpointer, gpointer, GType};",
    ];

    write_vec(w, &v)
}

pub fn only_for_glib<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"pub type gboolean = c_int;",
"pub const GFALSE:  c_int = 0;",
"pub const GTRUE:   c_int = 1;",
"",
"pub type gconstpointer = *const c_void;",
"pub type gpointer = *mut c_void;",
    ];

    write_vec(w, &v)
}

pub fn only_for_gtk<W: Write>(w: &mut W) -> Result<()>{
    let v = vec![
"",
"pub const GTK_ENTRY_BUFFER_MAX_SIZE: u16 = ::std::u16::MAX;",
"",
"//pub type GtkTreeModelForeachFunc = fn(model: *mut GtkTreeModel, path: *mut GtkTreePath, iter: *mut GtkTreeIter, data: gpointer) -> gboolean;",
"",
"pub const GTK_STYLE_PROVIDER_PRIORITY_FALLBACK: u32 = 1;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_THEME: u32 = 200;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_SETTINGS: u32 = 400;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_APPLICATION: u32 = 600;",
"pub const GTK_STYLE_PROVIDER_PRIORITY_USER: u32 = 800;",
    ];

    write_vec(w, &v)
}
<|endoftext|>


<commit_before>//! [PUT /_matrix/app/v1/transactions/{txnId}](https://matrix.org/docs/spec/application_service/r0.1.2#put-matrix-app-v1-transactions-txnid)

use ruma_api::ruma_api;
use ruma_events::AnyEvent;
use ruma_serde::Raw;

ruma_api! {
    metadata: {
        description: "This API is called by the homeserver when it wants to push an event (or batch of events) to the application service.",
        method: PUT,
        name: "push_events",
        path: "/_matrix/app/v1/transactions/:txn_id",
        rate_limited: false,
        authentication: QueryOnlyAccessToken,
    }

    request: {
        /// The transaction ID for this set of events.
        ///
        /// Homeservers generate these IDs and they are used to ensure idempotency of results.
        #[ruma_api(path)]
        pub txn_id: &'a str,

        /// A list of events.
        pub events: &'a [Raw<AnyEvent>],
    }

    #[derive(Default)]
    response: {}
}

impl<'a> Request<'a> {
    /// Creates a new `Request` with the given transaction ID and list of events.
    pub fn new(txn_id: &'a str, events: &'a [Raw<AnyEvent>]) -> Self {
        Self { txn_id, events }
    }
}

impl Response {
    /// Creates an empty `Response`.
    pub fn new() -> Self {
        Self
    }
}

#[cfg(feature = "server")]
#[cfg(test)]
mod tests {
    use ruma_api::{exports::http, OutgoingRequest};
    use ruma_events::AnyEvent;
    use ruma_serde::Raw;
    use serde_json::json;

    use super::Request;

    #[test]
    fn decode_request_contains_events_field() {
        let dummy_event: AnyEvent = serde_json::from_value(json!({
            "content": {},
            "type": "m.dummy"
        }))
        .unwrap();
        let dummy_event = Raw::from(dummy_event);
        let events = vec![dummy_event];

        let req: http::Request<Vec<u8>> = Request { events: &events, txn_id: "any_txn_id" }
            .try_into_http_request("https://homeserver.tld", Some("auth_tok"))
            .unwrap();
        let json_body: serde_json::Value = serde_json::from_slice(&req.body()).unwrap();

        assert_eq!(
            1,
            json_body.as_object().unwrap().get("events").unwrap().as_array().unwrap().len()
        );
    }
}
<commit_msg>appservice-api: Add constructor for push_events::v1::IncomingRequest<commit_after>//! [PUT /_matrix/app/v1/transactions/{txnId}](https://matrix.org/docs/spec/application_service/r0.1.2#put-matrix-app-v1-transactions-txnid)

use ruma_api::ruma_api;
use ruma_events::AnyEvent;
use ruma_serde::Raw;

ruma_api! {
    metadata: {
        description: "This API is called by the homeserver when it wants to push an event (or batch of events) to the application service.",
        method: PUT,
        name: "push_events",
        path: "/_matrix/app/v1/transactions/:txn_id",
        rate_limited: false,
        authentication: QueryOnlyAccessToken,
    }

    request: {
        /// The transaction ID for this set of events.
        ///
        /// Homeservers generate these IDs and they are used to ensure idempotency of results.
        #[ruma_api(path)]
        pub txn_id: &'a str,

        /// A list of events.
        pub events: &'a [Raw<AnyEvent>],
    }

    #[derive(Default)]
    response: {}
}

impl<'a> Request<'a> {
    /// Creates a new `Request` with the given transaction ID and list of events.
    pub fn new(txn_id: &'a str, events: &'a [Raw<AnyEvent>]) -> Self {
        Self { txn_id, events }
    }
}

impl IncomingRequest {
    /// Creates an `IncomingRequest` with the given transaction ID and list of events.
    pub fn new(txn_id: String, events: Vec<Raw<AnyEvent>>) -> IncomingRequest {
        IncomingRequest { txn_id, events }
    }
}

impl Response {
    /// Creates an empty `Response`.
    pub fn new() -> Self {
        Self
    }
}

#[cfg(feature = "server")]
#[cfg(test)]
mod tests {
    use ruma_api::{exports::http, OutgoingRequest};
    use ruma_events::AnyEvent;
    use ruma_serde::Raw;
    use serde_json::json;

    use super::Request;

    #[test]
    fn decode_request_contains_events_field() {
        let dummy_event: AnyEvent = serde_json::from_value(json!({
            "content": {},
            "type": "m.dummy"
        }))
        .unwrap();
        let dummy_event = Raw::from(dummy_event);
        let events = vec![dummy_event];

        let req: http::Request<Vec<u8>> = Request { events: &events, txn_id: "any_txn_id" }
            .try_into_http_request("https://homeserver.tld", Some("auth_tok"))
            .unwrap();
        let json_body: serde_json::Value = serde_json::from_slice(&req.body()).unwrap();

        assert_eq!(
            1,
            json_body.as_object().unwrap().get("events").unwrap().as_array().unwrap().len()
        );
    }
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// This explicit "extern crate" is needed to make the linker aware of the
// `libtock_alloc_init` symbol.
extern crate lang_items;

use arrayref::array_ref;
use core::convert::TryFrom;
use crypto::rng256::ThreadRng256;
use ctap2::ctap::command::{ AuthenticatorMakeCredentialParameters,
    AuthenticatorGetAssertionParameters, AuthenticatorClientPinParameters};
use ctap2::ctap::hid::receive::MessageAssembler;
use ctap2::ctap::hid::send::HidPacketIterator;
use ctap2::ctap::hid::{HidPacket, ChannelID, CtapHid, Message};
use ctap2::ctap::status_code::Ctap2StatusCode;
use ctap2::ctap::CtapState;
use libtock_drivers::timer::{ClockValue, Timestamp};

const COMMAND_INIT: u8 = 0x06;
const CHANNEL_BROADCAST: ChannelID = [0xFF, 0xFF, 0xFF, 0xFF];
const PACKET_TYPE_MASK: u8 = 0x80;

const CLOCK_FREQUENCY_HZ: usize = 32768;
const DUMMY_TIMESTAMP: Timestamp<isize> = Timestamp::from_ms(0);
const DUMMY_CLOCK_VALUE: ClockValue = ClockValue::new(0, CLOCK_FREQUENCY_HZ);

#[derive(Clone, PartialEq)]
pub enum InputType {
    CborMakeCredentialParameter,
    CborGetAssertionParameter,
    CborClientPinParameter,
    Ctap1,
}

// Converts a byte slice into Message
fn raw_to_message(data: &[u8]) -> Message {
    if data.len() <= 4 {
        let mut cid = [0; 4];
        cid[..data.len()].copy_from_slice(data);
        Message {
            cid,
            cmd: 0,
            payload: vec![],
        }
    } else {
        Message {
            cid: array_ref!(data, 0, 4).clone(),
            cmd: data[4],
            payload: data[5..].to_vec(),
        }
    }
}

// Returns an initialized ctap state, hid and the allocated cid
// after processing the init command.
fn initialize<CheckUserPresence>(
    ctap_state: &mut CtapState<ThreadRng256, CheckUserPresence>,
    ctap_hid: &mut CtapHid,
) -> ChannelID
where
    CheckUserPresence: Fn(ChannelID) -> Result<(), Ctap2StatusCode>,
{
    let nonce = vec![0x12, 0x34, 0x56, 0x78, 0x9A, 0xBC, 0xDE, 0xF0];
    let message = Message {
        cid: CHANNEL_BROADCAST,
        cmd: COMMAND_INIT,
        payload: nonce.clone(),
    };
    let mut assembler_reply = MessageAssembler::new();
    let mut result_cid: ChannelID = Default::default();
    for pkt_request in HidPacketIterator::new(message.clone()).unwrap() {
        for pkt_reply in ctap_hid.process_hid_packet(&pkt_request, DUMMY_CLOCK_VALUE, ctap_state) {
            if let Ok(Some(result)) = assembler_reply.parse_packet(&pkt_reply, DUMMY_TIMESTAMP) {
                result_cid.copy_from_slice(&result.payload[8..12]);
            }
        }
    }
    result_cid
}

// Checks whether the given data can be interpreted as the given type.
fn is_type(data: &[u8], input_type: InputType) -> bool {
    if input_type.clone() == InputType::Ctap1 {
        return true;
    }
    match cbor::read(data) {
        Err(_) => false,
        Ok(decoded_cbor) => {
            match input_type {
                InputType::CborMakeCredentialParameter => {
                    AuthenticatorMakeCredentialParameters::try_from(decoded_cbor).is_ok()
                }
                InputType::CborGetAssertionParameter => {
                    AuthenticatorGetAssertionParameters::try_from(decoded_cbor).is_ok()
                }
                InputType::CborClientPinParameter => {
                    AuthenticatorClientPinParameters::try_from(decoded_cbor).is_ok()
                }
                _ => true
            }
        }
    }
}

// Interprets the raw data as of the given input type and
// invokes message splitting, assembling and packet processing at CTAP HID level.
pub fn process_input(data: &[u8], input_type: InputType) {
    if !is_type(data, input_type.clone()) {
        return ();
    }
    // Initialize ctap state and hid and get the allocated cid.
    let mut rng = ThreadRng256 {};
    let user_immediately_present = |_| Ok(());
    let mut ctap_state = CtapState::new(&mut rng, user_immediately_present);
    let mut ctap_hid = CtapHid::new();
    let cid = initialize(&mut ctap_state, &mut ctap_hid);
    // Wrap input as message with allocated cid and command type.
    let mut command = cid.to_vec();
    match input_type {
        InputType::CborMakeCredentialParameter => {
            command.extend(&[0x10, 0x01]);
        }
        InputType::CborGetAssertionParameter => {
            command.extend(&[0x10, 0x02]);
        }
        InputType::CborClientPinParameter => {
            command.extend(&[0x10, 0x06]);
        }
        InputType::Ctap1 => {
            command.extend(&[0x03]);
        }
    }
    command.extend(data);
    // Process message at HID level.
    let message = raw_to_message(&command);
    if let Some(hid_packet_iterator) = HidPacketIterator::new(message) {
        let mut assembler_reply = MessageAssembler::new();
        for pkt_request in hid_packet_iterator {
            for pkt_reply in
                ctap_hid.process_hid_packet(&pkt_request, DUMMY_CLOCK_VALUE, &mut ctap_state)
            {
                // Only checks for assembling crashes, not for semantics.
                let _ = assembler_reply.parse_packet(&pkt_reply, DUMMY_TIMESTAMP);
            }
        }
    }
}

// Splits and reassembles the given data as HID packets.
pub fn split_assemble(data: &[u8]) {
    let mut message = raw_to_message(data);
    if let Some(hid_packet_iterator) = HidPacketIterator::new(message.clone()) {
        let mut assembler = MessageAssembler::new();
        let packets: Vec<HidPacket> = hid_packet_iterator.collect();
        if let Some((last_packet, first_packets)) = packets.split_last() {
            for packet in first_packets {
                assert_eq!(assembler.parse_packet(packet, DUMMY_TIMESTAMP), Ok(None));
            }
            message.cmd &= !PACKET_TYPE_MASK;
            assert_eq!(
                assembler.parse_packet(last_packet, DUMMY_TIMESTAMP),
                Ok(Some(message))
            );
        }
    }
}
<commit_msg>Formatting<commit_after>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// This explicit "extern crate" is needed to make the linker aware of the
// `libtock_alloc_init` symbol.
extern crate lang_items;

use arrayref::array_ref;
use core::convert::TryFrom;
use crypto::rng256::ThreadRng256;
use ctap2::ctap::command::{
    AuthenticatorClientPinParameters, AuthenticatorGetAssertionParameters,
    AuthenticatorMakeCredentialParameters,
};
use ctap2::ctap::hid::receive::MessageAssembler;
use ctap2::ctap::hid::send::HidPacketIterator;
use ctap2::ctap::hid::{ChannelID, CtapHid, HidPacket, Message};
use ctap2::ctap::status_code::Ctap2StatusCode;
use ctap2::ctap::CtapState;
use libtock_drivers::timer::{ClockValue, Timestamp};

const COMMAND_INIT: u8 = 0x06;
const CHANNEL_BROADCAST: ChannelID = [0xFF, 0xFF, 0xFF, 0xFF];
const PACKET_TYPE_MASK: u8 = 0x80;

const CLOCK_FREQUENCY_HZ: usize = 32768;
const DUMMY_TIMESTAMP: Timestamp<isize> = Timestamp::from_ms(0);
const DUMMY_CLOCK_VALUE: ClockValue = ClockValue::new(0, CLOCK_FREQUENCY_HZ);

#[derive(Clone, PartialEq)]
pub enum InputType {
    CborMakeCredentialParameter,
    CborGetAssertionParameter,
    CborClientPinParameter,
    Ctap1,
}

// Converts a byte slice into Message
fn raw_to_message(data: &[u8]) -> Message {
    if data.len() <= 4 {
        let mut cid = [0; 4];
        cid[..data.len()].copy_from_slice(data);
        Message {
            cid,
            cmd: 0,
            payload: vec![],
        }
    } else {
        Message {
            cid: array_ref!(data, 0, 4).clone(),
            cmd: data[4],
            payload: data[5..].to_vec(),
        }
    }
}

// Returns an initialized ctap state, hid and the allocated cid
// after processing the init command.
fn initialize<CheckUserPresence>(
    ctap_state: &mut CtapState<ThreadRng256, CheckUserPresence>,
    ctap_hid: &mut CtapHid,
) -> ChannelID
where
    CheckUserPresence: Fn(ChannelID) -> Result<(), Ctap2StatusCode>,
{
    let nonce = vec![0x12, 0x34, 0x56, 0x78, 0x9A, 0xBC, 0xDE, 0xF0];
    let message = Message {
        cid: CHANNEL_BROADCAST,
        cmd: COMMAND_INIT,
        payload: nonce.clone(),
    };
    let mut assembler_reply = MessageAssembler::new();
    let mut result_cid: ChannelID = Default::default();
    for pkt_request in HidPacketIterator::new(message.clone()).unwrap() {
        for pkt_reply in ctap_hid.process_hid_packet(&pkt_request, DUMMY_CLOCK_VALUE, ctap_state) {
            if let Ok(Some(result)) = assembler_reply.parse_packet(&pkt_reply, DUMMY_TIMESTAMP) {
                result_cid.copy_from_slice(&result.payload[8..12]);
            }
        }
    }
    result_cid
}

// Checks whether the given data can be interpreted as the given type.
fn is_type(data: &[u8], input_type: InputType) -> bool {
    if input_type.clone() == InputType::Ctap1 {
        return true;
    }
    match cbor::read(data) {
        Err(_) => false,
        Ok(decoded_cbor) => match input_type {
            InputType::CborMakeCredentialParameter => {
                AuthenticatorMakeCredentialParameters::try_from(decoded_cbor).is_ok()
            }
            InputType::CborGetAssertionParameter => {
                AuthenticatorGetAssertionParameters::try_from(decoded_cbor).is_ok()
            }
            InputType::CborClientPinParameter => {
                AuthenticatorClientPinParameters::try_from(decoded_cbor).is_ok()
            }
            _ => true,
        },
    }
}

// Interprets the raw data as of the given input type and
// invokes message splitting, assembling and packet processing at CTAP HID level.
pub fn process_input(data: &[u8], input_type: InputType) {
    if !is_type(data, input_type.clone()) {
        return ();
    }
    // Initialize ctap state and hid and get the allocated cid.
    let mut rng = ThreadRng256 {};
    let user_immediately_present = |_| Ok(());
    let mut ctap_state = CtapState::new(&mut rng, user_immediately_present);
    let mut ctap_hid = CtapHid::new();
    let cid = initialize(&mut ctap_state, &mut ctap_hid);
    // Wrap input as message with allocated cid and command type.
    let mut command = cid.to_vec();
    match input_type {
        InputType::CborMakeCredentialParameter => {
            command.extend(&[0x10, 0x01]);
        }
        InputType::CborGetAssertionParameter => {
            command.extend(&[0x10, 0x02]);
        }
        InputType::CborClientPinParameter => {
            command.extend(&[0x10, 0x06]);
        }
        InputType::Ctap1 => {
            command.extend(&[0x03]);
        }
    }
    command.extend(data);
    // Process message at HID level.
    let message = raw_to_message(&command);
    if let Some(hid_packet_iterator) = HidPacketIterator::new(message) {
        let mut assembler_reply = MessageAssembler::new();
        for pkt_request in hid_packet_iterator {
            for pkt_reply in
                ctap_hid.process_hid_packet(&pkt_request, DUMMY_CLOCK_VALUE, &mut ctap_state)
            {
                // Only checks for assembling crashes, not for semantics.
                let _ = assembler_reply.parse_packet(&pkt_reply, DUMMY_TIMESTAMP);
            }
        }
    }
}

// Splits and reassembles the given data as HID packets.
pub fn split_assemble(data: &[u8]) {
    let mut message = raw_to_message(data);
    if let Some(hid_packet_iterator) = HidPacketIterator::new(message.clone()) {
        let mut assembler = MessageAssembler::new();
        let packets: Vec<HidPacket> = hid_packet_iterator.collect();
        if let Some((last_packet, first_packets)) = packets.split_last() {
            for packet in first_packets {
                assert_eq!(assembler.parse_packet(packet, DUMMY_TIMESTAMP), Ok(None));
            }
            message.cmd &= !PACKET_TYPE_MASK;
            assert_eq!(
                assembler.parse_packet(last_packet, DUMMY_TIMESTAMP),
                Ok(Some(message))
            );
        }
    }
}
<|endoftext|>