<commit_before>use std::path::{Path, PathBuf};
use serde_json;
use directory::error::{OpenReadError, DeleteError, OpenWriteError};
use directory::{ReadOnlySource, WritePtr};
use std::result;
use std::io;
use Directory;
use std::sync::{Arc, RwLock};
use std::collections::HashSet;
use std::sync::RwLockWriteGuard;
use std::io::Write;
use core::MANAGED_FILEPATH;
use std::collections::HashMap;
use std::fmt;
use Result;
use Error;

/// Wrapper of directories that keeps track of files created by Tantivy.
///
/// A managed directory is just a wrapper of a directory
/// that keeps a (persisted) list of the files that
/// have been created (and not deleted) by tantivy so far.
///
/// Thanks to this list, it implements a `garbage_collect` method
/// that removes the files that were created by tantivy and are not
/// useful anymore.
#[derive(Debug)]
pub struct ManagedDirectory {
    directory: Box<Directory>,
    meta_informations: Arc<RwLock<MetaInformation>>,
}

#[derive(Debug, Default)]
struct MetaInformation {
    managed_paths: HashSet<PathBuf>,
    protected_files: HashMap<PathBuf, usize>,
}

/// A `FileProtection` prevents the garbage collection of a file.
///
/// See `ManagedDirectory.protect_file_from_delete`.
pub struct FileProtection {
    directory: ManagedDirectory,
    path: PathBuf,
}

fn unprotect_file_from_delete(directory: &ManagedDirectory, path: &Path) {
    let mut meta_informations_wlock = directory.meta_informations
        .write()
        .expect("Managed file lock poisoned");
    if let Some(counter_ref_mut) = meta_informations_wlock
        .protected_files
        .get_mut(path) {
        (*counter_ref_mut) -= 1;
    }
}

impl fmt::Debug for FileProtection {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
        write!(formatter, "FileProtectionFor({:?})", self.path)
    }
}

impl Drop for FileProtection {
    fn drop(&mut self) {
        unprotect_file_from_delete(&self.directory, &*self.path);
    }
}

/// Saves the file containing the list of existing files
/// that were created by tantivy.
fn save_managed_paths(directory: &mut Directory, wlock: &RwLockWriteGuard<MetaInformation>) -> io::Result<()> {
    let managed_paths = wlock.managed_paths.clone();
    let mut w = serde_json::to_vec(&managed_paths)?;
    write!(&mut w, "\n")?;
    directory.atomic_write(&MANAGED_FILEPATH, &w[..])?;
    Ok(())
}

impl ManagedDirectory {

    /// Wraps a directory as managed directory.
    pub fn new<Dir: Directory>(directory: Dir) -> Result<ManagedDirectory> {
        match directory.atomic_read(&MANAGED_FILEPATH) {
            Ok(data) => {
                let managed_files_json = String::from_utf8_lossy(&data);
                let managed_files: HashSet<PathBuf> = serde_json::from_str(&managed_files_json)
                    .map_err(|e| Error::CorruptedFile(MANAGED_FILEPATH.clone(), Box::new(e)))?;
                Ok(ManagedDirectory {
                    directory: box directory,
                    meta_informations: Arc::new(RwLock::new(
                        MetaInformation {
                            managed_paths: managed_files,
                            protected_files: HashMap::default()
                        })),
                })
            }
            Err(OpenReadError::FileDoesNotExist(_)) => {
                Ok(ManagedDirectory {
                    directory: box directory,
                    meta_informations: Arc::default(),
                })
            }
            Err(OpenReadError::IOError(e)) => {
                Err(From::from(e))
            }
        }
    }

    /// Garbage collect unused files.
    ///
    /// Removes the files that were created by `tantivy` and are not
    /// used by any segment anymore.
    ///
    /// * `living_files` - List of files that are still used by the index.
    ///
    /// This method does not panick nor returns errors.
    /// If a file cannot be deleted (for permission reasons for instance)
    /// an error is simply logged, and the file remains in the list of managed
    /// files.
    pub fn garbage_collect(&mut self, living_files: HashSet<PathBuf>) {
        let mut files_to_delete = vec!();
        {   // releasing the lock as .delete() will use it too.
            let meta_informations_rlock = self.meta_informations
                .read()
                .expect("Managed directory rlock poisoned in garbage collect.");
            for managed_path in &meta_informations_rlock.managed_paths {
                if !living_files.contains(managed_path) {
                    files_to_delete.push(managed_path.clone());
                }
            }
        }

        let mut deleted_files = vec!();
        {
            for file_to_delete in files_to_delete {
                match self.delete(&file_to_delete) {
                    Ok(_) => {
                        info!("Deleted {:?}", file_to_delete);
                        deleted_files.push(file_to_delete);
                    }
                    Err(file_error) => {
                        error!("Failed to delete {:?}", file_to_delete);
                        match file_error {
                            DeleteError::FileDoesNotExist(_) => {
                                deleted_files.push(file_to_delete);
                            }
                            DeleteError::IOError(_) => {
                                if !cfg!(target_os = "windows") {
                                    error!("Failed to delete {:?}", file_to_delete);
                                }
                            }
                            DeleteError::FileProtected(_) => {
                                // this is expected.
                            }
                        }

                    }
                }
            }
        }

        if !deleted_files.is_empty() {
            // update the list of managed files by removing
            // the file that were removed.
            let mut meta_informations_wlock = self.meta_informations
                .write()
                .expect("Managed directory wlock poisoned (2).");
            {
                let managed_paths_write = &mut meta_informations_wlock.managed_paths;
                for delete_file in &deleted_files {
                    managed_paths_write.remove(delete_file);
                }
            }
            if let Err(_) = save_managed_paths(self.directory.as_mut(), &meta_informations_wlock) {
                error!("Failed to save the list of managed files.");
            }
        }

    }

    /// Protects a file from being garbage collected.
    ///
    /// The method returns a `FileProtection` object.
    /// The file will not be garbage collected as long as the
    /// `FileProtection` object is kept alive.
    pub fn protect_file_from_delete(&self, path: &Path) -> FileProtection {
        let pathbuf = path.to_owned();
        {
            let mut meta_informations_wlock = self.meta_informations
                .write()
                .expect("Managed file lock poisoned on protect");
            *meta_informations_wlock
                .protected_files
                .entry(pathbuf.clone())
                .or_insert(0) += 1;
        }
        FileProtection {
            directory: self.clone(),
            path: pathbuf.clone(),
        }
    }

    /// Registers a file as managed
    ///
    /// This method must be called before the file is
    /// actually created to ensure that a failure between
    /// registering the filepath and creating the file
    /// will not lead to garbage files that will
    /// never get removed.
    fn register_file_as_managed(&mut self, filepath: &Path) -> io::Result<()> {
        let mut meta_wlock = self.meta_informations
                .write()
                .expect("Managed file lock poisoned");
        let has_changed = meta_wlock.managed_paths.insert(filepath.to_owned());
        if has_changed {
            save_managed_paths(self.directory.as_mut(), &meta_wlock)?;
        }
        Ok(())
    }
}

impl Directory for ManagedDirectory {

    fn open_read(&self, path: &Path) -> result::Result<ReadOnlySource, OpenReadError> {
        self.directory.open_read(path)
    }

    fn open_write(&mut self, path: &Path) -> result::Result<WritePtr, OpenWriteError> {
        self.register_file_as_managed(path)?;
        self.directory.open_write(path)
    }

    fn atomic_write(&mut self, path: &Path, data: &[u8]) -> io::Result<()> {
        self.register_file_as_managed(path)?;
        self.directory.atomic_write(path, data)
    }

    fn atomic_read(&self, path: &Path) -> result::Result<Vec<u8>, OpenReadError> {
        self.directory.atomic_read(path)
    }

    fn delete(&self, path: &Path) -> result::Result<(), DeleteError> {
        {
            let metas_rlock = self.meta_informations
                .read()
                .expect("poisoned lock in managed directory meta");
            if let Some(counter) = metas_rlock.protected_files.get(path) {
                if *counter > 0 {
                    return Err(DeleteError::FileProtected(path.to_owned()))
                }
            }
        }
        self.directory.delete(path)
    }

    fn exists(&self, path: &Path) -> bool {
        self.directory.exists(path)
    }

    fn box_clone(&self) -> Box<Directory> {
        box self.clone()
    }

}

