<commit_before>// Copyright 2013-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(all(unix, not(target_os = "ios"), not(target_os = "openbsd")))]
mod imp {
    use self::OsRngInner::*;

    use fs::File;
    use io;
    use libc;
    use mem;
    use rand::Rng;
    use rand::reader::ReaderRng;
    use sys::os::errno;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64",
                  target_arch = "x86",
                  target_arch = "arm",
                  target_arch = "aarch64",
                  target_arch = "powerpc",
                  target_arch = "powerpc64",
                  target_arch = "powerpc64le")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(any(target_arch = "arm", target_arch = "powerpc",
                  target_arch = "powerpc64", target_arch = "powerpc64le"))]
        const NR_GETRANDOM: libc::c_long = 384;
        #[cfg(target_arch = "aarch64")]
        const NR_GETRANDOM: libc::c_long = 278;

        unsafe {
            libc::syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64",
                      target_arch = "x86",
                      target_arch = "arm",
                      target_arch = "aarch64",
                      target_arch = "powerpc",
                      target_arch = "powerpc64",
                      target_arch = "powerpc64le"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        while read < v.len() {
            let result = getrandom(&mut v[read..]);
            if result == -1 {
                let err = errno() as libc::c_int;
                if err == libc::EINTR {
                    continue;
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as usize;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8; 4] = [0; 4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8; 4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8; 8] = [0; 8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8; 8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64",
                  target_arch = "x86",
                  target_arch = "arm",
                  target_arch = "aarch64",
                  target_arch = "powerpc",
                  target_arch = "powerpc64",
                  target_arch = "powerpc64le")))]
    fn is_getrandom_available() -> bool {
        use sync::atomic::{AtomicBool, Ordering};
        use sync::Once;

        static CHECKER: Once = Once::new();
        static AVAILABLE: AtomicBool = AtomicBool::new(false);

        CHECKER.call_once(|| {
            let mut buf: [u8; 0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = io::Error::last_os_error().raw_os_error();
                err != Some(libc::ENOSYS)
            } else {
                true
            };
            AVAILABLE.store(available, Ordering::Relaxed);
        });

        AVAILABLE.load(Ordering::Relaxed)
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64",
                      target_arch = "x86",
                      target_arch = "arm",
                      target_arch = "aarch64",
                      target_arch = "powerpc",
                      target_arch = "powerpc64",
                      target_arch = "powerpc64le"))))]
    fn is_getrandom_available() -> bool { false }

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    /// - OpenBSD: uses the `getentropy(2)` system call.
    ///
    /// This does not block.
    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReaderRng(ReaderRng<File>),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open("/dev/urandom"));
            let reader_rng = ReaderRng::new(reader);

            Ok(OsRng { inner: OsReaderRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReaderRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReaderRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReaderRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "openbsd")]
mod imp {
    use io;
    use libc;
    use mem;
    use sys::os::errno;
    use rand::Rng;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    /// - OpenBSD: uses the `getentropy(2)` system call.
    ///
    /// This does not block.
    pub struct OsRng {
        // dummy field to ensure that this struct cannot be constructed outside
        // of this module
        _dummy: (),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            Ok(OsRng { _dummy: () })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            // getentropy(2) permits a maximum buffer size of 256 bytes
            for s in v.chunks_mut(256) {
                let ret = unsafe {
                    libc::syscall(libc::NR_GETENTROPY, s.as_mut_ptr(), s.len())
                };
                if ret == -1 {
                    panic!("unexpected getentropy error: {}", errno());
                }
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    use io;
    use mem;
    use ptr;
    use rand::Rng;
    use libc::{c_int, size_t};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    /// - OpenBSD: uses the `getentropy(2)` system call.
    ///
    /// This does not block.
    pub struct OsRng {
        // dummy field to ensure that this struct cannot be constructed outside
        // of this module
        _dummy: (),
    }

    enum SecRandom {}

    #[allow(non_upper_case_globals)]
    const kSecRandomDefault: *const SecRandom = ptr::null();

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            Ok(OsRng { _dummy: () })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t,
                                   v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}",
                       io::Error::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    use io;
    use mem;
    use rand::Rng;
    use sys::c;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    /// - OpenBSD: uses the `getentropy(2)` system call.
    ///
    /// This does not block.
    pub struct OsRng {
        hcryptprov: c::HCRYPTPROV
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                c::CryptAcquireContextA(&mut hcp, 0 as c::LPCSTR, 0 as c::LPCSTR,
                                        c::PROV_RSA_FULL,
                                        c::CRYPT_VERIFYCONTEXT | c::CRYPT_SILENT)
            };

            if ret == 0 {
                Err(io::Error::last_os_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                c::CryptGenRandom(self.hcryptprov, v.len() as c::DWORD,
                                  v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}",
                       io::Error::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                c::CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}",
                       io::Error::last_os_error());
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use sync::mpsc::channel;
    use rand::Rng;
    use super::OsRng;
    use thread;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0; 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in 0..20 {
            let (tx, rx) = channel();
            txs.push(tx);

            thread::spawn(move|| {
                // wait until all the threads are ready to go.
                rx.recv().unwrap();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                thread::yield_now();
                let mut v = [0; 1000];

                for _ in 0..100 {
                    r.next_u32();
                    thread::yield_now();
                    r.next_u64();
                    thread::yield_now();
                    r.fill_bytes(&mut v);
                    thread::yield_now();
                }
            });
        }

        // start all the threads
        for tx in &txs {
            tx.send(()).unwrap();
        }
    }
}
<commit_msg>Incorrect getrandom() system call for PowerPC Linux<commit_after>// Copyright 2013-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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. This file may not be copied, modified, or distributed
// except according to those terms.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(all(unix, not(target_os = "ios"), not(target_os = "openbsd")))]
mod imp {
    use self::OsRngInner::*;

    use fs::File;
    use io;
    use libc;
    use mem;
    use rand::Rng;
    use rand::reader::ReaderRng;
    use sys::os::errno;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64",
                  target_arch = "x86",
                  target_arch = "arm",
                  target_arch = "aarch64",
                  target_arch = "powerpc",
                  target_arch = "powerpc64",
                  target_arch = "powerpc64le")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(target_arch = "arm")]
        const NR_GETRANDOM: libc::c_long = 384;
        #[cfg(any(target_arch = "powerpc", target_arch = "powerpc64",
                  target_arch = "powerpc64le"))]
        const NR_GETRANDOM: libc::c_long = 359;
        #[cfg(target_arch = "aarch64")]
        const NR_GETRANDOM: libc::c_long = 278;

        unsafe {
            libc::syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64",
                      target_arch = "x86",
                      target_arch = "arm",
                      target_arch = "aarch64",
                      target_arch = "powerpc",
                      target_arch = "powerpc64",
                      target_arch = "powerpc64le"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        while read < v.len() {
            let result = getrandom(&mut v[read..]);
            if result == -1 {
                let err = errno() as libc::c_int;
                if err == libc::EINTR {
                    continue;
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as usize;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8; 4] = [0; 4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8; 4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8; 8] = [0; 8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8; 8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64",
                  target_arch = "x86",
                  target_arch = "arm",
                  target_arch = "aarch64",
                  target_arch = "powerpc",
                  target_arch = "powerpc64",
                  target_arch = "powerpc64le")))]
    fn is_getrandom_available() -> bool {
        use sync::atomic::{AtomicBool, Ordering};
        use sync::Once;

        static CHECKER: Once = Once::new();
        static AVAILABLE: AtomicBool = AtomicBool::new(false);

        CHECKER.call_once(|| {
            let mut buf: [u8; 0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = io::Error::last_os_error().raw_os_error();
                err != Some(libc::ENOSYS)
            } else {
                true
            };
            AVAILABLE.store(available, Ordering::Relaxed);
        });

        AVAILABLE.load(Ordering::Relaxed)
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64",
                      target_arch = "x86",
                      target_arch = "arm",
                      target_arch = "aarch64",
                      target_arch = "powerpc",
                      target_arch = "powerpc64",
                      target_arch = "powerpc64le"))))]
    fn is_getrandom_available() -> bool { false }

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    /// - OpenBSD: uses the `getentropy(2)` system call.
    ///
    /// This does not block.
    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReaderRng(ReaderRng<File>),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open("/dev/urandom"));
            let reader_rng = ReaderRng::new(reader);

            Ok(OsRng { inner: OsReaderRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReaderRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReaderRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReaderRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "openbsd")]
