<commit_before>use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, AtomicUsize, ATOMIC_USIZE_INIT, AtomicIsize,
                        Ordering};
use std::sync::mpsc::{self, SyncSender, Receiver};
use std::sync::{Mutex, Arc};
use std::time::Duration;
use std::thread;
use std::fmt;
use std::error;

use {ManageConnection, CustomizeConnection, Pool, Config};

#[derive(Debug)]
pub struct Error;

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str("blammo")
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        "Error"
    }
}

#[derive(Debug, PartialEq)]
struct FakeConnection(bool);

struct OkManager;

impl ManageConnection for OkManager {
    type Connection = FakeConnection;
    type Error = Error;

    fn connect(&self) -> Result<FakeConnection, Error> {
        Ok(FakeConnection(true))
    }

    fn is_valid(&self, _: &mut FakeConnection) -> Result<(), Error> {
        Ok(())
    }

    fn has_broken(&self, _: &mut FakeConnection) -> bool {
        false
    }
}

struct NthConnectFailManager {
    n: Mutex<u32>,
}

impl ManageConnection for NthConnectFailManager {
    type Connection = FakeConnection;
    type Error = Error;

    fn connect(&self) -> Result<FakeConnection, Error> {
        let mut n = self.n.lock().unwrap();
        if *n > 0 {
            *n -= 1;
            Ok(FakeConnection(true))
        } else {
            Err(Error)
        }
    }

    fn is_valid(&self, _: &mut FakeConnection) -> Result<(), Error> {
        Ok(())
    }

    fn has_broken(&self, _: &mut FakeConnection) -> bool {
        false
    }
}

#[test]
fn test_pool_size_ok() {
    let config = Config::builder().pool_size(5).build();
    let manager = NthConnectFailManager { n: Mutex::new(5) };
    let pool = Pool::new(config, manager).unwrap();
    let mut conns = vec![];
    for _ in 0..5 {
        conns.push(pool.get().ok().unwrap());
    }
}

#[test]
fn test_acquire_release() {
    let config = Config::builder().pool_size(2).build();
    let pool = Pool::new(config, OkManager).unwrap();

    let conn1 = pool.get().ok().unwrap();
    let conn2 = pool.get().ok().unwrap();
    drop(conn1);
    let conn3 = pool.get().ok().unwrap();
    drop(conn2);
    drop(conn3);
}

#[test]
fn test_is_send_sync() {
    fn is_send_sync<T: Send + Sync>() {}
    is_send_sync::<Pool<OkManager>>();
}

#[test]
fn test_issue_2_unlocked_during_is_valid() {
    struct BlockingChecker {
        first: AtomicBool,
        s: Mutex<SyncSender<()>>,
        r: Mutex<Receiver<()>>,
    }

    impl ManageConnection for BlockingChecker {
        type Connection = FakeConnection;
        type Error = Error;

        fn connect(&self) -> Result<FakeConnection, Error> {
            Ok(FakeConnection(true))
        }

        fn is_valid(&self, _: &mut FakeConnection) -> Result<(), Error> {
            if self.first.compare_and_swap(true, false, Ordering::SeqCst) {
                self.s.lock().unwrap().send(()).unwrap();
                self.r.lock().unwrap().recv().unwrap();
            }
            Ok(())
        }

        fn has_broken(&self, _: &mut FakeConnection) -> bool {
            false
        }
    }

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

    let config = Config::builder()
                     .test_on_check_out(true)
                     .pool_size(2)
                     .build();
    let manager = BlockingChecker {
        first: AtomicBool::new(true),
        s: Mutex::new(s1),
        r: Mutex::new(r2),
    };
    let pool = Arc::new(Pool::new(config, manager).unwrap());

    let p2 = pool.clone();
    let t = thread::spawn(move || {
        p2.get().ok().unwrap();
    });

    r1.recv().unwrap();
    // get call by other task has triggered the health check
    pool.get().ok().unwrap();
    s2.send(()).ok().unwrap();

    t.join().ok().unwrap();
}

#[test]
fn test_drop_on_broken() {
    static DROPPED: AtomicBool = ATOMIC_BOOL_INIT;
    DROPPED.store(false, Ordering::SeqCst);

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.store(true, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = Error;

        fn connect(&self) -> Result<Connection, Error> {
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), Error> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            true
        }
    }

    let pool = Pool::new(Default::default(), Handler).unwrap();

    drop(pool.get().ok().unwrap());

    assert!(DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_initialization_failure() {
    let config = Config::builder()
                     .connection_timeout(Duration::from_secs(1))
                     .build();
    let manager = NthConnectFailManager { n: Mutex::new(0) };
    let err = Pool::new(config, manager).err().unwrap();
    assert!(err.to_string().contains("blammo"));
}

#[test]
fn test_lazy_initialization_failure() {
    let config = Config::builder()
                     .connection_timeout(Duration::from_secs(1))
                     .initialization_fail_fast(false)
                     .build();
    let manager = NthConnectFailManager { n: Mutex::new(0) };
    let pool = Pool::new(config, manager).unwrap();
    let err = pool.get().err().unwrap();
    assert!(err.to_string().contains("blammo"));
}

#[test]
fn test_get_timeout() {
    let config = Config::builder()
                     .pool_size(1)
                     .connection_timeout(Duration::from_secs(1))
                     .build();
    let pool = Pool::new(config, OkManager).unwrap();
    let _c = pool.get().unwrap();
    pool.get().err().unwrap();
}

#[test]
fn test_connection_customizer() {
    static DROPPED: AtomicBool = ATOMIC_BOOL_INIT;
    DROPPED.store(false, Ordering::SeqCst);

    struct Connection(i32);

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.store(true, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = Error;

        fn connect(&self) -> Result<Connection, Error> {
            Ok(Connection(0))
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), Error> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            true
        }
    }

    #[derive(Debug)]
    struct Customizer;

    impl CustomizeConnection<Connection, Error> for Customizer {
        fn on_acquire(&self, conn: &mut Connection) -> Result<(), Error> {
            if !DROPPED.load(Ordering::SeqCst) {
                Err(Error)
            } else {
                conn.0 = 1;
                Ok(())
            }
        }
    }

    let config = Config::builder()
                     .connection_customizer(Box::new(Customizer))
                     .build();
    let pool = Pool::new(config, Handler).unwrap();

    let conn = pool.get().unwrap();
    assert_eq!(1, conn.0);
    assert!(DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_idle_timeout() {
    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
    }

    struct Handler(AtomicIsize);

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = Error;

        fn connect(&self) -> Result<Connection, Error> {
            if self.0.fetch_sub(1, Ordering::SeqCst) > 0 {
                Ok(Connection)
            } else {
                Err(Error)
            }
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), Error> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                     .pool_size(5)
                     .idle_timeout(Some(Duration::from_secs(1)))
                     .build();
    let pool = Pool::new_inner(config, Handler(AtomicIsize::new(5)), Duration::from_secs(1))
                   .unwrap();
    let conn = pool.get().unwrap();
    thread::sleep(Duration::from_secs(2));
    assert_eq!(4, DROPPED.load(Ordering::SeqCst));
    drop(conn);
    assert_eq!(4, DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_max_lifetime() {
    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
    }

    struct Handler(AtomicIsize);

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = Error;

        fn connect(&self) -> Result<Connection, Error> {
            if self.0.fetch_sub(1, Ordering::SeqCst) > 0 {
                Ok(Connection)
            } else {
                Err(Error)
            }
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), Error> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                     .pool_size(5)
                     .max_lifetime(Some(Duration::from_secs(1)))
                     .connection_timeout(Duration::from_secs(1))
                     .build();
    let pool = Pool::new_inner(config, Handler(AtomicIsize::new(5)), Duration::from_secs(1))
                   .unwrap();
    let conn = pool.get().unwrap();
    thread::sleep(Duration::from_secs(2));
    assert_eq!(4, DROPPED.load(Ordering::SeqCst));
    drop(conn);
    thread::sleep(Duration::from_secs(2));
    assert_eq!(5, DROPPED.load(Ordering::SeqCst));
    assert!(pool.get().is_err());
}

#[test]
fn min_idle() {
    static CREATED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = Error;

        fn connect(&self) -> Result<Connection, Error> {
            CREATED.fetch_add(1, Ordering::SeqCst);
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), Error> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                     .pool_size(5)
                     .min_idle(Some(2))
                     .build();
    let pool = Pool::new(config, Handler).unwrap();
    assert_eq!(2, CREATED.load(Ordering::SeqCst));
    let _conns = (0..5).map(|_| pool.get().unwrap()).collect::<Vec<_>>();
    assert_eq!(5, CREATED.load(Ordering::SeqCst));
}

#[test]
fn conns_drop_on_pool_drop() {
    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = Error;

        fn connect(&self) -> Result<Connection, Error> {
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection, _: Duration) -> Result<(), Error> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                    .max_lifetime(Some(Duration::from_secs(10)))
                    .pool_size(10)
                    .build();
    let pool = Pool::new(config, Handler).unwrap();
    drop(pool);
    for _ in 0..10 {
        if DROPPED.load(Ordering::SeqCst) == 10 {
            return;
        }
        thread::sleep(Duration::from_secs(1));
    }
    panic!("timed out waiting for connections to drop");
}
<commit_msg>Fix build<commit_after>use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, AtomicUsize, ATOMIC_USIZE_INIT, AtomicIsize,
                        Ordering};
use std::sync::mpsc::{self, SyncSender, Receiver};
use std::sync::{Mutex, Arc};
use std::time::Duration;
use std::thread;
use std::fmt;
use std::error;

use {ManageConnection, CustomizeConnection, Pool, Config};

#[derive(Debug)]
pub struct Error;

impl fmt::Display for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str("blammo")
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        "Error"
    }
}

#[derive(Debug, PartialEq)]
struct FakeConnection(bool);

struct OkManager;

impl ManageConnection for OkManager {
    type Connection = FakeConnection;
    type Error = Error;

    fn connect(&self) -> Result<FakeConnection, Error> {
        Ok(FakeConnection(true))
    }

    fn is_valid(&self, _: &mut FakeConnection) -> Result<(), Error> {
        Ok(())
    }

    fn has_broken(&self, _: &mut FakeConnection) -> bool {
        false
    }
}

struct NthConnectFailManager {
    n: Mutex<u32>,
}

impl ManageConnection for NthConnectFailManager {
    type Connection = FakeConnection;
    type Error = Error;

    fn connect(&self) -> Result<FakeConnection, Error> {
        let mut n = self.n.lock().unwrap();
        if *n > 0 {
            *n -= 1;
            Ok(FakeConnection(true))
        } else {
            Err(Error)
        }
    }

    fn is_valid(&self, _: &mut FakeConnection) -> Result<(), Error> {
        Ok(())
    }

    fn has_broken(&self, _: &mut FakeConnection) -> bool {
        false
    }
}

#[test]
fn test_pool_size_ok() {
    let config = Config::builder().pool_size(5).build();
    let manager = NthConnectFailManager { n: Mutex::new(5) };
    let pool = Pool::new(config, manager).unwrap();
    let mut conns = vec![];
    for _ in 0..5 {
        conns.push(pool.get().ok().unwrap());
    }
}

#[test]
fn test_acquire_release() {
    let config = Config::builder().pool_size(2).build();
    let pool = Pool::new(config, OkManager).unwrap();

    let conn1 = pool.get().ok().unwrap();
    let conn2 = pool.get().ok().unwrap();
    drop(conn1);
    let conn3 = pool.get().ok().unwrap();
    drop(conn2);
    drop(conn3);
}

#[test]
fn test_is_send_sync() {
    fn is_send_sync<T: Send + Sync>() {}
    is_send_sync::<Pool<OkManager>>();
}

