<commit_before>//! Memory allocation APIs
//!
//! In a given program, the standard library has one “global” memory allocator
//! that is used for example by `Box<T>` and `Vec<T>`.
//!
//! Currently the default global allocator is unspecified. Libraries, however,
//! like `cdylib`s and `staticlib`s are guaranteed to use the [`System`] by
//! default.
//!
//! [`System`]: struct.System.html
//!
//! # The `#[global_allocator]` attribute
//!
//! This attribute allows configuring the choice of global allocator.
//! You can use this to implement a completely custom global allocator
//! to route all default allocation requests to a custom object.
//!
//! ```rust
//! use std::alloc::{GlobalAlloc, System, Layout};
//!
//! struct MyAllocator;
//!
//! unsafe impl GlobalAlloc for MyAllocator {
//!     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
//!         System.alloc(layout)
//!     }
//!
//!     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
//!         System.dealloc(ptr, layout)
//!     }
//! }
//!
//! #[global_allocator]
//! static GLOBAL: MyAllocator = MyAllocator;
//!
//! fn main() {
//!     // This `Vec` will allocate memory through `GLOBAL` above
//!     let mut v = Vec::new();
//!     v.push(1);
//! }
//! ```
//!
//! The attribute is used on a `static` item whose type implements the
//! [`GlobalAlloc`] trait. This type can be provided by an external library:
//!
//! [`GlobalAlloc`]: ../../core/alloc/trait.GlobalAlloc.html
//!
//! ```rust,ignore (demonstrates crates.io usage)
//! extern crate jemallocator;
//!
//! use jemallocator::Jemalloc;
//!
//! #[global_allocator]
//! static GLOBAL: Jemalloc = Jemalloc;
//!
//! fn main() {}
//! ```
//!
//! The `#[global_allocator]` can only be used once in a crate
//! or its recursive dependencies.

#![stable(feature = "alloc_module", since = "1.28.0")]

use core::sync::atomic::{AtomicPtr, Ordering};
use core::{mem, ptr};
use core::ptr::NonNull;

use crate::sys_common::util::dumb_print;

#[stable(feature = "alloc_module", since = "1.28.0")]
#[doc(inline)]
pub use alloc_crate::alloc::*;

/// The default memory allocator provided by the operating system.
///
/// This is based on `malloc` on Unix platforms and `HeapAlloc` on Windows,
/// plus related functions.
///
/// This type implements the `GlobalAlloc` trait and Rust programs by default
/// work as if they had this definition:
///
/// ```rust
/// use std::alloc::System;
///
/// #[global_allocator]
/// static A: System = System;
///
/// fn main() {
///     let a = Box::new(4); // Allocates from the system allocator.
///     println!("{}", a);
/// }
/// ```
///
/// You can also define your own wrapper around `System` if you'd like, such as
/// keeping track of the number of all bytes allocated:
///
/// ```rust
/// use std::alloc::{System, GlobalAlloc, Layout};
/// use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
///
/// struct Counter;
///
/// static ALLOCATED: AtomicUsize = AtomicUsize::new(0);
///
/// unsafe impl GlobalAlloc for Counter {
///     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
///         let ret = System.alloc(layout);
///         if !ret.is_null() {
///             ALLOCATED.fetch_add(layout.size(), SeqCst);
///         }
///         return ret
///     }
///
///     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
///         System.dealloc(ptr, layout);
///         ALLOCATED.fetch_sub(layout.size(), SeqCst);
///     }
/// }
///
/// #[global_allocator]
/// static A: Counter = Counter;
///
/// fn main() {
///     println!("allocated bytes before main: {}", ALLOCATED.load(SeqCst));
/// }
/// ```
///
/// It can also be used directly to allocate memory independently of whatever
/// global allocator has been selected for a Rust program. For example if a Rust
/// program opts in to using jemalloc as the global allocator, `System` will
/// still allocate memory using `malloc` and `HeapAlloc`.
#[stable(feature = "alloc_system_type", since = "1.28.0")]
#[derive(Debug, Copy, Clone)]
pub struct System;

// The Alloc impl just forwards to the GlobalAlloc impl, which is in `std::sys::*::alloc`.
#[unstable(feature = "allocator_api", issue = "32838")]
unsafe impl Alloc for System {
    #[inline]
    unsafe fn alloc(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
        NonNull::new(GlobalAlloc::alloc(self, layout)).ok_or(AllocErr)
    }

    #[inline]
    unsafe fn alloc_zeroed(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
        NonNull::new(GlobalAlloc::alloc_zeroed(self, layout)).ok_or(AllocErr)
    }

    #[inline]
    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        GlobalAlloc::dealloc(self, ptr.as_ptr(), layout)
    }

    #[inline]
    unsafe fn realloc(&mut self,
                      ptr: NonNull<u8>,
                      layout: Layout,
                      new_size: usize) -> Result<NonNull<u8>, AllocErr> {
        NonNull::new(GlobalAlloc::realloc(self, ptr.as_ptr(), layout, new_size)).ok_or(AllocErr)
    }
}

static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());

/// Registers a custom allocation error hook, replacing any that was previously registered.
///
/// The allocation error hook is invoked when an infallible memory allocation fails, before
/// the runtime aborts. The default hook prints a message to standard error,
/// but this behavior can be customized with the [`set_alloc_error_hook`] and
/// [`take_alloc_error_hook`] functions.
///
/// The hook is provided with a `Layout` struct which contains information
/// about the allocation that failed.
///
/// The allocation error hook is a global resource.
#[unstable(feature = "alloc_error_hook", issue = "51245")]
pub fn set_alloc_error_hook(hook: fn(Layout)) {
    HOOK.store(hook as *mut (), Ordering::SeqCst);
}

/// Unregisters the current allocation error hook, returning it.
///
/// *See also the function [`set_alloc_error_hook`].*
///
/// If no custom hook is registered, the default hook will be returned.
#[unstable(feature = "alloc_error_hook", issue = "51245")]
pub fn take_alloc_error_hook() -> fn(Layout) {
    let hook = HOOK.swap(ptr::null_mut(), Ordering::SeqCst);
    if hook.is_null() {
        default_alloc_error_hook
    } else {
        unsafe { mem::transmute(hook) }
    }
}

fn default_alloc_error_hook(layout: Layout) {
    dumb_print(format_args!("memory allocation of {} bytes failed", layout.size()));
}

#[cfg(not(test))]
#[doc(hidden)]
#[alloc_error_handler]
#[unstable(feature = "alloc_internals", issue = "0")]
pub fn rust_oom(layout: Layout) -> ! {
    let hook = HOOK.load(Ordering::SeqCst);
    let hook: fn(Layout) = if hook.is_null() {
        default_alloc_error_hook
    } else {
        unsafe { mem::transmute(hook) }
    };
    hook(layout);
    unsafe { crate::sys::abort_internal(); }
}

#[cfg(not(test))]
#[doc(hidden)]
#[allow(unused_attributes)]
#[unstable(feature = "alloc_internals", issue = "0")]
pub mod __default_lib_allocator {
    use super::{System, Layout, GlobalAlloc};
    // These magic symbol names are used as a fallback for implementing the
    // `__rust_alloc` etc symbols (see `src/liballoc/alloc.rs) when there is
    // no `#[global_allocator]` attribute.

    // for symbol names src/librustc/middle/allocator.rs
    // for signatures src/librustc_allocator/lib.rs

    // linkage directives are provided as part of the current compiler allocator
    // ABI

    #[rustc_std_internal_symbol]
    pub unsafe extern fn __rdl_alloc(size: usize, align: usize) -> *mut u8 {
        let layout = Layout::from_size_align_unchecked(size, align);
        System.alloc(layout)
    }

    #[rustc_std_internal_symbol]
    pub unsafe extern fn __rdl_dealloc(ptr: *mut u8,
                                       size: usize,
                                       align: usize) {
        System.dealloc(ptr, Layout::from_size_align_unchecked(size, align))
    }

    #[rustc_std_internal_symbol]
    pub unsafe extern fn __rdl_realloc(ptr: *mut u8,
                                       old_size: usize,
                                       align: usize,
                                       new_size: usize) -> *mut u8 {
        let old_layout = Layout::from_size_align_unchecked(old_size, align);
        System.realloc(ptr, old_layout, new_size)
    }

    #[rustc_std_internal_symbol]
    pub unsafe extern fn __rdl_alloc_zeroed(size: usize, align: usize) -> *mut u8 {
        let layout = Layout::from_size_align_unchecked(size, align);
        System.alloc_zeroed(layout)
    }
}
<commit_msg>Rollup merge of #59451 - TimDiekmann:patch-1, r=sfackler<commit_after>//! Memory allocation APIs
//!
//! In a given program, the standard library has one “global” memory allocator
//! that is used for example by `Box<T>` and `Vec<T>`.
//!
//! Currently the default global allocator is unspecified. Libraries, however,
//! like `cdylib`s and `staticlib`s are guaranteed to use the [`System`] by
//! default.
//!
//! [`System`]: struct.System.html
//!
//! # The `#[global_allocator]` attribute
//!
//! This attribute allows configuring the choice of global allocator.
//! You can use this to implement a completely custom global allocator
//! to route all default allocation requests to a custom object.
//!
//! ```rust
//! use std::alloc::{GlobalAlloc, System, Layout};
//!
//! struct MyAllocator;
//!
//! unsafe impl GlobalAlloc for MyAllocator {
//!     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
//!         System.alloc(layout)
//!     }
//!
//!     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
//!         System.dealloc(ptr, layout)
//!     }
//! }
//!
//! #[global_allocator]
//! static GLOBAL: MyAllocator = MyAllocator;
//!
//! fn main() {
//!     // This `Vec` will allocate memory through `GLOBAL` above
//!     let mut v = Vec::new();
//!     v.push(1);
//! }
//! ```
//!
//! The attribute is used on a `static` item whose type implements the
//! [`GlobalAlloc`] trait. This type can be provided by an external library:
//!
//! [`GlobalAlloc`]: ../../core/alloc/trait.GlobalAlloc.html
//!
//! ```rust,ignore (demonstrates crates.io usage)
//! extern crate jemallocator;
//!
//! use jemallocator::Jemalloc;
//!
//! #[global_allocator]
//! static GLOBAL: Jemalloc = Jemalloc;
//!
//! fn main() {}
//! ```
//!
//! The `#[global_allocator]` can only be used once in a crate
//! or its recursive dependencies.

#![stable(feature = "alloc_module", since = "1.28.0")]

use core::sync::atomic::{AtomicPtr, Ordering};
use core::{mem, ptr};
use core::ptr::NonNull;

use crate::sys_common::util::dumb_print;

#[stable(feature = "alloc_module", since = "1.28.0")]
#[doc(inline)]
pub use alloc_crate::alloc::*;

