<commit_before>//! Exposes the `chain` trait and `StackChain` type.

use super::response::Response;
use super::request::Request;
use super::alloy::Alloy;
use super::middleware::{Middleware, Status, Continue};
use super::mixin::Serve;

/// `chains` are the backbone of `Iron`. They coordinate `Middleware`
/// to ensure they are resolved and called in the right order,
/// create and distribute `Alloys`, and handle incoming requests.
///
/// `chains` are internal tools. Unless you want additional
/// or unusual behavior such as enhanced debug logging you
/// probably don't need to mess with `chain` internals.
///
/// That being said, custom `chains` are extremely powerful as they
/// allow you to completely control the resolution of `Middleware`.
pub trait Chain: Send + Clone {
    /// `dispatch` will be called once per `Request`, and may be
    /// called either with or without an existing `Alloy`. `dispatch` is responsible
    /// for delegating the request to the correct `Middleware` and in the correct
    /// order. Effectively, 99% of the work done by a `chain` is done here.
    fn dispatch(&mut self,
                request: &mut Request,
                response: &mut Response,
                opt_alloy: Option<&mut Alloy>) -> Status {
        let mut alloy = &mut Alloy::new();
        match opt_alloy {
            Some(a) => alloy = a,
            None => ()
        };

        let status = self.chain_enter(request, response, alloy);
        let _ = self.chain_exit(request, response, alloy);

        status
    }

    #[doc(hidden)]
    fn chain_enter(&mut self,
             request: &mut Request,
             response: &mut Response,
             alloy: &mut Alloy) -> Status;

    #[doc(hidden)]
    fn chain_exit(&mut self,
                  _request: &mut Request,
                  _response: &mut Response,
                  _alloy: &mut Alloy) -> Status {
        Continue
    }

    /// `link` is responsible for adding new `Middleware` to the `chain's` internal
    /// storage of `Middleware`. Different `chains` may implement different behavior
    /// for `link`, but - ideally - `Middleware` added here will be delegated to during
    /// `Requests`.
    fn link<M: Middleware>(&mut self, _middleware: M);

    /// Create a new instance of `chain`.
    fn new() -> Self;

    #[doc(hidden)]
    fn clone_box(&self) -> Box<Chain + Send> { box self.clone() as Box<Chain + Send> }
}

impl Clone for Box<Chain + Send> {
    fn clone(&self) -> Box<Chain + Send> { self.clone_box() }
}

/// The default `chain` used by `Iron`.
pub mod stackchain {
    use super::super::request::Request;
    use super::super::response::Response;
    use super::super::middleware::{Middleware, Continue, Unwind, Status};
    use super::super::alloy::Alloy;

    use super::Chain;

    /// The default `Chain` used by `Iron`.
    /// `StackChain` runs each `Request` through all `Middleware` in its stack.
    ///
    /// When it hits `Middleware` which returns `Unwind`, it passes
    /// the `Request` back up through all `Middleware` it has hit so far.
    ///
    /// If no `Middleware` return `Unwind` to indicate that they handled
    /// the request, then a 404 is automatically returned.
    pub struct StackChain {
        /// The storage used by `StackChain` to hold all `Middleware`
        /// that have been `linked` to it.
        stack: Vec<Box<Middleware + Send>>,
        unwind: Option<uint>
    }

    impl Clone for StackChain {
        fn clone(&self) -> StackChain {
            StackChain {
                stack: self.stack.clone(),
                unwind: self.unwind.clone()
            }
        }
    }

    /// `StackChain` is a `Chain`
    impl Chain for StackChain {
        fn dispatch(&mut self,
                    request: &mut Request,
                    response: &mut Response,
                    opt_alloy: Option<&mut Alloy>) -> Status {
            let mut alloy = &mut Alloy::new();
            match opt_alloy {
                Some(a) => alloy = a,
                None => ()
            };

            let status = self.chain_enter(request, response, alloy);

            match status {
                Unwind => (),
                Continue => {
                    // If no middleware returned unwind, then we send a 404.
                    // At least one middleware should return unwind when a
                    // terminal endpoint, such as a router, has been reached.
                    //
                    // We don't write to the body as other middleware may want
                    // to change headers.
                    response.status = ::http::status::NotFound;
                }
            }

            let _ = self.chain_exit(request, response, alloy);

            status
        }

        fn chain_enter(&mut self,
                 request: &mut Request,
                 response: &mut Response,
                 alloy: &mut Alloy) -> Status {
            // The `exit_stack` will hold all `Middleware` that are passed through
            // in the enter loop. This is so we know to take exactly the same
            // path through `Middleware` in reverse order than we did on the way in.
            self.unwind = None;

            'enter: for (i, middleware) in self.stack.mut_iter().enumerate() {
                match middleware.enter(request, response, alloy) {
                    Unwind   => {
                        self.unwind = Some(i);
                        return Unwind;
                    }
                    // Mark the middleware for traversal on exit.
                    Continue => ()
                }
            }

            Continue
        }