#[test]
fn test_issue_2_unlocked_during_is_valid() {
    struct BlockingChecker {
        first: AtomicBool,
        s: Mutex<SyncSender<()>>,
        r: Mutex<Receiver<()>>,
    }

    impl ManageConnection for BlockingChecker {
        type Connection = FakeConnection;
        type Error = Error;

        fn connect(&self) -> Result<FakeConnection, Error> {
            Ok(FakeConnection(true))
        }

        fn is_valid(&self, _: &mut FakeConnection) -> Result<(), Error> {
            if self.first.compare_and_swap(true, false, Ordering::SeqCst) {
                self.s.lock().unwrap().send(()).unwrap();
                self.r.lock().unwrap().recv().unwrap();
            }
            Ok(())
        }

        fn has_broken(&self, _: &mut FakeConnection) -> bool {
            false
        }
    }

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

    let config = Config::builder()
                     .test_on_check_out(true)
                     .pool_size(2)
                     .build();
    let manager = BlockingChecker {
        first: AtomicBool::new(true),
        s: Mutex::new(s1),
        r: Mutex::new(r2),
    };
    let pool = Arc::new(Pool::new(config, manager).unwrap());

    let p2 = pool.clone();
    let t = thread::spawn(move || {
        p2.get().ok().unwrap();
    });

    r1.recv().unwrap();
    // get call by other task has triggered the health check
    pool.get().ok().unwrap();
    s2.send(()).ok().unwrap();

    t.join().ok().unwrap();
}

#[test]
fn test_drop_on_broken() {
    static DROPPED: AtomicBool = ATOMIC_BOOL_INIT;
    DROPPED.store(false, Ordering::SeqCst);

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.store(true, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = Error;

        fn connect(&self) -> Result<Connection, Error> {
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), Error> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            true
        }
    }

    let pool = Pool::new(Default::default(), Handler).unwrap();

    drop(pool.get().ok().unwrap());

    assert!(DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_initialization_failure() {
    let config = Config::builder()
                     .connection_timeout(Duration::from_secs(1))
                     .build();
    let manager = NthConnectFailManager { n: Mutex::new(0) };
    let err = Pool::new(config, manager).err().unwrap();
    assert!(err.to_string().contains("blammo"));
}

#[test]
fn test_lazy_initialization_failure() {
    let config = Config::builder()
                     .connection_timeout(Duration::from_secs(1))
                     .initialization_fail_fast(false)
                     .build();
    let manager = NthConnectFailManager { n: Mutex::new(0) };
    let pool = Pool::new(config, manager).unwrap();
    let err = pool.get().err().unwrap();
    assert!(err.to_string().contains("blammo"));
}

#[test]
fn test_get_timeout() {
    let config = Config::builder()
                     .pool_size(1)
                     .connection_timeout(Duration::from_secs(1))
                     .build();
    let pool = Pool::new(config, OkManager).unwrap();
    let _c = pool.get().unwrap();
    pool.get().err().unwrap();
}

#[test]
fn test_connection_customizer() {
    static DROPPED: AtomicBool = ATOMIC_BOOL_INIT;
    DROPPED.store(false, Ordering::SeqCst);

    struct Connection(i32);

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.store(true, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = Error;

        fn connect(&self) -> Result<Connection, Error> {
            Ok(Connection(0))
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), Error> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            true
        }
    }

    #[derive(Debug)]
    struct Customizer;

    impl CustomizeConnection<Connection, Error> for Customizer {
        fn on_acquire(&self, conn: &mut Connection) -> Result<(), Error> {
            if !DROPPED.load(Ordering::SeqCst) {
                Err(Error)
            } else {
                conn.0 = 1;
                Ok(())
            }
        }
    }

    let config = Config::builder()
                     .connection_customizer(Box::new(Customizer))
                     .build();
    let pool = Pool::new(config, Handler).unwrap();

    let conn = pool.get().unwrap();
    assert_eq!(1, conn.0);
    assert!(DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_idle_timeout() {
    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
    }

    struct Handler(AtomicIsize);

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = Error;

        fn connect(&self) -> Result<Connection, Error> {
            if self.0.fetch_sub(1, Ordering::SeqCst) > 0 {
                Ok(Connection)
            } else {
                Err(Error)
            }
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), Error> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                     .pool_size(5)
                     .idle_timeout(Some(Duration::from_secs(1)))
                     .build();
    let pool = Pool::new_inner(config, Handler(AtomicIsize::new(5)), Duration::from_secs(1))
                   .unwrap();
    let conn = pool.get().unwrap();
    thread::sleep(Duration::from_secs(2));
    assert_eq!(4, DROPPED.load(Ordering::SeqCst));
    drop(conn);
    assert_eq!(4, DROPPED.load(Ordering::SeqCst));
}

#[test]
fn test_max_lifetime() {
    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
    }

    struct Handler(AtomicIsize);

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = Error;

        fn connect(&self) -> Result<Connection, Error> {
            if self.0.fetch_sub(1, Ordering::SeqCst) > 0 {
                Ok(Connection)
            } else {
                Err(Error)
            }
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), Error> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                     .pool_size(5)
                     .max_lifetime(Some(Duration::from_secs(1)))
                     .connection_timeout(Duration::from_secs(1))
                     .build();
    let pool = Pool::new_inner(config, Handler(AtomicIsize::new(5)), Duration::from_secs(1))
                   .unwrap();
    let conn = pool.get().unwrap();
    thread::sleep(Duration::from_secs(2));
    assert_eq!(4, DROPPED.load(Ordering::SeqCst));
    drop(conn);
    thread::sleep(Duration::from_secs(2));
    assert_eq!(5, DROPPED.load(Ordering::SeqCst));
    assert!(pool.get().is_err());
}

#[test]
fn min_idle() {
    static CREATED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = Error;

        fn connect(&self) -> Result<Connection, Error> {
            CREATED.fetch_add(1, Ordering::SeqCst);
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), Error> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                     .pool_size(5)
                     .min_idle(Some(2))
                     .build();
    let pool = Pool::new(config, Handler).unwrap();
    assert_eq!(2, CREATED.load(Ordering::SeqCst));
    let _conns = (0..5).map(|_| pool.get().unwrap()).collect::<Vec<_>>();
    assert_eq!(5, CREATED.load(Ordering::SeqCst));
}

#[test]
fn conns_drop_on_pool_drop() {
    static DROPPED: AtomicUsize = ATOMIC_USIZE_INIT;

    struct Connection;

    impl Drop for Connection {
        fn drop(&mut self) {
            DROPPED.fetch_add(1, Ordering::SeqCst);
        }
    }

    struct Handler;

    impl ManageConnection for Handler {
        type Connection = Connection;
        type Error = Error;

        fn connect(&self) -> Result<Connection, Error> {
            Ok(Connection)
        }

        fn is_valid(&self, _: &mut Connection) -> Result<(), Error> {
            Ok(())
        }

        fn has_broken(&self, _: &mut Connection) -> bool {
            false
        }
    }

    let config = Config::builder()
                    .max_lifetime(Some(Duration::from_secs(10)))
                    .pool_size(10)
                    .build();
    let pool = Pool::new(config, Handler).unwrap();
    drop(pool);
    for _ in 0..10 {
        if DROPPED.load(Ordering::SeqCst) == 10 {
            return;
        }
        thread::sleep(Duration::from_secs(1));
    }
    panic!("timed out waiting for connections to drop");
}
<|endoftext|>


<commit_before>//! Test support module; public to permit use from integration tests.
use std::collections::HashMap;
use std::env;
use std::ffi::OsStr;
use std::fmt;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use std::process::Command;

#[cfg(test)]
use anyhow::Result;

pub use crate::cli::self_update::test::{get_path, with_saved_path};
use crate::currentprocess;
use crate::dist::dist::TargetTriple;

// Things that can have environment variables applied to them.
pub trait Env {
    fn env<K, V>(&mut self, key: K, val: V)
    where
        K: AsRef<OsStr>,
        V: AsRef<OsStr>;
}

impl Env for Command {
    fn env<K, V>(&mut self, key: K, val: V)
    where
        K: AsRef<OsStr>,
        V: AsRef<OsStr>,
    {
        self.env(key, val);
    }
}

impl Env for HashMap<String, String> {
    fn env<K, V>(&mut self, key: K, val: V)
    where
        K: AsRef<OsStr>,
        V: AsRef<OsStr>,
    {
        let key = key.as_ref().to_os_string().into_string().unwrap();
        let val = val.as_ref().to_os_string().into_string().unwrap();
        self.insert(key, val);
    }
}

/// Returns a tempdir for running tests in
pub fn test_dir() -> io::Result<tempfile::TempDir> {
    let current_exe_path = env::current_exe().unwrap();
    let mut exe_dir = current_exe_path.parent().unwrap();
    if exe_dir.ends_with("deps") {
        exe_dir = exe_dir.parent().unwrap();
    }
    let test_dir = exe_dir.parent().unwrap().join("tests");
    fs::create_dir_all(&test_dir).unwrap();
    tempfile::Builder::new()
        .prefix("running-test-")
        .tempdir_in(test_dir)
}

/// Makes persistent unique directory inside path.
///
/// Should only be used with path=a tempdir that will be cleaned up, as the
/// directory tempdir_in_with_prefix creates won't be automatically cleaned up.
fn tempdir_in_with_prefix<P: AsRef<Path>>(path: P, prefix: &str) -> io::Result<PathBuf> {
    Ok(tempfile::Builder::new()
        .prefix(prefix)
        .tempdir_in(path.as_ref())?
        .into_path())
}

/// What is this host's triple - seems very redundant with from_host_or_build()
/// ... perhaps this is so that the test data we have is only exercised on known
/// triples?
///
/// NOTE: This *cannot* be called within a currentprocess context as it creates
/// its own context on Windows hosts. This is partly by chance but also partly
/// deliberate: If you need the host triple, or to call for_host(), you can do
/// so outside of calls to run() or unit test code that runs in a currentprocess
/// context.
///
/// IF it becomes very hard to workaround that, then we can either make a second
/// this_host_triple that doesn't make its own currentprocess or use
/// TargetTriple::from_host() from within the currentprocess context as needed.
pub fn this_host_triple() -> String {
    if cfg!(target_os = "windows") {
        // For windows, this host may be different to the target: we may be
        // building with i686 toolchain, but on an x86_64 host, so run the
        // actual detection logic and trust it.
        let tp = Box::new(currentprocess::TestProcess::default());
        return currentprocess::with(tp, || TargetTriple::from_host().unwrap().to_string());
    }
    let arch = if cfg!(target_arch = "x86") {
        "i686"
    } else if cfg!(target_arch = "x86_64") {
        "x86_64"
    } else if cfg!(target_arch = "riscv64") {
        "riscv64gc"
    } else if cfg!(target_arch = "aarch64") {
        "aarch64"
    } else {
        unimplemented!()
    };
    let os = if cfg!(target_os = "linux") {
        "unknown-linux"
    } else if cfg!(target_os = "macos") {
        "apple-darwin"
    } else if cfg!(target_os = "illumos") {
        "unknown-illumos"
    } else {
        unimplemented!()
    };
    let env = if cfg!(target_env = "gnu") {
        Some("gnu")
    } else {
        None
    };

    if let Some(env) = env {
        format!("{}-{}-{}", arch, os, env)
    } else {
        format!("{}-{}", arch, os)
    }
}

// Format a string with this host triple.
#[macro_export]
macro_rules! for_host {
    ($s: expr) => {
        &format!($s, $crate::test::this_host_triple())
    };
}

#[derive(Clone)]
/// The smallest form of test isolation: an isolated RUSTUP_HOME, for codepaths
/// that read and write config files but do not invoke processes, download data
/// etc.
pub struct RustupHome {
    pub rustupdir: PathBuf,
}

impl RustupHome {
    pub fn apply<E: Env>(&self, e: &mut E) {
        e.env("RUSTUP_HOME", self.rustupdir.to_string_lossy().to_string())
    }

    pub fn has<P: AsRef<Path>>(&self, path: P) -> bool {
        self.rustupdir.join(path).exists()
    }

    pub fn join<P: AsRef<Path>>(&self, path: P) -> PathBuf {
        self.rustupdir.join(path)
    }

    pub fn new_in<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        let rustupdir = tempdir_in_with_prefix(path, "rustup")?;
        Ok(RustupHome { rustupdir })
    }

    pub fn remove(&self) -> io::Result<()> {
        remove_dir_all::remove_dir_all(&self.rustupdir)
    }
}