/// The default memory allocator provided by the operating system.
///
/// This is based on `malloc` on Unix platforms and `HeapAlloc` on Windows,
/// plus related functions.
///
/// This type implements the `GlobalAlloc` trait and Rust programs by default
/// work as if they had this definition:
///
/// ```rust
/// use std::alloc::System;
///
/// #[global_allocator]
/// static A: System = System;
///
/// fn main() {
///     let a = Box::new(4); // Allocates from the system allocator.
///     println!("{}", a);
/// }
/// ```
///
/// You can also define your own wrapper around `System` if you'd like, such as
/// keeping track of the number of all bytes allocated:
///
/// ```rust
/// use std::alloc::{System, GlobalAlloc, Layout};
/// use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
///
/// struct Counter;
///
/// static ALLOCATED: AtomicUsize = AtomicUsize::new(0);
///
/// unsafe impl GlobalAlloc for Counter {
///     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
///         let ret = System.alloc(layout);
///         if !ret.is_null() {
///             ALLOCATED.fetch_add(layout.size(), SeqCst);
///         }
///         return ret
///     }
///
///     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
///         System.dealloc(ptr, layout);
///         ALLOCATED.fetch_sub(layout.size(), SeqCst);
///     }
/// }
///
/// #[global_allocator]
/// static A: Counter = Counter;
///
/// fn main() {
///     println!("allocated bytes before main: {}", ALLOCATED.load(SeqCst));
/// }
/// ```
///
/// It can also be used directly to allocate memory independently of whatever
/// global allocator has been selected for a Rust program. For example if a Rust
/// program opts in to using jemalloc as the global allocator, `System` will
/// still allocate memory using `malloc` and `HeapAlloc`.
#[stable(feature = "alloc_system_type", since = "1.28.0")]
#[derive(Debug, Default, Copy, Clone)]
pub struct System;

// The Alloc impl just forwards to the GlobalAlloc impl, which is in `std::sys::*::alloc`.
#[unstable(feature = "allocator_api", issue = "32838")]
unsafe impl Alloc for System {
    #[inline]
    unsafe fn alloc(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
        NonNull::new(GlobalAlloc::alloc(self, layout)).ok_or(AllocErr)
    }

    #[inline]
    unsafe fn alloc_zeroed(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
        NonNull::new(GlobalAlloc::alloc_zeroed(self, layout)).ok_or(AllocErr)
    }

    #[inline]
    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        GlobalAlloc::dealloc(self, ptr.as_ptr(), layout)
    }

    #[inline]
    unsafe fn realloc(&mut self,
                      ptr: NonNull<u8>,
                      layout: Layout,
                      new_size: usize) -> Result<NonNull<u8>, AllocErr> {
        NonNull::new(GlobalAlloc::realloc(self, ptr.as_ptr(), layout, new_size)).ok_or(AllocErr)
    }
}

static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());

/// Registers a custom allocation error hook, replacing any that was previously registered.
///
/// The allocation error hook is invoked when an infallible memory allocation fails, before
/// the runtime aborts. The default hook prints a message to standard error,
/// but this behavior can be customized with the [`set_alloc_error_hook`] and
/// [`take_alloc_error_hook`] functions.
///
/// The hook is provided with a `Layout` struct which contains information
/// about the allocation that failed.
///
/// The allocation error hook is a global resource.
#[unstable(feature = "alloc_error_hook", issue = "51245")]
pub fn set_alloc_error_hook(hook: fn(Layout)) {
    HOOK.store(hook as *mut (), Ordering::SeqCst);
}

/// Unregisters the current allocation error hook, returning it.
///
/// *See also the function [`set_alloc_error_hook`].*
///
/// If no custom hook is registered, the default hook will be returned.
#[unstable(feature = "alloc_error_hook", issue = "51245")]
pub fn take_alloc_error_hook() -> fn(Layout) {
    let hook = HOOK.swap(ptr::null_mut(), Ordering::SeqCst);
    if hook.is_null() {
        default_alloc_error_hook
    } else {
        unsafe { mem::transmute(hook) }
    }
}

fn default_alloc_error_hook(layout: Layout) {
    dumb_print(format_args!("memory allocation of {} bytes failed", layout.size()));
}

#[cfg(not(test))]
#[doc(hidden)]
#[alloc_error_handler]
#[unstable(feature = "alloc_internals", issue = "0")]
pub fn rust_oom(layout: Layout) -> ! {
    let hook = HOOK.load(Ordering::SeqCst);
    let hook: fn(Layout) = if hook.is_null() {
        default_alloc_error_hook
    } else {
        unsafe { mem::transmute(hook) }
    };
    hook(layout);
    unsafe { crate::sys::abort_internal(); }
}

#[cfg(not(test))]
#[doc(hidden)]
#[allow(unused_attributes)]
#[unstable(feature = "alloc_internals", issue = "0")]
pub mod __default_lib_allocator {
    use super::{System, Layout, GlobalAlloc};
    // These magic symbol names are used as a fallback for implementing the
    // `__rust_alloc` etc symbols (see `src/liballoc/alloc.rs) when there is
    // no `#[global_allocator]` attribute.

    // for symbol names src/librustc/middle/allocator.rs
    // for signatures src/librustc_allocator/lib.rs

    // linkage directives are provided as part of the current compiler allocator
    // ABI

    #[rustc_std_internal_symbol]
    pub unsafe extern fn __rdl_alloc(size: usize, align: usize) -> *mut u8 {
        let layout = Layout::from_size_align_unchecked(size, align);
        System.alloc(layout)
    }

    #[rustc_std_internal_symbol]
    pub unsafe extern fn __rdl_dealloc(ptr: *mut u8,
                                       size: usize,
                                       align: usize) {
        System.dealloc(ptr, Layout::from_size_align_unchecked(size, align))
    }

    #[rustc_std_internal_symbol]
    pub unsafe extern fn __rdl_realloc(ptr: *mut u8,
                                       old_size: usize,
                                       align: usize,
                                       new_size: usize) -> *mut u8 {
        let old_layout = Layout::from_size_align_unchecked(old_size, align);
        System.realloc(ptr, old_layout, new_size)
    }

    #[rustc_std_internal_symbol]
    pub unsafe extern fn __rdl_alloc_zeroed(size: usize, align: usize) -> *mut u8 {
        let layout = Layout::from_size_align_unchecked(size, align);
        System.alloc_zeroed(layout)
    }
}
<|endoftext|>


<commit_before>//!
//! Hold the `Game` struct and the `play()` function
//!

// Tcod
extern crate tcod;
use self::tcod::{console, Console};
use core::tcod::input;

// core::world
pub mod world;
use core::world::World;

// core::object
pub mod object;
// Game state depends on what actions the player can do
use core::object::actions::Actions;

// core::renderer
pub mod renderer;
use self::renderer::Renderer;
use self::renderer::RGB;

// core::log
pub mod log;
use self::log::GlobalLog;

// Use init here privately
// After all, what game function requires that something in the initalizer be ran?
mod init;

///
/// Enum representing the state of the game
///
pub enum State {
  // Game just created
  New,
  // Player acted
  Act(Actions),
  // Debug
  Debug
}

///
/// Engine struct. Holds a world, state, and renderer
///
pub struct Engine {
  world: World,
  state: State,
  ren: Renderer,
  root: console::Root,

  // Debug
  noclip: bool

}

impl Engine {

  ///
  /// Capture keyboard input from tcod and update player state
  ///
  fn process_keypress(&mut self, keypress: input::Key) {

    match keypress.code {

      // If the keycode isn't escape we continue checking for important keys
      input::KeyCode::Escape => panic!("Bye"),
      _ => {

        if keypress.printable != ' ' {

          let oldpos = self.world.player.actor.pos.clone();

          // Wipe the state
          self.state = State::New;

          match keypress.printable {

            // Movement keys are bound to vim-like controls (hjklyubn)
            'h' => {
              self.world.player.actor.move_cart(-1, 0);
              self.world.player.state = Actions::Move;
            },
            'j' => {
              self.world.player.actor.move_cart(0, 1);
              self.world.player.state = Actions::Move;
            },
            'k' => {
              self.world.player.actor.move_cart(0, -1);
              self.world.player.state = Actions::Move;
            },
            'l' => {
              self.world.player.actor.move_cart(1, 0);
              self.world.player.state = Actions::Move;
            },
            'y' => {
              self.world.player.actor.move_cart(-1, -1);
              self.world.player.state = Actions::Move;
            },
            'u' => {
              self.world.player.actor.move_cart(1, -1);
              self.world.player.state = Actions::Move;
            },
            'b' => {
              self.world.player.actor.move_cart(-1, 1);
              self.world.player.state = Actions::Move;
            },
            'n' => {
              self.world.player.actor.move_cart(1, 1);
              self.world.player.state = Actions::Move;
            },
            // "Regenerate" current level
            'w' => {
              self.world.test_traverse();
              self.state = State::Act(Actions::Unknown);
            },
            // Create an empty level for testing
            'q' => {
              self.world.test_empty();
              self.state = State::Act(Actions::Unknown);
            },
            // Wait
            '.' => {
              self.world.player.state = Actions::Wait;
            },
            // Go downstars (if possible)
            '>' => { self.world.player.state = Actions::DownStair },
            // Go upstairs (if possible)
            '<' => { self.world.player.state = Actions::UpStair },
            // Unbound key, so we just say we don't know what the player did

            // Debug keypresses

            // Toggle scent
            'r' => {
              self.ren.show_scent = !self.ren.show_scent;
              self.ren.draw_all(&mut self.root, &mut self.world);
              self.state = State::Debug;
            },

            // Toggle sound
            't' => {
              self.ren.show_sound = !self.ren.show_sound;
              self.ren.draw_all(&mut self.root, &mut self.world);
              self.state = State::Debug;
            },

            // Toggle FoV
            'f' => {
              self.ren.fov = !self.ren.fov;
              self.ren.draw_all(&mut self.root, &mut self.world);
              self.state = State::Debug;
            },

            // Toggle noclip
            'z' => {
              self.noclip = !self.noclip;
              self.state = State::Debug;
            },

            _ => { self.world.player.state = Actions::Unknown }

          }

          // If state is Debug, don't override
          match self.state {
            State::Debug => (),
            _ => {
              // Set game state to player state
              self.state = State::Act(self.world.player.state.clone());

              match self.world.player.state {
                Actions::Move => {
                  // Make sure player doesn't do anything dumb
                  if !self.world.is_valid(self.world.player.actor.pos.x, self.world.player.actor.pos.y) && !self.noclip {
                    self.world.player.actor.pos = oldpos;
                    self.world.player.state = Actions::Unknown;
                  }
                }
                _ => ()
              }

            }
          }

        }

        /*
        // Prints keycode to console in case if you're trying to find a key that isn't intutive
        else {
          println!("{:?}", keypress.code);
        }
        */

      }

    }

  }

  ///
  /// Return a new `Game`
  ///
  /// This function assumes you will just be passing in tcod::console::Root.width() and height(),
  /// so inputs are i32s instead of usizes (they get converted)
  ///
  pub fn new() -> Engine {

    // Get map height
    let map_dim = init::map_dimensions();

    // Get root console
    let root = init::root();

    Engine {
      world: World::new(map_dim),
      state: State::New,
      ren: Renderer::new(
        map_dim,
        (root.width() as isize, root.height() as isize),
        init::console_height(),
        init::panel_width()
      ),
      root: root,

      // Debug
      noclip: false
    }

  }

  ///
  /// Update the game state, then update depending on the new state
  ///
  fn update(&mut self) {

    match &self.state {
      // Moving or waiting prompts a world update
      &State::Act(Actions::Move) | &State::Act(Actions::Wait) => self.world.update(),

      // Trying to go up and downstairs prompts the respective response from world
      &State::Act(Actions::DownStair) => {
        if self.world.can_go_down() {
          self.world.go_down();
        }
      },
      &State::Act(Actions::UpStair) => {
        if self.world.can_go_up() {
          self.world.go_up();
        }
      }

      // All other variants are dropped
      _ => ()
    }

  }