mod imp {
    use io;
    use libc;
    use mem;
    use sys::os::errno;
    use rand::Rng;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    /// - OpenBSD: uses the `getentropy(2)` system call.
    ///
    /// This does not block.
    pub struct OsRng {
        // dummy field to ensure that this struct cannot be constructed outside
        // of this module
        _dummy: (),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            Ok(OsRng { _dummy: () })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            // getentropy(2) permits a maximum buffer size of 256 bytes
            for s in v.chunks_mut(256) {
                let ret = unsafe {
                    libc::syscall(libc::NR_GETENTROPY, s.as_mut_ptr(), s.len())
                };
                if ret == -1 {
                    panic!("unexpected getentropy error: {}", errno());
                }
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    use io;
    use mem;
    use ptr;
    use rand::Rng;
    use libc::{c_int, size_t};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    /// - OpenBSD: uses the `getentropy(2)` system call.
    ///
    /// This does not block.
    pub struct OsRng {
        // dummy field to ensure that this struct cannot be constructed outside
        // of this module
        _dummy: (),
    }

    enum SecRandom {}

    #[allow(non_upper_case_globals)]
    const kSecRandomDefault: *const SecRandom = ptr::null();

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            Ok(OsRng { _dummy: () })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t,
                                   v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}",
                       io::Error::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    use io;
    use mem;
    use rand::Rng;
    use sys::c;

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    /// - OpenBSD: uses the `getentropy(2)` system call.
    ///
    /// This does not block.
    pub struct OsRng {
        hcryptprov: c::HCRYPTPROV
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                c::CryptAcquireContextA(&mut hcp, 0 as c::LPCSTR, 0 as c::LPCSTR,
                                        c::PROV_RSA_FULL,
                                        c::CRYPT_VERIFYCONTEXT | c::CRYPT_SILENT)
            };

            if ret == 0 {
                Err(io::Error::last_os_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                c::CryptGenRandom(self.hcryptprov, v.len() as c::DWORD,
                                  v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}",
                       io::Error::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                c::CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}",
                       io::Error::last_os_error());
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use sync::mpsc::channel;
    use rand::Rng;
    use super::OsRng;
    use thread;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0; 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in 0..20 {
            let (tx, rx) = channel();
            txs.push(tx);

            thread::spawn(move|| {
                // wait until all the threads are ready to go.
                rx.recv().unwrap();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                thread::yield_now();
                let mut v = [0; 1000];

                for _ in 0..100 {
                    r.next_u32();
                    thread::yield_now();
                    r.next_u64();
                    thread::yield_now();
                    r.fill_bytes(&mut v);
                    thread::yield_now();
                }
            });
        }

        // start all the threads
        for tx in &txs {
            tx.send(()).unwrap();
        }
    }
}
<|endoftext|>


<commit_before>
use std::marker::PhantomData;
use crypto_stuff::*;

pub const MAXMSGLEN : usize = 65535;

#[derive(Copy, Clone)]
pub enum Token {E, S, Dhee, Dhes, Dhse, Dhss, Empty}

pub trait HandshakePattern {
    fn name(out: &mut [u8]) -> usize;
    fn get(premsg_pattern_i: &mut [Token],
           premsg_pattern_r: &mut [Token],
           msg_patterns: &mut [[Token; 8]; 5]);
}

#[derive(Debug)]
pub enum NoiseError {DecryptError}

struct SymmetricState<C: Cipher, H: Hash> {
    cipherstate : CipherState<C>,
    has_k: bool,
    has_psk: bool,
    h : [u8; MAXHASHLEN], /* Change once Rust has trait-associated consts */
    ck: [u8; MAXHASHLEN], /* Change once Rust has trait-associated consts */
    wtf : PhantomData<H>, /* So rust thinks I'm using H, this is ugly */
}

pub struct HandshakeState<P: HandshakePattern, D: Dh, C: Cipher, H: Hash, R: Random> {
    symmetricstate: SymmetricState<C, H>,
    s: Option<D>,
    e: Option<D>,
    rs: Option<[u8; DHLEN]>,
    re: Option<[u8; DHLEN]>,
    my_turn_to_send : bool,
    message_patterns : [[Token; 8]; 5],
    message_index: usize,
    initiator: bool,
    rng : R,
    wtf : PhantomData<P>, /* So rust thinks I'm using P, this is ugly */
}

impl <C: Cipher, H: Hash> SymmetricState<C, H> {

    fn new(handshake_name: &[u8]) -> SymmetricState<C, H> {
        let mut hname = [0u8; MAXHASHLEN];
        if handshake_name.len() <= H::hash_len() {
            copy_memory(handshake_name, &mut hname);
        } else {
            let mut hasher = H::new();
            hasher.input(handshake_name);
            hasher.result(&mut hname);
        }
        SymmetricState{
            cipherstate: CipherState::new(&[0u8; CIPHERKEYLEN], 0),
            has_k : false,
            has_psk: false,
            h: hname,
            ck : hname,
            wtf: PhantomData::<H>
        }
    }

    fn mix_key(&mut self, data: &[u8]) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], data, &mut hkdf_output.0, &mut hkdf_output.1);
        copy_memory(&hkdf_output.0, &mut self.ck);
        self.cipherstate = CipherState::new(&hkdf_output.1[..CIPHERKEYLEN], 0);
        self.has_k = true;
    }

    fn mix_hash(&mut self, data: &[u8]) {
        let mut hasher = H::new();
        hasher.input(&self.h[..H::hash_len()]);
        hasher.input(data);
        hasher.result(&mut self.h);
    }

    fn mix_preshared_key(&mut self, psk: &[u8]) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], psk, &mut hkdf_output.0, &mut hkdf_output.1);
        copy_memory(&hkdf_output.0, &mut self.ck);
        self.mix_hash(&hkdf_output.1[..H::hash_len()]);
        self.has_psk = true;
    }

    fn encrypt_and_hash(&mut self, plaintext: &[u8], out: &mut [u8]) -> usize {
        if self.has_k {
            self.cipherstate.encrypt_ad(&self.h[..H::hash_len()], plaintext, out);
            self.mix_hash(&out[..plaintext.len() + TAGLEN]);
            return plaintext.len() + TAGLEN;
        } else {
            copy_memory(plaintext, out);
            self.mix_hash(plaintext);
            return plaintext.len();
        }
    }

    fn decrypt_and_hash(&mut self, data: &[u8], out: &mut [u8]) -> bool {
        if self.has_k {
            if !self.cipherstate.decrypt_ad(&self.h[..H::hash_len()], data, out) {
                return false;
            }
        }
        else {
            copy_memory(data, out);
        }
        self.mix_hash(data);
        true
    }

    fn split(&self, initiator: bool) -> (CipherState<C>, CipherState<C>) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], &[0u8; 0], &mut hkdf_output.0, &mut hkdf_output.1);
        let c1 = CipherState::<C>::new(&hkdf_output.0[..CIPHERKEYLEN], 0);
        let c2 = CipherState::<C>::new(&hkdf_output.1[..CIPHERKEYLEN], 0);
        if initiator { (c1, c2) } else { (c2, c1) }
    }

}

impl <P: HandshakePattern, D: Dh, C: Cipher, H: Hash, R: Random> HandshakeState<P, D, C, H, R> {