impl fmt::Display for RustupHome {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.rustupdir.display())
    }
}

/// Create an isolated rustup home with no content, then call f with it, and
/// delete it afterwards.
#[cfg(test)]
pub(crate) fn with_rustup_home<F>(f: F) -> Result<()>
where
    F: FnOnce(&RustupHome) -> Result<()>,
{
    let test_dir = test_dir()?;
    let rustup_home = RustupHome::new_in(test_dir)?;
    f(&rustup_home)
}
<commit_msg>Fix host triple test for freebsd<commit_after>//! Test support module; public to permit use from integration tests.
use std::collections::HashMap;
use std::env;
use std::ffi::OsStr;
use std::fmt;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use std::process::Command;

#[cfg(test)]
use anyhow::Result;

pub use crate::cli::self_update::test::{get_path, with_saved_path};
use crate::currentprocess;
use crate::dist::dist::TargetTriple;

// Things that can have environment variables applied to them.
pub trait Env {
    fn env<K, V>(&mut self, key: K, val: V)
    where
        K: AsRef<OsStr>,
        V: AsRef<OsStr>;
}

impl Env for Command {
    fn env<K, V>(&mut self, key: K, val: V)
    where
        K: AsRef<OsStr>,
        V: AsRef<OsStr>,
    {
        self.env(key, val);
    }
}

impl Env for HashMap<String, String> {
    fn env<K, V>(&mut self, key: K, val: V)
    where
        K: AsRef<OsStr>,
        V: AsRef<OsStr>,
    {
        let key = key.as_ref().to_os_string().into_string().unwrap();
        let val = val.as_ref().to_os_string().into_string().unwrap();
        self.insert(key, val);
    }
}

/// Returns a tempdir for running tests in
pub fn test_dir() -> io::Result<tempfile::TempDir> {
    let current_exe_path = env::current_exe().unwrap();
    let mut exe_dir = current_exe_path.parent().unwrap();
    if exe_dir.ends_with("deps") {
        exe_dir = exe_dir.parent().unwrap();
    }
    let test_dir = exe_dir.parent().unwrap().join("tests");
    fs::create_dir_all(&test_dir).unwrap();
    tempfile::Builder::new()
        .prefix("running-test-")
        .tempdir_in(test_dir)
}

/// Makes persistent unique directory inside path.
///
/// Should only be used with path=a tempdir that will be cleaned up, as the
/// directory tempdir_in_with_prefix creates won't be automatically cleaned up.
fn tempdir_in_with_prefix<P: AsRef<Path>>(path: P, prefix: &str) -> io::Result<PathBuf> {
    Ok(tempfile::Builder::new()
        .prefix(prefix)
        .tempdir_in(path.as_ref())?
        .into_path())
}

/// What is this host's triple - seems very redundant with from_host_or_build()
/// ... perhaps this is so that the test data we have is only exercised on known
/// triples?
///
/// NOTE: This *cannot* be called within a currentprocess context as it creates
/// its own context on Windows hosts. This is partly by chance but also partly
/// deliberate: If you need the host triple, or to call for_host(), you can do
/// so outside of calls to run() or unit test code that runs in a currentprocess
/// context.
///
/// IF it becomes very hard to workaround that, then we can either make a second
/// this_host_triple that doesn't make its own currentprocess or use
/// TargetTriple::from_host() from within the currentprocess context as needed.
pub fn this_host_triple() -> String {
    if cfg!(target_os = "windows") {
        // For windows, this host may be different to the target: we may be
        // building with i686 toolchain, but on an x86_64 host, so run the
        // actual detection logic and trust it.
        let tp = Box::new(currentprocess::TestProcess::default());
        return currentprocess::with(tp, || TargetTriple::from_host().unwrap().to_string());
    }
    let arch = if cfg!(target_arch = "x86") {
        "i686"
    } else if cfg!(target_arch = "x86_64") {
        "x86_64"
    } else if cfg!(target_arch = "riscv64") {
        "riscv64gc"
    } else if cfg!(target_arch = "aarch64") {
        "aarch64"
    } else {
        unimplemented!()
    };
    let os = if cfg!(target_os = "linux") {
        "unknown-linux"
    } else if cfg!(target_os = "macos") {
        "apple-darwin"
    } else if cfg!(target_os = "illumos") {
        "unknown-illumos"
    } else if cfg!(target_os = "freebsd") {
        "unknown-freebsd"
    } else {
        unimplemented!()
    };
    let env = if cfg!(target_env = "gnu") {
        Some("gnu")
    } else {
        None
    };

    if let Some(env) = env {
        format!("{}-{}-{}", arch, os, env)
    } else {
        format!("{}-{}", arch, os)
    }
}

// Format a string with this host triple.
#[macro_export]
macro_rules! for_host {
    ($s: expr) => {
        &format!($s, $crate::test::this_host_triple())
    };
}

#[derive(Clone)]
/// The smallest form of test isolation: an isolated RUSTUP_HOME, for codepaths
/// that read and write config files but do not invoke processes, download data
/// etc.
pub struct RustupHome {
    pub rustupdir: PathBuf,
}

impl RustupHome {
    pub fn apply<E: Env>(&self, e: &mut E) {
        e.env("RUSTUP_HOME", self.rustupdir.to_string_lossy().to_string())
    }

    pub fn has<P: AsRef<Path>>(&self, path: P) -> bool {
        self.rustupdir.join(path).exists()
    }

    pub fn join<P: AsRef<Path>>(&self, path: P) -> PathBuf {
        self.rustupdir.join(path)
    }

    pub fn new_in<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        let rustupdir = tempdir_in_with_prefix(path, "rustup")?;
        Ok(RustupHome { rustupdir })
    }

    pub fn remove(&self) -> io::Result<()> {
        remove_dir_all::remove_dir_all(&self.rustupdir)
    }
}

impl fmt::Display for RustupHome {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.rustupdir.display())
    }
}

/// Create an isolated rustup home with no content, then call f with it, and
/// delete it afterwards.
#[cfg(test)]
pub(crate) fn with_rustup_home<F>(f: F) -> Result<()>
where
    F: FnOnce(&RustupHome) -> Result<()>,
{
    let test_dir = test_dir()?;
    let rustup_home = RustupHome::new_in(test_dir)?;
    f(&rustup_home)
}
<|endoftext|>


<commit_before>use std::slice;
use std::ffi::{CString,OsStr};
use std::path::Path;

use ::libc::{c_int,c_char};

use std::os::unix::prelude::*;

pub fn ptr_to_path<'a>(ptr: *const c_char) -> Option<&'a Path> {
    match ptr_to_os_str(ptr) {
        Some(path) => Some(Path::new(path)),
        None => None,
    }
}

pub fn ptr_to_os_str<'a>(ptr: *const c_char) -> Option<&'a OsStr> {
    if !ptr.is_null() {
        Some(unsafe { ptr_to_os_str_unchecked(ptr) })
    }
    else {
        None
    }
}

pub unsafe fn ptr_to_os_str_unchecked<'a>(ptr: *const c_char) -> &'a OsStr {
    OsStr::from_bytes(slice::from_raw_parts(ptr as *const u8, ::libc::strlen(ptr) as usize))
}

pub fn os_str_to_cstring<T: AsRef<OsStr>>(s: T) -> ::Result<CString> {
    match CString::new(s.as_ref().as_bytes()) {
        Ok(s) => Ok(s),
        Err(_) => return Err(::error::from_errno(::libc::EINVAL))
    }
}

pub fn errno_to_result(errno: c_int) -> ::Result<()> {
    match errno {
        0 => Ok(()),
        e => Err(::error::from_errno(e))
    }
}
<commit_msg>Inline util functions<commit_after>use std::slice;
use std::ffi::{CString,OsStr};
use std::path::Path;

use ::libc::{c_int,c_char};

use std::os::unix::prelude::*;

#[inline(always)]
pub fn ptr_to_path<'a>(ptr: *const c_char) -> Option<&'a Path> {
    match ptr_to_os_str(ptr) {
        Some(path) => Some(Path::new(path)),
        None => None,
    }
}

#[inline(always)]
pub fn ptr_to_os_str<'a>(ptr: *const c_char) -> Option<&'a OsStr> {
    if !ptr.is_null() {
        Some(unsafe { ptr_to_os_str_unchecked(ptr) })
    }
    else {
        None
    }
}

#[inline(always)]
pub unsafe fn ptr_to_os_str_unchecked<'a>(ptr: *const c_char) -> &'a OsStr {
    OsStr::from_bytes(slice::from_raw_parts(ptr as *const u8, ::libc::strlen(ptr) as usize))
}

#[inline(always)]
pub fn os_str_to_cstring<T: AsRef<OsStr>>(s: T) -> ::Result<CString> {
    match CString::new(s.as_ref().as_bytes()) {
        Ok(s) => Ok(s),
        Err(_) => return Err(::error::from_errno(::libc::EINVAL))
    }
}

#[inline(always)]
pub fn errno_to_result(errno: c_int) -> ::Result<()> {
    match errno {
        0 => Ok(()),
        e => Err(::error::from_errno(e))
    }
}
<|endoftext|>


<commit_before>use fastrand;
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};
use std::{io, iter::repeat_with};

use crate::error::IoResultExt;

fn tmpname(prefix: &OsStr, suffix: &OsStr, rand_len: usize) -> OsString {
    let mut buf = OsString::with_capacity(prefix.len() + suffix.len() + rand_len);
    buf.push(prefix);
    buf.push(
        repeat_with(fastrand::alphanumeric)
            .take(rand_len)
            .collect::<String>(),
    );
    buf.push(suffix);
    buf
}

pub fn create_helper<F, R>(
    base: &Path,
    prefix: &OsStr,
    suffix: &OsStr,
    random_len: usize,
    f: F,
) -> io::Result<R>
where
    F: Fn(PathBuf) -> io::Result<R>,
{
    let num_retries = if random_len != 0 {
        crate::NUM_RETRIES
    } else {
        1
    };

    for _ in 0..num_retries {
        let path = base.join(tmpname(prefix, suffix, random_len));
        return match f(path) {
            Err(ref e) if e.kind() == io::ErrorKind::AlreadyExists => continue,
            res => res,
        };
    }

    Err(io::Error::new(
        io::ErrorKind::AlreadyExists,
        "too many temporary files exist",
    ))
    .with_err_path(|| base)
}
<commit_msg>chore: avoid allocation when generating temp names<commit_after>use fastrand;
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};
use std::{io, iter::repeat_with};

use crate::error::IoResultExt;

fn tmpname(prefix: &OsStr, suffix: &OsStr, rand_len: usize) -> OsString {
    let mut buf = OsString::with_capacity(prefix.len() + suffix.len() + rand_len);
    buf.push(prefix);
    let mut char_buf = [0u8; 4];
    for c in repeat_with(fastrand::alphanumeric).take(rand_len) {
        buf.push(c.encode_utf8(&mut char_buf));
    }
    buf.push(suffix);
    buf
}

pub fn create_helper<F, R>(
    base: &Path,
    prefix: &OsStr,
    suffix: &OsStr,
    random_len: usize,
    f: F,
) -> io::Result<R>
where
    F: Fn(PathBuf) -> io::Result<R>,
{
    let num_retries = if random_len != 0 {
        crate::NUM_RETRIES
    } else {
        1
    };

    for _ in 0..num_retries {
        let path = base.join(tmpname(prefix, suffix, random_len));
        return match f(path) {
            Err(ref e) if e.kind() == io::ErrorKind::AlreadyExists => continue,
            res => res,
        };
    }

    Err(io::Error::new(
        io::ErrorKind::AlreadyExists,
        "too many temporary files exist",
    ))
    .with_err_path(|| base)
}
<|endoftext|>


<commit_before>use std::string::String;
use std::fs::File;
use std::io::{BufReader, Read};
use std::collections::HashMap;
use std::ptr;