  ///
  /// Play the game.
  ///
  pub fn play(&mut self) {

    // Some starting messages, will be removed in later versions
    let mut log = GlobalLog.lock().unwrap();
    log.push(("Welcome to Edgequest", RGB(255, 0, 255)));
    log.push(("Move with vim keys", RGB(255, 255, 255)));
    log.push(("esc to quit, w to regenerate", RGB(255, 255, 255)));
    log.push(("r to toggle scent, t to toggle sound", RGB(255, 255, 255)));
    log.push(("f to toggle FoV, z to toggle noclip", RGB(255, 255, 255)));
    drop(log);

    // Initial update
    self.update();

    // Draw all and capture keypresses
    while !self.root.window_closed() {

      // Draw what the camera sees
      self.ren.draw_all(&mut self.root, &mut self.world);

      // Capture game keys (Keys that change the state of the player)
      // This is what gives it the turn based nature, i.e. waits for player input before
      // doing anything
      let keypress = self.root.wait_for_keypress(true);
      self.process_keypress(keypress);

      // Update game
      self.update();

    }

  }

}<commit_msg>Update<commit_after>//!
//! Hold the `Game` struct and the `play()` function
//!

// Tcod
extern crate tcod;
use self::tcod::{console, Console};
use core::tcod::input;

// core::world
pub mod world;
use core::world::World;

// core::object
pub mod object;
// Game state depends on what actions the player can do
use core::object::actions::Actions;

// core::renderer
pub mod renderer;
use self::renderer::Renderer;
use self::renderer::RGB;

// core::log
pub mod log;
use self::log::GlobalLog;

// Use init here privately
// After all, what game function requires that something in the initalizer be ran?
mod init;

///
/// Enum representing the state of the game
///
pub enum State {
  // Game just created
  New,
  // Player acted
  Act(Actions),
  // Debug
  Debug
}

///
/// Engine struct. Holds a world, state, and renderer
///
pub struct Engine {
  world: World,
  state: State,
  ren: Renderer,
  root: console::Root,

  // Debug
  noclip: bool

}

impl Engine {

  ///
  /// Capture keyboard input from tcod and update player state
  ///
  fn process_keypress(&mut self, keypress: input::Key) {

    match keypress.code {

      // If the keycode isn't escape we continue checking for important keys
      input::KeyCode::Escape => panic!("Bye"),
      _ => {

        if keypress.printable != ' ' {

          let oldpos = self.world.player.actor.pos.clone();

          // Wipe the state
          self.state = State::New;

          match keypress.printable {

            // Movement keys are bound to vim-like controls (hjklyubn)
            'h' => {
              self.world.player.actor.move_cart(-1, 0);
              self.world.player.state = Actions::Move;
            },
            'j' => {
              self.world.player.actor.move_cart(0, 1);
              self.world.player.state = Actions::Move;
            },
            'k' => {
              self.world.player.actor.move_cart(0, -1);
              self.world.player.state = Actions::Move;
            },
            'l' => {
              self.world.player.actor.move_cart(1, 0);
              self.world.player.state = Actions::Move;
            },
            'y' => {
              self.world.player.actor.move_cart(-1, -1);
              self.world.player.state = Actions::Move;
            },
            'u' => {
              self.world.player.actor.move_cart(1, -1);
              self.world.player.state = Actions::Move;
            },
            'b' => {
              self.world.player.actor.move_cart(-1, 1);
              self.world.player.state = Actions::Move;
            },
            'n' => {
              self.world.player.actor.move_cart(1, 1);
              self.world.player.state = Actions::Move;
            },
            // "Regenerate" current level
            'w' => {
              self.world.test_traverse();
              self.state = State::Act(Actions::Unknown);
            },
            // Create an empty level for testing
            'q' => {
              self.world.test_empty();
              self.state = State::Act(Actions::Unknown);
            },
            // Wait
            '.' => {
              self.world.player.state = Actions::Wait;
            },
            // Go downstars (if possible)
            '>' => { self.world.player.state = Actions::DownStair },
            // Go upstairs (if possible)
            '<' => { self.world.player.state = Actions::UpStair },
            // Unbound key, so we just say we don't know what the player did

            // Debug keypresses

            // Toggle scent
            'r' => {
              self.ren.show_scent = !self.ren.show_scent;
              self.ren.draw_all(&mut self.root, &mut self.world);
              self.state = State::Debug;
            },

            // Toggle sound
            't' => {
              self.ren.show_sound = !self.ren.show_sound;
              self.ren.draw_all(&mut self.root, &mut self.world);
              self.state = State::Debug;
            },

            // Toggle FoV
            'f' => {
              self.ren.fov = !self.ren.fov;
              self.ren.draw_all(&mut self.root, &mut self.world);
              self.state = State::Debug;
            },

            // Toggle noclip
            'z' => {
              self.noclip = !self.noclip;
              self.state = State::Debug;
            },

            _ => { self.world.player.state = Actions::Unknown }

          }

          // If state is Debug, don't override
          match self.state {
            State::Debug => (),
            _ => {
              // Set game state to player state
              self.state = State::Act(self.world.player.state.clone());

              match self.world.player.state {
                Actions::Move => {
                  // Make sure player doesn't do anything dumb
                  if !self.world.is_valid(self.world.player.actor.pos.x, self.world.player.actor.pos.y) && !self.noclip {
                    self.world.player.actor.pos = oldpos;
                    self.world.player.state = Actions::Unknown;
                  }
                }
                _ => ()
              }

            }
          }

        }

        /*
        // Prints keycode to console in case if you're trying to find a key that isn't intutive
        else {
          println!("{:?}", keypress.code);
        }
        */

      }

    }

  }

  ///
  /// Return a new `Game`
  ///
  /// This function assumes you will just be passing in tcod::console::Root.width() and height(),
  /// so inputs are i32s instead of usizes (they get converted)
  ///
  pub fn new() -> Engine {

    // Get map height
    let map_dim = init::map_dimensions();

    // Get root console
    let root = init::root();

    Engine {
      world: World::new(map_dim),
      state: State::New,
      ren: Renderer::new(
        map_dim,
        (root.width() as isize, root.height() as isize),
        init::console_height(),
        init::panel_width()
      ),
      root: root,

      // Debug
      noclip: false
    }

  }

  ///
  /// Update the game state, then update depending on the new state
  ///
  fn update(&mut self) {

    match self.state {
      // Moving or waiting prompts a world update
      State::Act(Actions::Move) | State::Act(Actions::Wait) => self.world.update(),

      // Trying to go up and downstairs prompts the respective response from world
      State::Act(Actions::DownStair) => {
        if self.world.can_go_down() {
          self.world.go_down();
        }
      },
      State::Act(Actions::UpStair) => {
        if self.world.can_go_up() {
          self.world.go_up();
        }
      }

      // All other variants are dropped
      _ => ()
    }

  }

  ///
  /// Play the game.
  ///
  pub fn play(&mut self) {

    // Some starting messages, will be removed in later versions
    let mut log = GlobalLog.lock().unwrap();
    log.push(("Welcome to Edgequest", RGB(255, 0, 255)));
    log.push(("Move with vim keys", RGB(255, 255, 255)));
    log.push(("esc to quit, w to regenerate", RGB(255, 255, 255)));
    log.push(("r to toggle scent, t to toggle sound", RGB(255, 255, 255)));
    log.push(("f to toggle FoV, z to toggle noclip", RGB(255, 255, 255)));
    drop(log);

    // Initial update
    self.update();

    // Draw all and capture keypresses
    while !self.root.window_closed() {

      // Draw what the camera sees
      self.ren.draw_all(&mut self.root, &mut self.world);

      // Capture game keys (Keys that change the state of the player)
      // This is what gives it the turn based nature, i.e. waits for player input before
      // doing anything
      let keypress = self.root.wait_for_keypress(true);
      self.process_keypress(keypress);

      // Update game
      self.update();

    }

  }

}<|endoftext|>


<commit_before>use crate::Position;
use fixedbitset::FixedBitSet;

#[derive(Debug)]
pub struct Coverage {
    n: usize,
    /// A bit is set, if the row is covered.
    covered_rows: FixedBitSet,
    /// A bit is set, if the column is covered.
    covered_columns: FixedBitSet,
}

impl Coverage {
    #[inline]
    pub fn n(&self) -> usize {
        self.n
    }

    pub fn new(n: usize) -> Coverage {
        assert!(n > 0);
        Coverage {
            n: n,
            covered_rows: FixedBitSet::with_capacity(n),
            covered_columns: FixedBitSet::with_capacity(n),
        }
    }

    /// find a single uncovered (row, column) pair. Iterates in column, row order.
    #[inline]
    pub fn find_uncovered_column_row<F>(&self, mut f: F) -> Option<Position>
    where
        F: FnMut(Position) -> bool,
    {
        let n = self.n();

        for column in 0..n {
            if self.is_column_covered(column) {
                continue;
            }
            for row in 0..n {
                if self.is_row_covered(row) {
                    continue;
                }

                let pos = Position { row, column };
                if f(pos) {
                    return Some(pos);
                }
            }
        }

        return None;
    }

    /// iterates over all uncovered (row, column) pairs in row, column order
    #[inline]
    pub fn iter_uncovered_row_column<F>(&self, mut f: F)
    where
        F: FnMut(Position),
    {
        let n = self.n();

        for row in 0..n {
            if self.is_row_covered(row) {
                continue;
            }

            for column in 0..n {
                if self.is_column_covered(column) {
                    continue;
                }

                f(Position { row, column });
            }
        }
    }

    /// iterates over all uncovered (row, column) pairs in row, column order, and set covered if f returns true.
    #[inline]
    pub fn iter_uncovered_row_column_and_cover<F>(&mut self, mut f: F)
    where
        F: FnMut(Position) -> bool,
    {
        let n = self.n();

        for row in 0..n {
            if self.is_row_covered(row) {
                continue;
            }

            'column: for column in 0..n {
                if self.is_column_covered(column) {
                    continue;
                }

                let pos = Position { row, column };
                if f(pos) {
                    self.cover(pos);
                    // the complete row is covered! break the loop!
                    break 'column;
                }
            }
        }
    }

    #[inline]
    pub fn is_row_covered(&self, row: usize) -> bool {
        debug_assert!(row < self.n());
        self.covered_rows.contains(row)
    }

    #[inline]
    pub fn is_column_covered(&self, column: usize) -> bool {
        debug_assert!(column < self.n());
        self.covered_columns.contains(column)
    }

    #[inline]
    pub fn cover(&mut self, pos: Position) {
        self.cover_row(pos.row);
        self.cover_column(pos.column);
    }

    #[inline]
    pub fn cover_column(&mut self, column: usize) {
        debug_assert!(column < self.n());
        self.covered_columns.set(column, true);
    }

    #[inline]
    pub fn uncover_column(&mut self, column: usize) {
        debug_assert!(column < self.n());
        self.covered_columns.set(column, false);
    }

    #[inline]
    pub fn cover_row(&mut self, row: usize) {
        debug_assert!(row < self.n());
        self.covered_rows.set(row, true);
    }

    pub fn clear(&mut self) {
        self.covered_rows.clear();
        self.covered_columns.clear();
    }

    pub fn all_uncovered(&self) -> bool {
        is_bitset_clear(&self.covered_rows) && is_bitset_clear(&self.covered_columns)
    }
}