impl Clone for ManagedDirectory {
    fn clone(&self) -> ManagedDirectory {
        ManagedDirectory {
            directory: self.directory.box_clone(),
            meta_informations: self.meta_informations.clone(),
        }
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use directory::MmapDirectory;
    use std::path::Path;
    use std::io::Write;
    use tempdir::TempDir;

    lazy_static! {
        static ref TEST_PATH1: &'static Path = Path::new("some_path_for_test");
        static ref TEST_PATH2: &'static Path = Path::new("some_path_for_test2");
    }

    #[test]
    fn test_managed_directory() {
        let tempdir = TempDir::new("index").unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());
        {
            let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
            let mut managed_directory = ManagedDirectory::new(mmap_directory).unwrap();
            {
                let mut write_file = managed_directory.open_write(*TEST_PATH1).unwrap();
                write_file.flush().unwrap();
            }
            {
                managed_directory.atomic_write(*TEST_PATH2, &vec!(0u8,1u8)).unwrap();
            }
            {
                assert!(managed_directory.exists(*TEST_PATH1));
                assert!(managed_directory.exists(*TEST_PATH2));
            }
            {
                let living_files: HashSet<PathBuf> = [TEST_PATH1.to_owned()]
                    .into_iter()
                    .cloned()
                    .collect();
                managed_directory.garbage_collect(living_files);
            }
            {
                assert!(managed_directory.exists(*TEST_PATH1));
                assert!(!managed_directory.exists(*TEST_PATH2));
            }
        }
        {
            let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
            let mut managed_directory = ManagedDirectory::new(mmap_directory).unwrap();
            {
                assert!(managed_directory.exists(*TEST_PATH1));
                assert!(!managed_directory.exists(*TEST_PATH2));
            }
            {
                let living_files: HashSet<PathBuf> = HashSet::new();
                managed_directory.garbage_collect(living_files);
            }
            {
                assert!(!managed_directory.exists(*TEST_PATH1));
                assert!(!managed_directory.exists(*TEST_PATH2));
            }
        }
    }

    #[test]
    fn test_managed_directory_gc_while_mmapped() {
        let tempdir = TempDir::new("index").unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());
        let living_files = HashSet::new();

        let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
        let mut managed_directory = ManagedDirectory::new(mmap_directory).unwrap();
        managed_directory.atomic_write(*TEST_PATH1, &vec!(0u8,1u8)).unwrap();
        assert!(managed_directory.exists(*TEST_PATH1));

        let _mmap_read = managed_directory.open_read(*TEST_PATH1).unwrap();
        managed_directory.garbage_collect(living_files.clone());
        if cfg!(target_os = "windows") {
            // On Windows, gc should try and fail the file as it is mmapped.
            assert!(managed_directory.exists(*TEST_PATH1));
            // unmap should happen here.
            drop(_mmap_read);
            // The file should still be in the list of managed file and
            // eventually be deleted once mmap is released.
            managed_directory.garbage_collect(living_files);
            assert!(!managed_directory.exists(*TEST_PATH1));
        }
        else {
            assert!(!managed_directory.exists(*TEST_PATH1));
        }

    }

    #[test]
    fn test_managed_directory_protect() {
        let tempdir = TempDir::new("index").unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());
        let living_files = HashSet::new();

        let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
        let mut managed_directory = ManagedDirectory::new(mmap_directory).unwrap();
        managed_directory.atomic_write(*TEST_PATH1, &vec!(0u8,1u8)).unwrap();
        assert!(managed_directory.exists(*TEST_PATH1));

        {
            let _file_protection = managed_directory.protect_file_from_delete(*TEST_PATH1);
            managed_directory.garbage_collect(living_files.clone());
            assert!(managed_directory.exists(*TEST_PATH1));
        }

        managed_directory.garbage_collect(living_files.clone());
        assert!(!managed_directory.exists(*TEST_PATH1));

    }

}
<commit_msg>Avoid clone() call<commit_after>use std::path::{Path, PathBuf};
use serde_json;
use directory::error::{OpenReadError, DeleteError, OpenWriteError};
use directory::{ReadOnlySource, WritePtr};
use std::result;
use std::io;
use Directory;
use std::sync::{Arc, RwLock};
use std::collections::HashSet;
use std::sync::RwLockWriteGuard;
use std::io::Write;
use core::MANAGED_FILEPATH;
use std::collections::HashMap;
use std::fmt;
use Result;
use Error;

/// Wrapper of directories that keeps track of files created by Tantivy.
///
/// A managed directory is just a wrapper of a directory
/// that keeps a (persisted) list of the files that
/// have been created (and not deleted) by tantivy so far.
///
/// Thanks to this list, it implements a `garbage_collect` method
/// that removes the files that were created by tantivy and are not
/// useful anymore.
#[derive(Debug)]
pub struct ManagedDirectory {
    directory: Box<Directory>,
    meta_informations: Arc<RwLock<MetaInformation>>,
}

#[derive(Debug, Default)]
struct MetaInformation {
    managed_paths: HashSet<PathBuf>,
    protected_files: HashMap<PathBuf, usize>,
}

/// A `FileProtection` prevents the garbage collection of a file.
///
/// See `ManagedDirectory.protect_file_from_delete`.
pub struct FileProtection {
    directory: ManagedDirectory,
    path: PathBuf,
}

fn unprotect_file_from_delete(directory: &ManagedDirectory, path: &Path) {
    let mut meta_informations_wlock = directory.meta_informations
        .write()
        .expect("Managed file lock poisoned");
    if let Some(counter_ref_mut) = meta_informations_wlock
        .protected_files
        .get_mut(path) {
        (*counter_ref_mut) -= 1;
    }
}

impl fmt::Debug for FileProtection {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
        write!(formatter, "FileProtectionFor({:?})", self.path)
    }
}

impl Drop for FileProtection {
    fn drop(&mut self) {
        unprotect_file_from_delete(&self.directory, &*self.path);
    }
}

/// Saves the file containing the list of existing files
/// that were created by tantivy.
fn save_managed_paths(directory: &mut Directory, wlock: &RwLockWriteGuard<MetaInformation>) -> io::Result<()> {
    let mut w = serde_json::to_vec(&wlock.managed_paths)?;
    write!(&mut w, "\n")?;
    directory.atomic_write(&MANAGED_FILEPATH, &w[..])?;
    Ok(())
}

impl ManagedDirectory {

    /// Wraps a directory as managed directory.
    pub fn new<Dir: Directory>(directory: Dir) -> Result<ManagedDirectory> {
        match directory.atomic_read(&MANAGED_FILEPATH) {
            Ok(data) => {
                let managed_files_json = String::from_utf8_lossy(&data);
                let managed_files: HashSet<PathBuf> = serde_json::from_str(&managed_files_json)
                    .map_err(|e| Error::CorruptedFile(MANAGED_FILEPATH.clone(), Box::new(e)))?;
                Ok(ManagedDirectory {
                    directory: box directory,
                    meta_informations: Arc::new(RwLock::new(
                        MetaInformation {
                            managed_paths: managed_files,
                            protected_files: HashMap::default()
                        })),
                })
            }
            Err(OpenReadError::FileDoesNotExist(_)) => {
                Ok(ManagedDirectory {
                    directory: box directory,
                    meta_informations: Arc::default(),
                })
            }
            Err(OpenReadError::IOError(e)) => {
                Err(From::from(e))
            }
        }
    }

    /// Garbage collect unused files.
    ///
    /// Removes the files that were created by `tantivy` and are not
    /// used by any segment anymore.
    ///
    /// * `living_files` - List of files that are still used by the index.
    ///
    /// This method does not panick nor returns errors.
    /// If a file cannot be deleted (for permission reasons for instance)
    /// an error is simply logged, and the file remains in the list of managed
    /// files.
    pub fn garbage_collect(&mut self, living_files: HashSet<PathBuf>) {
        let mut files_to_delete = vec!();
        {   // releasing the lock as .delete() will use it too.
            let meta_informations_rlock = self.meta_informations
                .read()
                .expect("Managed directory rlock poisoned in garbage collect.");
            for managed_path in &meta_informations_rlock.managed_paths {
                if !living_files.contains(managed_path) {
                    files_to_delete.push(managed_path.clone());
                }
            }
        }

        let mut deleted_files = vec!();
        {
            for file_to_delete in files_to_delete {
                match self.delete(&file_to_delete) {
                    Ok(_) => {
                        info!("Deleted {:?}", file_to_delete);
                        deleted_files.push(file_to_delete);
                    }
                    Err(file_error) => {
                        error!("Failed to delete {:?}", file_to_delete);
                        match file_error {
                            DeleteError::FileDoesNotExist(_) => {
                                deleted_files.push(file_to_delete);
                            }
                            DeleteError::IOError(_) => {
                                if !cfg!(target_os = "windows") {
                                    error!("Failed to delete {:?}", file_to_delete);
                                }
                            }
                            DeleteError::FileProtected(_) => {
                                // this is expected.
                            }
                        }

                    }
                }
            }
        }

        if !deleted_files.is_empty() {
            // update the list of managed files by removing
            // the file that were removed.
            let mut meta_informations_wlock = self.meta_informations
                .write()
                .expect("Managed directory wlock poisoned (2).");
            {
                let managed_paths_write = &mut meta_informations_wlock.managed_paths;
                for delete_file in &deleted_files {
                    managed_paths_write.remove(delete_file);
                }
            }
            if let Err(_) = save_managed_paths(self.directory.as_mut(), &meta_informations_wlock) {
                error!("Failed to save the list of managed files.");
            }
        }

    }