use zip::read::{ZipFile, ZipArchive};
use zip::result::ZipError;
use byteorder::{LittleEndian, ReadBytesExt};
use quick_xml::{XmlReader, Event, AsStr};
use encoding::{Encoding, DecoderTrap};
use encoding::all::UTF_16LE;

use {DataType, ExcelReader, Range, CellErrorType};
use vba::VbaProject;
use utils;
use errors::*;

pub struct Xlsb {
    zip: ZipArchive<File>,
}

impl Xlsb {
    fn iter<'a>(&'a mut self, path: &str) -> Result<RecordIter<'a>> {
        match self.zip.by_name(path) {
            Ok(f) => Ok(RecordIter { r: BufReader::new(f) }),
            Err(ZipError::FileNotFound) => Err(format!("file {} does not exist", path).into()),
            Err(e) => Err(e.into()),
        }
    }
}

impl ExcelReader for Xlsb {

    fn new(f: File) -> Result<Self> {
        Ok(Xlsb { zip: try!(ZipArchive::new(f)) })
    }

    fn has_vba(&mut self) -> bool {
        self.zip.by_name("xl/vbaProject.bin").is_ok()
    }

    fn vba_project(&mut self) -> Result<VbaProject> {
        let mut f = try!(self.zip.by_name("xl/vbaProject.bin"));
        let len = f.size() as usize;
        VbaProject::new(&mut f, len)
    }

    /// MS-XLSB
    fn read_relationships(&mut self) -> Result<HashMap<Vec<u8>, String>> {
        let mut relationships = HashMap::new();
        match self.zip.by_name("xl/_rels/workbook.bin.rels") {
            Ok(f) => {
                let xml = XmlReader::from_reader(BufReader::new(f))
                    .with_check(false)
                    .trim_text(false);

                for res_event in xml {
                    match res_event {
                        Ok(Event::Start(ref e)) if e.name() == b"Relationship" => {
                            let mut id = Vec::new();
                            let mut target = String::new();
                            for a in e.attributes() {
                                match try!(a) {
                                    (b"Id", v) => id.extend_from_slice(v),
                                    (b"Target", v) => target = try!(v.as_str()).to_string(),
                                    _ => (),
                                }
                            }
                            relationships.insert(id, target);
                        }
                        Err(e) => return Err(e.into()),
                        _ => (),
                    }
                }
            },
            Err(ZipError::FileNotFound) => (),
            Err(e) => return Err(e.into()),
        }
        Ok(relationships)
    }

    /// MS-XLSB 2.1.7.45
    fn read_shared_strings(&mut self) -> Result<Vec<String>> {
        let mut iter = try!(self.iter("xl/sharedStrings.bin"));
        let mut buf = vec![0; 1024];

        let _ = try!(iter.fill_next(0x009F, &mut buf)); // BrtBeginSst
        let len = try!((&mut &*buf).read_u32::<LittleEndian>()) as usize;
        let mut strings = Vec::with_capacity(len);

        // BrtSSTItems
        for _ in 0..len {
            let _ = try!(iter.next_skip_blocks(0x0013, &[
                                               (0x0023, Some(0x0024)) // future
                                               ], &mut buf)); // BrtSSTItem
            let flags = buf[0];
            if flags & 0b11000000 != 0 {
                // suppose A and B == 0
                return Err("only regular shared strings are supported".into());
            }
            strings.push(try!(wide_str(&buf[1..])));
        }
        Ok(strings)
    }

    /// MS-XLSB 2.1.7.61
    fn read_sheets_names(&mut self, relationships: &HashMap<Vec<u8>, String>)
        -> Result<HashMap<String, String>>
    {
        let mut iter = try!(self.iter("xl/workbook.bin"));
        let mut buf = vec![0; 1024];

        let _ = try!(iter.fill_next(0x0083, &mut buf)); // BrtBeginBook

        // BrtBeginBundleShs
        let _ = try!(iter.next_skip_blocks(0x008F, &[
                                          (0x0080, None),         // BrtFileVersion
                                          (0x0099, None),         // BrtWbProp
                                          (0x02A4, Some(0x0224)), // File Sharing
                                          (0x0025, Some(0x0026)), // AC blocks
                                          (0x02A5, Some(0x0216)), // Book protection(iso)
                                          (0x0087, Some(0x0088)), // BOOKVIEWS
                                          ], &mut buf));
        let mut sheets = HashMap::new();
        loop {
            match try!(iter.read_type()) {
                0x0090 => return Ok(sheets), // BrtEndBundleShs
                0x009C => (), // BrtBundleSh
                typ => return Err(format!("Expecting end of sheet, got {:x}", typ).into()),
            }
            let len = try!(iter.fill_buffer(&mut buf));
            let rel_len = utils::start_u32(&buf[8..len]);
            if rel_len != 0xFFFFFFFF {
                let rel_len = rel_len as usize * 2;
                let relid = &buf[12..12 + rel_len];
                // converts utf16le to utf8 for HashMap search
                let relid = try!(UTF_16LE.decode(relid, DecoderTrap::Ignore).map_err(|e| e.to_string()));
                let path = format!("xl/{}", relationships[relid.as_bytes()]);
                let name = try!(wide_str(&buf[12 + rel_len..len]));
                sheets.insert(name, path);
            }
        }
    }

    /// MS-XLSB 2.1.7.62
    fn read_worksheet_range(&mut self, path: &str, strings: &[String]) -> Result<Range> {
        let mut iter = try!(self.iter(path));
        let mut buf = vec![0; 1024];

        let _ = try!(iter.fill_next(0x0081, &mut buf)); // BrtBeginSheet

        // BrtWsDim, skipping BrtWsProp
        let _ = try!(iter.next_skip_blocks(0x0094, &[(0x0093, None)], &mut buf));
        let (position, size) = unchecked_rfx(&buf[..16]);

        if size.0 == 0 || size.1 == 0 {
            return Ok(Range::default());
        }

        // BrtBeginSheetData
        let _ = try!(iter.next_skip_blocks(0x0091, &[
                                          (0x0085, Some(0x0086)), // Views
                                          (0x0025, Some(0x0026)), // AC blocks
                                          (0x01E5, None),         // BrtWsFmtInfo
                                          (0x0186, Some(0x0187)), // Col Infos
                                          ], &mut buf));

        let mut data = vec![DataType::Empty; (size.0 * size.1) as usize];

        // loop through all non empty rows
        loop {
            // BrtRowHdr
            let _ = try!(iter.next_skip_blocks(0x0000, &[
                                              (0x0025, Some(0x0026)), // AC blocks
                                              (0x0023, Some(0x0024)), // future
                                              ], &mut buf));
            let row = utils::start_u32(&buf);

            // get column indexes
            let cols = {
                let span_len = utils::start_u32(&buf[13..]) as usize;
                let mut span_iter = utils::to_u32(&buf[17..]).take(span_len * 2);
                let mut cols = Vec::with_capacity(size.1);
                for _ in 0..span_len {
                    cols.extend(span_iter.next().unwrap()..(span_iter.next().unwrap() + 1));
                }
                cols
            };

            // read all values for the row
            for idx in cols.iter()
                .map(|col| (row - position.0) as usize * size.1 + (*col - position.1) as usize) {
                loop {
                    // read record
                    let typ = try!(iter.read_type());
                    let _ = try!(iter.fill_buffer(&mut buf));

                    match typ {
                        0x0092 => return Err("Expecting cell, got BrtEndSheetData".into()),
                        0x0001 => (), // BrtCellBlank: nothing to do as it is the default value
                        0x0002 => { // BrtCellRk MS-XLSB 2.5.122
                            let d100 = (buf[8] & 1) != 0;
                            let is_int = (buf[8] & 2) != 0;
                            buf[8] &= 0xFC;
                            data[idx] = if is_int {
                                let v = unsafe {
                                    (ptr::read(&buf[8..12] as *const [u8] as *const i32) >> 2) as i64
                                };
                                DataType::Int( if d100 { v/100 } else { v })
                            } else {
                                let mut v = [0u8; 8];
                                v[4..].copy_from_slice(&buf[8..12]);
                                let v = unsafe {
                                    ptr::read(&v as *const [u8] as *const f64)
                                };
                                DataType::Float( if d100 { v/100.0 } else { v })
                            };
                        },
                        0x0003 => { // BrtCellError
                            data[idx] = DataType::Error(match buf[8] {
                                0x00 => CellErrorType::Null,
                                0x07 => CellErrorType::Div0,
                                0x0F => CellErrorType::Value,
                                0x17 => CellErrorType::Ref,
                                0x1D => CellErrorType::Name,
                                0x24 => CellErrorType::Num,
                                0x2A => CellErrorType::NA,
                                0x2B => CellErrorType::GettingData,
                                c => return Err(format!("Unrecognised cell error code 0x{:x}", c).into()),
                            });
                        },
                        0x0004 => { // BrtCellBool
                            data[idx] = DataType::Bool(buf[8] != 0);
                        },
                        0x0005 => { // BrtCellReal
                            let v = unsafe { ptr::read(&buf[8..16] as *const [u8] as *const f64) };
                            data[idx] = DataType::Float(v);
                        },
                        0x0006 => { // BrtCellSt
                            data[idx] = DataType::String(try!(wide_str(&buf[8..])));
                        },
                        0x0007 => { // BrtCellIsst
                            let isst = utils::start_u32(&buf[8..12]) as usize;
                            data[idx] = DataType::String(strings[isst].clone());
                        },
                        _ => continue, // anything else, ignore and try next
                    }

                    // return if the last cell has been set
                    if idx == data.len() - 1 {
                        return Ok(Range::new(position, size, data))
                    } else {
                        break;
                    }
                }
            }
        }
    }

}

struct RecordIter<'a> {
    r: BufReader<ZipFile<'a>>,
}

impl<'a> RecordIter<'a> {

    fn read_type(&mut self) -> Result<u16> {
        let b = try!(self.r.read_u8());
        let typ = if (b & 0x80) == 0x80 {
            (b & 0x7F) as u16 + (((try!(self.r.read_u8()) & 0x7F) as u16)<<7)
        } else {
            b as u16
        };
        Ok(typ)
    }

    fn fill_buffer(&mut self, buf: &mut Vec<u8>) -> Result<usize> {
        let mut b = try!(self.r.read_u8());
        let mut len = b as usize;
        for i in 0..3 {
            if (b & 0x80) == 0 { break; }
            b = try!(self.r.read_u8());
            len += (b as usize) << (7 * i);
        }
        if buf.len() < len { *buf = vec![0; len]; }

        let _ = try!(self.r.read_exact(&mut buf[..len]));
        Ok(len)
    }

    fn fill_next(&mut self, record_type: u16, buf: &mut Vec<u8>) -> Result<usize> {
        let typ = try!(self.read_type());
        if record_type != typ {
            Err(format!("Unexpected record: expecting 0x{:x}, found 0x{:x}",
                        record_type, typ).into())
        } else {
            self.fill_buffer(buf)
        }
    }

    /// Reads next type, and discard block between `start` and `end`
    fn next_skip_blocks(&mut self, record_type: u16, bounds: &[(u16, Option<u16>)],
                       buf: &mut Vec<u8>) -> Result<usize>
    {
        let mut end = None;
        loop {
            let typ = try!(self.read_type());
            match end {
                Some(e) if e == typ => end = None,
                Some(_) => (),
                None if typ == record_type => {
                    return self.fill_buffer(buf);
                },
                None => match bounds.iter().position(|b| b.0 == typ) {
                    Some(i) => {
                        end = bounds[i].1;
                    },
                    None => return Err(format!("Unexpected record after block: \
                        expecting 0x{:x} found 0x{:x}", record_type, typ).into())
                }
            }
            let _ = try!(self.fill_buffer(buf));
        }
    }

}

fn wide_str(buf: &[u8]) -> Result<String> {
    let len = utils::start_u32(buf);
    let s = &buf[4..4 + len as usize * 2];
    UTF_16LE.decode(s, DecoderTrap::Ignore).map_err(|e| e.to_string().into())
}