fn is_bitset_clear(bitset: &FixedBitSet) -> bool {
    for elm in bitset.as_slice().iter() {
        if *elm != 0 {
            return false;
        }
    }
    return true;
}
<commit_msg>Invert logic in Coverage "matrix"<commit_after>use crate::Position;
use fixedbitset::FixedBitSet;

#[derive(Debug)]
pub struct Coverage {
    n: usize,
    /// A bit is set, if the row is uncovered.
    uncovered_rows: FixedBitSet,
    /// A bit is set, if the column is uncovered.
    uncovered_columns: FixedBitSet,
}

impl Coverage {
    #[inline]
    pub fn n(&self) -> usize {
        self.n
    }

    pub fn new(n: usize) -> Coverage {
        assert!(n > 0);

        // We start with all bits set (all rows/columns) uncovered.
        let mut all_rows_uncovered = FixedBitSet::with_capacity(n);
        all_rows_uncovered.set_range(.., true);

        // We can simply clone the rows as we work on square matrices
        let all_columns_uncovered = all_rows_uncovered.clone();

        Coverage {
            n: n,
            uncovered_rows: all_rows_uncovered,
            uncovered_columns: all_columns_uncovered,
        }
    }

    /// find a single uncovered (row, column) pair. Iterates in column, row order.
    #[inline]
    pub fn find_uncovered_column_row<F>(&self, mut f: F) -> Option<Position>
    where
        F: FnMut(Position) -> bool,
    {
        let n = self.n();

        for column in 0..n {
            if self.is_column_covered(column) {
                continue;
            }
            for row in 0..n {
                if self.is_row_covered(row) {
                    continue;
                }

                let pos = Position { row, column };
                if f(pos) {
                    return Some(pos);
                }
            }
        }

        return None;
    }

    /// iterates over all uncovered (row, column) pairs in row, column order
    #[inline]
    pub fn iter_uncovered_row_column<F>(&self, mut f: F)
    where
        F: FnMut(Position),
    {
        let n = self.n();

        for row in 0..n {
            if self.is_row_covered(row) {
                continue;
            }

            for column in 0..n {
                if self.is_column_covered(column) {
                    continue;
                }

                f(Position { row, column });
            }
        }
    }

    /// iterates over all uncovered (row, column) pairs in row, column order, and set covered if f returns true.
    #[inline]
    pub fn iter_uncovered_row_column_and_cover<F>(&mut self, mut f: F)
    where
        F: FnMut(Position) -> bool,
    {
        let n = self.n();

        for row in 0..n {
            if self.is_row_covered(row) {
                continue;
            }

            'column: for column in 0..n {
                if self.is_column_covered(column) {
                    continue;
                }

                let pos = Position { row, column };
                if f(pos) {
                    self.cover(pos);
                    // the complete row is covered! break the loop!
                    break 'column;
                }
            }
        }
    }

    #[inline]
    pub fn is_row_covered(&self, row: usize) -> bool {
        debug_assert!(row < self.n());
        !self.uncovered_rows.contains(row)
    }

    #[inline]
    pub fn is_column_covered(&self, column: usize) -> bool {
        debug_assert!(column < self.n());
        !self.uncovered_columns.contains(column)
    }

    #[inline]
    pub fn cover(&mut self, pos: Position) {
        self.cover_row(pos.row);
        self.cover_column(pos.column);
    }

    #[inline]
    pub fn cover_column(&mut self, column: usize) {
        debug_assert!(column < self.n());
        self.uncovered_columns.set(column, false);
    }

    #[inline]
    pub fn uncover_column(&mut self, column: usize) {
        debug_assert!(column < self.n());
        self.uncovered_columns.set(column, true);
    }

    #[inline]
    pub fn cover_row(&mut self, row: usize) {
        debug_assert!(row < self.n());
        self.uncovered_rows.set(row, false);
    }

    pub fn clear(&mut self) {
        self.uncovered_rows.set_range(.., true);
        self.uncovered_columns.set_range(.., true);
    }

    pub fn all_uncovered(&self) -> bool {
        self.uncovered_rows.count_ones(..) + self.uncovered_columns.count_ones(..)
            == (self.n + self.n)
    }
}
<|endoftext|>


<commit_before>//! Plugin which issues vault tokens to services.

use compose_yml::v2 as dc;
use std::result;
use serde_yaml;
#[cfg(test)]
use std::cell::RefCell;
use std::collections::BTreeMap;
use std::env;
use std::fmt::Debug;
use std::fs;
use std::io::Read;
use std::path::PathBuf;
#[cfg(test)]
use std::rc::Rc;
use vault;
use vault::client::VaultDuration;

use errors::*;
use plugins;
use plugins::{Operation, PluginGenerate, PluginNew, PluginTransform};
use project::Project;
use util::err;

#[cfg(feature = "serde_derive")]
include!(concat!("vault_config.in.rs"));
#[cfg(feature = "serde_codegen")]
include!(concat!(env!("OUT_DIR"), "/plugins/transform/vault_config.rs"));

/// Load a vault token from `~/.vault-token`, where the command line client
/// puts it.
fn load_vault_token_from_file() -> Result<String> {
    let path = try!(env::home_dir()
            .ok_or_else(|| err("You do not appear to have a home directory")))
        .join(".vault-token");
    let mut f = try!(fs::File::open(&path)
        .map_err(|e| err!("Error opening {}: {}", path.display(), e)));
    let mut result = String::new();
    try!(f.read_to_string(&mut result)
        .map_err(|e| err!("Error reading {}: {}", path.display(), e)));
    Ok(result.trim().to_owned())
}

/// Find the vault token we'll use to generate new tokens.
fn find_vault_token() -> Result<String> {
    env::var("VAULT_MASTER_TOKEN")
        .or_else(|_| env::var("VAULT_TOKEN"))
        .or_else(|_| load_vault_token_from_file())
        .map_err(|e| {
            err!("{}.  You probably want to log in using the vault client or set \
                  VAULT_MASTER_TOKEN",
                 e)
        })
}

/// The "environment" in which to interpret a configuration file.  We don't
/// want to use the OS environment variables, but rather a fake environment
/// with a few carefully selected values.
#[derive(Debug)]
struct ConfigEnvironment<'a> {
    /// The context for our transformation, including the project, pod,
    /// etc.
    ctx: &'a plugins::Context<'a>,
    /// The name of the current service.
    service: &'a str,
}

impl<'a> dc::Environment for ConfigEnvironment<'a> {
    fn var(&self, key: &str) -> result::Result<String, env::VarError> {
        let result = match key {
            "PROJECT" => Ok(self.ctx.project.name()),
            "OVERRIDE" => Ok(self.ctx.ovr.name()),
            "POD" => Ok(self.ctx.pod.name()),
            "SERVICE" => Ok(self.service),
            _ => Err(env::VarError::NotPresent),
        };
        result.map(|s| s.to_owned())
    }
}

/// An abstract interface to Vault's token-generation capabilities.  We use
/// this to mock vault during tests.
trait GenerateToken: Debug {
    /// Get a `VAULT_ADDR` value to use along with this token.
    fn addr(&self) -> &str;
    /// Generate a token with the specified parameters.
    fn generate_token(&self,
                      display_name: &str,
                      policies: Vec<String>,
                      ttl: VaultDuration)
                      -> Result<String>;
}

/// A list of calls made to a `MockVault` instance.
#[cfg(test)]
type MockVaultCalls = Rc<RefCell<Vec<(String, Vec<String>, VaultDuration)>>>;

/// A fake interface to vault for testing purposes.
#[derive(Debug)]
#[cfg(test)]
struct MockVault {
    /// The tokens we were asked to generate.  We store these in a RefCell
    /// so that we can have "interior" mutability, because we don't want
    /// `generate_token` to be `&mut self` in the general case.
    calls: MockVaultCalls,
}

#[cfg(test)]
impl MockVault {
    /// Create a new MockVault.
    fn new() -> MockVault {
        MockVault { calls: Rc::new(RefCell::new(vec![])) }
    }

    /// Return a reference to record of calls made to our vault.
    fn calls(&self) -> MockVaultCalls {
        self.calls.clone()
    }
}

#[cfg(test)]
impl GenerateToken for MockVault {
    fn addr(&self) -> &str {
        "http://example.com:8200/"
    }

    fn generate_token(&self,
                      display_name: &str,
                      policies: Vec<String>,
                      ttl: VaultDuration)
                      -> Result<String> {
        self.calls.borrow_mut().push((display_name.to_owned(), policies, ttl));
        Ok("fake_token".to_owned())
    }
}

/// An interface to an actual vault server.
#[derive(Debug)]
struct Vault {
    /// The address of our vault server.
    addr: String,
    /// The master token that we'll use to issue new tokens.
    token: String,
}

impl Vault {
    /// Create a new vault client.
    fn new() -> Result<Vault> {
        let mut addr = try!(env::var("VAULT_ADDR").map_err(|_| {
            err("Please set the environment variable VAULT_ADDR to the URL of \
                 your vault server")
        }));
        // TODO MED: Temporary fix because of broken URL handling in
        // hashicorp_vault.  File an upstream bug.
        if addr.ends_with('/') {
            let new_len = addr.len() - 1;
            addr.truncate(new_len);
        }
        let token = try!(find_vault_token());
        Ok(Vault {
            addr: addr,
            token: token,
        })
    }
}

impl GenerateToken for Vault {
    fn addr(&self) -> &str {
        &self.addr
    }

    fn generate_token(&self,
                      display_name: &str,
                      policies: Vec<String>,
                      ttl: VaultDuration)
                      -> Result<String> {
        // We can't store `client` in `self`, because it has some obnoxious
        // lifetime parameters.  So we'll just recreate it.  This is
        // probably not the worst idea, because it uses `hyper` for HTTP,
        // and `hyper` HTTP connections used to have expiration issues that
        // were tricky for clients to deal with correctly.
        let client = try!(vault::Client::new(&self.addr, &self.token));
        let opts = vault::client::TokenOptions::default()
            .display_name(display_name)
            .renewable(false)
            .ttl(ttl)
            .policies(policies);
        let auth = try!(client.create_token(&opts));
        Ok(auth.client_token)
    }
}

/// Issues `VAULT_TOKEN` values to services.
#[derive(Debug)]
pub struct Plugin {
    /// Our `config/vault.yml` YAML file, parsed and read into memory.
    /// Optional because if we're being run as a `PluginGenerate`, we won't
    /// have it (but it's guaranteed otherwise).
    config: Option<Config>,
    /// Our source of tokens.
    generator: Option<Box<GenerateToken>>,
}

impl Plugin {
    /// Get the path to this plugin's config file.
    fn config_path(project: &Project) -> PathBuf {
        project.root_dir().join("config/vault.yml")
    }

    /// Create a new plugin, specifying an alternate source for tokens.
    fn new_with_generator<G>(project: &Project, generator: Option<G>) -> Result<Plugin>
        where G: GenerateToken + 'static
    {
        let path = Self::config_path(project);
        let config = if path.exists() {
            let f = try!(fs::File::open(&path));
            Some(try!(serde_yaml::from_reader(f)
                .map_err(|e| err!("Error reading {}: {}", path.display(), e))))
        } else {
            None
        };
        Ok(Plugin {
            config: config,
            generator: generator.map(|gen: G| -> Box<GenerateToken> { Box::new(gen) }),
        })
    }
}