    pub fn new(rng: R,
               initiator: bool,
               prologue: &[u8],
               new_preshared_key: Option<&[u8]>,
               new_s : Option<D>,
               new_e : Option<D>,
               new_rs: Option<[u8; DHLEN]>,
               new_re: Option<[u8; DHLEN]>) -> HandshakeState<P, D, C, H, R> {
        let mut handshake_name = [0u8; 128];
        let mut name_len: usize;
        if let Some(_) = new_preshared_key {
            copy_memory("NoisePSK_".as_bytes(), &mut handshake_name);
            name_len = 9;
        } else {
            copy_memory("Noise_".as_bytes(), &mut handshake_name);
            name_len = 6;
        }
        name_len += P::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += D::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += C::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += H::name(&mut handshake_name[name_len..]);

        let mut symmetricstate = SymmetricState::new(&handshake_name[..name_len]);

        symmetricstate.mix_hash(prologue);

        if let Some(preshared_key) = new_preshared_key {
            symmetricstate.mix_preshared_key(preshared_key);
        }

        let mut premsg_pattern_i = [Token::Empty; 2];
        let mut premsg_pattern_r = [Token::Empty; 2];
        let mut message_patterns = [[Token::Empty; 8]; 5];
        P::get(&mut premsg_pattern_i, &mut premsg_pattern_r, &mut message_patterns);
        if initiator {
            for token in &premsg_pattern_i {
                match *token {
                    Token::S => symmetricstate.mix_hash(new_s.as_ref().unwrap().pubkey()),
                    Token::E => symmetricstate.mix_hash(new_e.as_ref().unwrap().pubkey()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
            for token in &premsg_pattern_r {
                match *token {
                    Token::S => symmetricstate.mix_hash(&new_rs.unwrap()),
                    Token::E => symmetricstate.mix_hash(&new_re.unwrap()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
        } else {
            for token in &premsg_pattern_i {
                match *token {
                    Token::S => symmetricstate.mix_hash(&new_rs.unwrap()),
                    Token::E => symmetricstate.mix_hash(&new_re.unwrap()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
            for token in &premsg_pattern_r {
                match *token {
                    Token::S => symmetricstate.mix_hash(new_s.as_ref().unwrap().pubkey()),
                    Token::E => symmetricstate.mix_hash(new_e.as_ref().unwrap().pubkey()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
        }

        HandshakeState{
            symmetricstate: symmetricstate,
            s: new_s, e: new_e, rs: new_rs, re: new_re,
            my_turn_to_send: initiator,
            message_patterns: message_patterns,
            message_index: 0,
            initiator: initiator,
            rng: rng,
            wtf: PhantomData::<P>}
    }

    pub fn write_message(&mut self,
                         payload: &[u8],
                         message: &mut [u8]) -> (usize, Option<(CipherState<C>, CipherState<C>)>) {
        assert!(self.my_turn_to_send);
        let tokens = self.message_patterns[self.message_index];
        let mut last = false;
        if let Token::Empty = self.message_patterns[self.message_index+1][0] {
            last = true;
        }
        self.message_index += 1;

        let mut byte_index = 0;
        for token in &tokens {
            match *token {
                Token::E => {
                    self.e = Some(D::generate(&mut self.rng));
                    let pubkey = self.e.as_ref().unwrap().pubkey();
                    copy_memory(pubkey, &mut message[byte_index..]);
                    byte_index += DHLEN;
                    self.symmetricstate.mix_hash(&pubkey);
                    if self.symmetricstate.has_psk {
                        self.symmetricstate.mix_key(&pubkey);
                    }
                },
                Token::S => byte_index += self.symmetricstate.encrypt_and_hash(
                                &self.s.as_ref().unwrap().pubkey(), &mut message[byte_index..]),
                Token::Dhee => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhes => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Dhse => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhss => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Empty => break
            }
        }
        self.my_turn_to_send = false;
        byte_index += self.symmetricstate.encrypt_and_hash(payload, &mut message[byte_index..]);
        assert!(byte_index <= MAXMSGLEN);
        match last {
            true => (byte_index, Some(self.symmetricstate.split(self.initiator))),
            false => (byte_index, None)
        }
    }

    pub fn read_message(&mut self,
                        message: &[u8],
                        payload: &mut [u8]) ->
                            Result<(usize, Option<(CipherState<C>, CipherState<C>)>),
                                    NoiseError> {
        assert!(self.my_turn_to_send == false);
        assert!(message.len() <= MAXMSGLEN);

        let tokens = self.message_patterns[self.message_index];
        let mut last = false;
        if let Token::Empty = self.message_patterns[self.message_index+1][0] {
            last = true;
        }
        self.message_index += 1;

        let mut ptr = message;
        for token in &tokens {
            match *token {
                Token::E => {
                    let mut pubkey = [0u8; DHLEN];
                    copy_memory(&ptr[..DHLEN], &mut pubkey);
                    ptr = &ptr[DHLEN..];
                    self.re = Some(pubkey);
                    self.symmetricstate.mix_hash(&pubkey);
                    if self.symmetricstate.has_psk {
                        self.symmetricstate.mix_key(&pubkey);
                    }
                },
                Token::S => {
                    let data = match self.symmetricstate.has_k {
                        true =>  {
                            let temp = &ptr[..DHLEN + TAGLEN];
                            ptr = &ptr[DHLEN + TAGLEN..];
                            temp
                        }
                        false => {
                            let temp = &ptr[..DHLEN];
                            ptr = &ptr[DHLEN..];
                            temp
                        }
                    };
                    let mut pubkey = [0u8; DHLEN];
                    if !self.symmetricstate.decrypt_and_hash(data, &mut pubkey) {
                        return Err(NoiseError::DecryptError);
                    }
                    self.rs = Some(pubkey);
                },
                Token::Dhee => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhes => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhse => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Dhss => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Empty => break
            }
        }
        let payload_len = match self.symmetricstate.has_k {
            true => ptr.len() - TAGLEN,
            false => ptr.len()
        };
        if !self.symmetricstate.decrypt_and_hash(ptr, payload) {
            return Err(NoiseError::DecryptError);
        }
        self.my_turn_to_send = true;
        match last {
            true => Ok( (payload_len, Some(self.symmetricstate.split(self.initiator)) ) ),
            false => Ok( (payload_len, None) )
        }
    }

}

<commit_msg>tiny code formatting<commit_after>
use std::marker::PhantomData;
use crypto_stuff::*;

pub const MAXMSGLEN : usize = 65535;

#[derive(Copy, Clone)]
pub enum Token {E, S, Dhee, Dhes, Dhse, Dhss, Empty}

pub trait HandshakePattern {
    fn name(out: &mut [u8]) -> usize;
    fn get(premsg_pattern_i: &mut [Token],
           premsg_pattern_r: &mut [Token],
           msg_patterns: &mut [[Token; 8]; 5]);
}

#[derive(Debug)]
pub enum NoiseError {DecryptError}

struct SymmetricState<C: Cipher, H: Hash> {
    cipherstate : CipherState<C>,
    has_k: bool,
    has_psk: bool,
    h : [u8; MAXHASHLEN], /* Change once Rust has trait-associated consts */
    ck: [u8; MAXHASHLEN], /* Change once Rust has trait-associated consts */
    wtf : PhantomData<H>, /* So rust thinks I'm using H, this is ugly */
}

pub struct HandshakeState<P: HandshakePattern, D: Dh, C: Cipher, H: Hash, R: Random> {
    symmetricstate: SymmetricState<C, H>,
    s: Option<D>,
    e: Option<D>,
    rs: Option<[u8; DHLEN]>,
    re: Option<[u8; DHLEN]>,
    my_turn_to_send : bool,
    message_patterns : [[Token; 8]; 5],
    message_index: usize,
    initiator: bool,
    rng : R,
    wtf : PhantomData<P>, /* So rust thinks I'm using P, this is ugly */
}

impl <C: Cipher, H: Hash> SymmetricState<C, H> {

    fn new(handshake_name: &[u8]) -> SymmetricState<C, H> {
        let mut hname = [0u8; MAXHASHLEN];
        if handshake_name.len() <= H::hash_len() {
            copy_memory(handshake_name, &mut hname);
        } else {
            let mut hasher = H::new();
            hasher.input(handshake_name);
            hasher.result(&mut hname);
        }
        SymmetricState{
            cipherstate: CipherState::new(&[0u8; CIPHERKEYLEN], 0),
            has_k : false,
            has_psk: false,
            h: hname,
            ck : hname,
            wtf: PhantomData::<H>
        }
    }

    fn mix_key(&mut self, data: &[u8]) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], data, &mut hkdf_output.0, &mut hkdf_output.1);
        copy_memory(&hkdf_output.0, &mut self.ck);
        self.cipherstate = CipherState::new(&hkdf_output.1[..CIPHERKEYLEN], 0);
        self.has_k = true;
    }

    fn mix_hash(&mut self, data: &[u8]) {
        let mut hasher = H::new();
        hasher.input(&self.h[..H::hash_len()]);
        hasher.input(data);
        hasher.result(&mut self.h);
    }

    fn mix_preshared_key(&mut self, psk: &[u8]) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], psk, &mut hkdf_output.0, &mut hkdf_output.1);
        copy_memory(&hkdf_output.0, &mut self.ck);
        self.mix_hash(&hkdf_output.1[..H::hash_len()]);
        self.has_psk = true;
    }

    fn encrypt_and_hash(&mut self, plaintext: &[u8], out: &mut [u8]) -> usize {
        if self.has_k {
            self.cipherstate.encrypt_ad(&self.h[..H::hash_len()], plaintext, out);
            self.mix_hash(&out[..plaintext.len() + TAGLEN]);
            return plaintext.len() + TAGLEN;
        } else {
            copy_memory(plaintext, out);
            self.mix_hash(plaintext);
            return plaintext.len();
        }
    }

    fn decrypt_and_hash(&mut self, data: &[u8], out: &mut [u8]) -> bool {
        if self.has_k {
            if !self.cipherstate.decrypt_ad(&self.h[..H::hash_len()], data, out) {
                return false;
            }
        }
        else {
            copy_memory(data, out);
        }
        self.mix_hash(data);
        true
    }

    fn split(&self, initiator: bool) -> (CipherState<C>, CipherState<C>) {
        let mut hkdf_output = ([0u8; MAXHASHLEN], [0u8; MAXHASHLEN]);
        H::hkdf(&self.ck[..H::hash_len()], &[0u8; 0], &mut hkdf_output.0, &mut hkdf_output.1);
        let c1 = CipherState::<C>::new(&hkdf_output.0[..CIPHERKEYLEN], 0);
        let c2 = CipherState::<C>::new(&hkdf_output.1[..CIPHERKEYLEN], 0);
        if initiator { (c1, c2) } else { (c2, c1) }
    }

}

impl <P: HandshakePattern, D: Dh, C: Cipher, H: Hash, R: Random> HandshakeState<P, D, C, H, R> {

    pub fn new(rng: R,
               initiator: bool,
               prologue: &[u8],
               new_preshared_key: Option<&[u8]>,
               new_s : Option<D>,
               new_e : Option<D>,
               new_rs: Option<[u8; DHLEN]>,
               new_re: Option<[u8; DHLEN]>) -> HandshakeState<P, D, C, H, R> {
        let mut handshake_name = [0u8; 128];
        let mut name_len: usize;
        if let Some(_) = new_preshared_key {
            copy_memory("NoisePSK_".as_bytes(), &mut handshake_name);
            name_len = 9;
        } else {
            copy_memory("Noise_".as_bytes(), &mut handshake_name);
            name_len = 6;
        }
        name_len += P::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += D::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += C::name(&mut handshake_name[name_len..]);
        handshake_name[name_len] = '_' as u8;
        name_len += 1;
        name_len += H::name(&mut handshake_name[name_len..]);

        let mut symmetricstate = SymmetricState::new(&handshake_name[..name_len]);

        symmetricstate.mix_hash(prologue);

        if let Some(preshared_key) = new_preshared_key {
            symmetricstate.mix_preshared_key(preshared_key);
        }

        let mut premsg_pattern_i = [Token::Empty; 2];
        let mut premsg_pattern_r = [Token::Empty; 2];
        let mut message_patterns = [[Token::Empty; 8]; 5];
        P::get(&mut premsg_pattern_i, &mut premsg_pattern_r, &mut message_patterns);
        if initiator {
            for token in &premsg_pattern_i {
                match *token {
                    Token::S => symmetricstate.mix_hash(new_s.as_ref().unwrap().pubkey()),
                    Token::E => symmetricstate.mix_hash(new_e.as_ref().unwrap().pubkey()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
            for token in &premsg_pattern_r {
                match *token {
                    Token::S => symmetricstate.mix_hash(&new_rs.unwrap()),
                    Token::E => symmetricstate.mix_hash(&new_re.unwrap()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
        } else {
            for token in &premsg_pattern_i {
                match *token {
                    Token::S => symmetricstate.mix_hash(&new_rs.unwrap()),
                    Token::E => symmetricstate.mix_hash(&new_re.unwrap()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
            for token in &premsg_pattern_r {
                match *token {
                    Token::S => symmetricstate.mix_hash(new_s.as_ref().unwrap().pubkey()),
                    Token::E => symmetricstate.mix_hash(new_e.as_ref().unwrap().pubkey()),
                    Token::Empty => break,
                    _ => unreachable!()
                }
            }
        }

        HandshakeState{
            symmetricstate: symmetricstate,
            s: new_s, e: new_e, rs: new_rs, re: new_re,
            my_turn_to_send: initiator,
            message_patterns: message_patterns,
            message_index: 0,
            initiator: initiator,
            rng: rng,
            wtf: PhantomData::<P>}
    }

    pub fn write_message(&mut self,
                         payload: &[u8],
                         message: &mut [u8]) -> (usize, Option<(CipherState<C>, CipherState<C>)>) {
        assert!(self.my_turn_to_send);
        let tokens = self.message_patterns[self.message_index];
        let mut last = false;
        if let Token::Empty = self.message_patterns[self.message_index+1][0] {
            last = true;
        }
        self.message_index += 1;

        let mut byte_index = 0;
        for token in &tokens {
            match *token {
                Token::E => {
                    self.e = Some(D::generate(&mut self.rng));
                    let pubkey = self.e.as_ref().unwrap().pubkey();
                    copy_memory(pubkey, &mut message[byte_index..]);
                    byte_index += DHLEN;
                    self.symmetricstate.mix_hash(&pubkey);
                    if self.symmetricstate.has_psk {
                        self.symmetricstate.mix_key(&pubkey);
                    }
                },
                Token::S => {
                    byte_index += self.symmetricstate.encrypt_and_hash(
                                        &self.s.as_ref().unwrap().pubkey(),
                                        &mut message[byte_index..]);
                },
                Token::Dhee => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhes => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Dhse => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhss => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Empty => break
            }
        }
        self.my_turn_to_send = false;
        byte_index += self.symmetricstate.encrypt_and_hash(payload, &mut message[byte_index..]);
        assert!(byte_index <= MAXMSGLEN);
        match last {
            true => (byte_index, Some(self.symmetricstate.split(self.initiator))),
            false => (byte_index, None)
        }
    }

    pub fn read_message(&mut self,
                        message: &[u8],
                        payload: &mut [u8]) ->
                            Result<(usize, Option<(CipherState<C>, CipherState<C>)>),
                                    NoiseError> {
        assert!(self.my_turn_to_send == false);
        assert!(message.len() <= MAXMSGLEN);

        let tokens = self.message_patterns[self.message_index];
        let mut last = false;
        if let Token::Empty = self.message_patterns[self.message_index+1][0] {
            last = true;
        }
        self.message_index += 1;

        let mut ptr = message;
        for token in &tokens {
            match *token {
                Token::E => {
                    let mut pubkey = [0u8; DHLEN];
                    copy_memory(&ptr[..DHLEN], &mut pubkey);
                    ptr = &ptr[DHLEN..];
                    self.re = Some(pubkey);
                    self.symmetricstate.mix_hash(&pubkey);
                    if self.symmetricstate.has_psk {
                        self.symmetricstate.mix_key(&pubkey);
                    }
                },
                Token::S => {
                    let data = match self.symmetricstate.has_k {
                        true =>  {
                            let temp = &ptr[..DHLEN + TAGLEN];
                            ptr = &ptr[DHLEN + TAGLEN..];
                            temp
                        }
                        false => {
                            let temp = &ptr[..DHLEN];
                            ptr = &ptr[DHLEN..];
                            temp
                        }
                    };
                    let mut pubkey = [0u8; DHLEN];
                    if !self.symmetricstate.decrypt_and_hash(data, &mut pubkey) {
                        return Err(NoiseError::DecryptError);
                    }
                    self.rs = Some(pubkey);
                },
                Token::Dhee => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhes => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.re.unwrap())),
                Token::Dhse => self.symmetricstate.mix_key(&self.e.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Dhss => self.symmetricstate.mix_key(&self.s.as_ref().unwrap().dh(&self.rs.unwrap())),
                Token::Empty => break
            }
        }
        let payload_len = match self.symmetricstate.has_k {
            true => ptr.len() - TAGLEN,
            false => ptr.len()
        };
        if !self.symmetricstate.decrypt_and_hash(ptr, payload) {
            return Err(NoiseError::DecryptError);
        }
        self.my_turn_to_send = true;
        match last {
            true => Ok( (payload_len, Some(self.symmetricstate.split(self.initiator)) ) ),
            false => Ok( (payload_len, None) )
        }
    }

}

<|endoftext|>


<commit_before>use std::f32::consts::{PI, FRAC_1_PI};
use std::path::Path;
use std::cmp::min;

use img;
use na::{Norm, Inverse, origin};
use uuid::Uuid;

use ::{Vector3f, Point2i, Point2f, Point3f, Transform};
use geometry::{spherical_phi, spherical_theta};
use interaction::{Interaction, SurfaceInteraction};
use light::{Light, LightFlags, VisibilityTester, INFINITE};
use mipmap::{MIPMap, WrapMode};
use ray::Ray;
use sampling::Distribution2D;
use spectrum::Spectrum;

pub struct InfiniteAreaLight {
    id: Uuid,
    light_to_world: Transform,
    world_to_light: Transform,
    n_samples: u32,
    l_map: Box<MIPMap<Spectrum>>,
    world_center: Point3f,
    world_radius: f32,
    distribution: Box<Distribution2D>,
}

impl InfiniteAreaLight {
    pub fn new(l2w: Transform,
               n_samples: u32,
               power: Spectrum,
               texmap: &Path)
               -> InfiniteAreaLight {
        // Read texel data from texmap and initialise Lmap
        let (resolution, mut texels) = if let Ok(buf) = img::open(texmap) {
            info!("Loading environment map {} for infinite light",
                  texmap.display());
            let rgb = buf.to_rgb();
            let resolution = Point2i::new(rgb.width(), rgb.height());
            let mut pixels: Vec<Spectrum> = rgb.pixels()
                .map(|p| Spectrum::from_srgb(&p.data) * power)
                .collect();
            (resolution, pixels)
        } else {
            warn!("Environment map {} for infinite light not found! Using constant texture \
                   instead.",
                  texmap.display());
            (Point2i::new(1, 1), vec![power])
        };
        //
        let l_map = Box::new(MIPMap::new(&resolution, &texels[..], false, 0.0, WrapMode::Repeat));
        // initialize sampling PDFs for infinite area light
        // - compute scalar-valued image img from environment map
        let (width, height) = (2 * l_map.width(), 2 * l_map.height());
        let filter = 0.5 / min(width, height) as f32;
        let mut img = Vec::with_capacity(width * height);
        for v in 0..height {
            let vp = (v as f32 + 0.5) / height as f32;
            let sin_theta = (PI * (v as f32 + 0.5) / height as f32).sin();
            for u in 0..width {
                let up = (u as f32 + 0.5) / width as f32;
                img.push(l_map.lookup(&Point2f::new(up, vp), filter).y() * sin_theta);
            }
        }
        // - compute sampling distributions for rows and columns of image
        let distribution = Box::new(Distribution2D::new(&img[..], width, height));

        InfiniteAreaLight {
            id: Uuid::new_v4(),
            light_to_world: l2w,
            world_to_light: l2w.inverse().unwrap(),
            n_samples: n_samples,
            l_map: l_map,
            world_center: origin(),
            world_radius: 100.0, // TODO
            distribution: distribution,
        }
    }
}

impl Light for InfiniteAreaLight {
    fn id(&self) -> Uuid {
        self.id
    }

    fn sample_li(&self,
                 isect: &SurfaceInteraction,
                 u: &Point2f)
                 -> (Spectrum, Vector3f, f32, VisibilityTester) {
        // Find (u, v) sample coordinates in infinite light texture
        let (uv, map_pdf) = self.distribution.sample_continuous(u);
        if map_pdf == 0.0 {
            return (Spectrum::black(),
                    Vector3f::new(0.0, 0.0, 0.0),
                    0.0,
                    VisibilityTester::new(Interaction::from_point(&origin()),
                                          Interaction::from_point(&origin())));
        }
        // Convert infinite light sample point to direction
        let theta = uv[1] * PI;
        let phi = uv[0] * 2.0 * PI;
        let cos_theta = theta.cos();
        let sin_theta = theta.sin();
        let cos_phi = phi.cos();
        let sin_phi = phi.sin();
        let wi = self.light_to_world *
                 Vector3f::new(sin_theta * cos_phi, sin_theta * sin_phi, cos_theta);
        // Compute PDF for sampled infinite light direction
        let pdf = if sin_theta == 0.0 {
            0.0
        } else {
            map_pdf / (2.0 * PI * PI * sin_theta)
        };
        // Return radiance value for infinite light direction
        let target = isect.p + wi * (2.0 * self.world_radius);
        let vis = VisibilityTester::new(isect.into(), Interaction::from_point(&target));
        (self.l_map.lookup(&uv, 0.0), wi, pdf, vis)
    }

    fn pdf_li(&self, si: &SurfaceInteraction, w: &Vector3f) -> f32 {
        let wi = self.world_to_light * *w;
        let theta = spherical_theta(&wi);
        let phi = spherical_phi(&wi);
        let sin_theta = theta.sin();

        if sin_theta == 0.0 {
            0.0
        } else {
            self.distribution.pdf(&Point2f::new(phi * FRAC_1_PI * 0.5, theta * FRAC_1_PI)) /
            (2.0 * PI * PI * sin_theta)
        }
    }

    fn n_samples(&self) -> u32 {
        self.n_samples
    }

    fn flags(&self) -> LightFlags {
        INFINITE
    }

    fn power(&self) -> Spectrum {
        PI * self.world_radius * self.world_radius * self.l_map.lookup(&Point2f::new(0.5, 0.5), 0.5)
    }

    fn le(&self, ray: &Ray) -> Spectrum {
        let w = (self.world_to_light * ray.d).normalize();
        let st = Point2f::new(spherical_phi(&w) * FRAC_1_PI * 0.5,
                              spherical_theta(&w) * FRAC_1_PI);

        self.l_map.lookup(&st, 0.0)
    }
}
<commit_msg>Looks like TGA files don't need to be gamma corrected?<commit_after>use std::f32::consts::{PI, FRAC_1_PI};
use std::path::Path;
use std::cmp::min;

use img;
use na::{Norm, Inverse, origin};
use uuid::Uuid;

use ::{Vector3f, Point2i, Point2f, Point3f, Transform};
use geometry::{spherical_phi, spherical_theta};
use interaction::{Interaction, SurfaceInteraction};
use light::{Light, LightFlags, VisibilityTester, INFINITE};
use mipmap::{MIPMap, WrapMode};
use ray::Ray;
use sampling::Distribution2D;
use spectrum::Spectrum;

pub struct InfiniteAreaLight {
    id: Uuid,
    light_to_world: Transform,
    world_to_light: Transform,
    n_samples: u32,
    l_map: Box<MIPMap<Spectrum>>,
    world_center: Point3f,
    world_radius: f32,
    distribution: Box<Distribution2D>,
}

impl InfiniteAreaLight {
    pub fn new(l2w: Transform,
               n_samples: u32,
               power: Spectrum,
               texmap: &Path)
               -> InfiniteAreaLight {
        // Read texel data from texmap and initialise Lmap
        let (resolution, mut texels) = if let Ok(buf) = img::open(texmap) {
            info!("Loading environment map {} for infinite light",
                  texmap.display());
            let rgb = buf.to_rgb();
            let resolution = Point2i::new(rgb.width(), rgb.height());
            let mut pixels: Vec<Spectrum> = rgb.pixels()
                .map(|p| {
                    let r = p.data[0] as f32 / 255.0;
                    let g = p.data[1] as f32 / 255.0;
                    let b = p.data[2] as f32 / 255.0;
                    Spectrum::rgb(r, g, b) * power
                })
                .collect();
            (resolution, pixels)
        } else {
            warn!("Environment map {} for infinite light not found! Using constant texture \
                   instead.",
                  texmap.display());
            (Point2i::new(1, 1), vec![power])
        };
        //
        let l_map = Box::new(MIPMap::new(&resolution, &texels[..], false, 0.0, WrapMode::Repeat));
        // initialize sampling PDFs for infinite area light
        // - compute scalar-valued image img from environment map
        let (width, height) = (2 * l_map.width(), 2 * l_map.height());
        let filter = 0.5 / min(width, height) as f32;
        let mut img = Vec::with_capacity(width * height);
        for v in 0..height {
            let vp = (v as f32 + 0.5) / height as f32;
            let sin_theta = (PI * (v as f32 + 0.5) / height as f32).sin();
            for u in 0..width {
                let up = (u as f32 + 0.5) / width as f32;
                img.push(l_map.lookup(&Point2f::new(up, vp), filter).y() * sin_theta);
            }
        }
        // - compute sampling distributions for rows and columns of image
        let distribution = Box::new(Distribution2D::new(&img[..], width, height));

        InfiniteAreaLight {
            id: Uuid::new_v4(),
            light_to_world: l2w,
            world_to_light: l2w.inverse().unwrap(),
            n_samples: n_samples,
            l_map: l_map,
            world_center: origin(),
            world_radius: 28.0, // TODO
            distribution: distribution,
        }
    }
}

impl Light for InfiniteAreaLight {
    fn id(&self) -> Uuid {
        self.id
    }

    fn sample_li(&self,
                 isect: &SurfaceInteraction,
                 u: &Point2f)
                 -> (Spectrum, Vector3f, f32, VisibilityTester) {
        // Find (u, v) sample coordinates in infinite light texture
        let (uv, map_pdf) = self.distribution.sample_continuous(u);
        if map_pdf == 0.0 {
            return (Spectrum::black(),
                    Vector3f::new(0.0, 0.0, 0.0),
                    0.0,
                    VisibilityTester::new(Interaction::from_point(&origin()),
                                          Interaction::from_point(&origin())));
        }
        // Convert infinite light sample point to direction
        let theta = uv[1] * PI;
        let phi = uv[0] * 2.0 * PI;
        let cos_theta = theta.cos();
        let sin_theta = theta.sin();
        let cos_phi = phi.cos();
        let sin_phi = phi.sin();
        let wi = self.light_to_world *
                 Vector3f::new(sin_theta * cos_phi, sin_theta * sin_phi, cos_theta);
        // Compute PDF for sampled infinite light direction
        let pdf = if sin_theta == 0.0 {
            0.0
        } else {
            map_pdf / (2.0 * PI * PI * sin_theta)
        };
        // Return radiance value for infinite light direction
        let target = isect.p + wi * (2.0 * self.world_radius);
        let vis = VisibilityTester::new(isect.into(), Interaction::from_point(&target));
        (self.l_map.lookup(&uv, 0.0), wi, pdf, vis)
    }

    fn pdf_li(&self, si: &SurfaceInteraction, w: &Vector3f) -> f32 {
        let wi = self.world_to_light * *w;
        let theta = spherical_theta(&wi);
        let phi = spherical_phi(&wi);
        let sin_theta = theta.sin();

        if sin_theta == 0.0 {
            0.0
        } else {
            self.distribution.pdf(&Point2f::new(phi * FRAC_1_PI * 0.5, theta * FRAC_1_PI)) /
            (2.0 * PI * PI * sin_theta)
        }
    }

    fn n_samples(&self) -> u32 {
        self.n_samples
    }

    fn flags(&self) -> LightFlags {
        INFINITE
    }

    fn power(&self) -> Spectrum {
        PI * self.world_radius * self.world_radius * self.l_map.lookup(&Point2f::new(0.5, 0.5), 0.5)
    }

    fn le(&self, ray: &Ray) -> Spectrum {
        let w = (self.world_to_light * ray.d).normalize();
        let st = Point2f::new(spherical_phi(&w) * FRAC_1_PI * 0.5,
                              spherical_theta(&w) * FRAC_1_PI);

        self.l_map.lookup(&st, 0.0)
    }
}
<|endoftext|>


<commit_before>#![deny(missing_docs)]
#![warn(rust_2018_idioms)]
#![doc(html_root_url = "https://docs.rs/amq-protocol-tcp/6.0.0-alpha3/")]

//! # AMQP URI TCP connection handling
//!
//! amq-protocol-tcp is a library aiming at providing tools to help
//! connecting to an AMQP URI

use amq_protocol_uri::{AMQPScheme, AMQPUri};
use log::trace;
use mio::net::TcpStream as MioTcpStream;
use std::{error::Error, fmt, io};
use tcp_stream::{
    HandshakeError as TcpStreamHandshakeError,
    MidHandshakeTlsStream as TcpStreamMidHandshakeTlsStream,
};

/// Re-export TcpStream
pub use tcp_stream::{Identity, TcpStream};

/// Trait providing a method to connect to a TcpStream
pub trait AMQPUriTcpExt {
    /// connect to a TcpStream
    fn connect<S, F: FnOnce(TcpStream, AMQPUri) -> S>(self, f: F) -> Result<S, HandshakeError<S, F>>
    where
        Self: Sized,
    {
        self.connect_full(f, None)
    }
    /// connect to a TcpStream, registering it to the given Poll with the given Token to handle the
    /// handshake process. You should reregister it afterwards to better fit your needs
    fn connect_full<S, F: FnOnce(TcpStream, AMQPUri) -> S>(
        self,
        f: F,
        identity: Option<Identity<'_, '_>>,
    ) -> Result<S, HandshakeError<S, F>>;
}

impl AMQPUriTcpExt for AMQPUri {
    fn connect_full<S, F: FnOnce(TcpStream, AMQPUri) -> S>(
        self,
        f: F,
        identity: Option<Identity<'_, '_>>,
    ) -> Result<S, HandshakeError<S, F>> {
        let uri = format!("{}:{}", self.authority.host, self.authority.port);
        trace!("Connecting to {}", uri);
        let stream = TcpStream::connect(uri)?;

        match self.scheme {
            AMQPScheme::AMQP => Ok(f(stream, self)),
            AMQPScheme::AMQPS => match stream.into_tls(&self.authority.host, identity) {
                Ok(s) => Ok(f(s, self)),
                Err(err) => Err(HandshakeError::from(err, f, self)),
            },
        }
    }
}

impl AMQPUriTcpExt for &str {
    fn connect_full<S, F: FnOnce(TcpStream, AMQPUri) -> S>(
        self,
        f: F,
        identity: Option<Identity<'_, '_>>,
    ) -> Result<S, HandshakeError<S, F>> {
        self.parse::<AMQPUri>()
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?
            .connect_full(f, identity)
    }
}

/// A TLS stream which has been interrupted during the handshake
pub struct MidHandshakeTlsStream<S, F: FnOnce(TcpStream, AMQPUri) -> S>(
    TcpStreamMidHandshakeTlsStream,
    F,
    AMQPUri,
);

impl<S, F: FnOnce(TcpStream, AMQPUri) -> S> MidHandshakeTlsStream<S, F> {
    /// Get a reference to the inner stream
    pub fn get_ref(&self) -> &MioTcpStream {
        self.0.get_ref()
    }

    /// Get a mutablereference to the inner stream
    pub fn get_mut(&mut self) -> &mut MioTcpStream {
        self.0.get_mut()
    }

    /// Retry the handshake
    pub fn handshake(self) -> Result<S, HandshakeError<S, F>> {
        let MidHandshakeTlsStream(mid, f, uri) = self;
        match mid.handshake() {
            Ok(s) => Ok(f(s, uri)),
            Err(err) => Err(HandshakeError::from(err, f, uri)),
        }
    }
}

/// An error returned while performing the handshake
#[allow(clippy::large_enum_variant)]
pub enum HandshakeError<S, F: FnOnce(TcpStream, AMQPUri) -> S> {
    /// We hit WouldBlock during handshake
    WouldBlock(MidHandshakeTlsStream<S, F>),
    /// We hit a critical failure
    Failure(io::Error),
}

impl<S, F: FnOnce(TcpStream, AMQPUri) -> S> HandshakeError<S, F> {
    fn from(error: TcpStreamHandshakeError, f: F, uri: AMQPUri) -> Self {
        match error {
            TcpStreamHandshakeError::WouldBlock(mid) => {
                Self::WouldBlock(MidHandshakeTlsStream(mid, f, uri))
            }
            TcpStreamHandshakeError::Failure(error) => Self::Failure(error),
        }
    }
}

impl<S, F: FnOnce(TcpStream, AMQPUri) -> S> From<io::Error> for HandshakeError<S, F> {
    fn from(error: io::Error) -> Self {
        Self::Failure(error)
    }
}

impl<S, F: FnOnce(TcpStream, AMQPUri) -> S> fmt::Display for HandshakeError<S, F> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            HandshakeError::WouldBlock(_) => f.write_str("WouldBlock hit during handshake"),
            HandshakeError::Failure(err) => f.write_fmt(format_args!("IO error: {}", err)),
        }
    }
}

impl<S, F: FnOnce(TcpStream, AMQPUri) -> S> fmt::Debug for HandshakeError<S, F> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            HandshakeError::WouldBlock(_) => f.debug_tuple("HandshakeError::WouldBlock").finish(),
            HandshakeError::Failure(err) => {
                f.debug_tuple("HandshakeError::Failure").field(err).finish()
            }
        }
    }
}

impl<S, F: FnOnce(TcpStream, AMQPUri) -> S> Error for HandshakeError<S, F> {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            HandshakeError::Failure(err) => Some(err),
            _ => None,
        }
    }
}
<commit_msg>tcp: add HandshakeResult<commit_after>#![deny(missing_docs)]
#![warn(rust_2018_idioms)]
#![doc(html_root_url = "https://docs.rs/amq-protocol-tcp/6.0.0-alpha3/")]

//! # AMQP URI TCP connection handling
//!
//! amq-protocol-tcp is a library aiming at providing tools to help
//! connecting to an AMQP URI

use amq_protocol_uri::{AMQPScheme, AMQPUri};
use log::trace;
use mio::net::TcpStream as MioTcpStream;
use std::{error::Error, fmt, io};
use tcp_stream::{
    HandshakeError as TcpStreamHandshakeError,
    MidHandshakeTlsStream as TcpStreamMidHandshakeTlsStream,
};

/// Re-export TcpStream
pub use tcp_stream::{Identity, TcpStream};

/// Type holding the result or handshake state
pub type HandshakeResult<S, F> = Result<S, HandshakeError<S, F>>;

/// Trait providing a method to connect to a TcpStream
pub trait AMQPUriTcpExt {
    /// connect to a TcpStream
    fn connect<S, F: FnOnce(TcpStream, AMQPUri) -> S>(self, f: F) -> HandshakeResult<S, F>
    where
        Self: Sized,
    {
        self.connect_full(f, None)
    }
    /// connect to a TcpStream, registering it to the given Poll with the given Token to handle the
    /// handshake process. You should reregister it afterwards to better fit your needs
    fn connect_full<S, F: FnOnce(TcpStream, AMQPUri) -> S>(
        self,
        f: F,
        identity: Option<Identity<'_, '_>>,
    ) -> HandshakeResult<S, F>;
}

impl AMQPUriTcpExt for AMQPUri {
    fn connect_full<S, F: FnOnce(TcpStream, AMQPUri) -> S>(
        self,
        f: F,
        identity: Option<Identity<'_, '_>>,
    ) -> Result<S, HandshakeError<S, F>> {
        let uri = format!("{}:{}", self.authority.host, self.authority.port);
        trace!("Connecting to {}", uri);
        let stream = TcpStream::connect(uri)?;

        match self.scheme {
            AMQPScheme::AMQP => Ok(f(stream, self)),
            AMQPScheme::AMQPS => match stream.into_tls(&self.authority.host, identity) {
                Ok(s) => Ok(f(s, self)),
                Err(err) => Err(HandshakeError::from(err, f, self)),
            },
        }
    }
}

impl AMQPUriTcpExt for &str {
    fn connect_full<S, F: FnOnce(TcpStream, AMQPUri) -> S>(
        self,
        f: F,
        identity: Option<Identity<'_, '_>>,
    ) -> HandshakeResult<S, F> {
        self.parse::<AMQPUri>()
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?
            .connect_full(f, identity)
    }
}

/// A TLS stream which has been interrupted during the handshake
pub struct MidHandshakeTlsStream<S, F: FnOnce(TcpStream, AMQPUri) -> S>(
    TcpStreamMidHandshakeTlsStream,
    F,
    AMQPUri,
);

impl<S, F: FnOnce(TcpStream, AMQPUri) -> S> MidHandshakeTlsStream<S, F> {
    /// Get a reference to the inner stream
    pub fn get_ref(&self) -> &MioTcpStream {
        self.0.get_ref()
    }

    /// Get a mutablereference to the inner stream
    pub fn get_mut(&mut self) -> &mut MioTcpStream {
        self.0.get_mut()
    }

    /// Retry the handshake
    pub fn handshake(self) -> HandshakeResult<S, F> {
        let MidHandshakeTlsStream(mid, f, uri) = self;
        match mid.handshake() {
            Ok(s) => Ok(f(s, uri)),
            Err(err) => Err(HandshakeError::from(err, f, uri)),
        }
    }
}

/// An error returned while performing the handshake
#[allow(clippy::large_enum_variant)]
pub enum HandshakeError<S, F: FnOnce(TcpStream, AMQPUri) -> S> {
    /// We hit WouldBlock during handshake
    WouldBlock(MidHandshakeTlsStream<S, F>),
    /// We hit a critical failure
    Failure(io::Error),
}

impl<S, F: FnOnce(TcpStream, AMQPUri) -> S> HandshakeError<S, F> {
    fn from(error: TcpStreamHandshakeError, f: F, uri: AMQPUri) -> Self {
        match error {
            TcpStreamHandshakeError::WouldBlock(mid) => {
                Self::WouldBlock(MidHandshakeTlsStream(mid, f, uri))
            }
            TcpStreamHandshakeError::Failure(error) => Self::Failure(error),
        }
    }
}

impl<S, F: FnOnce(TcpStream, AMQPUri) -> S> From<io::Error> for HandshakeError<S, F> {
    fn from(error: io::Error) -> Self {
        Self::Failure(error)
    }
}

impl<S, F: FnOnce(TcpStream, AMQPUri) -> S> fmt::Display for HandshakeError<S, F> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            HandshakeError::WouldBlock(_) => f.write_str("WouldBlock hit during handshake"),
            HandshakeError::Failure(err) => f.write_fmt(format_args!("IO error: {}", err)),
        }
    }
}

impl<S, F: FnOnce(TcpStream, AMQPUri) -> S> fmt::Debug for HandshakeError<S, F> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            HandshakeError::WouldBlock(_) => f.debug_tuple("HandshakeError::WouldBlock").finish(),
            HandshakeError::Failure(err) => {
                f.debug_tuple("HandshakeError::Failure").field(err).finish()
            }
        }
    }
}