fn unchecked_rfx(buf: &[u8]) -> ((u32, u32), (usize, usize)) {
    let mut iter = utils::to_u32(&buf[..16]);
    let rw_first = iter.next().unwrap();
    let rw_last = iter.next().unwrap();
    let col_first = iter.next().unwrap();
    let col_last = iter.next().unwrap();

    ((rw_first, col_first),
     ((rw_last - rw_first + 1) as usize, (col_last - col_first + 1) as usize))
}
<commit_msg>xlsb: use unique shared list count<commit_after>use std::string::String;
use std::fs::File;
use std::io::{BufReader, Read};
use std::collections::HashMap;
use std::ptr;

use zip::read::{ZipFile, ZipArchive};
use zip::result::ZipError;
use byteorder::ReadBytesExt;
use quick_xml::{XmlReader, Event, AsStr};
use encoding::{Encoding, DecoderTrap};
use encoding::all::UTF_16LE;

use {DataType, ExcelReader, Range, CellErrorType};
use vba::VbaProject;
use utils;
use errors::*;

pub struct Xlsb {
    zip: ZipArchive<File>,
}

impl Xlsb {
    fn iter<'a>(&'a mut self, path: &str) -> Result<RecordIter<'a>> {
        match self.zip.by_name(path) {
            Ok(f) => Ok(RecordIter { r: BufReader::new(f) }),
            Err(ZipError::FileNotFound) => Err(format!("file {} does not exist", path).into()),
            Err(e) => Err(e.into()),
        }
    }
}

impl ExcelReader for Xlsb {

    fn new(f: File) -> Result<Self> {
        Ok(Xlsb { zip: try!(ZipArchive::new(f)) })
    }

    fn has_vba(&mut self) -> bool {
        self.zip.by_name("xl/vbaProject.bin").is_ok()
    }

    fn vba_project(&mut self) -> Result<VbaProject> {
        let mut f = try!(self.zip.by_name("xl/vbaProject.bin"));
        let len = f.size() as usize;
        VbaProject::new(&mut f, len)
    }

    /// MS-XLSB
    fn read_relationships(&mut self) -> Result<HashMap<Vec<u8>, String>> {
        let mut relationships = HashMap::new();
        match self.zip.by_name("xl/_rels/workbook.bin.rels") {
            Ok(f) => {
                let xml = XmlReader::from_reader(BufReader::new(f))
                    .with_check(false)
                    .trim_text(false);

                for res_event in xml {
                    match res_event {
                        Ok(Event::Start(ref e)) if e.name() == b"Relationship" => {
                            let mut id = Vec::new();
                            let mut target = String::new();
                            for a in e.attributes() {
                                match try!(a) {
                                    (b"Id", v) => id.extend_from_slice(v),
                                    (b"Target", v) => target = try!(v.as_str()).to_string(),
                                    _ => (),
                                }
                            }
                            relationships.insert(id, target);
                        }
                        Err(e) => return Err(e.into()),
                        _ => (),
                    }
                }
            },
            Err(ZipError::FileNotFound) => (),
            Err(e) => return Err(e.into()),
        }
        Ok(relationships)
    }

    /// MS-XLSB 2.1.7.45
    fn read_shared_strings(&mut self) -> Result<Vec<String>> {
        let mut iter = try!(self.iter("xl/sharedStrings.bin"));
        let mut buf = vec![0; 1024];

        let _ = try!(iter.fill_next(0x009F, &mut buf)); // BrtBeginSst
        let len = utils::start_u32(&buf[4..8]) as usize;
        let mut strings = Vec::with_capacity(len);

        // BrtSSTItems
        for _ in 0..len {
            let _ = try!(iter.next_skip_blocks(0x0013, &[
                                               (0x0023, Some(0x0024)) // future
                                               ], &mut buf)); // BrtSSTItem
            let flags = buf[0];
            if flags & 0b11000000 != 0 {
                // suppose A and B == 0
                return Err("only regular shared strings are supported".into());
            }
            strings.push(try!(wide_str(&buf[1..])));
        }
        Ok(strings)
    }

    /// MS-XLSB 2.1.7.61
    fn read_sheets_names(&mut self, relationships: &HashMap<Vec<u8>, String>)
        -> Result<HashMap<String, String>>
    {
        let mut iter = try!(self.iter("xl/workbook.bin"));
        let mut buf = vec![0; 1024];

        let _ = try!(iter.fill_next(0x0083, &mut buf)); // BrtBeginBook

        // BrtBeginBundleShs
        let _ = try!(iter.next_skip_blocks(0x008F, &[
                                          (0x0080, None),         // BrtFileVersion
                                          (0x0099, None),         // BrtWbProp
                                          (0x02A4, Some(0x0224)), // File Sharing
                                          (0x0025, Some(0x0026)), // AC blocks
                                          (0x02A5, Some(0x0216)), // Book protection(iso)
                                          (0x0087, Some(0x0088)), // BOOKVIEWS
                                          ], &mut buf));
        let mut sheets = HashMap::new();
        loop {
            match try!(iter.read_type()) {
                0x0090 => return Ok(sheets), // BrtEndBundleShs
                0x009C => (), // BrtBundleSh
                typ => return Err(format!("Expecting end of sheet, got {:x}", typ).into()),
            }
            let len = try!(iter.fill_buffer(&mut buf));
            let rel_len = utils::start_u32(&buf[8..len]);
            if rel_len != 0xFFFFFFFF {
                let rel_len = rel_len as usize * 2;
                let relid = &buf[12..12 + rel_len];
                // converts utf16le to utf8 for HashMap search
                let relid = try!(UTF_16LE.decode(relid, DecoderTrap::Ignore).map_err(|e| e.to_string()));
                let path = format!("xl/{}", relationships[relid.as_bytes()]);
                let name = try!(wide_str(&buf[12 + rel_len..len]));
                sheets.insert(name, path);
            }
        }
    }

    /// MS-XLSB 2.1.7.62
    fn read_worksheet_range(&mut self, path: &str, strings: &[String]) -> Result<Range> {
        let mut iter = try!(self.iter(path));
        let mut buf = vec![0; 1024];

        let _ = try!(iter.fill_next(0x0081, &mut buf)); // BrtBeginSheet

        // BrtWsDim, skipping BrtWsProp
        let _ = try!(iter.next_skip_blocks(0x0094, &[(0x0093, None)], &mut buf));
        let (position, size) = unchecked_rfx(&buf[..16]);

        if size.0 == 0 || size.1 == 0 {
            return Ok(Range::default());
        }

        // BrtBeginSheetData
        let _ = try!(iter.next_skip_blocks(0x0091, &[
                                          (0x0085, Some(0x0086)), // Views
                                          (0x0025, Some(0x0026)), // AC blocks
                                          (0x01E5, None),         // BrtWsFmtInfo
                                          (0x0186, Some(0x0187)), // Col Infos
                                          ], &mut buf));

        let mut data = vec![DataType::Empty; (size.0 * size.1) as usize];

        // loop through all non empty rows
        loop {
            // BrtRowHdr
            let _ = try!(iter.next_skip_blocks(0x0000, &[
                                              (0x0025, Some(0x0026)), // AC blocks
                                              (0x0023, Some(0x0024)), // future
                                              ], &mut buf));
            let row = utils::start_u32(&buf);

            // get column indexes
            let cols = {
                let span_len = utils::start_u32(&buf[13..]) as usize;
                let mut span_iter = utils::to_u32(&buf[17..]).take(span_len * 2);
                let mut cols = Vec::with_capacity(size.1);
                for _ in 0..span_len {
                    cols.extend(span_iter.next().unwrap()..(span_iter.next().unwrap() + 1));
                }
                cols
            };

            // read all values for the row
            for idx in cols.iter()
                .map(|col| (row - position.0) as usize * size.1 + (*col - position.1) as usize) {
                loop {
                    // read record
                    let typ = try!(iter.read_type());
                    let _ = try!(iter.fill_buffer(&mut buf));

                    match typ {
                        0x0092 => return Err("Expecting cell, got BrtEndSheetData".into()),
                        0x0001 => (), // BrtCellBlank: nothing to do as it is the default value
                        0x0002 => { // BrtCellRk MS-XLSB 2.5.122
                            let d100 = (buf[8] & 1) != 0;
                            let is_int = (buf[8] & 2) != 0;
                            buf[8] &= 0xFC;
                            data[idx] = if is_int {
                                let v = unsafe {
                                    (ptr::read(&buf[8..12] as *const [u8] as *const i32) >> 2) as i64
                                };
                                DataType::Int( if d100 { v/100 } else { v })
                            } else {
                                let mut v = [0u8; 8];
                                v[4..].copy_from_slice(&buf[8..12]);
                                let v = unsafe {
                                    ptr::read(&v as *const [u8] as *const f64)
                                };
                                DataType::Float( if d100 { v/100.0 } else { v })
                            };
                        },
                        0x0003 => { // BrtCellError
                            data[idx] = DataType::Error(match buf[8] {
                                0x00 => CellErrorType::Null,
                                0x07 => CellErrorType::Div0,
                                0x0F => CellErrorType::Value,
                                0x17 => CellErrorType::Ref,
                                0x1D => CellErrorType::Name,
                                0x24 => CellErrorType::Num,
                                0x2A => CellErrorType::NA,
                                0x2B => CellErrorType::GettingData,
                                c => return Err(format!("Unrecognised cell error code 0x{:x}", c).into()),
                            });
                        },
                        0x0004 => { // BrtCellBool
                            data[idx] = DataType::Bool(buf[8] != 0);
                        },
                        0x0005 => { // BrtCellReal
                            let v = unsafe { ptr::read(&buf[8..16] as *const [u8] as *const f64) };
                            data[idx] = DataType::Float(v);
                        },
                        0x0006 => { // BrtCellSt
                            data[idx] = DataType::String(try!(wide_str(&buf[8..])));
                        },
                        0x0007 => { // BrtCellIsst
                            let isst = utils::start_u32(&buf[8..12]) as usize;
                            data[idx] = DataType::String(strings[isst].clone());
                        },
                        _ => continue, // anything else, ignore and try next
                    }

                    // return if the last cell has been set
                    if idx == data.len() - 1 {
                        return Ok(Range::new(position, size, data))
                    } else {
                        break;
                    }
                }
            }
        }
    }

}

struct RecordIter<'a> {
    r: BufReader<ZipFile<'a>>,
}

impl<'a> RecordIter<'a> {

    fn read_type(&mut self) -> Result<u16> {
        let b = try!(self.r.read_u8());
        let typ = if (b & 0x80) == 0x80 {
            (b & 0x7F) as u16 + (((try!(self.r.read_u8()) & 0x7F) as u16)<<7)
        } else {
            b as u16
        };
        Ok(typ)
    }

    fn fill_buffer(&mut self, buf: &mut Vec<u8>) -> Result<usize> {
        let mut b = try!(self.r.read_u8());
        let mut len = b as usize;
        for i in 0..3 {
            if (b & 0x80) == 0 { break; }
            b = try!(self.r.read_u8());
            len += (b as usize) << (7 * i);
        }
        if buf.len() < len { *buf = vec![0; len]; }

        let _ = try!(self.r.read_exact(&mut buf[..len]));
        Ok(len)
    }

    fn fill_next(&mut self, record_type: u16, buf: &mut Vec<u8>) -> Result<usize> {
        let typ = try!(self.read_type());
        if record_type != typ {
            Err(format!("Unexpected record: expecting 0x{:x}, found 0x{:x}",
                        record_type, typ).into())
        } else {
            self.fill_buffer(buf)
        }
    }

    /// Reads next type, and discard block between `start` and `end`
    fn next_skip_blocks(&mut self, record_type: u16, bounds: &[(u16, Option<u16>)],
                       buf: &mut Vec<u8>) -> Result<usize>
    {
        let mut end = None;
        loop {
            let typ = try!(self.read_type());
            match end {
                Some(e) if e == typ => end = None,
                Some(_) => (),
                None if typ == record_type => {
                    return self.fill_buffer(buf);
                },
                None => match bounds.iter().position(|b| b.0 == typ) {
                    Some(i) => {
                        end = bounds[i].1;
                    },
                    None => return Err(format!("Unexpected record after block: \
                        expecting 0x{:x} found 0x{:x}", record_type, typ).into())
                }
            }
            let _ = try!(self.fill_buffer(buf));
        }
    }

}