impl plugins::Plugin for Plugin {
    fn name(&self) -> &'static str {
        Self::plugin_name()
    }
}

impl PluginNew for Plugin {
    fn plugin_name() -> &'static str {
        "vault"
    }

    fn is_configured_for(project: &Project) -> Result<bool> {
        let path = Self::config_path(project);
        Ok(path.exists())
    }

    fn new(project: &Project) -> Result<Self> {
        // An annoying special case.  We may be called as a code generator,
        // in which case we don't want to try to create a `GenerateToken`
        // instance.
        let token_gen = if try!(Self::is_configured_for(project)) {
            Some(try!(Vault::new()))
        } else {
            None
        };
        Self::new_with_generator(project, token_gen)
    }
}

impl PluginGenerate for Plugin {
    fn generator_description(&self) -> &'static str {
        "Get passwords & other secrets from a Vault server"
    }
}

impl PluginTransform for Plugin {
    fn transform(&self,
                 _op: Operation,
                 ctx: &plugins::Context,
                 file: &mut dc::File)
                 -> Result<()> {

        // Get our plugin config.
        let config = self.config
            .as_ref()
            .expect("config should always be present for transform");
        let generator = self.generator
            .as_ref()
            .expect("generator should always be present for transform");

        // Should this plugin be excluded in this override?
        if !ctx.ovr.is_enabled_by(&config.enable_in_overrides) {
            return Ok(());
        }

        // Apply to each service.
        for (name, mut service) in &mut file.services {
            // Set up a ConfigEnvironment that we can use to perform
            // interpolations of values like `$SERVICE` in our config file.
            let env = ConfigEnvironment {
                ctx: ctx,
                service: name,
            };

            // Define a local helper function to interpolate
            // `RawOr<String>` values using `env`.
            let interpolated = |raw_val: &dc::RawOr<String>| -> Result<String> {
                let mut val = raw_val.to_owned();
                Ok(try!(val.interpolate_env(&env)).to_owned())
            };

            // Insert our VAULT_ADDR value into the generated files.
            service.environment
                .insert("VAULT_ADDR".to_owned(), generator.addr().to_owned());

            // Get a list of policy "patterns" that apply to this service.
            let mut raw_policies = config.default_policies.clone();
            raw_policies.extend(config.pods
                .get(ctx.pod.name())
                .and_then(|pod| pod.get(name))
                .map_or_else(|| vec![], |s| s.policies.clone()));

            // Interpolate the variables found in our policy patterns.
            let mut policies = vec![];
            for result in raw_policies.iter().map(|p| interpolated(p)) {
                // We'd like to use std::result::fold here but it's unstable.
                policies.push(try!(result));
            }
            debug!("Generating token for '{}' with policies {:?}",
                   name,
                   &policies);

            // Generate a VAULT_TOKEN.
            let display_name = format!("{}_{}_{}_{}",
                                       ctx.project.name(),
                                       ctx.ovr.name(),
                                       ctx.pod.name(),
                                       name);
            let ttl = VaultDuration::seconds(config.default_ttl);
            let token = try!(generator.generate_token(&display_name, policies, ttl)
                .chain_err(|| format!("could not generate token for '{}'", name)));
            service.environment.insert("VAULT_TOKEN".to_owned(), token);

            // Add in any extra environment variables.
            for (var, val) in &config.extra_environment {
                service.environment.insert(var.to_owned(), try!(interpolated(val)));
            }
        }
        Ok(())
    }
}

#[test]
fn interpolates_policies() {
    use env_logger;
    let _ = env_logger::init();

    env::set_var("VAULT_ADDR", "http://example.com:8200/");
    env::set_var("VAULT_MASTER_TOKEN", "fake master token");

    let proj = Project::from_example("vault_integration").unwrap();
    let ovr = proj.ovr("production").unwrap();

    let vault = MockVault::new();
    let calls = vault.calls();
    let plugin = Plugin::new_with_generator(&proj, Some(vault)).unwrap();

    let frontend = proj.pod("frontend").unwrap();
    let ctx = plugins::Context::new(&proj, ovr, frontend);
    let mut file = frontend.merged_file(ovr).unwrap();
    plugin.transform(Operation::Output, &ctx, &mut file).unwrap();
    let web = file.services.get("web").unwrap();
    assert_eq!(web.environment.get("VAULT_ADDR").expect("has VAULT_ADDR"),
               "http://example.com:8200/");
    assert_eq!(web.environment.get("VAULT_TOKEN").expect("has VAULT_TOKEN"),
               "fake_token");
    assert_eq!(web.environment.get("VAULT_ENV").expect("has VAULT_ENV"),
               "production");

    let calls = calls.borrow();
    assert_eq!(calls.len(), 1);
    let (ref display_name, ref policies, ref ttl) = calls[0];
    assert_eq!(display_name, "vault_integration_production_frontend_web");
    assert_eq!(policies,
               &["vault_integration-production".to_owned(),
                 "vault_integration-production-frontend-web".to_owned(),
                 "vault_integration-production-ssl".to_owned()]);
    assert_eq!(ttl, &VaultDuration::seconds(2592000));
}

#[test]
fn only_applied_in_specified_overrides() {
    use env_logger;
    let _ = env_logger::init();

    let proj = Project::from_example("vault_integration").unwrap();
    let ovr = proj.ovr("test").unwrap();

    let vault = MockVault::new();
    let plugin = Plugin::new_with_generator(&proj, Some(vault)).unwrap();

    let frontend = proj.pod("frontend").unwrap();
    let ctx = plugins::Context::new(&proj, ovr, frontend);
    let mut file = frontend.merged_file(ovr).unwrap();
    plugin.transform(Operation::Output, &ctx, &mut file).unwrap();
    let web = file.services.get("web").unwrap();
    assert_eq!(web.environment.get("VAULT_ADDR"), None);
}
<commit_msg>vault: Make sure our generated tokens are renewable<commit_after>//! Plugin which issues vault tokens to services.

use compose_yml::v2 as dc;
use std::result;
use serde_yaml;
#[cfg(test)]
use std::cell::RefCell;
use std::collections::BTreeMap;
use std::env;
use std::fmt::Debug;
use std::fs;
use std::io::Read;
use std::path::PathBuf;
#[cfg(test)]
use std::rc::Rc;
use vault;
use vault::client::VaultDuration;

use errors::*;
use plugins;
use plugins::{Operation, PluginGenerate, PluginNew, PluginTransform};
use project::Project;
use util::err;

#[cfg(feature = "serde_derive")]
include!(concat!("vault_config.in.rs"));
#[cfg(feature = "serde_codegen")]
include!(concat!(env!("OUT_DIR"), "/plugins/transform/vault_config.rs"));

/// Load a vault token from `~/.vault-token`, where the command line client
/// puts it.
fn load_vault_token_from_file() -> Result<String> {
    let path = try!(env::home_dir()
            .ok_or_else(|| err("You do not appear to have a home directory")))
        .join(".vault-token");
    let mut f = try!(fs::File::open(&path)
        .map_err(|e| err!("Error opening {}: {}", path.display(), e)));
    let mut result = String::new();
    try!(f.read_to_string(&mut result)
        .map_err(|e| err!("Error reading {}: {}", path.display(), e)));
    Ok(result.trim().to_owned())
}

/// Find the vault token we'll use to generate new tokens.
fn find_vault_token() -> Result<String> {
    env::var("VAULT_MASTER_TOKEN")
        .or_else(|_| env::var("VAULT_TOKEN"))
        .or_else(|_| load_vault_token_from_file())
        .map_err(|e| {
            err!("{}.  You probably want to log in using the vault client or set \
                  VAULT_MASTER_TOKEN",
                 e)
        })
}

/// The "environment" in which to interpret a configuration file.  We don't
/// want to use the OS environment variables, but rather a fake environment
/// with a few carefully selected values.
#[derive(Debug)]
struct ConfigEnvironment<'a> {
    /// The context for our transformation, including the project, pod,
    /// etc.
    ctx: &'a plugins::Context<'a>,
    /// The name of the current service.
    service: &'a str,
}

impl<'a> dc::Environment for ConfigEnvironment<'a> {
    fn var(&self, key: &str) -> result::Result<String, env::VarError> {
        let result = match key {
            "PROJECT" => Ok(self.ctx.project.name()),
            "OVERRIDE" => Ok(self.ctx.ovr.name()),
            "POD" => Ok(self.ctx.pod.name()),
            "SERVICE" => Ok(self.service),
            _ => Err(env::VarError::NotPresent),
        };
        result.map(|s| s.to_owned())
    }
}

/// An abstract interface to Vault's token-generation capabilities.  We use
/// this to mock vault during tests.
trait GenerateToken: Debug {
    /// Get a `VAULT_ADDR` value to use along with this token.
    fn addr(&self) -> &str;
    /// Generate a token with the specified parameters.
    fn generate_token(&self,
                      display_name: &str,
                      policies: Vec<String>,
                      ttl: VaultDuration)
                      -> Result<String>;
}

/// A list of calls made to a `MockVault` instance.
#[cfg(test)]
type MockVaultCalls = Rc<RefCell<Vec<(String, Vec<String>, VaultDuration)>>>;

/// A fake interface to vault for testing purposes.
#[derive(Debug)]
#[cfg(test)]
struct MockVault {
    /// The tokens we were asked to generate.  We store these in a RefCell
    /// so that we can have "interior" mutability, because we don't want
    /// `generate_token` to be `&mut self` in the general case.
    calls: MockVaultCalls,
}

#[cfg(test)]
impl MockVault {
    /// Create a new MockVault.
    fn new() -> MockVault {
        MockVault { calls: Rc::new(RefCell::new(vec![])) }
    }

    /// Return a reference to record of calls made to our vault.
    fn calls(&self) -> MockVaultCalls {
        self.calls.clone()
    }
}

#[cfg(test)]
impl GenerateToken for MockVault {
    fn addr(&self) -> &str {
        "http://example.com:8200/"
    }

    fn generate_token(&self,
                      display_name: &str,
                      policies: Vec<String>,
                      ttl: VaultDuration)
                      -> Result<String> {
        self.calls.borrow_mut().push((display_name.to_owned(), policies, ttl));
        Ok("fake_token".to_owned())
    }
}

/// An interface to an actual vault server.
#[derive(Debug)]
struct Vault {
    /// The address of our vault server.
    addr: String,
    /// The master token that we'll use to issue new tokens.
    token: String,
}

impl Vault {
    /// Create a new vault client.
    fn new() -> Result<Vault> {
        let mut addr = try!(env::var("VAULT_ADDR").map_err(|_| {
            err("Please set the environment variable VAULT_ADDR to the URL of \
                 your vault server")
        }));
        // TODO MED: Temporary fix because of broken URL handling in
        // hashicorp_vault.  File an upstream bug.
        if addr.ends_with('/') {
            let new_len = addr.len() - 1;
            addr.truncate(new_len);
        }
        let token = try!(find_vault_token());
        Ok(Vault {
            addr: addr,
            token: token,
        })
    }
}

impl GenerateToken for Vault {
    fn addr(&self) -> &str {
        &self.addr
    }