    /// Protects a file from being garbage collected.
    ///
    /// The method returns a `FileProtection` object.
    /// The file will not be garbage collected as long as the
    /// `FileProtection` object is kept alive.
    pub fn protect_file_from_delete(&self, path: &Path) -> FileProtection {
        let pathbuf = path.to_owned();
        {
            let mut meta_informations_wlock = self.meta_informations
                .write()
                .expect("Managed file lock poisoned on protect");
            *meta_informations_wlock
                .protected_files
                .entry(pathbuf.clone())
                .or_insert(0) += 1;
        }
        FileProtection {
            directory: self.clone(),
            path: pathbuf.clone(),
        }
    }

    /// Registers a file as managed
    ///
    /// This method must be called before the file is
    /// actually created to ensure that a failure between
    /// registering the filepath and creating the file
    /// will not lead to garbage files that will
    /// never get removed.
    fn register_file_as_managed(&mut self, filepath: &Path) -> io::Result<()> {
        let mut meta_wlock = self.meta_informations
                .write()
                .expect("Managed file lock poisoned");
        let has_changed = meta_wlock.managed_paths.insert(filepath.to_owned());
        if has_changed {
            save_managed_paths(self.directory.as_mut(), &meta_wlock)?;
        }
        Ok(())
    }
}

impl Directory for ManagedDirectory {

    fn open_read(&self, path: &Path) -> result::Result<ReadOnlySource, OpenReadError> {
        self.directory.open_read(path)
    }

    fn open_write(&mut self, path: &Path) -> result::Result<WritePtr, OpenWriteError> {
        self.register_file_as_managed(path)?;
        self.directory.open_write(path)
    }

    fn atomic_write(&mut self, path: &Path, data: &[u8]) -> io::Result<()> {
        self.register_file_as_managed(path)?;
        self.directory.atomic_write(path, data)
    }

    fn atomic_read(&self, path: &Path) -> result::Result<Vec<u8>, OpenReadError> {
        self.directory.atomic_read(path)
    }

    fn delete(&self, path: &Path) -> result::Result<(), DeleteError> {
        {
            let metas_rlock = self.meta_informations
                .read()
                .expect("poisoned lock in managed directory meta");
            if let Some(counter) = metas_rlock.protected_files.get(path) {
                if *counter > 0 {
                    return Err(DeleteError::FileProtected(path.to_owned()))
                }
            }
        }
        self.directory.delete(path)
    }

    fn exists(&self, path: &Path) -> bool {
        self.directory.exists(path)
    }

    fn box_clone(&self) -> Box<Directory> {
        box self.clone()
    }

}

impl Clone for ManagedDirectory {
    fn clone(&self) -> ManagedDirectory {
        ManagedDirectory {
            directory: self.directory.box_clone(),
            meta_informations: self.meta_informations.clone(),
        }
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use directory::MmapDirectory;
    use std::path::Path;
    use std::io::Write;
    use tempdir::TempDir;

    lazy_static! {
        static ref TEST_PATH1: &'static Path = Path::new("some_path_for_test");
        static ref TEST_PATH2: &'static Path = Path::new("some_path_for_test2");
    }

    #[test]
    fn test_managed_directory() {
        let tempdir = TempDir::new("index").unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());
        {
            let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
            let mut managed_directory = ManagedDirectory::new(mmap_directory).unwrap();
            {
                let mut write_file = managed_directory.open_write(*TEST_PATH1).unwrap();
                write_file.flush().unwrap();
            }
            {
                managed_directory.atomic_write(*TEST_PATH2, &vec!(0u8,1u8)).unwrap();
            }
            {
                assert!(managed_directory.exists(*TEST_PATH1));
                assert!(managed_directory.exists(*TEST_PATH2));
            }
            {
                let living_files: HashSet<PathBuf> = [TEST_PATH1.to_owned()]
                    .into_iter()
                    .cloned()
                    .collect();
                managed_directory.garbage_collect(living_files);
            }
            {
                assert!(managed_directory.exists(*TEST_PATH1));
                assert!(!managed_directory.exists(*TEST_PATH2));
            }
        }
        {
            let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
            let mut managed_directory = ManagedDirectory::new(mmap_directory).unwrap();
            {
                assert!(managed_directory.exists(*TEST_PATH1));
                assert!(!managed_directory.exists(*TEST_PATH2));
            }
            {
                let living_files: HashSet<PathBuf> = HashSet::new();
                managed_directory.garbage_collect(living_files);
            }
            {
                assert!(!managed_directory.exists(*TEST_PATH1));
                assert!(!managed_directory.exists(*TEST_PATH2));
            }
        }
    }

    #[test]
    fn test_managed_directory_gc_while_mmapped() {
        let tempdir = TempDir::new("index").unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());
        let living_files = HashSet::new();

        let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
        let mut managed_directory = ManagedDirectory::new(mmap_directory).unwrap();
        managed_directory.atomic_write(*TEST_PATH1, &vec!(0u8,1u8)).unwrap();
        assert!(managed_directory.exists(*TEST_PATH1));

        let _mmap_read = managed_directory.open_read(*TEST_PATH1).unwrap();
        managed_directory.garbage_collect(living_files.clone());
        if cfg!(target_os = "windows") {
            // On Windows, gc should try and fail the file as it is mmapped.
            assert!(managed_directory.exists(*TEST_PATH1));
            // unmap should happen here.
            drop(_mmap_read);
            // The file should still be in the list of managed file and
            // eventually be deleted once mmap is released.
            managed_directory.garbage_collect(living_files);
            assert!(!managed_directory.exists(*TEST_PATH1));
        }
        else {
            assert!(!managed_directory.exists(*TEST_PATH1));
        }

    }

    #[test]
    fn test_managed_directory_protect() {
        let tempdir = TempDir::new("index").unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());
        let living_files = HashSet::new();

        let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
        let mut managed_directory = ManagedDirectory::new(mmap_directory).unwrap();
        managed_directory.atomic_write(*TEST_PATH1, &vec!(0u8,1u8)).unwrap();
        assert!(managed_directory.exists(*TEST_PATH1));

        {
            let _file_protection = managed_directory.protect_file_from_delete(*TEST_PATH1);
            managed_directory.garbage_collect(living_files.clone());
            assert!(managed_directory.exists(*TEST_PATH1));
        }

        managed_directory.garbage_collect(living_files.clone());
        assert!(!managed_directory.exists(*TEST_PATH1));

    }

}
<|endoftext|>


<commit_before>//
// FlightVars
// Copyright (c) 2015, 2016 Alvaro Polo
//
// 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/.

use std::io;
use std::thread;

use fsuipc;
use fsuipc::{Handle, Session};
use mio;

use domain::types::*;
use util::Consume;

pub mod types;
pub use self::types::*;

const POLLING_DELAY_MS: u64 = 50;

pub enum Envelope {
    Cmd(Command),
    Shutdown
}

pub struct Domain {
    worker: thread::JoinHandle<()>,
    tx: mio::Sender<Envelope>
}

impl Domain {
    pub fn new() -> Domain {
        let (worker, tx) = spawn_worker();
        Domain { worker: worker, tx: tx }
    }

    pub fn shutdown(self) {
        self.tx.send(Envelope::Shutdown).unwrap_or_else(|e| {
            warn!("unexpected error while sending shutdown message to fsuipc domain: {}", e);
        });
        self.worker.join().unwrap_or_else(|_| {
            warn!("unexpected error while waiting for fsuipc domain worker thread");
        });
    }

    pub fn consumer(&self) -> Consumer {
        Consumer { tx: self.tx.clone() }
    }
}

#[derive(Clone)]
pub struct Consumer {
    tx: mio::Sender<Envelope>
}

impl Consume for Consumer {
    type Item = Command;
    type Error = mio::NotifyError<Envelope>;
    fn consume(&mut self, cmd: Command) -> Result<(), mio::NotifyError<Envelope>> {
        self.tx.send(Envelope::Cmd(cmd))
    }
}

struct Observer {
    offset: Offset,
    client: Client,
    retain: Option<Value>,
    buffer: [u8; 4],
}

impl Observer {
    pub fn read(&mut self, session: &mut fsuipc::local::LocalSession) {
        session.read_bytes(
            u16::from(self.offset.addr()),
            &mut self.buffer as *mut [u8; 4] as *mut u8,
            usize::from(self.offset.len())).unwrap();
    }