        fn chain_exit(&mut self,
                 request: &mut Request,
                 response: &mut Response,
                 alloy: &mut Alloy) -> Status {
            match self.unwind {
                Some(i) => {
                    'exit: for middleware in self.stack.mut_slice_to(i).mut_iter().rev() {
                        let _ = middleware.exit(request, response, alloy);
                    }
                },
                None => {
                    'exit: for middleware in self.stack.mut_iter().rev() {
                        let _ = middleware.exit(request, response, alloy);
                    }
                }
            }

            Continue
        }

        /// Add `Middleware` to the `Chain`.
        fn link<M: Middleware>(&mut self, middleware: M) {
            self.stack.push(box middleware);
        }

        /// Create a new instance of `StackChain`.
        fn new() -> StackChain {
            StackChain {
                stack: vec![],
                unwind: None
            }
        }
    }

    impl FromIterator<Box<Middleware + Send>> for StackChain {
        fn from_iter<T: Iterator<Box<Middleware + Send>>>(mut iterator: T) -> StackChain {
            StackChain {
                stack: iterator.collect(),
                unwind: None
            }
        }
    }

    #[cfg(test)]
    mod test {
        pub use super::*;
        pub use super::super::super::request::Request;
        pub use super::super::super::response::Response;
        pub use super::super::super::alloy::Alloy;
        pub use super::super::super::middleware::{Middleware, Status, Continue, Unwind};
        pub use std::sync::{Arc, Mutex};

        #[deriving(Clone)]
        pub struct CallCount {
            enter: Arc<Mutex<u64>>,
            exit: Arc<Mutex<u64>>
        }

        impl Middleware for CallCount {
            fn enter(&mut self, _req: &mut Request,
                     _res: &mut Response, _alloy: &mut Alloy) -> Status {
                let mut enter = self.enter.lock();
                *enter += 1;
                Continue
            }

            fn exit(&mut self, _req: &mut Request,
                    _res: &mut Response, _alloy: &mut Alloy) -> Status {
                let mut exit = self.exit.lock();
                *exit += 1;
                Continue
            }
        }

        #[deriving(Clone)]
        pub struct Stopper;

        impl Middleware for Stopper {
            fn enter(&mut self, _req: &mut Request,
                     _res: &mut Response, _alloy: &mut Alloy) -> Status {
                Unwind // Stop .status from being accessed, which fails.
            }
        }

        mod dispatch {
            use super::{CallCount, Arc, Mutex, Stopper};
            use super::super::StackChain;
            use super::super::super::Chain;
            use std::mem::{uninitialized};

            #[test]
            fn calls_middleware_enter() {
                let mut testchain: StackChain = Chain::new();
                let enter = Arc::new(Mutex::new(0));
                let exit = Arc::new(Mutex::new(0));
                testchain.link(CallCount { enter: enter.clone(), exit: exit.clone() });
                testchain.link(Stopper);
                unsafe {
                    let _ = testchain.dispatch(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }
                assert_eq!(*enter.lock(), 1);
            }

            #[test]
            fn calls_middleware_exit() {
                let mut testchain: StackChain = Chain::new();
                let enter = Arc::new(Mutex::new(0));
                let exit = Arc::new(Mutex::new(0));
                testchain.link(CallCount { enter: enter.clone(), exit: exit.clone() });
                testchain.link(Stopper);
                unsafe {
                    let _ = testchain.dispatch(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }
                assert_eq!(*exit.lock(), 1);
            }

            #[test]
            fn calls_all_middleware_enter_exit() {
                let mut testchain: StackChain = Chain::new();
                let enter_exits: Vec<(Arc<Mutex<u64>>, Arc<Mutex<u64>>)> = vec![];

                for _ in range(0u8, 10) {
                    let (enter, exit) = (Arc::new(Mutex::new(0)), Arc::new(Mutex::new(0)));
                    testchain.link(CallCount { enter: enter.clone(), exit: exit.clone() });
                }

                testchain.link(Stopper);
                unsafe {
                    let _ = testchain.dispatch(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }

                for (enter, exit) in enter_exits.move_iter() {
                    assert_eq!(*enter.lock(), 1);
                    assert_eq!(*exit.lock(), 1);
                }
            }
        }

        mod chain_enter {
            use super::{CallCount, Arc, Mutex, Stopper};
            use super::super::StackChain;
            use super::super::super::Chain;
            use std::mem::{uninitialized};

            #[test]
            fn calls_middleware_enter() {
                let mut testchain: StackChain = Chain::new();
                let enter = Arc::new(Mutex::new(0));
                let exit = Arc::new(Mutex::new(0));
                testchain.link(CallCount { enter: enter.clone(), exit: exit.clone() });
                testchain.link(Stopper);
                unsafe {
                    let _ = testchain.chain_enter(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }
                assert_eq!(*enter.lock(), 1);
            }

            #[test]
            fn doesnt_call_middleware_exit() {
                let mut testchain: StackChain = Chain::new();
                let enter = Arc::new(Mutex::new(0));
                let exit = Arc::new(Mutex::new(0));
                testchain.link(CallCount { enter: enter.clone(), exit: exit.clone() });
                testchain.link(Stopper);
                unsafe {
                    let _ = testchain.chain_enter(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }
                assert_eq!(*exit.lock(), 0);
            }
        }

        mod chain_exit {
            use super::{CallCount, Arc, Mutex, Stopper};
            use super::super::StackChain;
            use super::super::super::Chain;
            use std::mem::uninitialized;

            #[test]
            fn calls_middleware_exit() {
                let mut testchain: StackChain = Chain::new();
                let enter = Arc::new(Mutex::new(0));
                let exit = Arc::new(Mutex::new(0));
                testchain.link(CallCount {
                    enter: enter.clone(), exit: exit.clone()
                });
                testchain.link(Stopper);
                unsafe {
                    let _  = testchain.chain_enter(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );

                    let _  = testchain.chain_exit(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }
                assert_eq!(*exit.lock(), 1);
            }

            #[test]
            fn doesnt_call_middleware_enter() {
                let mut testchain: StackChain = Chain::new();
                let enter = Arc::new(Mutex::new(0));
                let exit = Arc::new(Mutex::new(0));
                testchain.link(CallCount {
                    enter: enter.clone(), exit: exit.clone()
                });
                testchain.unwind = Some(1);
                unsafe {
                    let _  = testchain.chain_exit(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }
                assert_eq!(*enter.lock(), 0);
            }
        }

        mod bench {
            use super::super::super::super::middleware::Middleware;
            pub use super::Stopper;

            #[deriving(Clone)]
            struct Noop;

            impl Middleware for Noop {}

            macro_rules! bench_noop_x (
                ($name:ident, $num:expr, $method:ident) => {
                    #[bench]
                    fn $name(b: &mut Bencher) {
                        let mut testchain: StackChain = Chain::new();
                        for _ in range(0, $num) {
                            testchain.link(Noop);
                        }
                        testchain.link(Stopper);
                        b.iter(|| {
                            black_box(unsafe {
                                let _ = testchain.$method(
                                    uninitialized(),
                                    uninitialized(),
                                    uninitialized()
                                );
                            })
                        });
                    }
                }
            )

            macro_rules! bench_method (
                ($method:ident) => {
                    mod $method {
                        use std::mem::uninitialized;
                        use test::{Bencher, black_box};
                        use super::{Noop, Stopper};
                        use super::super::super::StackChain;
                        use super::super::super::super::Chain;

                        bench_noop_x!(bench_empty, 0u8, $method)
                        bench_noop_x!(bench_1, 1u8, $method)
                        bench_noop_x!(bench_2, 2u8, $method)
                        bench_noop_x!(bench_3, 3u8, $method)
                        bench_noop_x!(bench_4, 4u8, $method)
                        bench_noop_x!(bench_10, 10u8, $method)
                        bench_noop_x!(bench_100, 100u8, $method)
                    }
                }
            )

            bench_method!(dispatch)
            bench_method!(chain_enter)
        }
    }
}
<commit_msg>(feat) Added chain_error to Chain and StackChain.<commit_after>//! Exposes the `chain` trait and `StackChain` type.

use std::fmt::Show;

use super::response::Response;
use super::request::Request;
use super::alloy::Alloy;
use super::middleware::{Middleware, Status, Continue, Error};

/// `chains` are the backbone of `Iron`. They coordinate `Middleware`
/// to ensure they are resolved and called in the right order,
/// create and distribute `Alloys`, and handle incoming requests.
///
/// `chains` are internal tools. Unless you want additional
/// or unusual behavior such as enhanced debug logging you
/// probably don't need to mess with `chain` internals.
///
/// That being said, custom `chains` are extremely powerful as they
/// allow you to completely control the resolution of `Middleware`.
pub trait Chain: Send + Clone {
    /// `dispatch` will be called once per `Request`, and may be
    /// called either with or without an existing `Alloy`. `dispatch` is responsible
    /// for delegating the request to the correct `Middleware` and in the correct
    /// order. Effectively, 99% of the work done by a `chain` is done here.
    fn dispatch(&mut self,
                request: &mut Request,
                response: &mut Response,
                opt_alloy: Option<&mut Alloy>) -> Status {
        let mut alloy = &mut Alloy::new();
        match opt_alloy {
            Some(a) => alloy = a,
            None => ()
        };

        let mut status = self.chain_enter(request, response, alloy);
        match status {
            Error(ref mut e) => {
                let error: &mut Show = *e;
                let _ = self.chain_error(request, response, alloy, error);
            },
            _ => {
                let _ = self.chain_exit(request, response, alloy);
            }
        };

        status
    }

    #[doc(hidden)]
    fn chain_enter(&mut self,
             request: &mut Request,
             response: &mut Response,
             alloy: &mut Alloy) -> Status;

    #[doc(hidden)]
    fn chain_exit(&mut self,
                  _request: &mut Request,
                  _response: &mut Response,
                  _alloy: &mut Alloy) -> Status {
        Continue
    }

    #[doc(hidden)]
    fn chain_error(&mut self,
                  _request: &mut Request,
                  _response: &mut Response,
                  _alloy: &mut Alloy,
                  _error: &mut Show) { () }

    /// `link` is responsible for adding new `Middleware` to the `chain's` internal
    /// storage of `Middleware`. Different `chains` may implement different behavior
    /// for `link`, but - ideally - `Middleware` added here will be delegated to during
    /// `Requests`.
    fn link<M: Middleware>(&mut self, _middleware: M);

    /// Create a new instance of `chain`.
    fn new() -> Self;

    #[doc(hidden)]
    fn clone_box(&self) -> Box<Chain + Send> { box self.clone() as Box<Chain + Send> }
}

impl Clone for Box<Chain + Send> {
    fn clone(&self) -> Box<Chain + Send> { self.clone_box() }
}

/// The default `chain` used by `Iron`.
pub mod stackchain {
    use std::fmt::Show;

    use super::super::request::Request;
    use super::super::response::Response;
    use super::super::middleware::{Middleware, Continue, Unwind, Error, Status};
    use super::super::alloy::Alloy;

    use super::Chain;

    /// The default `Chain` used by `Iron`.
    /// `StackChain` runs each `Request` through all `Middleware` in its stack.
    ///
    /// When it hits `Middleware` which returns `Unwind`, it passes
    /// the `Request` back up through all `Middleware` it has hit so far.
<<<<<<< HEAD
    ///
    /// If no `Middleware` return `Unwind` to indicate that they handled
    /// the request, then a 404 is automatically returned.
=======
    #[deriving(Clone)]
>>>>>>> (feat) Added chain_error to Chain and StackChain.
    pub struct StackChain {
        /// The storage used by `StackChain` to hold all `Middleware`
        /// that have been `linked` to it.
        stack: Vec<Box<Middleware + Send>>,
        status: ChainStatus,
    }

    #[deriving(Clone)]
    enum ChainStatus { Unwound(uint), Errored(uint), Unhandled }

    /// `StackChain` is a `Chain`
    impl Chain for StackChain {
        fn dispatch(&mut self,
                    request: &mut Request,
                    response: &mut Response,
                    opt_alloy: Option<&mut Alloy>) -> Status {
            let mut alloy = &mut Alloy::new();
            match opt_alloy {
                Some(a) => alloy = a,
                None => ()
            };

            let status = self.chain_enter(request, response, alloy);

            match status {
                Unwind => (),
                Continue => {
                    // If no middleware returned unwind, then we send a 404.
                    // At least one middleware should return unwind when a
                    // terminal endpoint, such as a router, has been reached.
                    //
                    // We don't write to the body as other middleware may want
                    // to change headers.
                    response.status = ::http::status::NotFound;
                }
            }

            let _ = self.chain_exit(request, response, alloy);

            status
        }

        fn chain_enter(&mut self,
                 request: &mut Request,
                 response: &mut Response,
                 alloy: &mut Alloy) -> Status {
            // The `exit_stack` will hold all `Middleware` that are passed through
            // in the enter loop. This is so we know to take exactly the same
            // path through `Middleware` in reverse order than we did on the way in.
            self.status = Unhandled;

            'enter: for (i, middleware) in self.stack.mut_iter().enumerate() {
                match middleware.enter(request, response, alloy) {
                    Unwind   => {
                        self.status = Unwound(i);
                        return Unwind;
                    },
                    e @ Error(_) => {
                        self.status = Errored(i);
                        return e;
                    }
                    // Mark the middleware for traversal on exit.
                    Continue => ()
                }
            }

            self.status = Unhandled;
            Continue
        }

        fn chain_exit(&mut self,
                 request: &mut Request,
                 response: &mut Response,
                 alloy: &mut Alloy) -> Status {
            match self.status {
                Unwound(i) => {
                    for middleware in self.stack.mut_slice_to(i).mut_iter().rev() {
                        let _ = middleware.exit(request, response, alloy);
                    }
                },
                Unhandled => {
                    for middleware in self.stack.mut_iter().rev() {
                        let _ = middleware.exit(request, response, alloy);
                    }
                },
                Errored(_) => fail!("chain_exit called on a StackChain which Errored.")
            }

            Continue
        }

        fn chain_error(&mut self,
                      request: &mut Request,
                      response: &mut Response,
                      alloy: &mut Alloy,
                      error: &mut Show) {
            match self.status {
                Errored(i) => {
                    for middleware in self.stack.mut_slice_to(i).mut_iter().rev() {
                        let _ = middleware.on_error(request, response, alloy, error);
                    }
                },
                _ => fail!("chain_error called on a chain which did not error.")
            }
        }

        /// Add `Middleware` to the `Chain`.
        fn link<M: Middleware>(&mut self, middleware: M) {
            self.stack.push(box middleware);
        }

        /// Create a new instance of `StackChain`.
        fn new() -> StackChain {
            StackChain {
                stack: vec![],
                status: Unhandled
            }
        }
    }

    impl FromIterator<Box<Middleware + Send>> for StackChain {
        fn from_iter<T: Iterator<Box<Middleware + Send>>>(mut iterator: T) -> StackChain {
            StackChain {
                stack: iterator.collect(),
                status: Unhandled
            }
        }
    }

    #[cfg(test)]
    mod test {
        pub use super::*;
        pub use super::super::super::request::Request;
        pub use super::super::super::response::Response;
        pub use super::super::super::alloy::Alloy;
        pub use super::super::super::middleware::{Middleware, Status, Continue, Unwind};
        pub use std::sync::{Arc, Mutex};

        #[deriving(Clone)]
        pub struct CallCount {
            enter: Arc<Mutex<u64>>,
            exit: Arc<Mutex<u64>>
        }

        impl Middleware for CallCount {
            fn enter(&mut self, _req: &mut Request,
                     _res: &mut Response, _alloy: &mut Alloy) -> Status {
                let mut enter = self.enter.lock();
                *enter += 1;
                Continue
            }

            fn exit(&mut self, _req: &mut Request,
                    _res: &mut Response, _alloy: &mut Alloy) -> Status {
                let mut exit = self.exit.lock();
                *exit += 1;
                Continue
            }
        }

        #[deriving(Clone)]
        pub struct Stopper;

        impl Middleware for Stopper {
            fn enter(&mut self, _req: &mut Request,
                     _res: &mut Response, _alloy: &mut Alloy) -> Status {
                Unwind // Stop .status from being accessed, which fails.
            }
        }

        mod dispatch {
            use super::{CallCount, Arc, Mutex, Stopper};
            use super::super::StackChain;
            use super::super::super::Chain;
            use std::mem::{uninitialized};

            #[test]
            fn calls_middleware_enter() {
                let mut testchain: StackChain = Chain::new();
                let enter = Arc::new(Mutex::new(0));
                let exit = Arc::new(Mutex::new(0));
                testchain.link(CallCount { enter: enter.clone(), exit: exit.clone() });
                testchain.link(Stopper);
                unsafe {
                    let _ = testchain.dispatch(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }
                assert_eq!(*enter.lock(), 1);
            }

            #[test]
            fn calls_middleware_exit() {
                let mut testchain: StackChain = Chain::new();
                let enter = Arc::new(Mutex::new(0));
                let exit = Arc::new(Mutex::new(0));
                testchain.link(CallCount { enter: enter.clone(), exit: exit.clone() });
                testchain.link(Stopper);
                unsafe {
                    let _ = testchain.dispatch(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }
                assert_eq!(*exit.lock(), 1);
            }

            #[test]
            fn calls_all_middleware_enter_exit() {
                let mut testchain: StackChain = Chain::new();
                let enter_exits: Vec<(Arc<Mutex<u64>>, Arc<Mutex<u64>>)> = vec![];

                for _ in range(0u8, 10) {
                    let (enter, exit) = (Arc::new(Mutex::new(0)), Arc::new(Mutex::new(0)));
                    testchain.link(CallCount { enter: enter.clone(), exit: exit.clone() });
                }

                testchain.link(Stopper);
                unsafe {
                    let _ = testchain.dispatch(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }

                for (enter, exit) in enter_exits.move_iter() {
                    assert_eq!(*enter.lock(), 1);
                    assert_eq!(*exit.lock(), 1);
                }
            }
        }

        mod chain_enter {
            use super::{CallCount, Arc, Mutex, Stopper};
            use super::super::StackChain;
            use super::super::super::Chain;
            use std::mem::{uninitialized};

            #[test]
            fn calls_middleware_enter() {
                let mut testchain: StackChain = Chain::new();
                let enter = Arc::new(Mutex::new(0));
                let exit = Arc::new(Mutex::new(0));
                testchain.link(CallCount { enter: enter.clone(), exit: exit.clone() });
                testchain.link(Stopper);
                unsafe {
                    let _ = testchain.chain_enter(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }
                assert_eq!(*enter.lock(), 1);
            }

            #[test]
            fn doesnt_call_middleware_exit() {
                let mut testchain: StackChain = Chain::new();
                let enter = Arc::new(Mutex::new(0));
                let exit = Arc::new(Mutex::new(0));
                testchain.link(CallCount { enter: enter.clone(), exit: exit.clone() });
                testchain.link(Stopper);
                unsafe {
                    let _ = testchain.chain_enter(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }
                assert_eq!(*exit.lock(), 0);
            }
        }

        mod chain_exit {
            use super::{CallCount, Arc, Mutex, Stopper};
            use super::super::StackChain;
            use super::super::super::Chain;
            use std::mem::uninitialized;

            #[test]
            fn calls_middleware_exit() {
                let mut testchain: StackChain = Chain::new();
                let enter = Arc::new(Mutex::new(0));
                let exit = Arc::new(Mutex::new(0));
                testchain.link(CallCount {
                    enter: enter.clone(), exit: exit.clone()
                });
                testchain.link(Stopper);
                unsafe {
                    let _  = testchain.chain_enter(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );

                    let _  = testchain.chain_exit(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }
                assert_eq!(*exit.lock(), 1);
            }

            #[test]
            fn doesnt_call_middleware_enter() {
                let mut testchain: StackChain = Chain::new();
                let enter = Arc::new(Mutex::new(0));
                let exit = Arc::new(Mutex::new(0));
                testchain.link(CallCount {
                    enter: enter.clone(), exit: exit.clone()
                });
                testchain.status = Unwound(1);
                unsafe {
                    let _  = testchain.chain_exit(
                        uninitialized(),
                        uninitialized(),
                        uninitialized()
                    );
                }
                assert_eq!(*enter.lock(), 0);
            }
        }

        mod bench {
            use super::super::super::super::middleware::Middleware;
            pub use super::Stopper;

            #[deriving(Clone)]
            struct Noop;

            impl Middleware for Noop {}

            macro_rules! bench_noop_x (
                ($name:ident, $num:expr, $method:ident) => {
                    #[bench]
                    fn $name(b: &mut Bencher) {
                        let mut testchain: StackChain = Chain::new();
                        for _ in range(0, $num) {
                            testchain.link(Noop);
                        }
                        testchain.link(Stopper);
                        b.iter(|| {
                            black_box(unsafe {
                                let _ = testchain.$method(
                                    uninitialized(),
                                    uninitialized(),
                                    uninitialized()
                                );
                            })
                        });
                    }
                }
            )

            macro_rules! bench_method (
                ($method:ident) => {
                    mod $method {
                        use std::mem::uninitialized;
                        use test::{Bencher, black_box};
                        use super::{Noop, Stopper};
                        use super::super::super::StackChain;
                        use super::super::super::super::Chain;

                        bench_noop_x!(bench_empty, 0u8, $method)
                        bench_noop_x!(bench_1, 1u8, $method)
                        bench_noop_x!(bench_2, 2u8, $method)
                        bench_noop_x!(bench_3, 3u8, $method)
                        bench_noop_x!(bench_4, 4u8, $method)
                        bench_noop_x!(bench_10, 10u8, $method)
                        bench_noop_x!(bench_100, 100u8, $method)
                    }
                }
            )

            bench_method!(dispatch)
            bench_method!(chain_enter)
        }
    }
}
<|endoftext|>


<commit_before>extern crate serde_json;

use self::super::{ERR_HOST, ERR_HTSI, ERR_INVALID, MAP, OK};
use self::super::super::http;

use base64::{encode, decode};
use std::slice;
use std::u16;
use std::u8;

#[derive(Serialize)]
struct RequestData {
    dad: u8,
    sad: u8,
    lenc: usize,
    command: String,
    lenr: usize,
}

#[allow(non_snake_case)]
#[derive(Deserialize)]
struct ResponseData {
    dad: u8,
    sad: u8,
    lenr: usize,
    response: String,
    responseCode: i8,
}

pub fn data(ctn: u16,
            dad: *mut u8,
            sad: *mut u8,
            lenc: usize,
            command: *const u8,
            lenr: *mut usize,
            response: *mut u8)
            -> i8 {

    debug!("ctn: {}", ctn);
    if !(ctn >= u16::MIN && ctn <= u16::MAX) {
        error!("ctn is not an u16. Returning {}", ERR_INVALID);
        return ERR_INVALID;
    }

    let safe_dad: &mut u8 = unsafe { &mut *dad };
    debug!("dad: {}", safe_dad);
    if !(safe_dad >= &mut u8::MIN && safe_dad <= &mut u8::MAX) {
        error!("dad is not an u8. Returning {}", ERR_INVALID);
        return ERR_INVALID;
    }

    let safe_sad: &mut u8 = unsafe { &mut *sad };
    debug!("sad: {}", safe_sad);
    if !(safe_sad >= &mut u8::MIN && safe_sad <= &mut u8::MAX) {
        error!("sad is not an u8. Returning {}", ERR_INVALID);
        return ERR_INVALID;
    }

    debug!("lenc: {}", lenc);
    if !(lenc > u16::MIN as usize && lenc < u16::MAX as usize) {
        error!("lenc is not an u16. Returning {}", ERR_INVALID);
        return ERR_INVALID;
    }

    let safe_command = unsafe { slice::from_raw_parts(command, lenc) };
    debug!("command: {:?}", safe_command);

    let safe_lenr: &mut usize = unsafe { &mut *lenr };
    debug!("lenr: {}", safe_lenr);
    sanitize_lenr(&mut *safe_lenr);

    let safe_response = unsafe { slice::from_raw_parts_mut(response, *safe_lenr) };
    debug!("response with {} slices formed", safe_response.len());

    if !MAP.lock().unwrap().contains_key(&ctn) {
        debug!("Card terminal has not been opened. Returning {}",
               ERR_INVALID);
        return ERR_INVALID;
    }

    let request_data = RequestData {
        dad: *safe_dad,
        sad: *safe_sad,
        lenc,
        command: encode(safe_command),
        lenr: *safe_lenr,
    };

    let path = get_request_path(ctn);
    let http_response = match http::post(path, &request_data) {
        Ok(http_response) => http_response,
        Err(error) => {
            debug!("Error: {:?}", error);
            error!("Request failed! Returning {}", ERR_HTSI);
            return ERR_HTSI;
        }
    };

    let (http_status, response_body) = http::extract_response(http_response);
    match http_status {
        http::HttpStatus::Ok => {
            let data: ResponseData = serde_json::from_str(&response_body).unwrap();

            if data.responseCode == OK {
                *safe_dad = data.dad;
                *safe_sad = data.sad;
                *safe_lenr = data.lenr;

                let decoded = decode(&data.response).unwrap();
                debug!("Decoded response field {:?}", decoded);

                for (place, element) in safe_response.iter_mut().zip(decoded.iter()) {
                    *place = *element;
                }
            }
            debug!("Returning {}", data.responseCode);
            return data.responseCode;
        }
        _ => {
            error!("Response not OK! Returning {}", ERR_HOST);
            ERR_HOST
        }
    }
}

fn get_request_path(ctn: u16) -> String {
    let pn = MAP.lock().unwrap();
    let pn = pn.get(&ctn).unwrap();

    let mut path = String::from("ct_data");
    path.push_str("/");
    path.push_str(&ctn.to_string());
    path.push_str("/");
    path.push_str(&pn.to_string());

    path
}

fn sanitize_lenr(lenr: &mut usize) {
    let max_usize = u16::MAX as usize;
    let min_usize = u16::MIN as usize;
    if *lenr < min_usize || *lenr > max_usize {
        debug!("... sanitize lenr to {}", u16::MAX);
        *lenr = max_usize;
    }
}
<commit_msg>adjust range test<commit_after>extern crate serde_json;

use self::super::{ERR_HOST, ERR_HTSI, ERR_INVALID, MAP, OK};
use self::super::super::http;

use base64::{encode, decode};
use std::slice;
use std::u16;
use std::u8;

#[derive(Serialize)]
struct RequestData {
    dad: u8,
    sad: u8,
    lenc: usize,
    command: String,
    lenr: usize,
}

#[allow(non_snake_case)]
#[derive(Deserialize)]
struct ResponseData {
    dad: u8,
    sad: u8,
    lenr: usize,
    response: String,
    responseCode: i8,
}

pub fn data(ctn: u16,
            dad: *mut u8,
            sad: *mut u8,
            lenc: usize,
            command: *const u8,
            lenr: *mut usize,
            response: *mut u8)
            -> i8 {

    debug!("ctn: {}", ctn);
    if ctn < u16::MIN && ctn > u16::MAX {
        error!("ctn is not an u16. Returning {}", ERR_INVALID);
        return ERR_INVALID;
    }

    let safe_dad: &mut u8 = unsafe { &mut *dad };
    debug!("dad: {}", safe_dad);
    if safe_dad < &mut u8::MIN && safe_dad > &mut u8::MAX {
        error!("dad is not an u8. Returning {}", ERR_INVALID);
        return ERR_INVALID;
    }

    let safe_sad: &mut u8 = unsafe { &mut *sad };
    debug!("sad: {}", safe_sad);
    if safe_sad < &mut u8::MIN && safe_sad > &mut u8::MAX {
        error!("sad is not an u8. Returning {}", ERR_INVALID);
        return ERR_INVALID;
    }

    debug!("lenc: {}", lenc);
    if lenc < u16::MIN as usize && lenc > u16::MAX as usize {
        error!("lenc is not an u16. Returning {}", ERR_INVALID);
        return ERR_INVALID;
    }

    let safe_command = unsafe { slice::from_raw_parts(command, lenc) };
    debug!("command: {:?}", safe_command);

    let safe_lenr: &mut usize = unsafe { &mut *lenr };
    debug!("lenr: {}", safe_lenr);
    sanitize_lenr(&mut *safe_lenr);

    let safe_response = unsafe { slice::from_raw_parts_mut(response, *safe_lenr) };
    debug!("response with {} slices formed", safe_response.len());

    if !MAP.lock().unwrap().contains_key(&ctn) {
        debug!("Card terminal has not been opened. Returning {}",
               ERR_INVALID);
        return ERR_INVALID;
    }

    let request_data = RequestData {
        dad: *safe_dad,
        sad: *safe_sad,
        lenc,
        command: encode(safe_command),
        lenr: *safe_lenr,
    };

    let path = get_request_path(ctn);
    let http_response = match http::post(path, &request_data) {
        Ok(http_response) => http_response,
        Err(error) => {
            debug!("Error: {:?}", error);
            error!("Request failed! Returning {}", ERR_HTSI);
            return ERR_HTSI;
        }
    };

    let (http_status, response_body) = http::extract_response(http_response);
    match http_status {
        http::HttpStatus::Ok => {
            let data: ResponseData = serde_json::from_str(&response_body).unwrap();

            if data.responseCode == OK {
                *safe_dad = data.dad;
                *safe_sad = data.sad;
                *safe_lenr = data.lenr;

                let decoded = decode(&data.response).unwrap();
                debug!("Decoded response field {:?}", decoded);

                for (place, element) in safe_response.iter_mut().zip(decoded.iter()) {
                    *place = *element;
                }
            }
            debug!("Returning {}", data.responseCode);
            return data.responseCode;
        }
        _ => {
            error!("Response not OK! Returning {}", ERR_HOST);
            ERR_HOST
        }
    }
}

fn get_request_path(ctn: u16) -> String {
    let pn = MAP.lock().unwrap();
    let pn = pn.get(&ctn).unwrap();

    let mut path = String::from("ct_data");
    path.push_str("/");
    path.push_str(&ctn.to_string());
    path.push_str("/");
    path.push_str(&pn.to_string());

    path
}

fn sanitize_lenr(lenr: &mut usize) {
    let max_usize = u16::MAX as usize;
    let min_usize = u16::MIN as usize;
    if *lenr < min_usize || *lenr > max_usize {
        debug!("... sanitize lenr to {}", u16::MAX);
        *lenr = max_usize;
    }
}
<|endoftext|>


<commit_before>use super::algorithm::ECIES;
use crate::{errors::ECIESError, transport::Transport, types::PeerId};
use anyhow::{bail, Context as _};
use bytes::{Bytes, BytesMut};
use futures::{ready, Sink, SinkExt};
use secp256k1::SecretKey;
use std::{
    fmt::Debug,
    io,
    pin::Pin,
    task::{Context, Poll},
};
use tokio::stream::*;
use tokio_util::codec::*;
use tracing::*;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
/// Current ECIES state of a connection
pub enum ECIESState {
    Auth,
    Ack,
    Header,
    Body,
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// Raw egress values for an ECIES protocol
pub enum EgressECIESValue {
    Auth,
    Ack,
    Message(Bytes),
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// Raw ingress values for an ECIES protocol
pub enum IngressECIESValue {
    AuthReceive(PeerId),
    Ack,
    Message(Vec<u8>),
}

/// Tokio codec for ECIES
#[derive(Debug)]
pub struct ECIESCodec {
    ecies: ECIES,
    state: ECIESState,
}

impl ECIESCodec {
    /// Create a new server codec using the given secret key
    pub fn new_server(secret_key: SecretKey) -> Result<Self, ECIESError> {
        Ok(Self {
            ecies: ECIES::new_server(secret_key)?,
            state: ECIESState::Auth,
        })
    }

    /// Create a new client codec using the given secret key and the server's public id
    pub fn new_client(secret_key: SecretKey, remote_id: PeerId) -> Result<Self, ECIESError> {
        Ok(Self {
            ecies: ECIES::new_client(secret_key, remote_id)?,
            state: ECIESState::Auth,
        })
    }
}

impl Decoder for ECIESCodec {
    type Item = IngressECIESValue;
    type Error = io::Error;

    #[instrument(level = "trace", skip(self, buf), fields(peer=&*format!("{:?}", self.ecies.remote_id.map(|s| s.to_string())), state=&*format!("{:?}", self.state)))]
    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        loop {
            match self.state {
                ECIESState::Auth => {
                    trace!("parsing auth");
                    if buf.len() < 2 {
                        return Ok(None);
                    }

                    let payload_size = u16::from_be_bytes([buf[0], buf[1]]) as usize;
                    let total_size = payload_size + 2;

                    if buf.len() < total_size {
                        trace!("current len {}, need {}", buf.len(), total_size);
                        return Ok(None);
                    }

                    self.ecies.read_auth(&mut *buf.split_to(total_size))?;

                    self.state = ECIESState::Header;
                    return Ok(Some(IngressECIESValue::AuthReceive(self.ecies.remote_id())));
                }
                ECIESState::Ack => {
                    trace!("parsing ack with len {}", buf.len());
                    if buf.len() < 2 {
                        return Ok(None);
                    }

                    let payload_size = u16::from_be_bytes([buf[0], buf[1]]) as usize;
                    let total_size = payload_size + 2;

                    if buf.len() < total_size {
                        trace!("current len {}, need {}", buf.len(), total_size);
                        return Ok(None);
                    }

                    self.ecies.read_ack(&mut *buf.split_to(total_size))?;

                    self.state = ECIESState::Header;
                    return Ok(Some(IngressECIESValue::Ack));
                }
                ECIESState::Header => {
                    if buf.len() < ECIES::header_len() {
                        return Ok(None);
                    }

                    self.ecies
                        .read_header(&mut *buf.split_to(ECIES::header_len()))?;

                    self.state = ECIESState::Body;
                }
                ECIESState::Body => {
                    if buf.len() < self.ecies.body_len() {
                        return Ok(None);
                    }

                    let mut data = buf.split_to(self.ecies.body_len());
                    let ret = self.ecies.read_body(&mut *data)?;

                    self.state = ECIESState::Header;
                    return Ok(Some(IngressECIESValue::Message(ret.to_vec())));
                }
            }
        }
    }
}

impl Encoder<EgressECIESValue> for ECIESCodec {
    type Error = io::Error;

    #[instrument(level = "trace", skip(self, buf), fields(peer=&*format!("{:?}", self.ecies.remote_id.map(|s| s.to_string())), state=&*format!("{:?}", self.state)))]
    fn encode(&mut self, item: EgressECIESValue, buf: &mut BytesMut) -> Result<(), Self::Error> {
        match item {
            EgressECIESValue::Auth => {
                self.state = ECIESState::Ack;
                self.ecies.write_auth(buf);
                Ok(())
            }
            EgressECIESValue::Ack => {
                self.state = ECIESState::Header;
                self.ecies.write_ack(buf);
                Ok(())
            }
            EgressECIESValue::Message(data) => {
                self.ecies.write_header(buf, data.len());
                self.ecies.write_body(buf, &data);
                Ok(())
            }
        }
    }
}

/// `ECIES` stream over TCP exchanging raw bytes
#[derive(Debug)]
pub struct ECIESStream<Io> {
    stream: Framed<Io, ECIESCodec>,
    remote_id: PeerId,
}

impl<Io> ECIESStream<Io>
where
    Io: Transport,
{
    /// Connect to an `ECIES` server
    #[instrument(skip(transport, secret_key), fields(peer=&*format!("{:?}", transport.remote_addr())))]
    pub async fn connect(
        transport: Io,
        secret_key: SecretKey,
        remote_id: PeerId,
    ) -> anyhow::Result<Self> {
        let ecies = ECIESCodec::new_client(secret_key, remote_id)
            .map_err(|_| io::Error::new(io::ErrorKind::Other, "invalid handshake"))?;

        let mut transport = ecies.framed(transport);

        trace!("sending ecies auth ...");
        transport.send(EgressECIESValue::Auth).await?;

        trace!("waiting for ecies ack ...");
        let ack = transport.try_next().await?;

        trace!("parsing ecies ack ...");
        if matches!(ack, Some(IngressECIESValue::Ack)) {
            Ok(Self {
                stream: transport,
                remote_id,
            })
        } else {
            bail!("invalid handshake: expected ack, got {:?} instead", ack)
        }
    }

    /// Listen on a just connected ECIES client
    #[instrument(skip(transport, secret_key), fields(peer=&*format!("{:?}", transport.remote_addr())))]
    pub async fn incoming(transport: Io, secret_key: SecretKey) -> anyhow::Result<Self> {
        let ecies = ECIESCodec::new_server(secret_key).context("handshake error")?;

        debug!("incoming ecies stream ...");
        let mut transport = ecies.framed(transport);
        let ack = transport.try_next().await?;

        debug!("receiving ecies auth");
        let remote_id = match ack {
            Some(IngressECIESValue::AuthReceive(remote_id)) => remote_id,
            other => {
                debug!("expected auth, got {:?} instead", other);
                bail!("invalid handshake");
            }
        };

        debug!("sending ecies ack ...");
        transport
            .send(EgressECIESValue::Ack)
            .await
            .context("failed to send ECIES auth")?;

        Ok(Self {
            stream: transport,
            remote_id,
        })
    }

    /// Get the remote id
    pub fn remote_id(&self) -> PeerId {
        self.remote_id
    }
}

impl<Io> Stream for ECIESStream<Io>
where
    Io: Transport,
{
    type Item = Result<Bytes, io::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match ready!(Pin::new(&mut self.get_mut().stream).poll_next(cx)) {
            Some(Ok(IngressECIESValue::Message(body))) => Poll::Ready(Some(Ok(body.into()))),
            Some(other) => Poll::Ready(Some(Err(io::Error::new(
                io::ErrorKind::Other,
                format!(
                    "ECIES stream protocol error: expected message, received {:?}",
                    other
                ),
            )))),
            None => Poll::Ready(None),
        }
    }
}

impl<Io> Sink<Bytes> for ECIESStream<Io>
where
    Io: Transport,
{
    type Error = io::Error;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Pin::new(&mut self.get_mut().stream).poll_ready(cx)
    }

    fn start_send(self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
        let this = self.get_mut();
        Pin::new(&mut this.stream).start_send(EgressECIESValue::Message(item))?;

        Ok(())
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Pin::new(&mut self.get_mut().stream).poll_flush(cx)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Pin::new(&mut self.get_mut().stream).poll_close(cx)
    }
}
<commit_msg>Remove useless conversion<commit_after>use super::algorithm::ECIES;
use crate::{errors::ECIESError, transport::Transport, types::PeerId};
use anyhow::{bail, Context as _};
use bytes::{Bytes, BytesMut};
use futures::{ready, Sink, SinkExt};
use secp256k1::SecretKey;
use std::{
    fmt::Debug,
    io,
    pin::Pin,
    task::{Context, Poll},
};
use tokio::stream::*;
use tokio_util::codec::*;
use tracing::*;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
/// Current ECIES state of a connection
pub enum ECIESState {
    Auth,
    Ack,
    Header,
    Body,
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// Raw egress values for an ECIES protocol
pub enum EgressECIESValue {
    Auth,
    Ack,
    Message(Bytes),
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// Raw ingress values for an ECIES protocol
pub enum IngressECIESValue {
    AuthReceive(PeerId),
    Ack,
    Message(Bytes),
}

/// Tokio codec for ECIES
#[derive(Debug)]
pub struct ECIESCodec {
    ecies: ECIES,
    state: ECIESState,
}

impl ECIESCodec {
    /// Create a new server codec using the given secret key
    pub fn new_server(secret_key: SecretKey) -> Result<Self, ECIESError> {
        Ok(Self {
            ecies: ECIES::new_server(secret_key)?,
            state: ECIESState::Auth,
        })
    }

    /// Create a new client codec using the given secret key and the server's public id
    pub fn new_client(secret_key: SecretKey, remote_id: PeerId) -> Result<Self, ECIESError> {
        Ok(Self {
            ecies: ECIES::new_client(secret_key, remote_id)?,
            state: ECIESState::Auth,
        })
    }
}

impl Decoder for ECIESCodec {
    type Item = IngressECIESValue;
    type Error = io::Error;

    #[instrument(level = "trace", skip(self, buf), fields(peer=&*format!("{:?}", self.ecies.remote_id.map(|s| s.to_string())), state=&*format!("{:?}", self.state)))]
    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        loop {
            match self.state {
                ECIESState::Auth => {
                    trace!("parsing auth");
                    if buf.len() < 2 {
                        return Ok(None);
                    }

                    let payload_size = u16::from_be_bytes([buf[0], buf[1]]) as usize;
                    let total_size = payload_size + 2;

                    if buf.len() < total_size {
                        trace!("current len {}, need {}", buf.len(), total_size);
                        return Ok(None);
                    }

                    self.ecies.read_auth(&mut *buf.split_to(total_size))?;

                    self.state = ECIESState::Header;
                    return Ok(Some(IngressECIESValue::AuthReceive(self.ecies.remote_id())));
                }
                ECIESState::Ack => {
                    trace!("parsing ack with len {}", buf.len());
                    if buf.len() < 2 {
                        return Ok(None);
                    }

                    let payload_size = u16::from_be_bytes([buf[0], buf[1]]) as usize;
                    let total_size = payload_size + 2;

                    if buf.len() < total_size {
                        trace!("current len {}, need {}", buf.len(), total_size);
                        return Ok(None);
                    }

                    self.ecies.read_ack(&mut *buf.split_to(total_size))?;

                    self.state = ECIESState::Header;
                    return Ok(Some(IngressECIESValue::Ack));
                }
                ECIESState::Header => {
                    if buf.len() < ECIES::header_len() {
                        return Ok(None);
                    }

                    self.ecies
                        .read_header(&mut *buf.split_to(ECIES::header_len()))?;

                    self.state = ECIESState::Body;
                }
                ECIESState::Body => {
                    if buf.len() < self.ecies.body_len() {
                        return Ok(None);
                    }

                    let mut data = buf.split_to(self.ecies.body_len());
                    let ret = Bytes::copy_from_slice(&self.ecies.read_body(&mut *data)?);

                    self.state = ECIESState::Header;
                    return Ok(Some(IngressECIESValue::Message(ret)));
                }
            }
        }
    }
}

impl Encoder<EgressECIESValue> for ECIESCodec {
    type Error = io::Error;

    #[instrument(level = "trace", skip(self, buf), fields(peer=&*format!("{:?}", self.ecies.remote_id.map(|s| s.to_string())), state=&*format!("{:?}", self.state)))]
    fn encode(&mut self, item: EgressECIESValue, buf: &mut BytesMut) -> Result<(), Self::Error> {
        match item {
            EgressECIESValue::Auth => {
                self.state = ECIESState::Ack;
                self.ecies.write_auth(buf);
                Ok(())
            }
            EgressECIESValue::Ack => {
                self.state = ECIESState::Header;
                self.ecies.write_ack(buf);
                Ok(())
            }
            EgressECIESValue::Message(data) => {
                self.ecies.write_header(buf, data.len());
                self.ecies.write_body(buf, &data);
                Ok(())
            }
        }
    }
}

/// `ECIES` stream over TCP exchanging raw bytes
#[derive(Debug)]
pub struct ECIESStream<Io> {
    stream: Framed<Io, ECIESCodec>,
    remote_id: PeerId,
}

impl<Io> ECIESStream<Io>
where
    Io: Transport,
{
    /// Connect to an `ECIES` server
    #[instrument(skip(transport, secret_key), fields(peer=&*format!("{:?}", transport.remote_addr())))]
    pub async fn connect(
        transport: Io,
        secret_key: SecretKey,
        remote_id: PeerId,
    ) -> anyhow::Result<Self> {
        let ecies = ECIESCodec::new_client(secret_key, remote_id)
            .map_err(|_| io::Error::new(io::ErrorKind::Other, "invalid handshake"))?;

        let mut transport = ecies.framed(transport);

        trace!("sending ecies auth ...");
        transport.send(EgressECIESValue::Auth).await?;

        trace!("waiting for ecies ack ...");
        let ack = transport.try_next().await?;

        trace!("parsing ecies ack ...");
        if matches!(ack, Some(IngressECIESValue::Ack)) {
            Ok(Self {
                stream: transport,
                remote_id,
            })
        } else {
            bail!("invalid handshake: expected ack, got {:?} instead", ack)
        }
    }

    /// Listen on a just connected ECIES client
    #[instrument(skip(transport, secret_key), fields(peer=&*format!("{:?}", transport.remote_addr())))]
    pub async fn incoming(transport: Io, secret_key: SecretKey) -> anyhow::Result<Self> {
        let ecies = ECIESCodec::new_server(secret_key).context("handshake error")?;

        debug!("incoming ecies stream ...");
        let mut transport = ecies.framed(transport);
        let ack = transport.try_next().await?;

        debug!("receiving ecies auth");
        let remote_id = match ack {
            Some(IngressECIESValue::AuthReceive(remote_id)) => remote_id,
            other => {
                debug!("expected auth, got {:?} instead", other);
                bail!("invalid handshake");
            }
        };

        debug!("sending ecies ack ...");
        transport
            .send(EgressECIESValue::Ack)
            .await
            .context("failed to send ECIES auth")?;

        Ok(Self {
            stream: transport,
            remote_id,
        })
    }

    /// Get the remote id
    pub fn remote_id(&self) -> PeerId {
        self.remote_id
    }
}

impl<Io> Stream for ECIESStream<Io>
where
    Io: Transport,
{
    type Item = Result<Bytes, io::Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match ready!(Pin::new(&mut self.get_mut().stream).poll_next(cx)) {
            Some(Ok(IngressECIESValue::Message(body))) => Poll::Ready(Some(Ok(body))),
            Some(other) => Poll::Ready(Some(Err(io::Error::new(
                io::ErrorKind::Other,
                format!(
                    "ECIES stream protocol error: expected message, received {:?}",
                    other
                ),
            )))),
            None => Poll::Ready(None),
        }
    }
}

impl<Io> Sink<Bytes> for ECIESStream<Io>
where
    Io: Transport,
{
    type Error = io::Error;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Pin::new(&mut self.get_mut().stream).poll_ready(cx)
    }

    fn start_send(self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
        let this = self.get_mut();
        Pin::new(&mut this.stream).start_send(EgressECIESValue::Message(item))?;

        Ok(())
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Pin::new(&mut self.get_mut().stream).poll_flush(cx)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Pin::new(&mut self.get_mut().stream).poll_close(cx)
    }
}
<|endoftext|>


<commit_before>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! gralloc: Cross-platform, Rust-based, Vulkan centric GPU allocation and
//! mapping.

use std::collections::BTreeMap as Map;

use base::round_up_to_page_size;

use crate::rutabaga_gralloc::formats::*;
use crate::rutabaga_gralloc::system_gralloc::SystemGralloc;
use crate::rutabaga_utils::*;

#[cfg(feature = "minigbm")]
use crate::rutabaga_gralloc::minigbm::MinigbmDevice;

#[cfg(feature = "vulkano")]
use crate::rutabaga_gralloc::vulkano_gralloc::VulkanoGralloc;

/*
 * Rutabaga gralloc flags are copied from minigbm, but redundant legacy flags are left out.
 * For example, USE_WRITE / USE_CURSOR_64X64 / USE_CURSOR don't add much value.
 */
const RUTABAGA_GRALLOC_USE_SCANOUT: u32 = 1 << 0;
const RUTABAGA_GRALLOC_USE_RENDERING: u32 = 1 << 2;
const RUTABAGA_GRALLOC_USE_LINEAR: u32 = 1 << 4;
const RUTABAGA_GRALLOC_USE_TEXTURING: u32 = 1 << 5;
const RUTABAGA_GRALLOC_USE_CAMERA_WRITE: u32 = 1 << 6;
const RUTABAGA_GRALLOC_USE_CAMERA_READ: u32 = 1 << 7;
#[allow(dead_code)]
const RUTABAGA_GRALLOC_USE_PROTECTED: u32 = 1 << 8;

/* SW_{WRITE,READ}_RARELY omitted since not even Android uses this much. */
const RUTABAGA_GRALLOC_USE_SW_READ_OFTEN: u32 = 1 << 9;
const RUTABAGA_GRALLOC_USE_SW_WRITE_OFTEN: u32 = 1 << 11;

#[allow(dead_code)]
const RUTABAGA_GRALLOC_VIDEO_DECODER: u32 = 1 << 13;
#[allow(dead_code)]
const RUTABAGA_GRALLOC_VIDEO_ENCODER: u32 = 1 << 14;

/// Usage flags for constructing a buffer object.
#[derive(Copy, Clone, Eq, PartialEq, Default)]
pub struct RutabagaGrallocFlags(pub u32);

impl RutabagaGrallocFlags {
    /// Returns empty set of flags.
    #[inline(always)]
    pub fn empty() -> RutabagaGrallocFlags {
        RutabagaGrallocFlags(0)
    }

    /// Returns the given set of raw `RUTABAGA_GRALLOC` flags wrapped in a RutabagaGrallocFlags
    /// struct.
    #[inline(always)]
    pub fn new(raw: u32) -> RutabagaGrallocFlags {
        RutabagaGrallocFlags(raw)
    }

    /// Sets the scanout flag's presence.
    #[inline(always)]
    pub fn use_scanout(self, e: bool) -> RutabagaGrallocFlags {
        if e {
            RutabagaGrallocFlags(self.0 | RUTABAGA_GRALLOC_USE_SCANOUT)
        } else {
            RutabagaGrallocFlags(self.0 & !RUTABAGA_GRALLOC_USE_SCANOUT)
        }
    }

    /// Sets the rendering flag's presence.
    #[inline(always)]
    pub fn use_rendering(self, e: bool) -> RutabagaGrallocFlags {
        if e {
            RutabagaGrallocFlags(self.0 | RUTABAGA_GRALLOC_USE_RENDERING)
        } else {
            RutabagaGrallocFlags(self.0 & !RUTABAGA_GRALLOC_USE_RENDERING)
        }
    }

    /// Sets the linear flag's presence.
    #[inline(always)]
    pub fn use_linear(self, e: bool) -> RutabagaGrallocFlags {
        if e {
            RutabagaGrallocFlags(self.0 | RUTABAGA_GRALLOC_USE_LINEAR)
        } else {
            RutabagaGrallocFlags(self.0 & !RUTABAGA_GRALLOC_USE_LINEAR)
        }
    }

    /// Returns true if the texturing flag is set.
    #[inline(always)]
    pub fn uses_texturing(self) -> bool {
        self.0 & RUTABAGA_GRALLOC_USE_TEXTURING != 0
    }

    /// Returns true if the rendering flag is set.
    #[inline(always)]
    pub fn uses_rendering(self) -> bool {
        self.0 & RUTABAGA_GRALLOC_USE_RENDERING != 0
    }

    /// Returns true if the memory will accessed by the CPU or an IP block that prefers host
    /// visible allocations (i.e, camera).
    #[inline(always)]
    pub fn host_visible(self) -> bool {
        self.0 & RUTABAGA_GRALLOC_USE_SW_READ_OFTEN != 0
            || self.0 & RUTABAGA_GRALLOC_USE_SW_WRITE_OFTEN != 0
            || self.0 & RUTABAGA_GRALLOC_USE_CAMERA_WRITE != 0
            || self.0 & RUTABAGA_GRALLOC_USE_CAMERA_READ != 0
    }

    /// Returns true if the memory will read by the CPU or an IP block that prefers cached
    /// allocations (i.e, camera).
    #[inline(always)]
    pub fn host_cached(self) -> bool {
        self.0 & RUTABAGA_GRALLOC_USE_CAMERA_READ != 0
            || self.0 & RUTABAGA_GRALLOC_USE_SW_READ_OFTEN != 0
    }
}

/// Information required to allocate a swapchain image.
#[derive(Copy, Clone, Default)]
pub struct ImageAllocationInfo {
    pub width: u32,
    pub height: u32,
    pub drm_format: DrmFormat,
    pub flags: RutabagaGrallocFlags,
}

/// The memory requirements, compression and layout of a swapchain image.
#[derive(Copy, Clone, Default)]
pub struct ImageMemoryRequirements {
    pub info: ImageAllocationInfo,
    pub map_info: u32,
    pub strides: [u32; 4],
    pub offsets: [u32; 4],
    pub modifier: u64,
    pub size: u64,
    pub vulkan_info: Option<VulkanInfo>,
}

/// Trait that needs to be implemented to service graphics memory requests.  Two step allocation
/// process:
///
///   (1) Get memory requirements for a given allocation request.
///   (2) Allocate using those requirements.
pub trait Gralloc {
    /// This function must return true if the implementation can:
    ///
    ///   (1) allocate GPU memory and
    ///   (2) {export to}/{import from} into a OS-specific RutabagaHandle.
    fn supports_external_gpu_memory(&self) -> bool;

    /// This function must return true the implementation can {export to}/{import from} a Linux
    /// dma-buf.  This often used for sharing with the scanout engine or multimedia subsystems.
    fn supports_dmabuf(&self) -> bool;

    /// Implementations must return the resource layout, compression, and caching properties of
    /// an allocation request.
    fn get_image_memory_requirements(
        &mut self,
        info: ImageAllocationInfo,
    ) -> RutabagaResult<ImageMemoryRequirements>;

    /// Implementations must allocate memory given the requirements and return a RutabagaHandle
    /// upon success.
    fn allocate_memory(&mut self, reqs: ImageMemoryRequirements) -> RutabagaResult<RutabagaHandle>;
}

/// Enumeration of possible allocation backends.
#[derive(Copy, Clone, Eq, PartialEq, PartialOrd, Ord)]
pub enum GrallocBackend {
    #[allow(dead_code)]
    Vulkano,
    #[allow(dead_code)]
    Minigbm,
    System,
}

/// A container for a variety of allocation backends.
pub struct RutabagaGralloc {
    grallocs: Map<GrallocBackend, Box<dyn Gralloc>>,
}

impl RutabagaGralloc {
    /// Returns a new RutabagaGralloc instance upon success.  All allocation backends that have
    /// been built are initialized.  The default system allocator is always initialized.
    pub fn new() -> RutabagaResult<RutabagaGralloc> {
        let mut grallocs: Map<GrallocBackend, Box<dyn Gralloc>> = Default::default();

        let system = SystemGralloc::init()?;
        grallocs.insert(GrallocBackend::System, system);

        #[cfg(feature = "minigbm")]
        {
            // crosvm integration tests build with the "wl-dmabuf" feature, which translates in
            // rutabaga to the "minigbm" feature.  These tests run on hosts where a rendernode is
            // not present, and minigbm can not be initialized.
            //
            // Thus, to keep kokoro happy, allow minigbm initialization to fail silently for now.
            if let Ok(gbm_device) = MinigbmDevice::init() {
                grallocs.insert(GrallocBackend::Minigbm, gbm_device);
            }
        }

        #[cfg(feature = "vulkano")]
        {
            let vulkano = VulkanoGralloc::init()?;
            grallocs.insert(GrallocBackend::Vulkano, vulkano);
        }

        Ok(RutabagaGralloc { grallocs })
    }

    /// Returns true if one of the allocation backends supports GPU external memory.
    pub fn supports_external_gpu_memory(&self) -> bool {
        for gralloc in self.grallocs.values() {
            if gralloc.supports_external_gpu_memory() {
                return true;
            }
        }

        false
    }

    /// Returns true if one of the allocation backends supports dma_buf.
    pub fn supports_dmabuf(&self) -> bool {
        for gralloc in self.grallocs.values() {
            if gralloc.supports_dmabuf() {
                return true;
            }
        }

        false
    }

    /// Returns the best allocation backend to service a particular request.
    fn determine_optimal_backend(&self, _info: ImageAllocationInfo) -> GrallocBackend {
        // This function could be more sophisticated and consider the allocation info.  For example,
        // nobody has ever tried Mali allocated memory + a mediatek/rockchip display and as such it
        // probably doesn't work.  In addition, YUV calculations in minigbm have yet to make it
        // towards the Vulkan api.  This function allows for a variety of quirks, but for now just
        // choose the most shiny backend that the user has built.  The rationale is "why would you
        // build it if you don't want to use it".
        let mut _backend = GrallocBackend::System;
        #[cfg(feature = "minigbm")]
        {
            _backend = GrallocBackend::Minigbm;
        }

        #[cfg(feature = "vulkano")]
        {
            _backend = GrallocBackend::Vulkano;
        }

        _backend
    }

    /// Returns a image memory requirements for the given `info` upon success.
    pub fn get_image_memory_requirements(
        &mut self,
        info: ImageAllocationInfo,
    ) -> RutabagaResult<ImageMemoryRequirements> {
        let backend = self.determine_optimal_backend(info);

        let gralloc = self
            .grallocs
            .get_mut(&backend)
            .ok_or(RutabagaError::Unsupported)?;

        let mut reqs = gralloc.get_image_memory_requirements(info)?;
        reqs.size = round_up_to_page_size(reqs.size as usize) as u64;
        Ok(reqs)
    }

    /// Allocates memory given the particular `reqs` upon success.
    pub fn allocate_memory(
        &mut self,
        reqs: ImageMemoryRequirements,
    ) -> RutabagaResult<RutabagaHandle> {
        let backend = self.determine_optimal_backend(reqs.info);

        let gralloc = self
            .grallocs
            .get_mut(&backend)
            .ok_or(RutabagaError::Unsupported)?;

        gralloc.allocate_memory(reqs)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn create_render_target() {
        let gralloc_result = RutabagaGralloc::new();
        if gralloc_result.is_err() {
            return;
        }

        let mut gralloc = gralloc_result.unwrap();

        let info = ImageAllocationInfo {
            width: 512,
            height: 1024,
            drm_format: DrmFormat::new(b'X', b'R', b'2', b'4'),
            flags: RutabagaGrallocFlags::empty().use_scanout(true),
        };

        let reqs = gralloc.get_image_memory_requirements(info).unwrap();
        let min_reqs = canonical_image_requirements(info).unwrap();

        assert_eq!(reqs.strides[0] >= min_reqs.strides[0], true);
        assert_eq!(reqs.size >= min_reqs.size, true);

        let _handle = gralloc.allocate_memory(reqs).unwrap();

        // Reallocate with same requirements
        let _handle2 = gralloc.allocate_memory(reqs).unwrap();
    }

    #[test]
    fn create_video_buffer() {
        let gralloc_result = RutabagaGralloc::new();
        if gralloc_result.is_err() {
            return;
        }

        let mut gralloc = gralloc_result.unwrap();

        let info = ImageAllocationInfo {
            width: 512,
            height: 1024,
            drm_format: DrmFormat::new(b'N', b'V', b'1', b'2'),
            flags: RutabagaGrallocFlags::empty().use_linear(true),
        };

        let reqs = gralloc.get_image_memory_requirements(info).unwrap();
        let min_reqs = canonical_image_requirements(info).unwrap();

        assert_eq!(reqs.strides[0] >= min_reqs.strides[0], true);
        assert_eq!(reqs.strides[1] >= min_reqs.strides[1], true);
        assert_eq!(reqs.strides[2], 0);
        assert_eq!(reqs.strides[3], 0);

        assert_eq!(reqs.offsets[0] >= min_reqs.offsets[0], true);
        assert_eq!(reqs.offsets[1] >= min_reqs.offsets[1], true);
        assert_eq!(reqs.offsets[2], 0);
        assert_eq!(reqs.offsets[3], 0);

        assert_eq!(reqs.size >= min_reqs.size, true);

        let _handle = gralloc.allocate_memory(reqs).unwrap();

        // Reallocate with same requirements
        let _handle2 = gralloc.allocate_memory(reqs).unwrap();
    }
}
<commit_msg>rutabaga_gfx: rutabaga_gralloc: fix Kokoro unhappiness<commit_after>// Copyright 2021 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! gralloc: Cross-platform, Rust-based, Vulkan centric GPU allocation and
//! mapping.

use std::collections::BTreeMap as Map;

use base::round_up_to_page_size;

use crate::rutabaga_gralloc::formats::*;
use crate::rutabaga_gralloc::system_gralloc::SystemGralloc;
use crate::rutabaga_utils::*;

#[cfg(feature = "minigbm")]
use crate::rutabaga_gralloc::minigbm::MinigbmDevice;

#[cfg(feature = "vulkano")]
use crate::rutabaga_gralloc::vulkano_gralloc::VulkanoGralloc;

/*
 * Rutabaga gralloc flags are copied from minigbm, but redundant legacy flags are left out.
 * For example, USE_WRITE / USE_CURSOR_64X64 / USE_CURSOR don't add much value.
 */
const RUTABAGA_GRALLOC_USE_SCANOUT: u32 = 1 << 0;
const RUTABAGA_GRALLOC_USE_RENDERING: u32 = 1 << 2;
const RUTABAGA_GRALLOC_USE_LINEAR: u32 = 1 << 4;
const RUTABAGA_GRALLOC_USE_TEXTURING: u32 = 1 << 5;
const RUTABAGA_GRALLOC_USE_CAMERA_WRITE: u32 = 1 << 6;
const RUTABAGA_GRALLOC_USE_CAMERA_READ: u32 = 1 << 7;
#[allow(dead_code)]
const RUTABAGA_GRALLOC_USE_PROTECTED: u32 = 1 << 8;

/* SW_{WRITE,READ}_RARELY omitted since not even Android uses this much. */
const RUTABAGA_GRALLOC_USE_SW_READ_OFTEN: u32 = 1 << 9;
const RUTABAGA_GRALLOC_USE_SW_WRITE_OFTEN: u32 = 1 << 11;

#[allow(dead_code)]
const RUTABAGA_GRALLOC_VIDEO_DECODER: u32 = 1 << 13;
#[allow(dead_code)]
const RUTABAGA_GRALLOC_VIDEO_ENCODER: u32 = 1 << 14;

/// Usage flags for constructing a buffer object.
#[derive(Copy, Clone, Eq, PartialEq, Default)]
pub struct RutabagaGrallocFlags(pub u32);

impl RutabagaGrallocFlags {
    /// Returns empty set of flags.
    #[inline(always)]
    pub fn empty() -> RutabagaGrallocFlags {
        RutabagaGrallocFlags(0)
    }

    /// Returns the given set of raw `RUTABAGA_GRALLOC` flags wrapped in a RutabagaGrallocFlags
    /// struct.
    #[inline(always)]
    pub fn new(raw: u32) -> RutabagaGrallocFlags {
        RutabagaGrallocFlags(raw)
    }

    /// Sets the scanout flag's presence.
    #[inline(always)]
    pub fn use_scanout(self, e: bool) -> RutabagaGrallocFlags {
        if e {
            RutabagaGrallocFlags(self.0 | RUTABAGA_GRALLOC_USE_SCANOUT)
        } else {
            RutabagaGrallocFlags(self.0 & !RUTABAGA_GRALLOC_USE_SCANOUT)
        }
    }

    /// Sets the rendering flag's presence.
    #[inline(always)]
    pub fn use_rendering(self, e: bool) -> RutabagaGrallocFlags {
        if e {
            RutabagaGrallocFlags(self.0 | RUTABAGA_GRALLOC_USE_RENDERING)
        } else {
            RutabagaGrallocFlags(self.0 & !RUTABAGA_GRALLOC_USE_RENDERING)
        }
    }

    /// Sets the linear flag's presence.
    #[inline(always)]
    pub fn use_linear(self, e: bool) -> RutabagaGrallocFlags {
        if e {
            RutabagaGrallocFlags(self.0 | RUTABAGA_GRALLOC_USE_LINEAR)
        } else {
            RutabagaGrallocFlags(self.0 & !RUTABAGA_GRALLOC_USE_LINEAR)
        }
    }

    /// Returns true if the texturing flag is set.
    #[inline(always)]
    pub fn uses_texturing(self) -> bool {
        self.0 & RUTABAGA_GRALLOC_USE_TEXTURING != 0
    }

    /// Returns true if the rendering flag is set.
    #[inline(always)]
    pub fn uses_rendering(self) -> bool {
        self.0 & RUTABAGA_GRALLOC_USE_RENDERING != 0
    }

    /// Returns true if the memory will accessed by the CPU or an IP block that prefers host
    /// visible allocations (i.e, camera).
    #[inline(always)]
    pub fn host_visible(self) -> bool {
        self.0 & RUTABAGA_GRALLOC_USE_SW_READ_OFTEN != 0
            || self.0 & RUTABAGA_GRALLOC_USE_SW_WRITE_OFTEN != 0
            || self.0 & RUTABAGA_GRALLOC_USE_CAMERA_WRITE != 0
            || self.0 & RUTABAGA_GRALLOC_USE_CAMERA_READ != 0
    }

    /// Returns true if the memory will read by the CPU or an IP block that prefers cached
    /// allocations (i.e, camera).
    #[inline(always)]
    pub fn host_cached(self) -> bool {
        self.0 & RUTABAGA_GRALLOC_USE_CAMERA_READ != 0
            || self.0 & RUTABAGA_GRALLOC_USE_SW_READ_OFTEN != 0
    }
}

/// Information required to allocate a swapchain image.
#[derive(Copy, Clone, Default)]
pub struct ImageAllocationInfo {
    pub width: u32,
    pub height: u32,
    pub drm_format: DrmFormat,
    pub flags: RutabagaGrallocFlags,
}

/// The memory requirements, compression and layout of a swapchain image.
#[derive(Copy, Clone, Default)]
pub struct ImageMemoryRequirements {
    pub info: ImageAllocationInfo,
    pub map_info: u32,
    pub strides: [u32; 4],
    pub offsets: [u32; 4],
    pub modifier: u64,
    pub size: u64,
    pub vulkan_info: Option<VulkanInfo>,
}

/// Trait that needs to be implemented to service graphics memory requests.  Two step allocation
/// process:
///
///   (1) Get memory requirements for a given allocation request.
///   (2) Allocate using those requirements.
pub trait Gralloc {
    /// This function must return true if the implementation can:
    ///
    ///   (1) allocate GPU memory and
    ///   (2) {export to}/{import from} into a OS-specific RutabagaHandle.
    fn supports_external_gpu_memory(&self) -> bool;

    /// This function must return true the implementation can {export to}/{import from} a Linux
    /// dma-buf.  This often used for sharing with the scanout engine or multimedia subsystems.
    fn supports_dmabuf(&self) -> bool;

    /// Implementations must return the resource layout, compression, and caching properties of
    /// an allocation request.
    fn get_image_memory_requirements(
        &mut self,
        info: ImageAllocationInfo,
    ) -> RutabagaResult<ImageMemoryRequirements>;

    /// Implementations must allocate memory given the requirements and return a RutabagaHandle
    /// upon success.
    fn allocate_memory(&mut self, reqs: ImageMemoryRequirements) -> RutabagaResult<RutabagaHandle>;
}

/// Enumeration of possible allocation backends.
#[derive(Copy, Clone, Eq, PartialEq, PartialOrd, Ord)]
pub enum GrallocBackend {
    #[allow(dead_code)]
    Vulkano,
    #[allow(dead_code)]
    Minigbm,
    System,
}

/// A container for a variety of allocation backends.
pub struct RutabagaGralloc {
    grallocs: Map<GrallocBackend, Box<dyn Gralloc>>,
}

impl RutabagaGralloc {
    /// Returns a new RutabagaGralloc instance upon success.  All allocation backends that have
    /// been built are initialized.  The default system allocator is always initialized.
    pub fn new() -> RutabagaResult<RutabagaGralloc> {
        let mut grallocs: Map<GrallocBackend, Box<dyn Gralloc>> = Default::default();

        let system = SystemGralloc::init()?;
        grallocs.insert(GrallocBackend::System, system);

        #[cfg(feature = "minigbm")]
        {
            // crosvm integration tests build with the "wl-dmabuf" feature, which translates in
            // rutabaga to the "minigbm" feature.  These tests run on hosts where a rendernode is
            // not present, and minigbm can not be initialized.
            //
            // Thus, to keep kokoro happy, allow minigbm initialization to fail silently for now.
            if let Ok(gbm_device) = MinigbmDevice::init() {
                grallocs.insert(GrallocBackend::Minigbm, gbm_device);
            }
        }

        #[cfg(feature = "vulkano")]
        {
            let vulkano = VulkanoGralloc::init()?;
            grallocs.insert(GrallocBackend::Vulkano, vulkano);
        }

        Ok(RutabagaGralloc { grallocs })
    }

    /// Returns true if one of the allocation backends supports GPU external memory.
    pub fn supports_external_gpu_memory(&self) -> bool {
        for gralloc in self.grallocs.values() {
            if gralloc.supports_external_gpu_memory() {
                return true;
            }
        }

        false
    }

    /// Returns true if one of the allocation backends supports dma_buf.
    pub fn supports_dmabuf(&self) -> bool {
        for gralloc in self.grallocs.values() {
            if gralloc.supports_dmabuf() {
                return true;
            }
        }

        false
    }

    /// Returns the best allocation backend to service a particular request.
    fn determine_optimal_backend(&self, _info: ImageAllocationInfo) -> GrallocBackend {
        // This function could be more sophisticated and consider the allocation info.  For example,
        // nobody has ever tried Mali allocated memory + a mediatek/rockchip display and as such it
        // probably doesn't work.  In addition, YUV calculations in minigbm have yet to make it
        // towards the Vulkan api.  This function allows for a variety of quirks, but for now just
        // choose the most shiny backend that the user has built.  The rationale is "why would you
        // build it if you don't want to use it".
        let mut _backend = GrallocBackend::System;

        #[cfg(feature = "minigbm")]
        {
            // See note on "wl-dmabuf" and Kokoro in Gralloc::new().
            if self.grallocs.contains_key(&GrallocBackend::Minigbm) {
                _backend = GrallocBackend::Minigbm;
            }
        }

        #[cfg(feature = "vulkano")]
        {
            _backend = GrallocBackend::Vulkano;
        }

        _backend
    }

    /// Returns a image memory requirements for the given `info` upon success.
    pub fn get_image_memory_requirements(
        &mut self,
        info: ImageAllocationInfo,
    ) -> RutabagaResult<ImageMemoryRequirements> {
        let backend = self.determine_optimal_backend(info);

        let gralloc = self
            .grallocs
            .get_mut(&backend)
            .ok_or(RutabagaError::Unsupported)?;

        let mut reqs = gralloc.get_image_memory_requirements(info)?;
        reqs.size = round_up_to_page_size(reqs.size as usize) as u64;
        Ok(reqs)
    }

    /// Allocates memory given the particular `reqs` upon success.
    pub fn allocate_memory(
        &mut self,
        reqs: ImageMemoryRequirements,
    ) -> RutabagaResult<RutabagaHandle> {
        let backend = self.determine_optimal_backend(reqs.info);

        let gralloc = self
            .grallocs
            .get_mut(&backend)
            .ok_or(RutabagaError::Unsupported)?;

        gralloc.allocate_memory(reqs)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn create_render_target() {
        let gralloc_result = RutabagaGralloc::new();
        if gralloc_result.is_err() {
            return;
        }

        let mut gralloc = gralloc_result.unwrap();

        let info = ImageAllocationInfo {
            width: 512,
            height: 1024,
            drm_format: DrmFormat::new(b'X', b'R', b'2', b'4'),
            flags: RutabagaGrallocFlags::empty().use_scanout(true),
        };

        let reqs = gralloc.get_image_memory_requirements(info).unwrap();
        let min_reqs = canonical_image_requirements(info).unwrap();

        assert_eq!(reqs.strides[0] >= min_reqs.strides[0], true);
        assert_eq!(reqs.size >= min_reqs.size, true);

        let _handle = gralloc.allocate_memory(reqs).unwrap();

        // Reallocate with same requirements
        let _handle2 = gralloc.allocate_memory(reqs).unwrap();
    }

    #[test]
    fn create_video_buffer() {
        let gralloc_result = RutabagaGralloc::new();
        if gralloc_result.is_err() {
            return;
        }

        let mut gralloc = gralloc_result.unwrap();

        let info = ImageAllocationInfo {
            width: 512,
            height: 1024,
            drm_format: DrmFormat::new(b'N', b'V', b'1', b'2'),
            flags: RutabagaGrallocFlags::empty().use_linear(true),
        };

        let reqs = gralloc.get_image_memory_requirements(info).unwrap();
        let min_reqs = canonical_image_requirements(info).unwrap();

        assert_eq!(reqs.strides[0] >= min_reqs.strides[0], true);
        assert_eq!(reqs.strides[1] >= min_reqs.strides[1], true);
        assert_eq!(reqs.strides[2], 0);
        assert_eq!(reqs.strides[3], 0);

        assert_eq!(reqs.offsets[0] >= min_reqs.offsets[0], true);
        assert_eq!(reqs.offsets[1] >= min_reqs.offsets[1], true);
        assert_eq!(reqs.offsets[2], 0);
        assert_eq!(reqs.offsets[3], 0);

        assert_eq!(reqs.size >= min_reqs.size, true);

        let _handle = gralloc.allocate_memory(reqs).unwrap();

        // Reallocate with same requirements
        let _handle2 = gralloc.allocate_memory(reqs).unwrap();
    }
}
<|endoftext|>