    fn generate_token(&self,
                      display_name: &str,
                      policies: Vec<String>,
                      ttl: VaultDuration)
                      -> Result<String> {
        // We can't store `client` in `self`, because it has some obnoxious
        // lifetime parameters.  So we'll just recreate it.  This is
        // probably not the worst idea, because it uses `hyper` for HTTP,
        // and `hyper` HTTP connections used to have expiration issues that
        // were tricky for clients to deal with correctly.
        let client = try!(vault::Client::new(&self.addr, &self.token));
        let opts = vault::client::TokenOptions::default()
            .display_name(display_name)
            .renewable(true)
            .ttl(ttl)
            .policies(policies);
        let auth = try!(client.create_token(&opts));
        Ok(auth.client_token)
    }
}

/// Issues `VAULT_TOKEN` values to services.
#[derive(Debug)]
pub struct Plugin {
    /// Our `config/vault.yml` YAML file, parsed and read into memory.
    /// Optional because if we're being run as a `PluginGenerate`, we won't
    /// have it (but it's guaranteed otherwise).
    config: Option<Config>,
    /// Our source of tokens.
    generator: Option<Box<GenerateToken>>,
}

impl Plugin {
    /// Get the path to this plugin's config file.
    fn config_path(project: &Project) -> PathBuf {
        project.root_dir().join("config/vault.yml")
    }

    /// Create a new plugin, specifying an alternate source for tokens.
    fn new_with_generator<G>(project: &Project, generator: Option<G>) -> Result<Plugin>
        where G: GenerateToken + 'static
    {
        let path = Self::config_path(project);
        let config = if path.exists() {
            let f = try!(fs::File::open(&path));
            Some(try!(serde_yaml::from_reader(f)
                .map_err(|e| err!("Error reading {}: {}", path.display(), e))))
        } else {
            None
        };
        Ok(Plugin {
            config: config,
            generator: generator.map(|gen: G| -> Box<GenerateToken> { Box::new(gen) }),
        })
    }
}

impl plugins::Plugin for Plugin {
    fn name(&self) -> &'static str {
        Self::plugin_name()
    }
}

impl PluginNew for Plugin {
    fn plugin_name() -> &'static str {
        "vault"
    }

    fn is_configured_for(project: &Project) -> Result<bool> {
        let path = Self::config_path(project);
        Ok(path.exists())
    }

    fn new(project: &Project) -> Result<Self> {
        // An annoying special case.  We may be called as a code generator,
        // in which case we don't want to try to create a `GenerateToken`
        // instance.
        let token_gen = if try!(Self::is_configured_for(project)) {
            Some(try!(Vault::new()))
        } else {
            None
        };
        Self::new_with_generator(project, token_gen)
    }
}

impl PluginGenerate for Plugin {
    fn generator_description(&self) -> &'static str {
        "Get passwords & other secrets from a Vault server"
    }
}

impl PluginTransform for Plugin {
    fn transform(&self,
                 _op: Operation,
                 ctx: &plugins::Context,
                 file: &mut dc::File)
                 -> Result<()> {

        // Get our plugin config.
        let config = self.config
            .as_ref()
            .expect("config should always be present for transform");
        let generator = self.generator
            .as_ref()
            .expect("generator should always be present for transform");

        // Should this plugin be excluded in this override?
        if !ctx.ovr.is_enabled_by(&config.enable_in_overrides) {
            return Ok(());
        }

        // Apply to each service.
        for (name, mut service) in &mut file.services {
            // Set up a ConfigEnvironment that we can use to perform
            // interpolations of values like `$SERVICE` in our config file.
            let env = ConfigEnvironment {
                ctx: ctx,
                service: name,
            };

            // Define a local helper function to interpolate
            // `RawOr<String>` values using `env`.
            let interpolated = |raw_val: &dc::RawOr<String>| -> Result<String> {
                let mut val = raw_val.to_owned();
                Ok(try!(val.interpolate_env(&env)).to_owned())
            };

            // Insert our VAULT_ADDR value into the generated files.
            service.environment
                .insert("VAULT_ADDR".to_owned(), generator.addr().to_owned());

            // Get a list of policy "patterns" that apply to this service.
            let mut raw_policies = config.default_policies.clone();
            raw_policies.extend(config.pods
                .get(ctx.pod.name())
                .and_then(|pod| pod.get(name))
                .map_or_else(|| vec![], |s| s.policies.clone()));

            // Interpolate the variables found in our policy patterns.
            let mut policies = vec![];
            for result in raw_policies.iter().map(|p| interpolated(p)) {
                // We'd like to use std::result::fold here but it's unstable.
                policies.push(try!(result));
            }
            debug!("Generating token for '{}' with policies {:?}",
                   name,
                   &policies);

            // Generate a VAULT_TOKEN.
            let display_name = format!("{}_{}_{}_{}",
                                       ctx.project.name(),
                                       ctx.ovr.name(),
                                       ctx.pod.name(),
                                       name);
            let ttl = VaultDuration::seconds(config.default_ttl);
            let token = try!(generator.generate_token(&display_name, policies, ttl)
                .chain_err(|| format!("could not generate token for '{}'", name)));
            service.environment.insert("VAULT_TOKEN".to_owned(), token);

            // Add in any extra environment variables.
            for (var, val) in &config.extra_environment {
                service.environment.insert(var.to_owned(), try!(interpolated(val)));
            }
        }
        Ok(())
    }
}

#[test]
fn interpolates_policies() {
    use env_logger;
    let _ = env_logger::init();

    env::set_var("VAULT_ADDR", "http://example.com:8200/");
    env::set_var("VAULT_MASTER_TOKEN", "fake master token");

    let proj = Project::from_example("vault_integration").unwrap();
    let ovr = proj.ovr("production").unwrap();

    let vault = MockVault::new();
    let calls = vault.calls();
    let plugin = Plugin::new_with_generator(&proj, Some(vault)).unwrap();

    let frontend = proj.pod("frontend").unwrap();
    let ctx = plugins::Context::new(&proj, ovr, frontend);
    let mut file = frontend.merged_file(ovr).unwrap();
    plugin.transform(Operation::Output, &ctx, &mut file).unwrap();
    let web = file.services.get("web").unwrap();
    assert_eq!(web.environment.get("VAULT_ADDR").expect("has VAULT_ADDR"),
               "http://example.com:8200/");
    assert_eq!(web.environment.get("VAULT_TOKEN").expect("has VAULT_TOKEN"),
               "fake_token");
    assert_eq!(web.environment.get("VAULT_ENV").expect("has VAULT_ENV"),
               "production");

    let calls = calls.borrow();
    assert_eq!(calls.len(), 1);
    let (ref display_name, ref policies, ref ttl) = calls[0];
    assert_eq!(display_name, "vault_integration_production_frontend_web");
    assert_eq!(policies,
               &["vault_integration-production".to_owned(),
                 "vault_integration-production-frontend-web".to_owned(),
                 "vault_integration-production-ssl".to_owned()]);
    assert_eq!(ttl, &VaultDuration::seconds(2592000));
}

#[test]
fn only_applied_in_specified_overrides() {
    use env_logger;
    let _ = env_logger::init();

    let proj = Project::from_example("vault_integration").unwrap();
    let ovr = proj.ovr("test").unwrap();

    let vault = MockVault::new();
    let plugin = Plugin::new_with_generator(&proj, Some(vault)).unwrap();

    let frontend = proj.pod("frontend").unwrap();
    let ctx = plugins::Context::new(&proj, ovr, frontend);
    let mut file = frontend.merged_file(ovr).unwrap();
    plugin.transform(Operation::Output, &ctx, &mut file).unwrap();
    let web = file.services.get("web").unwrap();
    assert_eq!(web.environment.get("VAULT_ADDR"), None);
}
<|endoftext|>


<commit_before>use std::{fs::File, io::BufReader, io::prelude::*, path::Path, string::String, vec::Vec};

use json;
use json::JsonValue;
use tiled;
use tiled::Map;

use crate::critter::CritterData;
use crate::game::constants::{CHARACTER_BUF_LENGTH, CHARACTER_JSON_PATH, ZOMBIE_JSON_PATH};

pub fn load_map_file(filename: &str) -> Map {
  let file = match File::open(&Path::new(&filename)) {
    Ok(f) => f,
    Err(e) => panic!("File {} not found: {}", filename, e),
  };
  let reader = BufReader::new(file);
  match tiled::parse(reader) {
    Ok(m) => m,
    Err(e) => panic!("Map parse error {:?}", e)
  }
}

pub fn get_map_tile(map: &Map, layer_index: usize, x: usize, y: usize) -> u32 {
  let layer = match map.layers.get(layer_index) {
    None => panic!("Layer_index value out of index {:?}", map.layers),
    Some(l) => l
  };
  let y_index = match layer.tiles.iter().rev().nth(y) {
    None => panic!("X value out of index {:?}", map.layers[0]),
    Some(x) => x
  };
  match y_index.get(x) {
    None => panic!("Y value out of index {:?}", layer.tiles[y]),
    Some(val) => *val
  }
}

fn read_sprite_file(filename: &str) -> String {
  let path = Path::new(&filename);
  let mut file = match File::open(&path) {
    Ok(f) => f,
    Err(e) => panic!("File {} not found: {}", filename, e),
  };
  let mut buf = String::new();
  match file.read_to_string(&mut buf) {
    Ok(_) => buf,
    Err(e) => panic!("read file {} error {}", filename, e),
  }
}

pub fn load_character() -> Vec<CritterData> {
  let mut sprites = Vec::with_capacity(CHARACTER_BUF_LENGTH + 64);
  let character_json = read_sprite_file(CHARACTER_JSON_PATH);
  let character = match json::parse(&character_json) {
    Ok(res) => res,
    Err(e) => panic!("Character {} parse error {:?}", CHARACTER_JSON_PATH, e),
  };

  let frame_data = |character: &JsonValue, key: &String| -> [f32; 4] {
    [
      character["frames"][key]["frame"]["x"].as_f32().unwrap(),
      character["frames"][key]["frame"]["y"].as_f32().unwrap(),
      character["frames"][key]["frame"]["w"].as_f32().unwrap(),
      character["frames"][key]["frame"]["h"].as_f32().unwrap(),
    ]
  };

  for x in 0..16 {
    for y in 0..14 {
      let key = &format!("run_{}_{}", x, y);
      sprites.push(CritterData::new(frame_data(&character, key)));
    }
  }

  for x in 0..15 {
    for y in 0..4 {
      let key = &format!("fire_{}_{}", x, y);
      sprites.push(CritterData::new(frame_data(&character, key)));
    }
  }

  sprites
}