    pub fn trigger_event(&mut self) {
        let buf_value = self.offset.len().decode_value(&self.buffer);
        let must_trigger = self.retain.as_ref().map(|v| *v != buf_value).unwrap_or(true);
        if must_trigger {
            let event = Event::Update(Var::FsuipcOffset(self.offset), buf_value);
            debug!("triggering event {:?} to client {}", event, self.client.name());
            self.client.sender().send(event).unwrap();
            self.retain = Some(buf_value);
        }
    }
}

struct Context {
    handle: fsuipc::local::LocalHandle,
    observers: Vec<Observer>,
}

impl Context {
    pub fn new()  -> io::Result<Context> {
        Ok(Context {
            handle: try!(fsuipc::local::LocalHandle::new()),
            observers: Vec::new(),
        })
    }

    fn process_write(&mut self, offset: Offset, value: Value) {
        debug!("writing value {} to offset {}", value, offset);
        match offset.len() {
            OffsetLen::UnsignedByte => self.process_write_of(offset.addr(), u8::from(value)),
            OffsetLen::SignedByte => self.process_write_of(offset.addr(), i8::from(value)),
            OffsetLen::UnsignedWord => self.process_write_of(offset.addr(), u16::from(value)),
            OffsetLen::SignedWord => self.process_write_of(offset.addr(), i16::from(value)),
            OffsetLen::UnsignedDouble => self.process_write_of(offset.addr(), u32::from(value)),
            OffsetLen::SignedDouble => self.process_write_of(offset.addr(), i32::from(value)),
        }
    }

    fn process_obs(&mut self, offset: Offset, client: Client) {
        debug!("client {} observing offset {}", client.name(), offset);
        self.observers.push(Observer {
            offset: offset,
            client: client,
            retain: None,
            buffer: [0; 4]
        });
    }

    fn process_write_of<T>(&mut self, offset: OffsetAddr, value: T) {
        let mut session = self.handle.session();
        session.write(u16::from(offset), &value).unwrap();
        session.process().unwrap();
    }
}

fn spawn_worker() -> (thread::JoinHandle<()>, mio::Sender<Envelope>) {
    let event_loop = mio::EventLoop::new().unwrap();
    let tx = event_loop.channel();
    let worker = thread::spawn(move || {
        let mut event_loop = event_loop;
        let mut ctx = match Context::new() {
            Ok(ctx) => ctx,
            Err(e) => {
                error!("cannot create FSUIPC context (is FSUIPC installed & running?): {}", e);
                return;
            },
        };
        event_loop.timeout_ms((), POLLING_DELAY_MS).unwrap();
        event_loop.run(&mut ctx).unwrap();
    });
    (worker, tx)
}

impl mio::Handler for Context {
    type Timeout = ();
    type Message = Envelope;

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Context>, _: ()) {
        let mut session = self.handle.session();
        for obs in self.observers.iter_mut() {
            obs.read(&mut session);
        }
        session.process().unwrap();
        for obs in self.observers.iter_mut() {
            obs.trigger_event();
        }
        event_loop.timeout_ms((), POLLING_DELAY_MS).unwrap();
    }

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Context>, msg: Envelope) {
        match msg {
            Envelope::Cmd(Command::Write(Var::FsuipcOffset(offset), value)) => {
                self.process_write(offset, value)
            },
            Envelope::Cmd(Command::Observe(Var::FsuipcOffset(offset), client)) => {
                self.process_obs(offset, client)
            },
            Envelope::Shutdown => event_loop.shutdown(),
            _ => {},
        }
    }
}

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

    #[test]
    fn should_init_and_shutdown() {
        let mut domain = Domain::new();
        domain.shutdown();
    }
}
<commit_msg>Improve error handling in `fsuipc` domain<commit_after>//
// FlightVars
// Copyright (c) 2015, 2016 Alvaro Polo
//
// 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/.

use std::io;
use std::thread;

use fsuipc;
use fsuipc::{Handle, Session};
use mio;

use domain::types::*;
use util::Consume;

pub mod types;
pub use self::types::*;

const POLLING_DELAY_MS: u64 = 50;

#[derive(Debug)]
pub enum Envelope {
    Cmd(Command),
    Shutdown
}

pub struct Domain {
    worker: thread::JoinHandle<()>,
    tx: mio::Sender<Envelope>
}

impl Domain {
    pub fn new() -> Domain {
        let (worker, tx) = spawn_worker();
        Domain { worker: worker, tx: tx }
    }

    pub fn shutdown(self) {
        self.tx.send(Envelope::Shutdown).unwrap_or_else(|e| {
            warn!("unexpected error while sending shutdown message to fsuipc domain: {}", e);
        });
        self.worker.join().unwrap_or_else(|_| {
            warn!("unexpected error while waiting for fsuipc domain worker thread");
        });
    }

    pub fn consumer(&self) -> Consumer {
        Consumer { tx: self.tx.clone() }
    }
}

#[derive(Clone)]
pub struct Consumer {
    tx: mio::Sender<Envelope>
}

impl Consume for Consumer {
    type Item = Command;
    type Error = mio::NotifyError<Envelope>;
    fn consume(&mut self, cmd: Command) -> Result<(), mio::NotifyError<Envelope>> {
        self.tx.send(Envelope::Cmd(cmd))
    }
}

struct Observer {
    offset: Offset,
    client: Client,
    retain: Option<Value>,
    buffer: [u8; 4],
}

impl Observer {
    pub fn read(&mut self, session: &mut fsuipc::local::LocalSession) {
        let offset = u16::from(self.offset.addr());
        let buffer = &mut self.buffer as *mut [u8; 4] as *mut u8;
        let nbytes = usize::from(self.offset.len());
        if let Err(e) = session.read_bytes(offset, buffer, nbytes) {
            error!("unexpected error while reading bytes from FSUIPC session: {}", e);
        }
    }

    pub fn trigger_event(&mut self) {
        let buf_value = self.offset.len().decode_value(&self.buffer);
        let must_trigger = self.retain.as_ref().map(|v| *v != buf_value).unwrap_or(true);
        if must_trigger {
            let event = Event::Update(Var::FsuipcOffset(self.offset), buf_value);
            debug!("triggering event {:?} to client {}", event, self.client.name());
            if let Err(e) = self.client.sender().send(event) {
                error!("expected error while sending event to client {}: {}",
                    self.client.name(), e);
            }
            self.retain = Some(buf_value);
        }
    }
}

struct Context {
    handle: fsuipc::local::LocalHandle,
    observers: Vec<Observer>,
}

impl Context {
    pub fn new()  -> io::Result<Context> {
        Ok(Context {
            handle: try!(fsuipc::local::LocalHandle::new()),
            observers: Vec::new(),
        })
    }

    fn process_write(&mut self, offset: Offset, value: Value) -> io::Result<()> {
        debug!("writing value {} to offset {}", value, offset);
        match offset.len() {
            OffsetLen::UnsignedByte => self.process_write_of(offset.addr(), u8::from(value)),
            OffsetLen::SignedByte => self.process_write_of(offset.addr(), i8::from(value)),
            OffsetLen::UnsignedWord => self.process_write_of(offset.addr(), u16::from(value)),
            OffsetLen::SignedWord => self.process_write_of(offset.addr(), i16::from(value)),
            OffsetLen::UnsignedDouble => self.process_write_of(offset.addr(), u32::from(value)),
            OffsetLen::SignedDouble => self.process_write_of(offset.addr(), i32::from(value)),
        }
    }

    fn process_obs(&mut self, offset: Offset, client: Client) {
        debug!("client {} observing offset {}", client.name(), offset);
        self.observers.push(Observer {
            offset: offset,
            client: client,
            retain: None,
            buffer: [0; 4]
        });
    }

    fn process_write_of<T>(&mut self, offset: OffsetAddr, value: T) -> io::Result<()> {
        let mut session = self.handle.session();
        try!(session.write(u16::from(offset), &value));
        try!(session.process());
        Ok(())
    }
}

fn spawn_worker() -> (thread::JoinHandle<()>, mio::Sender<Envelope>) {
    let event_loop = mio::EventLoop::new().unwrap();
    let tx = event_loop.channel();
    let worker = thread::spawn(move || {
        let mut event_loop = event_loop;
        let mut ctx = match Context::new() {
            Ok(ctx) => ctx,
            Err(e) => {
                error!("cannot create FSUIPC context (is FSUIPC installed & running?): {}", e);
                return;
            },
        };
        event_loop.timeout_ms((), POLLING_DELAY_MS).unwrap();
        event_loop.run(&mut ctx).unwrap();
    });
    (worker, tx)
}