impl<S, F: FnOnce(TcpStream, AMQPUri) -> S> Error for HandshakeError<S, F> {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            HandshakeError::Failure(err) => Some(err),
            _ => None,
        }
    }
}
<|endoftext|>


<commit_before>use uuid::Uuid;

use common::err::{Error, TResult};
use common::types::Type;

// let ctx = TokamakContext::new();
// ctx.from(RandomGenerator).select(...);

use super::TokamakContext;

pub struct DataFrame<'a> {
  pub ctx : &'a TokamakContext,
  pub decl: DataFrameDecl<'a>
}

pub enum DataFrameDecl<'a> {
  From(Box<DataSet>),
  Select(Box<DataFrame<'a>>, Vec<Expr>)
}

pub enum Expr {
  Plus,
  Field(String)
}

impl<'a> DataFrame<'a> {
  pub fn kind(&self) -> &'static str {
    match self.decl {
      DataFrameDecl::From  (_) => "from",
      DataFrameDecl::Select(_,_) => "select"
    }
  }

  pub fn select(self, exprs: Vec<Expr>) -> DataFrame<'a> {
    DataFrame {ctx: self.ctx, decl: DataFrameDecl::Select(Box::new(self), exprs)}
  }

  fn count(&self) -> TResult<usize> {
    Ok(0)
  }
}