fn wide_str(buf: &[u8]) -> Result<String> {
    let len = utils::start_u32(buf);
    let s = &buf[4..4 + len as usize * 2];
    UTF_16LE.decode(s, DecoderTrap::Ignore).map_err(|e| e.to_string().into())
}

fn unchecked_rfx(buf: &[u8]) -> ((u32, u32), (usize, usize)) {
    let mut iter = utils::to_u32(&buf[..16]);
    let rw_first = iter.next().unwrap();
    let rw_last = iter.next().unwrap();
    let col_first = iter.next().unwrap();
    let col_last = iter.next().unwrap();

    ((rw_first, col_first),
     ((rw_last - rw_first + 1) as usize, (col_last - col_first + 1) as usize))
}
<|endoftext|>


<commit_before>// Copyright 2017 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.

use std;
use std::fs::File;
use std::io;
use std::os::unix::io::{AsRawFd, RawFd};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread;

use sys_util::{EventFd, GuestMemory, Poller};

use super::{VirtioDevice, Queue, INTERRUPT_STATUS_USED_RING, TYPE_RNG};

const QUEUE_SIZE: u16 = 256;
const QUEUE_SIZES: &'static [u16] = &[QUEUE_SIZE];

#[derive(Debug)]
pub enum RngError {
    /// Can't access /dev/random
    AccessingRandomDev(io::Error),
}
pub type Result<T> = std::result::Result<T, RngError>;

struct Worker {
    queue: Queue,
    mem: GuestMemory,
    random_file: File,
    interrupt_status: Arc<AtomicUsize>,
    interrupt_evt: EventFd,
}

impl Worker {
    fn process_queue(&mut self) -> bool {
        let queue = &mut self.queue;

        let mut used_desc_heads = [(0, 0); QUEUE_SIZE as usize];
        let mut used_count = 0;
        for avail_desc in queue.iter(&self.mem) {
            let mut len = 0;

            // Drivers can only read from the random device.
            if avail_desc.is_write_only() {
                // Fill the read with data from the random device on the host.
                if self.mem.read_to_memory(avail_desc.addr,
                                           &mut self.random_file,
                                           avail_desc.len as usize)
                        .is_ok() {
                    len = avail_desc.len;
                }
            }

            used_desc_heads[used_count] = (avail_desc.index, len);
            used_count += 1;
        }

        for &(desc_index, len) in &used_desc_heads[..used_count] {
            queue.add_used(&self.mem, desc_index, len);
        }
        used_count > 0
    }

    fn signal_used_queue(&self) {
        self.interrupt_status
            .fetch_or(INTERRUPT_STATUS_USED_RING as usize, Ordering::SeqCst);
        self.interrupt_evt.write(1).unwrap();
    }

    fn run(&mut self, queue_evt: EventFd, kill_evt: EventFd) {
        const Q_AVAIL: u32 = 0;
        const KILL: u32 = 1;

        let mut poller = Poller::new(2);
        'poll: loop {
            let tokens = match poller.poll(&[(Q_AVAIL, &queue_evt), (KILL, &kill_evt)]) {
                Ok(v) => v,
                Err(e) => {
                    error!("failed polling for events: {:?}", e);
                    break;
                }
            };

            let mut needs_interrupt = false;
            for &token in tokens {
                match token {
                    Q_AVAIL => {
                        if let Err(e) = queue_evt.read() {
                            error!("failed reading queue EventFd: {:?}", e);
                            break 'poll;
                        }
                        needs_interrupt |= self.process_queue();
                    }
                    KILL => break 'poll,
                    _ => unreachable!(),
                }
            }
            if needs_interrupt {
                self.signal_used_queue();
            }
        }
    }
}

/// Virtio device for exposing entropy to the guest OS through virtio.
pub struct Rng {
    kill_evt: Option<EventFd>,
    random_file: Option<File>,
}

impl Rng {
    /// Create a new virtio rng device that gets random data from /dev/random.
    pub fn new() -> Result<Rng> {
        let random_file = File::open("/dev/random")
            .map_err(RngError::AccessingRandomDev)?;
        Ok(Rng {
               kill_evt: None,
               random_file: Some(random_file),
           })
    }
}

impl Drop for Rng {
    fn drop(&mut self) {
        if let Some(kill_evt) = self.kill_evt.take() {
            // Ignore the result because there is nothing we can do about it.
            let _ = kill_evt.write(1);
        }
    }
}

impl VirtioDevice for Rng {
    fn keep_fds(&self) -> Vec<RawFd> {
        let mut keep_fds = Vec::new();

        if let Some(ref random_file) = self.random_file {
            keep_fds.push(random_file.as_raw_fd());
        }

        keep_fds
    }

    fn device_type(&self) -> u32 {
        TYPE_RNG
    }

    fn queue_max_sizes(&self) -> &[u16] {
        QUEUE_SIZES
    }

    fn activate(&mut self,
                mem: GuestMemory,
                interrupt_evt: EventFd,
                status: Arc<AtomicUsize>,
                mut queues: Vec<Queue>,
                mut queue_evts: Vec<EventFd>) {
        if queues.len() != 1 || queue_evts.len() != 1 {
            return;
        }

        let (self_kill_evt, kill_evt) =
            match EventFd::new().and_then(|e| Ok((e.try_clone()?, e))) {
                Ok(v) => v,
                Err(e) => {
                    error!("failed to create kill EventFd pair: {:?}", e);
                    return;
                }
            };
        self.kill_evt = Some(self_kill_evt);

        let queue = queues.remove(0);

        if let Some(random_file) = self.random_file.take() {
            let worker_result = thread::Builder::new()
                .name("virtio_rng".to_string())
                .spawn(move || {
                    let mut worker = Worker {
                        queue: queue,
                        mem: mem,
                        random_file: random_file,
                        interrupt_status: status,
                        interrupt_evt: interrupt_evt,
                    };
                    worker.run(queue_evts.remove(0), kill_evt);
                });

            if let Err(e) = worker_result {
                error!("failed to spawn virtio_rng worker: {}", e);
                return;
            }
        }
    }
}
<commit_msg>rng: switch to /dev/urandom<commit_after>// Copyright 2017 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.

use std;
use std::fs::File;
use std::io;
use std::os::unix::io::{AsRawFd, RawFd};
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread;

use sys_util::{EventFd, GuestMemory, Poller};

use super::{VirtioDevice, Queue, INTERRUPT_STATUS_USED_RING, TYPE_RNG};

const QUEUE_SIZE: u16 = 256;
const QUEUE_SIZES: &'static [u16] = &[QUEUE_SIZE];

#[derive(Debug)]
pub enum RngError {
    /// Can't access /dev/urandom
    AccessingRandomDev(io::Error),
}
pub type Result<T> = std::result::Result<T, RngError>;

struct Worker {
    queue: Queue,
    mem: GuestMemory,
    random_file: File,
    interrupt_status: Arc<AtomicUsize>,
    interrupt_evt: EventFd,
}

impl Worker {
    fn process_queue(&mut self) -> bool {
        let queue = &mut self.queue;

        let mut used_desc_heads = [(0, 0); QUEUE_SIZE as usize];
        let mut used_count = 0;
        for avail_desc in queue.iter(&self.mem) {
            let mut len = 0;

            // Drivers can only read from the random device.
            if avail_desc.is_write_only() {
                // Fill the read with data from the random device on the host.
                if self.mem.read_to_memory(avail_desc.addr,
                                           &mut self.random_file,
                                           avail_desc.len as usize)
                        .is_ok() {
                    len = avail_desc.len;
                }
            }

            used_desc_heads[used_count] = (avail_desc.index, len);
            used_count += 1;
        }

        for &(desc_index, len) in &used_desc_heads[..used_count] {
            queue.add_used(&self.mem, desc_index, len);
        }
        used_count > 0
    }

    fn signal_used_queue(&self) {
        self.interrupt_status
            .fetch_or(INTERRUPT_STATUS_USED_RING as usize, Ordering::SeqCst);
        self.interrupt_evt.write(1).unwrap();
    }

    fn run(&mut self, queue_evt: EventFd, kill_evt: EventFd) {
        const Q_AVAIL: u32 = 0;
        const KILL: u32 = 1;

        let mut poller = Poller::new(2);
        'poll: loop {
            let tokens = match poller.poll(&[(Q_AVAIL, &queue_evt), (KILL, &kill_evt)]) {
                Ok(v) => v,
                Err(e) => {
                    error!("failed polling for events: {:?}", e);
                    break;
                }
            };

            let mut needs_interrupt = false;
            for &token in tokens {
                match token {
                    Q_AVAIL => {
                        if let Err(e) = queue_evt.read() {
                            error!("failed reading queue EventFd: {:?}", e);
                            break 'poll;
                        }
                        needs_interrupt |= self.process_queue();
                    }
                    KILL => break 'poll,
                    _ => unreachable!(),
                }
            }
            if needs_interrupt {
                self.signal_used_queue();
            }
        }
    }
}

/// Virtio device for exposing entropy to the guest OS through virtio.
pub struct Rng {
    kill_evt: Option<EventFd>,
    random_file: Option<File>,
}

impl Rng {
    /// Create a new virtio rng device that gets random data from /dev/urandom.
    pub fn new() -> Result<Rng> {
        let random_file = File::open("/dev/urandom")
            .map_err(RngError::AccessingRandomDev)?;
        Ok(Rng {
               kill_evt: None,
               random_file: Some(random_file),
           })
    }
}

impl Drop for Rng {
    fn drop(&mut self) {
        if let Some(kill_evt) = self.kill_evt.take() {
            // Ignore the result because there is nothing we can do about it.
            let _ = kill_evt.write(1);
        }
    }
}

impl VirtioDevice for Rng {
    fn keep_fds(&self) -> Vec<RawFd> {
        let mut keep_fds = Vec::new();

        if let Some(ref random_file) = self.random_file {
            keep_fds.push(random_file.as_raw_fd());
        }

        keep_fds
    }

    fn device_type(&self) -> u32 {
        TYPE_RNG
    }

    fn queue_max_sizes(&self) -> &[u16] {
        QUEUE_SIZES
    }

    fn activate(&mut self,
                mem: GuestMemory,
                interrupt_evt: EventFd,
                status: Arc<AtomicUsize>,
                mut queues: Vec<Queue>,
                mut queue_evts: Vec<EventFd>) {
        if queues.len() != 1 || queue_evts.len() != 1 {
            return;
        }

        let (self_kill_evt, kill_evt) =
            match EventFd::new().and_then(|e| Ok((e.try_clone()?, e))) {
                Ok(v) => v,
                Err(e) => {
                    error!("failed to create kill EventFd pair: {:?}", e);
                    return;
                }
            };
        self.kill_evt = Some(self_kill_evt);

        let queue = queues.remove(0);

        if let Some(random_file) = self.random_file.take() {
            let worker_result = thread::Builder::new()
                .name("virtio_rng".to_string())
                .spawn(move || {
                    let mut worker = Worker {
                        queue: queue,
                        mem: mem,
                        random_file: random_file,
                        interrupt_status: status,
                        interrupt_evt: interrupt_evt,
                    };
                    worker.run(queue_evts.remove(0), kill_evt);
                });

            if let Err(e) = worker_result {
                error!("failed to spawn virtio_rng worker: {}", e);
                return;
            }
        }
    }
}
<|endoftext|>


<commit_before>#![allow(non_snake_case)]
#![allow(non_camel_case_types)]

extern crate libc;

use std::ffi::CStr;

pub const PKG_VERSION: &'static str = env!("CARGO_PKG_VERSION");

pub fn check_version() {
    let raw_version = unsafe { CStr::from_ptr(HELIX_RUNTIME_VERSION) };
    let version = raw_version.to_str().expect("HELIX_RUNTIME_VERSION must be defined");

    if PKG_VERSION != version {
        panic!("libcsys-ruby version ({}) doesn't match helix_runtime version ({}).", PKG_VERSION, version);
    }
}