impl mio::Handler for Context {
    type Timeout = ();
    type Message = Envelope;

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Context>, _: ()) {
        let mut session = self.handle.session();
        for obs in self.observers.iter_mut() {
            obs.read(&mut session);
        }
        session.process().unwrap();
        for obs in self.observers.iter_mut() {
            obs.trigger_event();
        }
        event_loop.timeout_ms((), POLLING_DELAY_MS).unwrap();
    }

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Context>, msg: Envelope) {
        match msg {
            Envelope::Cmd(Command::Write(Var::FsuipcOffset(offset), value)) => {
                if let Err(e) = self.process_write(offset, value) {
                    error!("unexpected IO error while processing write command: {}", e);
                }
            },
            Envelope::Cmd(Command::Observe(Var::FsuipcOffset(offset), client)) => {
                self.process_obs(offset, client)
            },
            Envelope::Shutdown => event_loop.shutdown(),
            other => {
                warn!("fsuipc domain received an unexpected message: {:?}", other);
            },
        }
    }
}

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

    #[test]
    fn should_init_and_shutdown() {
        let mut domain = Domain::new();
        domain.shutdown();
    }
}
<|endoftext|>


<commit_before>//! Unimplemented.

use super::boolean::*;
#[repr(C)]
#[derive(Copy, Clone, Debug)]
pub enum SA_TYPE {
    SA_NORESTORE = 0,
    SA_RESTORE,
    SA_DEFAULT, // was === SA_RESTORE
    SA_NOSAVE,
    SA_SAVE,
    SA_SAVEASK,
    SA_SUICIDE,
}
#[allow(non_snake_case)]
#[repr(C)]
#[derive(Copy)]
pub struct structRstart {
    pub R_Quiet: Rboolean,
    pub R_Slave: Rboolean,
    pub R_Interactive: Rboolean,
    pub R_Verbose: Rboolean,
    pub LoadSiteFile: Rboolean,
    pub LoadInitFile: Rboolean,
    pub DebugInitFile: Rboolean,
    pub RestoreAction: SA_TYPE,
    pub SaveAction: SA_TYPE,
    pub vsize: ::libc::size_t,
    pub nsize: ::libc::size_t,
    pub max_vsize: ::libc::size_t,
    pub max_nsize: ::libc::size_t,
    pub ppsize: ::libc::size_t,
    pub NoRenviron: ::libc::c_int,
}
impl ::std::clone::Clone for structRstart {
    fn clone(&self) -> Self {
        *self
    }
}
impl ::std::default::Default for structRstart {
    fn default() -> Self {
        unsafe { ::std::mem::zeroed() }
    }
}
pub type Rstart = *mut structRstart;

#[link(name = "R")]
extern {
    pub fn R_DefParams(arg1: Rstart) -> ();
    pub fn R_SetParams(arg1: Rstart) -> ();
    pub fn R_SetWin32(arg1: Rstart) -> ();
    pub fn R_SizeFromEnv(arg1: Rstart) -> ();
    pub fn R_common_command_line(arg1: *mut ::libc::c_int,
                                 arg2: *mut *mut ::libc::c_char,
                                 arg3: Rstart)
                                 -> ();
    pub fn R_set_command_line_arguments(argc: ::libc::c_int, argv: *mut *mut ::libc::c_char) -> ();
    pub fn setup_Rmainloop() -> ();
}
<commit_msg>Remove unimplemented document mark for ext/r_startup<commit_after>//! R_ext/RStartup.h
//!
//! C functions to be called from alternative front-ends.
//!
//! Part of the API for such front-ends but not for packages.
//!
use super::boolean::*;

#[repr(C)]
#[derive(Copy, Clone, Debug)]
pub enum SA_TYPE {
    SA_NORESTORE = 0,
    SA_RESTORE,
    SA_DEFAULT, // was === SA_RESTORE
    SA_NOSAVE,
    SA_SAVE,
    SA_SAVEASK,
    SA_SUICIDE,
}
#[allow(non_snake_case)]
#[repr(C)]
#[derive(Copy)]
pub struct structRstart {
    pub R_Quiet: Rboolean,
    pub R_Slave: Rboolean,
    pub R_Interactive: Rboolean,
    pub R_Verbose: Rboolean,
    pub LoadSiteFile: Rboolean,
    pub LoadInitFile: Rboolean,
    pub DebugInitFile: Rboolean,
    pub RestoreAction: SA_TYPE,
    pub SaveAction: SA_TYPE,
    pub vsize: ::libc::size_t,
    pub nsize: ::libc::size_t,
    pub max_vsize: ::libc::size_t,
    pub max_nsize: ::libc::size_t,
    pub ppsize: ::libc::size_t,
    pub NoRenviron: ::libc::c_int,
}
impl ::std::clone::Clone for structRstart {
    fn clone(&self) -> Self {
        *self
    }
}
impl ::std::default::Default for structRstart {
    fn default() -> Self {
        unsafe { ::std::mem::zeroed() }
    }
}
pub type Rstart = *mut structRstart;

#[link(name = "R")]
extern {
    pub fn R_DefParams(arg1: Rstart) -> ();
    pub fn R_SetParams(arg1: Rstart) -> ();
    pub fn R_SetWin32(arg1: Rstart) -> ();
    pub fn R_SizeFromEnv(arg1: Rstart) -> ();
    pub fn R_common_command_line(arg1: *mut ::libc::c_int,
                                 arg2: *mut *mut ::libc::c_char,
                                 arg3: Rstart)
                                 -> ();
    pub fn R_set_command_line_arguments(argc: ::libc::c_int, argv: *mut *mut ::libc::c_char) -> ();
    pub fn setup_Rmainloop() -> ();
}
<|endoftext|>


<commit_before>
<commit_msg>clean example<commit_after><|endoftext|>


<commit_before>// Copyright (C) 2019, Cloudflare, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#[macro_use]
extern crate log;

use std::net::ToSocketAddrs;

use ring::rand::*;

const MAX_DATAGRAM_SIZE: usize = 1350;

const USAGE: &str = "Usage:
  http3-client [options] URL
  http3-client -h | --help

Options:
  --method METHOD          Use the given HTTP request method [default: GET].
  --body FILE              Send the given file as request body.
  --max-data BYTES         Connection-wide flow control limit [default: 10000000].
  --max-stream-data BYTES  Per-stream flow control limit [default: 1000000].
  --wire-version VERSION   The version number to send to the server [default: babababa].
  --no-verify              Don't verify server's certificate.
  --no-grease              Don't send GREASE.
  -h --help                Show this screen.
";