pub trait DataSet {
  fn name(&self) -> &str;

  //fn schema(&self) -> &Vec<Box<Type>>;
}

pub struct CustomDataSource {
  src_type : String,
  schema   : Vec<String>,
  props    : Vec<(String, String)>
}

impl DataSet for CustomDataSource {
  fn name(&self) -> &str {
    "aaa"
  }
}

pub fn RandomGenerator(types: Vec<&str>) -> Box<DataSet>
{
  Box::new(CustomDataSource {
    src_type: "random".to_string(),
    schema  : types.into_iter().map(|s| s.to_string()).collect::<Vec<String>>(),
    props   : Vec::new()
  })
}

fn typestr_to_schema(ctx: &TokamakContext, types: Vec<&str>) -> Vec<Box<Type>>
{
  types.iter()
    .map( |s| ctx.get_type(s).unwrap().clone_box() )
    .collect::<Vec<Box<Type>>>()
}<commit_msg>Implement name field to RandomGenerator.<commit_after>use uuid::Uuid;

use common::err::{Error, TResult};
use common::types::Type;

// let ctx = TokamakContext::new();
// ctx.from(RandomGenerator).select(...);

use super::TokamakContext;

pub struct DataFrame<'a> {
  pub ctx : &'a TokamakContext,
  pub decl: DataFrameDecl<'a>
}