pub type void = libc::c_void;
pub type c_func = *const void;
pub type c_string = *const libc::c_char;
// pub type c_func = extern "C" fn(...);

#[repr(C)]
#[derive(Eq, PartialEq, Hash, Copy, Clone, Debug)]
pub struct ID(*mut void);

#[repr(C)]
#[derive(Eq, PartialEq, Copy, Clone, Debug)]
pub struct VALUE(*mut void);

#[repr(C)]
#[derive(Eq, PartialEq, Copy, Clone, Debug)]
pub struct RubyException(isize);

impl RubyException {
    pub fn new() -> RubyException {
        RubyException(0)
    }

    pub fn empty() -> RubyException {
        RubyException(0)
    }

    pub fn for_tag(tag: isize) -> RubyException {
        RubyException(tag)
    }
}

pub const EMPTY_EXCEPTION: RubyException = RubyException(0);

#[repr(C)]
pub enum st_retval {
    ST_CONTINUE,
    ST_STOP,
    ST_DELETE,
    // ST_CHECK
}

#[cfg_attr(windows, link(name="helix-runtime"))]
extern "C" {
    #[link_name = "HELIX_RUNTIME_VERSION"]
    pub static HELIX_RUNTIME_VERSION: c_string;

    #[link_name = "HELIX_Qfalse"]
    pub static Qfalse: VALUE;

    #[link_name = "HELIX_Qtrue"]
    pub static Qtrue: VALUE;

    #[link_name = "HELIX_Qnil"]
    pub static Qnil: VALUE;

    #[link_name = "HELIX_PRIsVALUE"]
    pub static PRIsVALUE: c_string;

    #[link_name = "HELIX_SPRINTF_TO_S"]
    pub static SPRINTF_TO_S: c_string;

    #[link_name = "HELIX_SPRINTF_INSPECT"]
    pub static SPRINTF_INSPECT: c_string;

    #[link_name = "rb_cObject"]
    pub static rb_cObject: VALUE;

    #[link_name = "rb_eRuntimeError"]
    pub static rb_eRuntimeError: VALUE;

    #[link_name = "rb_eTypeError"]
    pub static rb_eTypeError: VALUE;

    #[link_name = "HELIX_RSTRING_LEN"]
    pub fn RSTRING_LEN(string: VALUE) -> isize;

    #[link_name = "HELIX_RSTRING_PTR"]
    pub fn RSTRING_PTR(string: VALUE) -> c_string;

    #[link_name = "HELIX_rb_utf8_str_new"]
    pub fn rb_utf8_str_new(string: c_string, len: libc::c_long) -> VALUE;

    #[link_name = "HELIX_RARRAY_LEN"]
    pub fn RARRAY_LEN(array: VALUE) -> isize;

    #[link_name = "HELIX_RARRAY_PTR"]
    pub fn RARRAY_PTR(array: VALUE) -> *mut VALUE;

    #[link_name = "HELIX_RARRAY_CONST_PTR"]
    pub fn RARRAY_CONST_PTR(array: VALUE) -> *const VALUE;

    #[link_name = "HELIX_RHASH_SIZE"]
    pub fn RHASH_SIZE(hash: VALUE) -> isize;

    #[link_name = "HELIX_RB_TYPE_P"]
    pub fn RB_TYPE_P(val: VALUE, rb_type: isize) -> bool;

    #[link_name = "HELIX_TYPE"]
    pub fn TYPE(val: VALUE) -> isize;

    pub fn rb_check_type(v: VALUE, rb_type: isize);

    #[link_name = "HELIX_NUM2U64"]
    pub fn NUM2U64(v: VALUE) -> u64;

    #[link_name = "HELIX_U642NUM"]
    pub fn U642NUM(num: u64) -> VALUE;

    #[link_name = "HELIX_NUM2I64"]
    pub fn NUM2I64(v: VALUE) -> i64;

    #[link_name = "HELIX_I642NUM"]
    pub fn I642NUM(num: i64) -> VALUE;

    #[link_name = "HELIX_NUM2U32"]
    pub fn NUM2U32(v: VALUE) -> u32;

    #[link_name = "HELIX_U322NUM"]
    pub fn U322NUM(num: u32) -> VALUE;

    #[link_name = "HELIX_NUM2I32"]
    pub fn NUM2I32(v: VALUE) -> i32;

    #[link_name = "HELIX_I322NUM"]
    pub fn I322NUM(num: i32) -> VALUE;

    #[link_name = "HELIX_NUM2F64"]
    pub fn NUM2F64(v: VALUE) -> f64;

    #[link_name = "HELIX_F642NUM"]
    pub fn F642NUM(num: f64) -> VALUE;

    #[link_name = "HELIX_T_STRING"]
    pub static T_STRING: isize;

    #[link_name = "HELIX_T_ARRAY"]
    pub static T_ARRAY: isize;

    #[link_name = "HELIX_T_HASH"]
    pub static T_HASH: isize;

    #[link_name = "HELIX_T_TRUE"]
    pub static T_TRUE: isize;

    #[link_name = "HELIX_T_FALSE"]
    pub static T_FALSE: isize;

    #[link_name = "HELIX_T_SYMBOL"]
    pub static T_SYMBOL: isize;

    #[link_name = "HELIX_T_FIXNUM"]
    pub static T_FIXNUM: isize;

    #[link_name = "HELIX_T_FLOAT"]
    pub static T_FLOAT: isize;

    #[link_name = "HELIX_T_BIGNUM"]
    pub static T_BIGNUM: isize;

    // unknown if working?
    // fn rb_define_variable(name: c_string, value: *const VALUE);
    pub fn rb_obj_class(obj: VALUE) -> VALUE;
    pub fn rb_obj_classname(obj: VALUE) -> c_string;
    pub fn rb_const_get(class: VALUE, name: ID) -> VALUE;
    pub fn rb_define_global_const(name: c_string, value: VALUE);
    pub fn rb_define_module(name: c_string) -> VALUE;
    pub fn rb_define_module_under(namespace: VALUE, name: c_string) -> VALUE;
    pub fn rb_define_class(name: c_string, superclass: VALUE) -> VALUE;
    pub fn rb_define_class_under(namespace: VALUE, name: c_string, superclass: VALUE) -> VALUE;
    pub fn rb_define_alloc_func(klass: VALUE, func: extern "C" fn(klass: VALUE) -> VALUE);
    pub fn rb_define_method(class: VALUE, name: c_string, func: c_func, arity: isize);
    pub fn rb_define_singleton_method(class: VALUE, name: c_string, func: c_func, arity: isize);
    pub fn rb_sprintf(specifier: c_string, ...) -> VALUE;
    pub fn rb_inspect(value: VALUE) -> VALUE;
    pub fn rb_intern(string: c_string) -> ID;
    pub fn rb_intern_str(string: VALUE) -> ID;
    pub fn rb_sym2id(symbol: VALUE) -> ID;
    pub fn rb_id2sym(id: ID) -> VALUE;
    pub fn rb_id2str(id: ID) -> VALUE;
    pub fn rb_ary_new_capa(capa: isize) -> VALUE;
    pub fn rb_ary_entry(ary: VALUE, offset: isize) -> VALUE;
    pub fn rb_ary_push(ary: VALUE, item: VALUE) -> VALUE;
    pub fn rb_hash_new() -> VALUE;
    pub fn rb_hash_aset(hash: VALUE, key: VALUE, value: VALUE) -> VALUE;
    pub fn rb_hash_foreach(hash: VALUE, f: extern "C" fn(key: VALUE, value: VALUE, farg: *mut void) -> st_retval, farg: *mut void);

    pub fn rb_raise(exc: VALUE, string: c_string, ...) -> !;
    pub fn rb_jump_tag(state: RubyException) -> !;
    pub fn rb_protect(try: extern "C" fn(v: *mut void) -> VALUE,
                      arg: *mut void,
                      state: *mut RubyException)
                      -> VALUE;

    #[link_name = "HELIX_Data_Wrap_Struct"]
    pub fn Data_Wrap_Struct(klass: VALUE, mark: extern "C" fn(*mut void), free: extern "C" fn(*mut void), data: *mut void) -> VALUE;

    #[link_name = "HELIX_Data_Get_Struct_Value"]
    pub fn Data_Get_Struct_Value(obj: VALUE) -> *mut void;

    #[link_name = "HELIX_Data_Set_Struct_Value"]
    pub fn Data_Set_Struct_Value(obj: VALUE, data: *mut void);
}
<commit_msg>Expose a smattering of functions in libcruby<commit_after>#![allow(non_snake_case)]
#![allow(non_camel_case_types)]

extern crate libc;

use std::ffi::CStr;

pub const PKG_VERSION: &'static str = env!("CARGO_PKG_VERSION");

pub fn check_version() {
    let raw_version = unsafe { CStr::from_ptr(HELIX_RUNTIME_VERSION) };
    let version = raw_version.to_str().expect("HELIX_RUNTIME_VERSION must be defined");

    if PKG_VERSION != version {
        panic!("libcsys-ruby version ({}) doesn't match helix_runtime version ({}).", PKG_VERSION, version);
    }
}

pub type void = libc::c_void;
pub type c_func = *const void;
pub type c_string = *const libc::c_char;
// pub type c_func = extern "C" fn(...);

#[repr(C)]
#[derive(Eq, PartialEq, Hash, Copy, Clone, Debug)]
pub struct ID(*mut void);

#[repr(C)]
#[derive(Eq, PartialEq, Copy, Clone, Debug)]
pub struct VALUE(*mut void);

#[repr(C)]
#[derive(Eq, PartialEq, Copy, Clone, Debug)]
pub struct RubyException(isize);

impl RubyException {
    pub fn new() -> RubyException {
        RubyException(0)
    }

    pub fn empty() -> RubyException {
        RubyException(0)
    }

    pub fn for_tag(tag: isize) -> RubyException {
        RubyException(tag)
    }
}

pub const EMPTY_EXCEPTION: RubyException = RubyException(0);

#[repr(C)]
pub enum st_retval {
    ST_CONTINUE,
    ST_STOP,
    ST_DELETE,
    // ST_CHECK
}