fn main() {
    let mut buf = [0; 65535];
    let mut out = [0; MAX_DATAGRAM_SIZE];

    env_logger::builder()
        .default_format_timestamp_nanos(true)
        .init();

    let args = docopt::Docopt::new(USAGE)
        .and_then(|dopt| dopt.parse())
        .unwrap_or_else(|e| e.exit());

    let max_data = args.get_str("--max-data");
    let max_data = u64::from_str_radix(max_data, 10).unwrap();

    let max_stream_data = args.get_str("--max-stream-data");
    let max_stream_data = u64::from_str_radix(max_stream_data, 10).unwrap();

    let version = args.get_str("--wire-version");
    let version = u32::from_str_radix(version, 16).unwrap();

    let url = url::Url::parse(args.get_str("URL")).unwrap();

    // Setup the event loop.
    let poll = mio::Poll::new().unwrap();
    let mut events = mio::Events::with_capacity(1024);

    // Resolve server address.
    let peer_addr = url.to_socket_addrs().unwrap().next().unwrap();
    info!("connecting to {:}", peer_addr);

    // Bind to INADDR_ANY or IN6ADDR_ANY depending on the IP family of the
    // server address. This is needed on macOS and BSD variants that don't
    // support binding to IN6ADDR_ANY for both v4 and v6.
    let bind_addr = match peer_addr {
        std::net::SocketAddr::V4(_) => "0.0.0.0:0",
        std::net::SocketAddr::V6(_) => "[::]:0",
    };

    // Create the UDP socket backing the QUIC connection, and register it with
    // the event loop.
    let socket = std::net::UdpSocket::bind(bind_addr).unwrap();
    socket.connect(peer_addr).unwrap();

    let socket = mio::net::UdpSocket::from_socket(socket).unwrap();
    poll.register(
        &socket,
        mio::Token(0),
        mio::Ready::readable(),
        mio::PollOpt::edge(),
    )
    .unwrap();

    // Create the configuration for the QUIC connection.
    let mut config = quiche::Config::new(version).unwrap();

    config.verify_peer(true);

    config
        .set_application_protos(quiche::h3::APPLICATION_PROTOCOL)
        .unwrap();

    config.set_idle_timeout(5000);
    config.set_max_packet_size(MAX_DATAGRAM_SIZE as u64);
    config.set_initial_max_data(max_data);
    config.set_initial_max_stream_data_bidi_local(max_stream_data);
    config.set_initial_max_stream_data_bidi_remote(max_stream_data);
    config.set_initial_max_stream_data_uni(max_stream_data);
    config.set_initial_max_streams_bidi(100);
    config.set_initial_max_streams_uni(100);
    config.set_disable_migration(true);

    let mut http3_conn = None;

    if args.get_bool("--no-verify") {
        config.verify_peer(false);
    }

    if args.get_bool("--no-grease") {
        config.grease(false);
    }

    if std::env::var_os("SSLKEYLOGFILE").is_some() {
        config.log_keys();
    }

    // Generate a random source connection ID for the connection.
    let mut scid = [0; quiche::MAX_CONN_ID_LEN];
    SystemRandom::new().fill(&mut scid[..]).unwrap();

    // Create a QUIC connection and initiate handshake.
    let mut conn = quiche::connect(url.domain(), &scid, &mut config).unwrap();

    let write = match conn.send(&mut out) {
        Ok(v) => v,

        Err(e) => panic!("initial send failed: {:?}", e),
    };

    while let Err(e) = socket.send(&out[..write]) {
        if e.kind() == std::io::ErrorKind::WouldBlock {
            debug!("send() would block");
            continue;
        }

        panic!("send() failed: {:?}", e);
    }

    debug!("written {}", write);

    let req_start = std::time::Instant::now();

    loop {
        poll.poll(&mut events, conn.timeout()).unwrap();

        // Read incoming UDP packets from the socket and feed them to quiche,
        // until there are no more packets to read.
        'read: loop {
            // If the event loop reported no events, it means that the timeout
            // has expired, so handle it without attempting to read packets. We
            // will then proceed with the send loop.
            if events.is_empty() {
                debug!("timed out");

                conn.on_timeout();

                break 'read;
            }

            let len = match socket.recv(&mut buf) {
                Ok(v) => v,

                Err(e) => {
                    // There are no more UDP packets to read, so end the read
                    // loop.
                    if e.kind() == std::io::ErrorKind::WouldBlock {
                        debug!("recv() would block");
                        break 'read;
                    }

                    panic!("recv() failed: {:?}", e);
                },
            };

            debug!("got {} bytes", len);

            // Process potentially coalesced packets.
            let read = match conn.recv(&mut buf[..len]) {
                Ok(v) => v,

                Err(quiche::Error::Done) => {
                    debug!("done reading");
                    break;
                },

                Err(e) => {
                    error!("recv failed: {:?}", e);
                    break 'read;
                },
            };

            debug!("processed {} bytes", read);
        }

        if conn.is_closed() {
            info!("connection closed, {:?}", conn.stats());
            break;
        }

        // Create a new HTTP/3 connection and end an HTTP request as soon as
        // the QUIC connection is established.
        if conn.is_established() && http3_conn.is_none() {
            let h3_config = quiche::h3::Config::new(0, 1024, 0, 0).unwrap();

            let mut h3_conn =
                quiche::h3::Connection::with_transport(&mut conn, &h3_config)
                    .unwrap();

            let mut path = String::from(url.path());

            if let Some(query) = url.query() {
                path.push('?');
                path.push_str(query);
            }

            let mut req = vec![
                quiche::h3::Header::new(":method", args.get_str("--method")),
                quiche::h3::Header::new(":scheme", url.scheme()),
                quiche::h3::Header::new(":authority", url.host_str().unwrap()),
                quiche::h3::Header::new(":path", &path),
                quiche::h3::Header::new("user-agent", "quiche"),
            ];

            info!("sending HTTP request {:?}", req);

            let body = if args.get_bool("--body") {
                std::fs::read(args.get_str("--body")).ok()
            } else {
                None
            };

            if body.is_some() {
                req.push(quiche::h3::Header::new(
                    "content-length",
                    &body.as_ref().unwrap().len().to_string(),
                ));
            }

            let s = match h3_conn.send_request(&mut conn, &req, body.is_none()) {
                Ok(stream_id) => stream_id,

                Err(e) => {
                    error!("failed to send request {:?}", e);
                    break;
                },
            };

            if let Some(body) = body {
                if let Err(e) = h3_conn.send_body(&mut conn, s, &body, true) {
                    error!("failed to send request body {:?}", e);
                    break;
                }
            }

            http3_conn = Some(h3_conn);
        }

        if let Some(http3_conn) = &mut http3_conn {
            // Process HTTP/3 events.
            loop {
                match http3_conn.poll(&mut conn) {
                    Ok((stream_id, quiche::h3::Event::Headers(headers))) => {
                        info!(
                            "got response headers {:?} on stream id {}",
                            headers, stream_id
                        );
                    },

                    Ok((stream_id, quiche::h3::Event::Data)) => {
                        if let Ok(read) =
                            http3_conn.recv_body(&mut conn, stream_id, &mut buf)
                        {
                            debug!(
                                "got {} bytes of response data on stream {}",
                                read, stream_id
                            );

                            print!("{}", unsafe {
                                std::str::from_utf8_unchecked(&buf[..read])
                            });
                        }
                    },

                    Ok((_stream_id, quiche::h3::Event::Finished)) => {
                        info!(
                            "response received in {:?}, closing...",
                            req_start.elapsed()
                        );

                        match conn.close(true, 0x00, b"kthxbye") {
                            // Already closed.
                            Ok(_) | Err(quiche::Error::Done) => (),

                            Err(e) => panic!("error closing conn: {:?}", e),
                        }

                        break;
                    },

                    Err(quiche::h3::Error::Done) => {
                        break;
                    },

                    Err(e) => {
                        error!("HTTP/3 processing failed: {:?}", e);

                        break;
                    },
                }
            }
        }

        // Generate outgoing QUIC packets and send them on the UDP socket, until
        // quiche reports that there are no more packets to be sent.
        loop {
            let write = match conn.send(&mut out) {
                Ok(v) => v,

                Err(quiche::Error::Done) => {
                    debug!("done writing");
                    break;
                },

                Err(e) => {
                    error!("send failed: {:?}", e);
                    conn.close(false, 0x1, b"fail").ok();
                    break;
                },
            };

            if let Err(e) = socket.send(&out[..write]) {
                if e.kind() == std::io::ErrorKind::WouldBlock {
                    debug!("send() would block");
                    break;
                }

                panic!("send() failed: {:?}", e);
            }

            debug!("written {}", write);
        }

        if conn.is_closed() {
            info!("connection closed, {:?}", conn.stats());
            break;
        }
    }
}
<commit_msg>http3-client: add request headers<commit_after>// Copyright (C) 2019, Cloudflare, Inc.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright notice,
//       this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#[macro_use]
extern crate log;

use std::net::ToSocketAddrs;

use ring::rand::*;

const MAX_DATAGRAM_SIZE: usize = 1350;

const USAGE: &str = "Usage:
  http3-client [options] URL
  http3-client -h | --help

Options:
  --method METHOD          Use the given HTTP request method [default: GET].
  --body FILE              Send the given file as request body.
  --max-data BYTES         Connection-wide flow control limit [default: 10000000].
  --max-stream-data BYTES  Per-stream flow control limit [default: 1000000].
  --wire-version VERSION   The version number to send to the server [default: babababa].
  --no-verify              Don't verify server's certificate.
  --no-grease              Don't send GREASE.
  -H --header HEADER ...   Add a request header.
  -h --help                Show this screen.
";