pub enum DataFrameDecl<'a> {
  From(Box<DataSet>),
  Select(Box<DataFrame<'a>>, Vec<Expr>)
}

pub enum Expr {
  Plus,
  Field(String)
}

impl<'a> DataFrame<'a> {
  pub fn kind(&self) -> &'static str {
    match self.decl {
      DataFrameDecl::From  (_) => "from",
      DataFrameDecl::Select(_,_) => "select"
    }
  }

  pub fn select(self, exprs: Vec<Expr>) -> DataFrame<'a> {
    DataFrame {ctx: self.ctx, decl: DataFrameDecl::Select(Box::new(self), exprs)}
  }

  fn count(&self) -> TResult<usize> {
    Ok(0)
  }
}

pub trait DataSet {
  fn name(&self) -> &str;

  //fn schema(&self) -> &Vec<Box<Type>>;
}

pub struct CustomDataSource {
  name     : String,
  src_type : String,
  schema   : Vec<String>,
  props    : Vec<(String, String)>
}

impl DataSet for CustomDataSource {
  fn name(&self) -> &str {
    &self.name
  }
}

pub fn RandomGenerator(types: Vec<&str>) -> Box<DataSet>
{
  Box::new(CustomDataSource {
    name    :  Uuid::new_v4().to_hyphenated_string(),
    src_type: "random".to_string(),
    schema  : types.into_iter().map(|s| s.to_string()).collect::<Vec<String>>(),
    props   : Vec::new()
  })
}