pub fn load_zombie() -> Vec<CritterData> {
  let mut sprites = Vec::with_capacity(256);
  let zombie_json = read_sprite_file(ZOMBIE_JSON_PATH);
  let zombie = match json::parse(&zombie_json) {
    Ok(res) => res,
    Err(e) => panic!("Zombie {} parse error {:?}", ZOMBIE_JSON_PATH, e),
  };

  let frame_data = |zombie: &JsonValue, key: &String| -> [f32; 4] {
    [
      zombie["frames"][key]["frame"]["x"].as_f32().unwrap(),
      zombie["frames"][key]["frame"]["y"].as_f32().unwrap(),
      zombie["frames"][key]["frame"]["w"].as_f32().unwrap(),
      zombie["frames"][key]["frame"]["h"].as_f32().unwrap(),
    ]
  };

  for x in 0..7 {
    for y in 0..7 {
      let key = &format!("critical_{}_{}", x, y);
      sprites.push(CritterData::new(frame_data(&zombie, key)));
    }
  }
  for x in 0..7 {
    for y in 0..5 {
      let key = &format!("normal_{}_{}", x, y);
      sprites.push(CritterData::new(frame_data(&zombie, key)));
    }
  }
  for x in 0..7 {
    for y in 0..4 {
      let key = &format!("still_{}_{}", x, y);
      sprites.push(CritterData::new(frame_data(&zombie, key)));
    }
  }

  for x in 0..7 {
    for y in 0..7 {
      let key = &format!("walk_{}_{}", x, y);
      sprites.push(CritterData::new(frame_data(&zombie, key)));
    }
  }
  sprites
}
<commit_msg>Refactor sprite file loading<commit_after>use std::{fs::File, io::BufReader, io::prelude::*, path::Path, string::String, vec::Vec};

use json;
use json::JsonValue;
use tiled;
use tiled::Map;

use crate::critter::CritterData;
use crate::game::constants::{CHARACTER_BUF_LENGTH, CHARACTER_JSON_PATH, ZOMBIE_JSON_PATH};

pub fn load_map_file(filename: &str) -> Map {
  let file = match File::open(&Path::new(&filename)) {
    Ok(f) => f,
    Err(e) => panic!("File {} not found: {}", filename, e),
  };
  let reader = BufReader::new(file);
  match tiled::parse(reader) {
    Ok(m) => m,
    Err(e) => panic!("Map parse error {:?}", e)
  }
}

pub fn get_map_tile(map: &Map, layer_index: usize, x: usize, y: usize) -> u32 {
  let layer = match map.layers.get(layer_index) {
    None => panic!("Layer_index value out of index {:?}", map.layers),
    Some(l) => l
  };
  let y_index = match layer.tiles.iter().rev().nth(y) {
    None => panic!("X value out of index {:?}", map.layers[0]),
    Some(x) => x
  };
  match y_index.get(x) {
    None => panic!("Y value out of index {:?}", layer.tiles[y]),
    Some(val) => *val
  }
}

fn read_sprite_file(filename: &str) -> String {
  let path = Path::new(&filename);
  let mut file = match File::open(&path) {
    Ok(f) => f,
    Err(e) => panic!("File {} not found: {}", filename, e),
  };
  let mut buf = String::new();
  match file.read_to_string(&mut buf) {
    Ok(_) => buf,
    Err(e) => panic!("read file {} error {}", filename, e),
  }
}

fn get_frame_data(character: &JsonValue, key: &String) -> CritterData {
  CritterData::new([
    character["frames"][key]["frame"]["x"].as_f32().unwrap(),
    character["frames"][key]["frame"]["y"].as_f32().unwrap(),
    character["frames"][key]["frame"]["w"].as_f32().unwrap(),
    character["frames"][key]["frame"]["h"].as_f32().unwrap(),
  ])
}

pub fn load_character() -> Vec<CritterData> {
  let mut sprites = Vec::with_capacity(CHARACTER_BUF_LENGTH + 64);
  let character_json = read_sprite_file(CHARACTER_JSON_PATH);
  let character = match json::parse(&character_json) {
    Ok(res) => res,
    Err(e) => panic!("Character {} parse error {:?}", CHARACTER_JSON_PATH, e),
  };

  for x in 0..16 {
    for y in 0..14 {
      let key = &format!("run_{}_{}", x, y);
      sprites.push(get_frame_data(&character, key));
    }
  }

  for x in 0..15 {
    for y in 0..4 {
      let key = &format!("fire_{}_{}", x, y);
      sprites.push(get_frame_data(&character, key));
    }
  }

  sprites
}

pub fn load_zombie() -> Vec<CritterData> {
  let mut sprites = Vec::with_capacity(256);
  let zombie_json = read_sprite_file(ZOMBIE_JSON_PATH);
  let zombie = match json::parse(&zombie_json) {
    Ok(res) => res,
    Err(e) => panic!("Zombie {} parse error {:?}", ZOMBIE_JSON_PATH, e),
  };

  for x in 0..7 {
    for y in 0..7 {
      let key = &format!("critical_{}_{}", x, y);
      sprites.push(get_frame_data(&zombie, key));
    }
  }
  for x in 0..7 {
    for y in 0..5 {
      let key = &format!("normal_{}_{}", x, y);
      sprites.push(get_frame_data(&zombie, key));
    }
  }
  for x in 0..7 {
    for y in 0..4 {
      let key = &format!("still_{}_{}", x, y);
      sprites.push(get_frame_data(&zombie, key));
    }
  }

  for x in 0..7 {
    for y in 0..7 {
      let key = &format!("walk_{}_{}", x, y);
      sprites.push(get_frame_data(&zombie, key));
    }
  }
  sprites
}
<|endoftext|>


<commit_before>use std::collections::BTreeMap;
use std::default::Default;
use std::fs::File;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};

use digest::Digest;
use ignore::WalkBuilder;
use serde_bytes;
use serde_cbor;
use serde_json;

use sha2;
use sha3;

use error;

#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct Database(Entry);

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum Entry {
    Directory(BTreeMap<PathBuf, Entry>),
    File(Metrics),
}