fn main() {
    let mut buf = [0; 65535];
    let mut out = [0; MAX_DATAGRAM_SIZE];

    env_logger::builder()
        .default_format_timestamp_nanos(true)
        .init();

    let args = docopt::Docopt::new(USAGE)
        .and_then(|dopt| dopt.parse())
        .unwrap_or_else(|e| e.exit());

    let max_data = args.get_str("--max-data");
    let max_data = u64::from_str_radix(max_data, 10).unwrap();

    let max_stream_data = args.get_str("--max-stream-data");
    let max_stream_data = u64::from_str_radix(max_stream_data, 10).unwrap();

    let version = args.get_str("--wire-version");
    let version = u32::from_str_radix(version, 16).unwrap();

    let url = url::Url::parse(args.get_str("URL")).unwrap();

    // Request headers (can be multiple).
    let req_headers = args.get_vec("--header");

    // Setup the event loop.
    let poll = mio::Poll::new().unwrap();
    let mut events = mio::Events::with_capacity(1024);

    // Resolve server address.
    let peer_addr = url.to_socket_addrs().unwrap().next().unwrap();
    info!("connecting to {:}", peer_addr);

    // Bind to INADDR_ANY or IN6ADDR_ANY depending on the IP family of the
    // server address. This is needed on macOS and BSD variants that don't
    // support binding to IN6ADDR_ANY for both v4 and v6.
    let bind_addr = match peer_addr {
        std::net::SocketAddr::V4(_) => "0.0.0.0:0",
        std::net::SocketAddr::V6(_) => "[::]:0",
    };

    // Create the UDP socket backing the QUIC connection, and register it with
    // the event loop.
    let socket = std::net::UdpSocket::bind(bind_addr).unwrap();
    socket.connect(peer_addr).unwrap();

    let socket = mio::net::UdpSocket::from_socket(socket).unwrap();
    poll.register(
        &socket,
        mio::Token(0),
        mio::Ready::readable(),
        mio::PollOpt::edge(),
    )
    .unwrap();

    // Create the configuration for the QUIC connection.
    let mut config = quiche::Config::new(version).unwrap();

    config.verify_peer(true);

    config
        .set_application_protos(quiche::h3::APPLICATION_PROTOCOL)
        .unwrap();

    config.set_idle_timeout(5000);
    config.set_max_packet_size(MAX_DATAGRAM_SIZE as u64);
    config.set_initial_max_data(max_data);
    config.set_initial_max_stream_data_bidi_local(max_stream_data);
    config.set_initial_max_stream_data_bidi_remote(max_stream_data);
    config.set_initial_max_stream_data_uni(max_stream_data);
    config.set_initial_max_streams_bidi(100);
    config.set_initial_max_streams_uni(100);
    config.set_disable_migration(true);

    let mut http3_conn = None;

    if args.get_bool("--no-verify") {
        config.verify_peer(false);
    }

    if args.get_bool("--no-grease") {
        config.grease(false);
    }

    if std::env::var_os("SSLKEYLOGFILE").is_some() {
        config.log_keys();
    }

    // Generate a random source connection ID for the connection.
    let mut scid = [0; quiche::MAX_CONN_ID_LEN];
    SystemRandom::new().fill(&mut scid[..]).unwrap();

    // Create a QUIC connection and initiate handshake.
    let mut conn = quiche::connect(url.domain(), &scid, &mut config).unwrap();

    let write = match conn.send(&mut out) {
        Ok(v) => v,

        Err(e) => panic!("initial send failed: {:?}", e),
    };

    while let Err(e) = socket.send(&out[..write]) {
        if e.kind() == std::io::ErrorKind::WouldBlock {
            debug!("send() would block");
            continue;
        }

        panic!("send() failed: {:?}", e);
    }

    debug!("written {}", write);

    let req_start = std::time::Instant::now();

    loop {
        poll.poll(&mut events, conn.timeout()).unwrap();

        // Read incoming UDP packets from the socket and feed them to quiche,
        // until there are no more packets to read.
        'read: loop {
            // If the event loop reported no events, it means that the timeout
            // has expired, so handle it without attempting to read packets. We
            // will then proceed with the send loop.
            if events.is_empty() {
                debug!("timed out");

                conn.on_timeout();

                break 'read;
            }

            let len = match socket.recv(&mut buf) {
                Ok(v) => v,

                Err(e) => {
                    // There are no more UDP packets to read, so end the read
                    // loop.
                    if e.kind() == std::io::ErrorKind::WouldBlock {
                        debug!("recv() would block");
                        break 'read;
                    }

                    panic!("recv() failed: {:?}", e);
                },
            };

            debug!("got {} bytes", len);

            // Process potentially coalesced packets.
            let read = match conn.recv(&mut buf[..len]) {
                Ok(v) => v,

                Err(quiche::Error::Done) => {
                    debug!("done reading");
                    break;
                },

                Err(e) => {
                    error!("recv failed: {:?}", e);
                    break 'read;
                },
            };

            debug!("processed {} bytes", read);
        }

        if conn.is_closed() {
            info!("connection closed, {:?}", conn.stats());
            break;
        }

        // Create a new HTTP/3 connection and end an HTTP request as soon as
        // the QUIC connection is established.
        if conn.is_established() && http3_conn.is_none() {
            let h3_config = quiche::h3::Config::new(0, 1024, 0, 0).unwrap();

            let mut h3_conn =
                quiche::h3::Connection::with_transport(&mut conn, &h3_config)
                    .unwrap();

            let mut path = String::from(url.path());

            if let Some(query) = url.query() {
                path.push('?');
                path.push_str(query);
            }

            let mut req = vec![
                quiche::h3::Header::new(":method", args.get_str("--method")),
                quiche::h3::Header::new(":scheme", url.scheme()),
                quiche::h3::Header::new(":authority", url.host_str().unwrap()),
                quiche::h3::Header::new(":path", &path),
                quiche::h3::Header::new("user-agent", "quiche"),
            ];

            // Add custom headers to the request.
            for header in &req_headers {
                let header_split: Vec<&str> = header.splitn(2, ": ").collect();
                // Convert header field to lowercase
                let (field, value) =
                    (&header_split[0].to_lowercase().to_string(), header_split[1]);
                req.push(quiche::h3::Header::new(field, value));
            }

            info!("sending HTTP request {:?}", req);

            let body = if args.get_bool("--body") {
                std::fs::read(args.get_str("--body")).ok()
            } else {
                None
            };

            if body.is_some() {
                req.push(quiche::h3::Header::new(
                    "content-length",
                    &body.as_ref().unwrap().len().to_string(),
                ));
            }

            let s = match h3_conn.send_request(&mut conn, &req, body.is_none()) {
                Ok(stream_id) => stream_id,

                Err(e) => {
                    error!("failed to send request {:?}", e);
                    break;
                },
            };

            if let Some(body) = body {
                if let Err(e) = h3_conn.send_body(&mut conn, s, &body, true) {
                    error!("failed to send request body {:?}", e);
                    break;
                }
            }

            http3_conn = Some(h3_conn);
        }

        if let Some(http3_conn) = &mut http3_conn {
            // Process HTTP/3 events.
            loop {
                match http3_conn.poll(&mut conn) {
                    Ok((stream_id, quiche::h3::Event::Headers(headers))) => {
                        info!(
                            "got response headers {:?} on stream id {}",
                            headers, stream_id
                        );
                    },

                    Ok((stream_id, quiche::h3::Event::Data)) => {
                        if let Ok(read) =
                            http3_conn.recv_body(&mut conn, stream_id, &mut buf)
                        {
                            debug!(
                                "got {} bytes of response data on stream {}",
                                read, stream_id
                            );

                            print!("{}", unsafe {
                                std::str::from_utf8_unchecked(&buf[..read])
                            });
                        }
                    },

                    Ok((_stream_id, quiche::h3::Event::Finished)) => {
                        info!(
                            "response received in {:?}, closing...",
                            req_start.elapsed()
                        );

                        match conn.close(true, 0x00, b"kthxbye") {
                            // Already closed.
                            Ok(_) | Err(quiche::Error::Done) => (),

                            Err(e) => panic!("error closing conn: {:?}", e),
                        }

                        break;
                    },

                    Err(quiche::h3::Error::Done) => {
                        break;
                    },

                    Err(e) => {
                        error!("HTTP/3 processing failed: {:?}", e);

                        break;
                    },
                }
            }
        }

        // Generate outgoing QUIC packets and send them on the UDP socket, until
        // quiche reports that there are no more packets to be sent.
        loop {
            let write = match conn.send(&mut out) {
                Ok(v) => v,

                Err(quiche::Error::Done) => {
                    debug!("done writing");
                    break;
                },

                Err(e) => {
                    error!("send failed: {:?}", e);
                    conn.close(false, 0x1, b"fail").ok();
                    break;
                },
            };

            if let Err(e) = socket.send(&out[..write]) {
                if e.kind() == std::io::ErrorKind::WouldBlock {
                    debug!("send() would block");
                    break;
                }

                panic!("send() failed: {:?}", e);
            }

            debug!("written {}", write);
        }

        if conn.is_closed() {
            info!("connection closed, {:?}", conn.stats());
            break;
        }
    }
}
<|endoftext|>


<commit_before>use crate::context::Context;
use crate::error::RenderError;
use crate::helpers::{HelperDef, HelperResult};
use crate::json::value::JsonTruthy;
use crate::output::Output;
use crate::registry::Registry;
use crate::render::{Helper, RenderContext, Renderable};

#[derive(Clone, Copy)]
pub struct IfHelper {
    positive: bool,
}

impl HelperDef for IfHelper {
    fn call<'reg: 'rc, 'rc>(
        &self,
        h: &Helper<'reg, 'rc>,
        r: &'reg Registry<'reg>,
        ctx: &'rc Context,
        rc: &mut RenderContext<'reg, 'rc>,
        out: &mut dyn Output,
    ) -> HelperResult {
        let param = h
            .param(0)
            .ok_or_else(|| RenderError::new("Param not found for helper \"if\""))?;
        let include_zero = h
            .hash_get("includeZero")
            .and_then(|v| v.value().as_bool())
            .unwrap_or(false);

        let mut value = param.value().is_truthy(include_zero);

        if !self.positive {
            value = !value;
        }

        let tmpl = if value { h.template() } else { h.inverse() };
        match tmpl {
            Some(ref t) => t.render(r, ctx, rc, out),
            None => Ok(()),
        }
    }
}