fn typestr_to_schema(ctx: &TokamakContext, types: Vec<&str>) -> Vec<Box<Type>>
{
  types.iter()
    .map( |s| ctx.get_type(s).unwrap().clone_box() )
    .collect::<Vec<Box<Type>>>()
}<|endoftext|>


<commit_before>mod idt;

macro_rules! handler {
    ($name: ident) => {{
        #[naked]
        extern "C" fn wrapper() -> ! {
            unsafe {
                asm!("mov rdi, rsp
                      sub rsp, 8 // align the stack pointer
                      call $0"
                      :: "i"($name as extern "C" fn(*const ExceptionStackFrame) -> !)
                      : "rdi" : "intel");
                ::core::intrinsics::unreachable();
            }
        }
        wrapper
    }}
}

macro_rules! handler_with_error_code {
    ($name: ident) => {{
        #[naked]
        extern "C" fn wrapper() -> ! {
            unsafe {
                asm!("pop rsi // pop error code into rsi
                      mov rdi, rsp
                      sub rsp, 8 // align the stack pointer
                      call $0"
                      :: "i"($name as extern "C" fn(*const ExceptionStackFrame, u64) -> !)
                      : "rdi" : "intel");
                ::core::intrinsics::unreachable();
            }
        }
        wrapper
    }}
}