impl Default for Entry {
    fn default() -> Entry {
        Entry::Directory(BTreeMap::default())
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Metrics {
    sha2: HashSum,
    sha3: HashSum,
    size: u64,      // File size
    nul: bool,      // Does the file contain a NUL byte?
    nonascii: bool, // Does the file contain non-ASCII bytes?
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct HashSum(#[serde(with = "serde_bytes")] Vec<u8>);

#[derive(Default)]
struct EngineSize(u64);
impl EngineSize {
    fn input(&mut self, input: &[u8]) {
        self.0 += input.len() as u64;
    }
    fn result(self) -> u64 {
        self.0
    }
}

#[derive(Default)]
struct EngineNul(bool);
impl EngineNul {
    fn input(&mut self, input: &[u8]) {
        self.0 = self.0 || input.iter().any(|x| *x == 0);
    }
    fn result(self) -> bool {
        self.0
    }
}

#[derive(Default)]
struct EngineNonascii(bool);
impl EngineNonascii {
    fn input(&mut self, input: &[u8]) {
        self.0 = self.0 || input.iter().any(|x| x & 0x80 != 0);
    }
    fn result(self) -> bool {
        self.0
    }
}

#[derive(Default)]
struct Engines {
    sha2: sha2::Sha256,
    sha3: sha3::Sha3_256,
    size: EngineSize,
    nul: EngineNul,
    nonascii: EngineNonascii,
}

impl Engines {
    fn input(&mut self, input: &[u8]) {
        self.sha2.input(input);
        self.sha3.input(input);
        self.size.input(input);
        self.nul.input(input);
        self.nonascii.input(input);
    }
    fn result(self) -> Metrics {
        Metrics {
            sha2: HashSum(Vec::from(self.sha2.result().as_slice())),
            sha3: HashSum(Vec::from(self.sha3.result().as_slice())),
            size: self.size.result(),
            nul: self.nul.result(),
            nonascii: self.nonascii.result(),
        }
    }
}

fn compute_metrics<P>(path: P) -> Result<Metrics, error::Error>
where
    P: AsRef<Path>
{
    let mut f = File::open(path)?;

    let mut engines = Engines::default();

    let mut buffer = [0; 4096];
    loop {
        let n = f.read(&mut buffer[..])?;
        if n == 0 { break }
        engines.input(&buffer[0..n]);
    }
    Ok(engines.result())
}

trait BTreeMapExt<K, V> where K: Ord, V: Default {
    fn get_mut_default(&mut self, key: K) -> &mut V;
}

impl<K, V> BTreeMapExt<K, V> for BTreeMap<K, V>
where
    K: Ord + Clone,
    V: Default,
{
    fn get_mut_default(&mut self, key: K) -> &mut V {
        if self.contains_key(&key) {
            match self.get_mut(&key) {
                Some(value) => value,
                None => unreachable!(), // Already checked that key exists
            }
        } else {
            {
                // FIXME: Would prefer to avoid the clone
                match self.insert(key.clone(), V::default()) {
                    Some(_) => unreachable!(), // Already checked for key
                    None => (),
                }
            }
            {
                match self.get_mut(&key) {
                    Some(value) => value,
                    None => unreachable!(), // Already checked that key exists
                }
            }
        }
    }
}

impl Entry {
    fn insert(&mut self, path: PathBuf, file: Entry) {
        // Inner nodes in the tree should always be directories. If
        // the node is not a directory, that means we are inserting a
        // duplicate file. However, this function is only called from
        // the directory walker, which makes it impossible to observe
        // any duplicates. (And the database, after construction, is
        // always immutable.)
        match self {
            &mut Entry::Directory(ref mut entries) => {
                let mut components = path.components();
                let count = components.clone().count();
                let first = Path::new(components.next().expect("unreachable").as_os_str()).to_owned();
                let rest = components.as_path().to_owned();
                if count > 1 {
                    let mut subentry = entries.get_mut_default(first);
                    subentry.insert(rest, file);
                } else {
                    match entries.insert(first, file) {
                        Some(_) => unreachable!(), // See above
                        None => (),
                    }
                }
            }
            &mut Entry::File(_) => unreachable!()
        }
    }

    fn lookup(&self, path: &PathBuf) -> Option<&Entry> {
        match *self {
            Entry::Directory(ref entries) => {
                let mut components = path.components();
                let count = components.clone().count();
                let first = Path::new(components.next().expect("unreachable").as_os_str()).to_owned();
                let rest = components.as_path().to_owned();
                if count > 1 {
                    entries.get(&first).and_then(
                        |subentry| subentry.lookup(&rest))
                } else {
                    entries.get(&first)
                }
            }
            Entry::File(_) => unreachable!()
        }
    }
}

fn summarize<P>(path: P, expected: &Metrics, actual: &Metrics)
where
    P: AsRef<Path>,
{
    let diff = expected.size != actual.size ||
        expected.sha2 != actual.sha2 ||
        expected.sha3 != actual.sha3;
    if diff {
        let suspicious_size = expected.size > 0 && actual.size == 0;
        let suspicious_nul = expected.nul != actual.nul;
        let suspicious_nonascii = expected.nonascii != actual.nonascii;
        println!("{} differs", path.as_ref().display());
        if suspicious_size {
            println!("  suspicious: file was truncated");
        }
        if suspicious_nul {
            println!("  suspicious: original had no NUL bytes, but now does");
        }
        if suspicious_nonascii {
            println!("  suspicious: original had no non-ASCII bytes, but now does");
        }
    }
}

impl Database {
    fn insert(&mut self, path: PathBuf, entry: Entry) {
        self.0.insert(path, entry);
    }

    fn lookup(&self, path: &PathBuf) -> Option<&Entry> {
        self.0.lookup(path)
    }

    pub fn build<P>(root: P) -> Result<Database, error::Error>
    where
        P: AsRef<Path>,
    {
        let mut database = Database::default();
        for entry in WalkBuilder::new(&root).build() {
            let entry = entry?;
            if entry.file_type().map_or(false, |t| t.is_file()) {
                let metrics = compute_metrics(entry.path())?;
                let result = Entry::File(metrics);
                let short_path = if entry.path() == root.as_ref() {
                    Path::new(entry.path().file_name().expect("unreachable"))
                } else {
                    entry.path().strip_prefix(&root)?
                };
                database.insert(short_path.to_owned(), result);
            }
        }
        Ok(database)
    }

    pub fn check<P>(&self, root: P) -> Result<(), error::Error>
    where
        P: AsRef<Path>,
    {
        // FIXME: Doesn't check for missing entries
        for entry in WalkBuilder::new(&root).build() {
            let entry = entry?;
            if entry.file_type().map_or(false, |t| t.is_file()) {
                let actual = compute_metrics(entry.path())?;

                let short_path = if entry.path() == root.as_ref() {
                    Path::new(entry.path().file_name().expect("unreachable"))
                } else {
                    entry.path().strip_prefix(&root)?
                };
                let expected = self.lookup(&short_path.to_owned());
                match expected {
                    Some(&Entry::File(ref metrics)) => summarize(
                        short_path, metrics, &actual),
                    Some(&Entry::Directory(_)) => unreachable!(),
                    None => println!("{} was created", short_path.display()),
                }
            }
        }
        Ok(())
    }

    pub fn load_json<P>(path: P) -> Result<Database, error::Error>
    where
        P: AsRef<Path>
    {
        let f = File::open(path)?;
        Ok(serde_json::from_reader(f)?)
    }

    pub fn dump_json<P>(&self, path: P) -> Result<(), error::Error>
    where
        P: AsRef<Path>
    {
        let json = serde_json::to_string(self)?;
        let mut f = File::create(path)?;
        write!(f, "{}", json)?;
        Ok(())
    }

    pub fn load_cbor<P>(path: P) -> Result<Database, error::Error>
    where
        P: AsRef<Path>
    {
        let f = File::open(path)?;
        Ok(serde_cbor::from_reader(f)?)
    }

    pub fn dump_cbor<P>(&self, path: P) -> Result<(), error::Error>
    where
        P: AsRef<Path>
    {
        let cbor = serde_cbor::to_vec(self)?;
        let mut f = File::create(path)?;
        f.write_all(cbor.as_slice())?;
        Ok(())
    }
}

// impl std::fmt::Display for Database {
//     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
//         for (path, entry) in self.0.iter() {
//             match entry {
//                 &Entry::File(ref hashes) => {
//                     let hash: Vec<_> = hashes.sha2.0.iter().map(
//                         |b| format!("{:02x}", b)).collect();
//                     writeln!(f, "{} {}", hash.join(""), Path::new(path).display())?
//                 }
//             }
//         }
//         Ok(())
//     }
// }
<commit_msg>Entry API makes it so much easier to write this.<commit_after>use std::collections::BTreeMap;
use std::default::Default;
use std::fs::File;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};

use digest::Digest;
use ignore::WalkBuilder;
use serde_bytes;
use serde_cbor;
use serde_json;

use sha2;
use sha3;

use error;

#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct Database(Entry);

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum Entry {
    Directory(BTreeMap<PathBuf, Entry>),
    File(Metrics),
}

impl Default for Entry {
    fn default() -> Entry {
        Entry::Directory(BTreeMap::default())
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Metrics {
    sha2: HashSum,
    sha3: HashSum,
    size: u64,      // File size
    nul: bool,      // Does the file contain a NUL byte?
    nonascii: bool, // Does the file contain non-ASCII bytes?
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct HashSum(#[serde(with = "serde_bytes")] Vec<u8>);

#[derive(Default)]
struct EngineSize(u64);
impl EngineSize {
    fn input(&mut self, input: &[u8]) {
        self.0 += input.len() as u64;
    }
    fn result(self) -> u64 {
        self.0
    }
}

#[derive(Default)]
struct EngineNul(bool);
impl EngineNul {
    fn input(&mut self, input: &[u8]) {
        self.0 = self.0 || input.iter().any(|x| *x == 0);
    }
    fn result(self) -> bool {
        self.0
    }
}

#[derive(Default)]
struct EngineNonascii(bool);
impl EngineNonascii {
    fn input(&mut self, input: &[u8]) {
        self.0 = self.0 || input.iter().any(|x| x & 0x80 != 0);
    }
    fn result(self) -> bool {
        self.0
    }
}

#[derive(Default)]
struct Engines {
    sha2: sha2::Sha256,
    sha3: sha3::Sha3_256,
    size: EngineSize,
    nul: EngineNul,
    nonascii: EngineNonascii,
}

impl Engines {
    fn input(&mut self, input: &[u8]) {
        self.sha2.input(input);
        self.sha3.input(input);
        self.size.input(input);
        self.nul.input(input);
        self.nonascii.input(input);
    }
    fn result(self) -> Metrics {
        Metrics {
            sha2: HashSum(Vec::from(self.sha2.result().as_slice())),
            sha3: HashSum(Vec::from(self.sha3.result().as_slice())),
            size: self.size.result(),
            nul: self.nul.result(),
            nonascii: self.nonascii.result(),
        }
    }
}

fn compute_metrics<P>(path: P) -> Result<Metrics, error::Error>
where
    P: AsRef<Path>
{
    let mut f = File::open(path)?;

    let mut engines = Engines::default();

    let mut buffer = [0; 4096];
    loop {
        let n = f.read(&mut buffer[..])?;
        if n == 0 { break }
        engines.input(&buffer[0..n]);
    }
    Ok(engines.result())
}

trait BTreeMapExt<K, V> where K: Ord, V: Default {
    fn get_mut_default(&mut self, key: K) -> &mut V;
}

impl<K, V> BTreeMapExt<K, V> for BTreeMap<K, V>
where
    K: Ord + Clone,
    V: Default,
{
    fn get_mut_default(&mut self, key: K) -> &mut V {
        self.entry(key).or_insert(V::default())
    }
}

impl Entry {
    fn insert(&mut self, path: PathBuf, file: Entry) {
        // Inner nodes in the tree should always be directories. If
        // the node is not a directory, that means we are inserting a
        // duplicate file. However, this function is only called from
        // the directory walker, which makes it impossible to observe
        // any duplicates. (And the database, after construction, is
        // always immutable.)
        match self {
            &mut Entry::Directory(ref mut entries) => {
                let mut components = path.components();
                let count = components.clone().count();
                let first = Path::new(components.next().expect("unreachable").as_os_str()).to_owned();
                let rest = components.as_path().to_owned();
                if count > 1 {
                    let mut subentry = entries.get_mut_default(first);
                    subentry.insert(rest, file);
                } else {
                    match entries.insert(first, file) {
                        Some(_) => unreachable!(), // See above
                        None => (),
                    }
                }
            }
            &mut Entry::File(_) => unreachable!()
        }
    }

    fn lookup(&self, path: &PathBuf) -> Option<&Entry> {
        match *self {
            Entry::Directory(ref entries) => {
                let mut components = path.components();
                let count = components.clone().count();
                let first = Path::new(components.next().expect("unreachable").as_os_str()).to_owned();
                let rest = components.as_path().to_owned();
                if count > 1 {
                    entries.get(&first).and_then(
                        |subentry| subentry.lookup(&rest))
                } else {
                    entries.get(&first)
                }
            }
            Entry::File(_) => unreachable!()
        }
    }
}

fn summarize<P>(path: P, expected: &Metrics, actual: &Metrics)
where
    P: AsRef<Path>,
{
    let diff = expected.size != actual.size ||
        expected.sha2 != actual.sha2 ||
        expected.sha3 != actual.sha3;
    if diff {
        let suspicious_size = expected.size > 0 && actual.size == 0;
        let suspicious_nul = expected.nul != actual.nul;
        let suspicious_nonascii = expected.nonascii != actual.nonascii;
        println!("{} differs", path.as_ref().display());
        if suspicious_size {
            println!("  suspicious: file was truncated");
        }
        if suspicious_nul {
            println!("  suspicious: original had no NUL bytes, but now does");
        }
        if suspicious_nonascii {
            println!("  suspicious: original had no non-ASCII bytes, but now does");
        }
    }
}

impl Database {
    fn insert(&mut self, path: PathBuf, entry: Entry) {
        self.0.insert(path, entry);
    }

    fn lookup(&self, path: &PathBuf) -> Option<&Entry> {
        self.0.lookup(path)
    }

    pub fn build<P>(root: P) -> Result<Database, error::Error>
    where
        P: AsRef<Path>,
    {
        let mut database = Database::default();
        for entry in WalkBuilder::new(&root).build() {
            let entry = entry?;
            if entry.file_type().map_or(false, |t| t.is_file()) {
                let metrics = compute_metrics(entry.path())?;
                let result = Entry::File(metrics);
                let short_path = if entry.path() == root.as_ref() {
                    Path::new(entry.path().file_name().expect("unreachable"))
                } else {
                    entry.path().strip_prefix(&root)?
                };
                database.insert(short_path.to_owned(), result);
            }
        }
        Ok(database)
    }

    pub fn check<P>(&self, root: P) -> Result<(), error::Error>
    where
        P: AsRef<Path>,
    {
        // FIXME: Doesn't check for missing entries
        for entry in WalkBuilder::new(&root).build() {
            let entry = entry?;
            if entry.file_type().map_or(false, |t| t.is_file()) {
                let actual = compute_metrics(entry.path())?;

                let short_path = if entry.path() == root.as_ref() {
                    Path::new(entry.path().file_name().expect("unreachable"))
                } else {
                    entry.path().strip_prefix(&root)?
                };
                let expected = self.lookup(&short_path.to_owned());
                match expected {
                    Some(&Entry::File(ref metrics)) => summarize(
                        short_path, metrics, &actual),
                    Some(&Entry::Directory(_)) => unreachable!(),
                    None => println!("{} was created", short_path.display()),
                }
            }
        }
        Ok(())
    }

    pub fn load_json<P>(path: P) -> Result<Database, error::Error>
    where
        P: AsRef<Path>
    {
        let f = File::open(path)?;
        Ok(serde_json::from_reader(f)?)
    }

    pub fn dump_json<P>(&self, path: P) -> Result<(), error::Error>
    where
        P: AsRef<Path>
    {
        let json = serde_json::to_string(self)?;
        let mut f = File::create(path)?;
        write!(f, "{}", json)?;
        Ok(())
    }

    pub fn load_cbor<P>(path: P) -> Result<Database, error::Error>
    where
        P: AsRef<Path>
    {
        let f = File::open(path)?;
        Ok(serde_cbor::from_reader(f)?)
    }

    pub fn dump_cbor<P>(&self, path: P) -> Result<(), error::Error>
    where
        P: AsRef<Path>
    {
        let cbor = serde_cbor::to_vec(self)?;
        let mut f = File::create(path)?;
        f.write_all(cbor.as_slice())?;
        Ok(())
    }
}

// impl std::fmt::Display for Database {
//     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
//         for (path, entry) in self.0.iter() {
//             match entry {
//                 &Entry::File(ref hashes) => {
//                     let hash: Vec<_> = hashes.sha2.0.iter().map(
//                         |b| format!("{:02x}", b)).collect();
//                     writeln!(f, "{} {}", hash.join(""), Path::new(path).display())?
//                 }
//             }
//         }
//         Ok(())
//     }
// }
<|endoftext|>