pub static IF_HELPER: IfHelper = IfHelper { positive: true };
pub static UNLESS_HELPER: IfHelper = IfHelper { positive: false };

#[cfg(test)]
mod test {
    use crate::helpers::WITH_HELPER;
    use crate::registry::Registry;
    use serde_json::value::Value as Json;
    use std::str::FromStr;

    #[test]
    fn test_if() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#if this}}hello{{/if}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t1", "{{#unless this}}hello{{else}}world{{/unless}}")
            .is_ok());

        let r0 = handlebars.render("t0", &true);
        assert_eq!(r0.ok().unwrap(), "hello".to_string());

        let r1 = handlebars.render("t1", &true);
        assert_eq!(r1.ok().unwrap(), "world".to_string());

        let r2 = handlebars.render("t0", &false);
        assert_eq!(r2.ok().unwrap(), "".to_string());
    }

    #[test]
    fn test_if_context() {
        let json_str = r#"{"a":{"b":99,"c":{"d": true}}}"#;
        let data = Json::from_str(json_str).unwrap();

        let mut handlebars = Registry::new();
        handlebars.register_helper("with", Box::new(WITH_HELPER));
        assert!(handlebars
            .register_template_string("t0", "{{#if a.c.d}}hello {{a.b}}{{/if}}")
            .is_ok());
        assert!(handlebars
            .register_template_string(
                "t1",
                "{{#with a}}{{#if c.d}}hello {{../a.b}}{{/if}}{{/with}}"
            )
            .is_ok());

        let r0 = handlebars.render("t0", &data);
        assert_eq!(r0.unwrap(), "hello 99".to_string());

        let r1 = handlebars.render("t1", &data);
        assert_eq!(r1.unwrap(), "hello 99".to_string());
    }

    #[test]
    fn test_if_include_zero() {
        use std::f64;
        let handlebars = Registry::new();

        assert_eq!(
            "0".to_owned(),
            handlebars
                .render_template("{{#if a}}1{{else}}0{{/if}}", &json!({"a": 0}))
                .unwrap()
        );
        assert_eq!(
            "1".to_owned(),
            handlebars
                .render_template(
                    "{{#if a includeZero=true}}1{{else}}0{{/if}}",
                    &json!({"a": 0})
                )
                .unwrap()
        );
        assert_eq!(
            "0".to_owned(),
            handlebars
                .render_template(
                    "{{#if a includeZero=true}}1{{else}}0{{/if}}",
                    &json!({ "a": f64::NAN })
                )
                .unwrap()
        );
    }

    #[test]
    fn test_invisible_line_stripping() {
        let hbs = Registry::new();
        assert_eq!(
            "yes\n",
            hbs.render_template("{{#if a}}\nyes\n{{/if}}\n", &json!({"a": true}))
                .unwrap()
        );

        assert_eq!(
            "x\ny",
            hbs.render_template("{{#if a}}x{{/if}}\ny", &json!({"a": true}))
                .unwrap()
        );

        assert_eq!(
            "y\nz",
            hbs.render_template("{{#if a}}\nx\n{{^}}\ny\n{{/if}}\nz", &json!({"a": false}))
                .unwrap()
        );

        assert_eq!(
            r#"  foo
  bar
  baz"#,
            hbs.render_template(
                r#"  {{#if true}}
  foo
  bar
  {{/if}}
  baz"#,
                &json!({})
            )
            .unwrap()
        );
    }
}
<commit_msg>(test) converge the root cause of leading whitespace issue<commit_after>use crate::context::Context;
use crate::error::RenderError;
use crate::helpers::{HelperDef, HelperResult};
use crate::json::value::JsonTruthy;
use crate::output::Output;
use crate::registry::Registry;
use crate::render::{Helper, RenderContext, Renderable};

#[derive(Clone, Copy)]
pub struct IfHelper {
    positive: bool,
}

impl HelperDef for IfHelper {
    fn call<'reg: 'rc, 'rc>(
        &self,
        h: &Helper<'reg, 'rc>,
        r: &'reg Registry<'reg>,
        ctx: &'rc Context,
        rc: &mut RenderContext<'reg, 'rc>,
        out: &mut dyn Output,
    ) -> HelperResult {
        let param = h
            .param(0)
            .ok_or_else(|| RenderError::new("Param not found for helper \"if\""))?;
        let include_zero = h
            .hash_get("includeZero")
            .and_then(|v| v.value().as_bool())
            .unwrap_or(false);

        let mut value = param.value().is_truthy(include_zero);

        if !self.positive {
            value = !value;
        }

        let tmpl = if value { h.template() } else { h.inverse() };
        match tmpl {
            Some(ref t) => t.render(r, ctx, rc, out),
            None => Ok(()),
        }
    }
}

pub static IF_HELPER: IfHelper = IfHelper { positive: true };
pub static UNLESS_HELPER: IfHelper = IfHelper { positive: false };

#[cfg(test)]
mod test {
    use crate::helpers::WITH_HELPER;
    use crate::registry::Registry;
    use serde_json::value::Value as Json;
    use std::str::FromStr;

    #[test]
    fn test_if() {
        let mut handlebars = Registry::new();
        assert!(handlebars
            .register_template_string("t0", "{{#if this}}hello{{/if}}")
            .is_ok());
        assert!(handlebars
            .register_template_string("t1", "{{#unless this}}hello{{else}}world{{/unless}}")
            .is_ok());

        let r0 = handlebars.render("t0", &true);
        assert_eq!(r0.ok().unwrap(), "hello".to_string());

        let r1 = handlebars.render("t1", &true);
        assert_eq!(r1.ok().unwrap(), "world".to_string());

        let r2 = handlebars.render("t0", &false);
        assert_eq!(r2.ok().unwrap(), "".to_string());
    }

    #[test]
    fn test_if_context() {
        let json_str = r#"{"a":{"b":99,"c":{"d": true}}}"#;
        let data = Json::from_str(json_str).unwrap();

        let mut handlebars = Registry::new();
        handlebars.register_helper("with", Box::new(WITH_HELPER));
        assert!(handlebars
            .register_template_string("t0", "{{#if a.c.d}}hello {{a.b}}{{/if}}")
            .is_ok());
        assert!(handlebars
            .register_template_string(
                "t1",
                "{{#with a}}{{#if c.d}}hello {{../a.b}}{{/if}}{{/with}}"
            )
            .is_ok());

        let r0 = handlebars.render("t0", &data);
        assert_eq!(r0.unwrap(), "hello 99".to_string());

        let r1 = handlebars.render("t1", &data);
        assert_eq!(r1.unwrap(), "hello 99".to_string());
    }

    #[test]
    fn test_if_include_zero() {
        use std::f64;
        let handlebars = Registry::new();

        assert_eq!(
            "0".to_owned(),
            handlebars
                .render_template("{{#if a}}1{{else}}0{{/if}}", &json!({"a": 0}))
                .unwrap()
        );
        assert_eq!(
            "1".to_owned(),
            handlebars
                .render_template(
                    "{{#if a includeZero=true}}1{{else}}0{{/if}}",
                    &json!({"a": 0})
                )
                .unwrap()
        );
        assert_eq!(
            "0".to_owned(),
            handlebars
                .render_template(
                    "{{#if a includeZero=true}}1{{else}}0{{/if}}",
                    &json!({ "a": f64::NAN })
                )
                .unwrap()
        );
    }

    #[test]
    fn test_invisible_line_stripping() {
        let hbs = Registry::new();
        assert_eq!(
            "yes\n",
            hbs.render_template("{{#if a}}\nyes\n{{/if}}\n", &json!({"a": true}))
                .unwrap()
        );

        assert_eq!(
            "x\ny",
            hbs.render_template("{{#if a}}x{{/if}}\ny", &json!({"a": true}))
                .unwrap()
        );

        assert_eq!(
            "y\nz",
            hbs.render_template("{{#if a}}\nx\n{{^}}\ny\n{{/if}}\nz", &json!({"a": false}))
                .unwrap()
        );

        assert_eq!(
            r#"yes
  foo
  bar
  baz"#,
            hbs.render_template(
                r#"yes
  {{#if true}}
  foo
  bar
  {{/if}}
  baz"#,
                &json!({})
            )
            .unwrap()
        );

        assert_eq!(
            r#"  foo
  bar
  baz"#,
            hbs.render_template(
                r#"  {{#if true}}
  foo
  bar
  {{/if}}
  baz"#,
                &json!({})
            )
            .unwrap()
        );
    }
}
<|endoftext|>