lazy_static! {
    static ref IDT: idt::Idt = {
        let mut idt = idt::Idt::new();

        idt.set_handler(0, handler!(divide_by_zero_handler));
        idt.set_handler(6, handler!(invalid_opcode_handler));
        idt.set_handler(8, double_fault_handler);
        idt.set_handler(14, handler_with_error_code!(page_fault_handler));

        idt
    };
}

pub fn init() {
    IDT.load();
}

use vga_buffer::print_error;

extern "C" fn divide_by_zero_handler(stack_frame: *const ExceptionStackFrame) -> ! {
    unsafe {
        print_error(format_args!("EXCEPTION: DIVIDE BY ZERO\n{:#?}", *stack_frame));
    }
    loop {}
}

extern "C" fn invalid_opcode_handler(stack_frame: *const ExceptionStackFrame)
    -> !
{
    unsafe {
        print_error(format_args!("EXCEPTION: INVALID OPCODE at {:#x}\n{:#?}",
            (*stack_frame).instruction_pointer, *stack_frame));
    }
    loop {}
}

extern "C" fn page_fault_handler(stack_frame: *const ExceptionStackFrame,
                                 error_code: u64) -> !
{
    unsafe {
        print_error(format_args!(
            "EXCEPTION: PAGE FAULT with error code {:?}\n{:#?}",
            PageFaultErrorCode::from_bits(error_code), *stack_frame));
    }
    loop {}
}

extern "C" fn double_fault_handler() -> ! {
    unsafe { print_error(format_args!("EXCEPTION: DOUBLE FAULT")) };
    loop {}
}

#[derive(Debug)]
#[repr(C)]
struct ExceptionStackFrame {
    instruction_pointer: u64,
    code_segment: u64,
    cpu_flags: u64,
    stack_pointer: u64,
    stack_segment: u64,
}

bitflags! {
    flags PageFaultErrorCode: u64 {
        const PROTECTION_VIOLATION = 1 << 0,
        const CAUSED_BY_WRITE = 1 << 1,
        const USER_MODE = 1 << 2,
        const MALFORMED_TABLE = 1 << 3,
        const INSTRUCTION_FETCH = 1 << 4,
    }
}
<commit_msg>Print cr2 in page_fault_handler<commit_after>mod idt;

macro_rules! handler {
    ($name: ident) => {{
        #[naked]
        extern "C" fn wrapper() -> ! {
            unsafe {
                asm!("mov rdi, rsp
                      sub rsp, 8 // align the stack pointer
                      call $0"
                      :: "i"($name as extern "C" fn(*const ExceptionStackFrame) -> !)
                      : "rdi" : "intel");
                ::core::intrinsics::unreachable();
            }
        }
        wrapper
    }}
}

macro_rules! handler_with_error_code {
    ($name: ident) => {{
        #[naked]
        extern "C" fn wrapper() -> ! {
            unsafe {
                asm!("pop rsi // pop error code into rsi
                      mov rdi, rsp
                      sub rsp, 8 // align the stack pointer
                      call $0"
                      :: "i"($name as extern "C" fn(*const ExceptionStackFrame, u64) -> !)
                      : "rdi" : "intel");
                ::core::intrinsics::unreachable();
            }
        }
        wrapper
    }}
}

lazy_static! {
    static ref IDT: idt::Idt = {
        let mut idt = idt::Idt::new();

        idt.set_handler(0, handler!(divide_by_zero_handler));
        idt.set_handler(6, handler!(invalid_opcode_handler));
        idt.set_handler(8, double_fault_handler);
        idt.set_handler(14, handler_with_error_code!(page_fault_handler));

        idt
    };
}

pub fn init() {
    IDT.load();
}

use vga_buffer::print_error;

extern "C" fn divide_by_zero_handler(stack_frame: *const ExceptionStackFrame) -> ! {
    unsafe {
        print_error(format_args!("EXCEPTION: DIVIDE BY ZERO\n{:#?}", *stack_frame));
    }
    loop {}
}

extern "C" fn invalid_opcode_handler(stack_frame: *const ExceptionStackFrame) -> ! {
    unsafe {
        print_error(format_args!("EXCEPTION: INVALID OPCODE at {:#x}\n{:#?}",
                                 (*stack_frame).instruction_pointer,
                                 *stack_frame));
    }
    loop {}
}

extern "C" fn page_fault_handler(stack_frame: *const ExceptionStackFrame, error_code: u64) -> ! {
    use x86::controlregs;
    unsafe {
        print_error(format_args!("EXCEPTION: PAGE FAULT while accessing {:#x}\nerror code: \
                                  {:?}\n{:#?}",
                                 controlregs::cr2(),
                                 PageFaultErrorCode::from_bits(error_code).unwrap(),
                                 *stack_frame));
    }
    loop {}
}

extern "C" fn double_fault_handler() -> ! {
    unsafe { print_error(format_args!("EXCEPTION: DOUBLE FAULT")) };
    loop {}
}

#[derive(Debug)]
#[repr(C)]
struct ExceptionStackFrame {
    instruction_pointer: u64,
    code_segment: u64,
    cpu_flags: u64,
    stack_pointer: u64,
    stack_segment: u64,
}

bitflags! {
    flags PageFaultErrorCode: u64 {
        const PROTECTION_VIOLATION = 1 << 0,
        const CAUSED_BY_WRITE = 1 << 1,
        const USER_MODE = 1 << 2,
        const MALFORMED_TABLE = 1 << 3,
        const INSTRUCTION_FETCH = 1 << 4,
    }
}
<|endoftext|>