#[cfg_attr(windows, link(name="helix-runtime"))]
extern "C" {
    #[link_name = "HELIX_RUNTIME_VERSION"]
    pub static HELIX_RUNTIME_VERSION: c_string;

    #[link_name = "HELIX_Qfalse"]
    pub static Qfalse: VALUE;

    #[link_name = "HELIX_Qtrue"]
    pub static Qtrue: VALUE;

    #[link_name = "HELIX_Qnil"]
    pub static Qnil: VALUE;

    #[link_name = "HELIX_PRIsVALUE"]
    pub static PRIsVALUE: c_string;

    #[link_name = "HELIX_SPRINTF_TO_S"]
    pub static SPRINTF_TO_S: c_string;

    #[link_name = "HELIX_SPRINTF_INSPECT"]
    pub static SPRINTF_INSPECT: c_string;

    #[link_name = "rb_cObject"]
    pub static rb_cObject: VALUE;

    #[link_name = "rb_eRuntimeError"]
    pub static rb_eRuntimeError: VALUE;

    #[link_name = "rb_eTypeError"]
    pub static rb_eTypeError: VALUE;

    #[link_name = "HELIX_RSTRING_LEN"]
    pub fn RSTRING_LEN(string: VALUE) -> isize;

    #[link_name = "HELIX_RSTRING_PTR"]
    pub fn RSTRING_PTR(string: VALUE) -> c_string;

    #[link_name = "HELIX_rb_utf8_str_new"]
    pub fn rb_utf8_str_new(string: c_string, len: libc::c_long) -> VALUE;

    #[link_name = "HELIX_RARRAY_LEN"]
    pub fn RARRAY_LEN(array: VALUE) -> isize;

    #[link_name = "HELIX_RARRAY_PTR"]
    pub fn RARRAY_PTR(array: VALUE) -> *mut VALUE;

    #[link_name = "HELIX_RARRAY_CONST_PTR"]
    pub fn RARRAY_CONST_PTR(array: VALUE) -> *const VALUE;

    #[link_name = "HELIX_RHASH_SIZE"]
    pub fn RHASH_SIZE(hash: VALUE) -> isize;

    #[link_name = "HELIX_RB_TYPE_P"]
    pub fn RB_TYPE_P(val: VALUE, rb_type: isize) -> bool;

    #[link_name = "HELIX_TYPE"]
    pub fn TYPE(val: VALUE) -> isize;

    pub fn rb_check_type(v: VALUE, rb_type: isize);

    #[link_name = "HELIX_NUM2U64"]
    pub fn NUM2U64(v: VALUE) -> u64;

    #[link_name = "HELIX_U642NUM"]
    pub fn U642NUM(num: u64) -> VALUE;

    #[link_name = "HELIX_NUM2I64"]
    pub fn NUM2I64(v: VALUE) -> i64;

    #[link_name = "HELIX_I642NUM"]
    pub fn I642NUM(num: i64) -> VALUE;

    #[link_name = "HELIX_NUM2U32"]
    pub fn NUM2U32(v: VALUE) -> u32;

    #[link_name = "HELIX_U322NUM"]
    pub fn U322NUM(num: u32) -> VALUE;

    #[link_name = "HELIX_NUM2I32"]
    pub fn NUM2I32(v: VALUE) -> i32;

    #[link_name = "HELIX_I322NUM"]
    pub fn I322NUM(num: i32) -> VALUE;

    #[link_name = "HELIX_NUM2F64"]
    pub fn NUM2F64(v: VALUE) -> f64;

    #[link_name = "HELIX_F642NUM"]
    pub fn F642NUM(num: f64) -> VALUE;

    #[link_name = "HELIX_T_STRING"]
    pub static T_STRING: isize;

    #[link_name = "HELIX_T_ARRAY"]
    pub static T_ARRAY: isize;

    #[link_name = "HELIX_T_HASH"]
    pub static T_HASH: isize;

    #[link_name = "HELIX_T_TRUE"]
    pub static T_TRUE: isize;

    #[link_name = "HELIX_T_FALSE"]
    pub static T_FALSE: isize;

    #[link_name = "HELIX_T_SYMBOL"]
    pub static T_SYMBOL: isize;

    #[link_name = "HELIX_T_FIXNUM"]
    pub static T_FIXNUM: isize;

    #[link_name = "HELIX_T_FLOAT"]
    pub static T_FLOAT: isize;

    #[link_name = "HELIX_T_BIGNUM"]
    pub static T_BIGNUM: isize;

    // unknown if working?
    // fn rb_define_variable(name: c_string, value: *const VALUE);
    pub fn rb_obj_class(obj: VALUE) -> VALUE;
    pub fn rb_obj_classname(obj: VALUE) -> c_string;
    pub fn rb_const_get(class: VALUE, name: ID) -> VALUE;
    pub fn rb_define_global_const(name: c_string, value: VALUE);
    pub fn rb_define_module(name: c_string) -> VALUE;
    pub fn rb_define_module_under(namespace: VALUE, name: c_string) -> VALUE;
    pub fn rb_define_class(name: c_string, superclass: VALUE) -> VALUE;
    pub fn rb_define_class_under(namespace: VALUE, name: c_string, superclass: VALUE) -> VALUE;
    pub fn rb_define_alloc_func(klass: VALUE, func: extern "C" fn(klass: VALUE) -> VALUE);
    pub fn rb_define_method(class: VALUE, name: c_string, func: c_func, arity: isize);
    pub fn rb_define_singleton_method(class: VALUE, name: c_string, func: c_func, arity: isize);
    pub fn rb_sprintf(specifier: c_string, ...) -> VALUE;
    pub fn rb_inspect(value: VALUE) -> VALUE;
    pub fn rb_intern(string: c_string) -> ID;
    pub fn rb_intern_str(string: VALUE) -> ID;
    pub fn rb_sym2id(symbol: VALUE) -> ID;
    pub fn rb_id2sym(id: ID) -> VALUE;
    pub fn rb_id2str(id: ID) -> VALUE;
    pub fn rb_ary_new() -> VALUE;
    pub fn rb_ary_new_capa(capa: isize) -> VALUE;
    pub fn rb_ary_entry(ary: VALUE, offset: isize) -> VALUE;
    pub fn rb_ary_push(ary: VALUE, item: VALUE) -> VALUE;
    pub fn rb_hash_new() -> VALUE;
    pub fn rb_hash_aref(hash: VALUE, key: VALUE) -> VALUE;
    pub fn rb_hash_aset(hash: VALUE, key: VALUE, value: VALUE) -> VALUE;
    pub fn rb_hash_foreach(hash: VALUE, f: extern "C" fn(key: VALUE, value: VALUE, farg: *mut void) -> st_retval, farg: *mut void);
    pub fn rb_gc_mark(value: VALUE);
    pub fn rb_funcall(value: VALUE, mid: ID, argc: libc::c_int, ...) -> VALUE;
    pub fn rb_scan_args(argc: libc::c_int, argv: *const VALUE, fmt: c_string, ...);
    pub fn rb_block_given_p() -> bool;
    pub fn rb_yield(value: VALUE) -> VALUE;
    pub fn rb_obj_dup(value: VALUE) -> VALUE;
    pub fn rb_obj_init_copy(value: VALUE, orig: VALUE) -> VALUE;

    pub fn rb_raise(exc: VALUE, string: c_string, ...) -> !;
    pub fn rb_jump_tag(state: RubyException) -> !;
    pub fn rb_protect(try: extern "C" fn(v: *mut void) -> VALUE,
                      arg: *mut void,
                      state: *mut RubyException)
                      -> VALUE;

    #[link_name = "HELIX_Data_Wrap_Struct"]
    pub fn Data_Wrap_Struct(klass: VALUE, mark: extern "C" fn(*mut void), free: extern "C" fn(*mut void), data: *mut void) -> VALUE;

    #[link_name = "HELIX_Data_Get_Struct_Value"]
    pub fn Data_Get_Struct_Value(obj: VALUE) -> *mut void;

    #[link_name = "HELIX_Data_Set_Struct_Value"]
    pub fn Data_Set_Struct_Value(obj: VALUE, data: *mut void);
}
<|endoftext|>


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

//! Functions to unify writing to devices

#[cfg(test)]
use std::io::Cursor;
use std::{
    fs::{File, OpenOptions},
    io::{self, BufWriter, Seek, SeekFrom, Write},
    path::Path,
};

use devicemapper::{Sectors, IEC, SECTOR_SIZE};

use crate::stratis::StratisResult;

/// The SyncAll trait unifies the File type with other types that do
/// not implement sync_all().
pub trait SyncAll: Write {
    fn sync_all(&mut self) -> io::Result<()>;
}

impl SyncAll for File {
    /// Invokes File::sync_all() thereby syncing all the data
    fn sync_all(&mut self) -> io::Result<()> {
        File::sync_all(self)
    }
}

#[cfg(test)]
impl<T> SyncAll for Cursor<T>
where
    Cursor<T>: Write,
{
    /// A no-op. No data need be synced, because it is already in the Cursor
    /// inner value, which has type T.
    fn sync_all(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl<T> SyncAll for BufWriter<T>
where
    T: SyncAll,
{
    fn sync_all(&mut self) -> io::Result<()> {
        self.flush()?;
        self.get_mut().sync_all()
    }
}

/// Write buf at offset length times.
fn write_sectors<P: AsRef<Path>>(
    path: P,
    offset: Sectors,
    length: Sectors,
    buf: &[u8; SECTOR_SIZE],
) -> StratisResult<()> {
    let mut f = BufWriter::with_capacity(
        convert_const!(IEC::Mi, u64, usize),
        OpenOptions::new().write(true).open(path)?,
    );

    f.seek(SeekFrom::Start(*offset.bytes()))?;
    for _ in 0..*length {
        f.write_all(buf)?;
    }

    f.sync_all()?;
    Ok(())
}

/// Zero sectors at the given offset for length sectors.
/// Note that this method buffers the zeros and syncs only when all are
/// written.
pub fn wipe_sectors<P: AsRef<Path>>(
    path: P,
    offset: Sectors,
    length: Sectors,
) -> StratisResult<()> {
    write_sectors(path, offset, length, &[0u8; SECTOR_SIZE])
}
<commit_msg>Do not allocate more space in the vector than will be used<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! Functions to unify writing to devices

#[cfg(test)]
use std::io::Cursor;
use std::{
    cmp::min,
    fs::{File, OpenOptions},
    io::{self, BufWriter, Seek, SeekFrom, Write},
    path::Path,
};

use devicemapper::{Sectors, IEC, SECTOR_SIZE};

use crate::stratis::StratisResult;

/// The SyncAll trait unifies the File type with other types that do
/// not implement sync_all().
pub trait SyncAll: Write {
    fn sync_all(&mut self) -> io::Result<()>;
}

impl SyncAll for File {
    /// Invokes File::sync_all() thereby syncing all the data
    fn sync_all(&mut self) -> io::Result<()> {
        File::sync_all(self)
    }
}

#[cfg(test)]
impl<T> SyncAll for Cursor<T>
where
    Cursor<T>: Write,
{
    /// A no-op. No data need be synced, because it is already in the Cursor
    /// inner value, which has type T.
    fn sync_all(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl<T> SyncAll for BufWriter<T>
where
    T: SyncAll,
{
    fn sync_all(&mut self) -> io::Result<()> {
        self.flush()?;
        self.get_mut().sync_all()
    }
}

/// Write buf at offset length times.
fn write_sectors<P: AsRef<Path>>(
    path: P,
    offset: Sectors,
    length: Sectors,
    buf: &[u8; SECTOR_SIZE],
) -> StratisResult<()> {
    let mut f = BufWriter::with_capacity(
        convert_const!(min(IEC::Mi, *(length.bytes())), u64, usize),
        OpenOptions::new().write(true).open(path)?,
    );

    f.seek(SeekFrom::Start(*offset.bytes()))?;
    for _ in 0..*length {
        f.write_all(buf)?;
    }

    f.sync_all()?;
    Ok(())
}

/// Zero sectors at the given offset for length sectors.
/// Note that this method buffers the zeros and syncs only when all are
/// written.
pub fn wipe_sectors<P: AsRef<Path>>(
    path: P,
    offset: Sectors,
    length: Sectors,
) -> StratisResult<()> {
    write_sectors(path, offset, length, &[0u8; SECTOR_SIZE])
}
<|endoftext|>


<commit_before>use editor::Editor;
use redraw::RedrawTask;
use buffer::Buffer;

impl Editor {
    /// Delete a character
    #[inline]
    pub fn delete(&mut self) {
        let (x, y) = self.pos();
        if x == 0 {
            if y != 0 {
                let s = self.buffer.remove_line(y);
                self.buffer[y - 1].push_str(&s);
                let len = self.buffer[y - 1].len();
                self.goto((len, y - 1));
                self.redraw_task = RedrawTask::Lines(y..y + 1);
            }
        } else if x < self.buffer[y].len() {
            self.buffer[y].remove(x);
            self.redraw_task = RedrawTask::LinesAfter(y);
        }

        self.hint();
    }
}
<commit_msg>Fix delete's newline behaviour<commit_after>use editor::Editor;
use redraw::RedrawTask;
use buffer::Buffer;
use cursor::Cursor;

impl Editor {
    /// Delete a character
    #[inline]
    pub fn delete(&mut self) {
        let &Cursor{ x, y, ..} = self.cursor();
        if x == self.buffer[y].len() {
            if y < self.buffer.len() - 1 {
                let s = self.buffer.remove_line(y + 1);
                self.buffer[y].push_str(&s);
                self.redraw_task = RedrawTask::Lines(y..y+1);
            }
        } else if x < self.buffer[y].len() {
            self.buffer[y].remove(x);
            self.redraw_task = RedrawTask::LinesAfter(y);
        }

        self.hint();
    }
}
<|endoftext|>