<commit_before>//! leveldb iterators
//!
//! Iteration is one of the most important parts of leveldb. This module provides
//! Iterators to iterate over key, values and pairs of both.
use leveldb_sys::{leveldb_iterator_t, leveldb_iter_seek_to_first, leveldb_iter_destroy,
                  leveldb_iter_seek_to_last, leveldb_create_iterator, leveldb_iter_valid,
                  leveldb_iter_next, leveldb_iter_key, leveldb_iter_value,
                  leveldb_readoptions_destroy, leveldb_iter_seek};
use libc::{size_t, c_char};
use std::iter;
use super::Database;
use super::options::{ReadOptions, c_readoptions};
use super::key::{Key, from_u8};
use std::slice::from_raw_parts;
use std::marker::PhantomData;

#[allow(missing_docs)]
struct RawIterator {
    ptr: *mut leveldb_iterator_t,
}

#[allow(missing_docs)]
impl Drop for RawIterator {
    fn drop(&mut self) {
        unsafe { leveldb_iter_destroy(self.ptr) }
    }
}

/// An iterator over the leveldb keyspace.
///
/// Returns key and value as a tuple.
pub struct Iterator<'a, K: Key + 'a> {
    start: bool,
    // Iterator accesses the Database through a leveldb_iter_t pointer
    // but needs to hold the reference for lifetime tracking
    #[allow(dead_code)]
    database: PhantomData<&'a Database<K>>,
    iter: RawIterator,
    from: Option<&'a K>,
    to: Option<&'a K>,
}

/// An iterator over the leveldb keyspace.
///
/// Returns just the keys.
pub struct KeyIterator<'a, K: Key + 'a> {
    start: bool,
    // Iterator accesses the Database through a leveldb_iter_t pointer
    // but needs to hold the reference for lifetime tracking
    #[allow(dead_code)]
    database: PhantomData<&'a Database<K>>,
    iter: RawIterator,
    from: Option<&'a K>,
    to: Option<&'a K>,
}

/// An iterator over the leveldb keyspace.
///
/// Returns just the value.
pub struct ValueIterator<'a, K: Key + 'a> {
    start: bool,
    // Iterator accesses the Database through a leveldb_iter_t pointer
    // but needs to hold the reference for lifetime tracking
    #[allow(dead_code)]
    database: PhantomData<&'a Database<K>>,
    iter: RawIterator,
    from: Option<&'a K>,
    to: Option<&'a K>,
}

/// A trait to allow access to the three main iteration styles of leveldb.
pub trait Iterable<'a, K: Key + 'a> {
    /// Return an Iterator iterating over (Key,Value) pairs
    fn iter(&'a self, options: ReadOptions<'a, K>) -> Iterator<K>;
    /// Returns an Iterator iterating over Keys only.
    fn keys_iter(&'a self, options: ReadOptions<'a, K>) -> KeyIterator<K>;
    /// Returns an Iterator iterating over Values only.
    fn value_iter(&'a self, options: ReadOptions<'a, K>) -> ValueIterator<K>;
}

impl<'a, K: Key + 'a> Iterable<'a, K> for Database<K> {
    fn iter(&'a self, options: ReadOptions<'a, K>) -> Iterator<K> {
        Iterator::new(self, options)
    }

    fn keys_iter(&'a self, options: ReadOptions<'a, K>) -> KeyIterator<K> {
        KeyIterator::new(self, options)
    }

    fn value_iter(&'a self, options: ReadOptions<'a, K>) -> ValueIterator<K> {
        ValueIterator::new(self, options)
    }
}

#[allow(missing_docs)]
pub trait LevelDBIterator<'a, K: Key> {
    #[inline]
    fn raw_iterator(&self) -> *mut leveldb_iterator_t;

    #[inline]
    fn start(&self) -> bool;

    #[inline]
    fn started(&mut self);

    fn from(self, key: &'a K) -> Self;
    fn to(self, key: &'a K) -> Self;

    fn from_key(&self) -> Option<&K>;
    fn to_key(&self) -> Option<&K>;

    fn valid(&self) -> bool {
        unsafe { leveldb_iter_valid(self.raw_iterator()) != 0 }
    }

    fn advance(&mut self) -> bool {
        unsafe {
            if !self.start() {

                leveldb_iter_next(self.raw_iterator());
            } else {
                if let Some(k) = self.from_key() {
                    self.seek(k)
                }
                self.started();
            }
        }
        self.valid()
    }

    fn key(&self) -> K {
        unsafe {
            let length: size_t = 0;
            let value = leveldb_iter_key(self.raw_iterator(), &length) as *const u8;
            from_u8(from_raw_parts(value, length as usize))
        }
    }

    fn value(&self) -> Vec<u8> {
        unsafe {
            let length: size_t = 0;
            let value = leveldb_iter_value(self.raw_iterator(), &length) as *const u8;
            from_raw_parts(value, length as usize).to_vec()
        }
    }

    fn seek_to_first(&self) {
        unsafe { leveldb_iter_seek_to_first(self.raw_iterator()) }
    }

    fn seek_to_last(&self) {
        if let Some(k) = self.to_key() {
            self.seek(k);
        } else {
            unsafe {
                leveldb_iter_seek_to_last(self.raw_iterator());
            }
        }
    }

    fn seek(&self, key: &K) {
        unsafe {
            key.as_slice(|k| {
                leveldb_iter_seek(self.raw_iterator(),
                                  k.as_ptr() as *mut c_char,
                                  k.len() as size_t);
            })
        }
    }
}

impl<'a, K: Key> Iterator<'a, K> {
    fn new(database: &'a Database<K>, options: ReadOptions<'a, K>) -> Iterator<'a, K> {
        unsafe {
            let c_readoptions = c_readoptions(&options);
            let ptr = leveldb_create_iterator(database.database.ptr, c_readoptions);
            leveldb_readoptions_destroy(c_readoptions);
            leveldb_iter_seek_to_first(ptr);
            Iterator {
                start: true,
                iter: RawIterator { ptr: ptr },
                database: PhantomData,
                from: None,
                to: None,
            }
        }
    }

    /// return the last element of the iterator
    pub fn last(self) -> Option<(K, Vec<u8>)> {
        self.seek_to_last();
        Some((self.key(), self.value()))
    }
}

impl<'a, K: Key> LevelDBIterator<'a, K> for Iterator<'a,K> {
    #[inline]
    fn raw_iterator(&self) -> *mut leveldb_iterator_t {
        self.iter.ptr
    }

    #[inline]
    fn start(&self) -> bool {
        self.start
    }

    #[inline]
    fn started(&mut self) {
        self.start = false
    }

    fn from(mut self, key: &'a K) -> Self {
        self.from = Some(key);
        self
    }

    fn to(mut self, key: &'a K) -> Self {
        self.to = Some(key);
        self
    }

    fn from_key(&self) -> Option<&K> {
        self.from
    }

    fn to_key(&self) -> Option<&K> {
        self.to
    }
}

impl<'a,K: Key> KeyIterator<'a,K> {
    fn new(database: &'a Database<K>, options: ReadOptions<'a, K>) -> KeyIterator<'a, K> {
        unsafe {
            let c_readoptions = c_readoptions(&options);
            let ptr = leveldb_create_iterator(database.database.ptr, c_readoptions);
            leveldb_readoptions_destroy(c_readoptions);
            leveldb_iter_seek_to_first(ptr);
            KeyIterator {
                start: true,
                iter: RawIterator { ptr: ptr },
                database: PhantomData,
                from: None,
                to: None,
            }
        }
    }

    /// return the last element of the iterator
    pub fn last(self) -> Option<K> {
        self.seek_to_last();
        Some(self.key())
    }
}

impl<'a,K: Key> LevelDBIterator<'a, K> for KeyIterator<'a,K> {
    #[inline]
    fn raw_iterator(&self) -> *mut leveldb_iterator_t {
        self.iter.ptr
    }

    #[inline]
    fn start(&self) -> bool {
        self.start
    }

    #[inline]
    fn started(&mut self) {
        self.start = false
    }

    fn from(mut self, key: &'a K) -> Self {
        self.from = Some(key);
        self
    }

    fn to(mut self, key: &'a K) -> Self {
        self.to = Some(key);
        self
    }

    fn from_key(&self) -> Option<&K> {
        self.from
    }

    fn to_key(&self) -> Option<&K> {
        self.to
    }
}

impl<'a,K: Key> ValueIterator<'a,K> {
    fn new(database: &'a Database<K>, options: ReadOptions<'a, K>) -> ValueIterator<'a, K> {
        unsafe {
            let c_readoptions = c_readoptions(&options);
            let ptr = leveldb_create_iterator(database.database.ptr, c_readoptions);
            leveldb_readoptions_destroy(c_readoptions);
            leveldb_iter_seek_to_first(ptr);
            ValueIterator {
                start: true,
                iter: RawIterator { ptr: ptr },
                database: PhantomData,
                from: None,
                to: None,
            }
        }
    }

    /// return the last element of the iterator
    pub fn last(self) -> Option<Vec<u8>> {
        self.seek_to_last();
        Some(self.value())
    }
}

impl<'a,K: Key> LevelDBIterator<'a, K> for ValueIterator<'a,K> {
    #[inline]
    fn raw_iterator(&self) -> *mut leveldb_iterator_t {
        self.iter.ptr
    }

    #[inline]
    fn start(&self) -> bool {
        self.start
    }

    #[inline]
    fn started(&mut self) {
        self.start = false
    }

    fn from(mut self, key: &'a K) -> Self {
        self.from = Some(key);
        self
    }

    fn to(mut self, key: &'a K) -> Self {
        self.to = Some(key);
        self
    }

    fn from_key(&self) -> Option<&K> {
        self.from
    }

    fn to_key(&self) -> Option<&K> {
        self.to
    }
}

impl<'a,K: Key> iter::Iterator for Iterator<'a,K> {
  type Item = (K,Vec<u8>);

    fn next(&mut self) -> Option<(K, Vec<u8>)> {
        if self.advance() {
            Some((self.key(), self.value()))
        } else {
            None
        }
    }
}

impl<'a, K: Key> iter::Iterator for KeyIterator<'a,K> {
  type Item = K;

    fn next(&mut self) -> Option<K> {
        if self.advance() {
            Some(self.key())
        } else {
            None
        }
    }
}

impl<'a, K: Key> iter::Iterator for ValueIterator<'a,K> {
  type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        if self.advance() {
            Some(self.value())
        } else {
            None
        }
    }
}
<commit_msg>Deduplicated iterator code.<commit_after>//! leveldb iterators
//!
//! Iteration is one of the most important parts of leveldb. This module provides
//! Iterators to iterate over key, values and pairs of both.
use leveldb_sys::{leveldb_iterator_t, leveldb_iter_seek_to_first, leveldb_iter_destroy,
                  leveldb_iter_seek_to_last, leveldb_create_iterator, leveldb_iter_valid,
                  leveldb_iter_next, leveldb_iter_key, leveldb_iter_value,
                  leveldb_readoptions_destroy, leveldb_iter_seek};
use libc::{size_t, c_char};
use std::iter;
use super::Database;
use super::options::{ReadOptions, c_readoptions};
use super::key::{Key, from_u8};
use std::slice::from_raw_parts;
use std::marker::PhantomData;

#[allow(missing_docs)]
struct RawIterator {
    ptr: *mut leveldb_iterator_t,
}

#[allow(missing_docs)]
impl Drop for RawIterator {
    fn drop(&mut self) {
        unsafe { leveldb_iter_destroy(self.ptr) }
    }
}

/// An iterator over the leveldb keyspace.
///
/// Returns key and value as a tuple.
pub struct Iterator<'a, K: Key + 'a> {
    start: bool,
    // Iterator accesses the Database through a leveldb_iter_t pointer
    // but needs to hold the reference for lifetime tracking
    #[allow(dead_code)]
    database: PhantomData<&'a Database<K>>,
    iter: RawIterator,
    from: Option<&'a K>,
    to: Option<&'a K>,
}

/// An iterator over the leveldb keyspace.
///
/// Returns just the keys.
pub struct KeyIterator<'a, K: Key + 'a> {
    inner: Iterator<'a, K>,
}

/// An iterator over the leveldb keyspace.
///
/// Returns just the value.
pub struct ValueIterator<'a, K: Key + 'a> {
    inner: Iterator<'a, K>,
}

/// A trait to allow access to the three main iteration styles of leveldb.
pub trait Iterable<'a, K: Key + 'a> {
    /// Return an Iterator iterating over (Key,Value) pairs
    fn iter(&'a self, options: ReadOptions<'a, K>) -> Iterator<K>;
    /// Returns an Iterator iterating over Keys only.
    fn keys_iter(&'a self, options: ReadOptions<'a, K>) -> KeyIterator<K>;
    /// Returns an Iterator iterating over Values only.
    fn value_iter(&'a self, options: ReadOptions<'a, K>) -> ValueIterator<K>;
}

impl<'a, K: Key + 'a> Iterable<'a, K> for Database<K> {
    fn iter(&'a self, options: ReadOptions<'a, K>) -> Iterator<K> {
        Iterator::new(self, options)
    }

    fn keys_iter(&'a self, options: ReadOptions<'a, K>) -> KeyIterator<K> {
        KeyIterator::new(self, options)
    }

    fn value_iter(&'a self, options: ReadOptions<'a, K>) -> ValueIterator<K> {
        ValueIterator::new(self, options)
    }
}

#[allow(missing_docs)]
pub trait LevelDBIterator<'a, K: Key> {
    #[inline]
    fn raw_iterator(&self) -> *mut leveldb_iterator_t;

    #[inline]
    fn start(&self) -> bool;

    #[inline]
    fn started(&mut self);

    fn from(self, key: &'a K) -> Self;
    fn to(self, key: &'a K) -> Self;

    fn from_key(&self) -> Option<&K>;
    fn to_key(&self) -> Option<&K>;

    fn valid(&self) -> bool {
        unsafe { leveldb_iter_valid(self.raw_iterator()) != 0 }
    }

    fn advance(&mut self) -> bool {
        unsafe {
            if !self.start() {

                leveldb_iter_next(self.raw_iterator());
            } else {
                if let Some(k) = self.from_key() {
                    self.seek(k)
                }
                self.started();
            }
        }
        self.valid()
    }

    fn key(&self) -> K {
        unsafe {
            let length: size_t = 0;
            let value = leveldb_iter_key(self.raw_iterator(), &length) as *const u8;
            from_u8(from_raw_parts(value, length as usize))
        }
    }

    fn value(&self) -> Vec<u8> {
        unsafe {
            let length: size_t = 0;
            let value = leveldb_iter_value(self.raw_iterator(), &length) as *const u8;
            from_raw_parts(value, length as usize).to_vec()
        }
    }

    fn seek_to_first(&self) {
        unsafe { leveldb_iter_seek_to_first(self.raw_iterator()) }
    }

    fn seek_to_last(&self) {
        if let Some(k) = self.to_key() {
            self.seek(k);
        } else {
            unsafe {
                leveldb_iter_seek_to_last(self.raw_iterator());
            }
        }
    }

    fn seek(&self, key: &K) {
        unsafe {
            key.as_slice(|k| {
                leveldb_iter_seek(self.raw_iterator(),
                                  k.as_ptr() as *mut c_char,
                                  k.len() as size_t);
            })
        }
    }
}

impl<'a, K: Key> Iterator<'a, K> {
    fn new(database: &'a Database<K>, options: ReadOptions<'a, K>) -> Iterator<'a, K> {
        unsafe {
            let c_readoptions = c_readoptions(&options);
            let ptr = leveldb_create_iterator(database.database.ptr, c_readoptions);
            leveldb_readoptions_destroy(c_readoptions);
            leveldb_iter_seek_to_first(ptr);
            Iterator {
                start: true,
                iter: RawIterator { ptr: ptr },
                database: PhantomData,
                from: None,
                to: None,
            }
        }
    }

    /// return the last element of the iterator
    pub fn last(self) -> Option<(K, Vec<u8>)> {
        self.seek_to_last();
        Some((self.key(), self.value()))
    }
}

impl<'a, K: Key> LevelDBIterator<'a, K> for Iterator<'a,K> {
    #[inline]
    fn raw_iterator(&self) -> *mut leveldb_iterator_t {
        self.iter.ptr
    }

    #[inline]
    fn start(&self) -> bool {
        self.start
    }

    #[inline]
    fn started(&mut self) {
        self.start = false
    }

    fn from(mut self, key: &'a K) -> Self {
        self.from = Some(key);
        self
    }

    fn to(mut self, key: &'a K) -> Self {
        self.to = Some(key);
        self
    }

    fn from_key(&self) -> Option<&K> {
        self.from
    }

    fn to_key(&self) -> Option<&K> {
        self.to
    }
}

impl<'a,K: Key> KeyIterator<'a,K> {
    fn new(database: &'a Database<K>, options: ReadOptions<'a, K>) -> KeyIterator<'a, K> {
        KeyIterator { inner: Iterator::new(database, options) }
    }

    /// return the last element of the iterator
    pub fn last(self) -> Option<K> {
        self.seek_to_last();
        Some(self.key())
    }
}

impl<'a,K: Key> LevelDBIterator<'a, K> for KeyIterator<'a,K> {
    #[inline]
    fn raw_iterator(&self) -> *mut leveldb_iterator_t {
        self.inner.iter.ptr
    }

    #[inline]
    fn start(&self) -> bool {
        self.inner.start
    }

    #[inline]
    fn started(&mut self) {
        self.inner.start = false
    }

    fn from(mut self, key: &'a K) -> Self {
        self.inner.from = Some(key);
        self
    }

    fn to(mut self, key: &'a K) -> Self {
        self.inner.to = Some(key);
        self
    }

    fn from_key(&self) -> Option<&K> {
        self.inner.from
    }

    fn to_key(&self) -> Option<&K> {
        self.inner.to
    }
}

impl<'a,K: Key> ValueIterator<'a,K> {
    fn new(database: &'a Database<K>, options: ReadOptions<'a, K>) -> ValueIterator<'a, K> {
        ValueIterator { inner: Iterator::new(database, options) }
    }

    /// return the last element of the iterator
    pub fn last(self) -> Option<Vec<u8>> {
        self.seek_to_last();
        Some(self.value())
    }
}

impl<'a,K: Key> LevelDBIterator<'a, K> for ValueIterator<'a,K> {
    #[inline]
    fn raw_iterator(&self) -> *mut leveldb_iterator_t {
        self.inner.iter.ptr
    }

    #[inline]
    fn start(&self) -> bool {
        self.inner.start
    }

    #[inline]
    fn started(&mut self) {
        self.inner.start = false
    }

    fn from(mut self, key: &'a K) -> Self {
        self.inner.from = Some(key);
        self
    }

    fn to(mut self, key: &'a K) -> Self {
        self.inner.to = Some(key);
        self
    }

    fn from_key(&self) -> Option<&K> {
        self.inner.from
    }

    fn to_key(&self) -> Option<&K> {
        self.inner.to
    }
}

impl<'a,K: Key> iter::Iterator for Iterator<'a,K> {
  type Item = (K,Vec<u8>);

    fn next(&mut self) -> Option<(K, Vec<u8>)> {
        if self.advance() {
            Some((self.key(), self.value()))
        } else {
            None
        }
    }
}

impl<'a, K: Key> iter::Iterator for KeyIterator<'a,K> {
  type Item = K;

    fn next(&mut self) -> Option<K> {
        if self.advance() {
            Some(self.key())
        } else {
            None
        }
    }
}

impl<'a, K: Key> iter::Iterator for ValueIterator<'a,K> {
  type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        if self.advance() {
            Some(self.value())
        } else {
            None
        }
    }
}
<|endoftext|>


<commit_before>//! Kata OS command line interface

// This brief bootstrap of Rust-in-Kata prototypes a minimal modular design
// for the DebugConsole CLI use case.
//
// * kata_io Read/Write interface (or move to std::, but that requires alloc)
// * kata_uart_client implementation of the kata_io interface
// * kata_line_reader
// * kata_shell
// * kata_debug_console main entry point fn run()

#![no_std]

use core::slice;
use kata_os_common::camkes::Camkes;

extern "C" {
    static cpio_archive: *const u8; // CPIO archive of built-in files
}

static mut CAMKES: Camkes = Camkes::new("DebugConsole");

#[no_mangle]
pub unsafe extern "C" fn pre_init() {
    // TODO(b/200946906): Review per-component heap allocations, including this one.
    const HEAP_SIZE: usize = 1 << 20;
    static mut HEAP_MEMORY: [u8; HEAP_SIZE] = [0; HEAP_SIZE];
    CAMKES.pre_init(
        log::LevelFilter::Debug,
        &mut HEAP_MEMORY,
    );
}

/// Entry point for DebugConsole. Runs the shell with UART IO.
#[no_mangle]
pub extern "C" fn run() -> ! {
    let mut tx = kata_uart_client::Tx::new();
    let mut rx = kata_io::BufReader::new(kata_uart_client::Rx::new());
    let cpio_archive_ref = unsafe {
        // XXX want begin-end or begin+size instead of a fixed-size block
        slice::from_raw_parts(cpio_archive, 16777216)
    };
    kata_shell::repl(&mut tx, &mut rx, cpio_archive_ref);
}
<commit_msg>DebugConsole: reduce heap size to 16KB<commit_after>//! Kata OS command line interface

// This brief bootstrap of Rust-in-Kata prototypes a minimal modular design
// for the DebugConsole CLI use case.
//
// * kata_io Read/Write interface (or move to std::, but that requires alloc)
// * kata_uart_client implementation of the kata_io interface
// * kata_line_reader
// * kata_shell
// * kata_debug_console main entry point fn run()

#![no_std]

use core::slice;
use kata_os_common::camkes::Camkes;

extern "C" {
    static cpio_archive: *const u8; // CPIO archive of built-in files
}

static mut CAMKES: Camkes = Camkes::new("DebugConsole");

#[no_mangle]
pub unsafe extern "C" fn pre_init() {
    const HEAP_SIZE: usize = 16 * 1024;
    static mut HEAP_MEMORY: [u8; HEAP_SIZE] = [0; HEAP_SIZE];
    CAMKES.pre_init(
        log::LevelFilter::Debug,
        &mut HEAP_MEMORY,
    );
}

/// Entry point for DebugConsole. Runs the shell with UART IO.
#[no_mangle]
pub extern "C" fn run() -> ! {
    let mut tx = kata_uart_client::Tx::new();
    let mut rx = kata_io::BufReader::new(kata_uart_client::Rx::new());
    let cpio_archive_ref = unsafe {
        // XXX want begin-end or begin+size instead of a fixed-size block
        slice::from_raw_parts(cpio_archive, 16777216)
    };
    kata_shell::repl(&mut tx, &mut rx, cpio_archive_ref);
}
<|endoftext|>


<commit_before>// TODO: Add License information.

//! ESIL (Evaluable Strings Intermediate Language) is the IL used by radare2.
//! For a complete documentation of ESIL please checkout:
//!  https://github.com/radare/radare2/wiki/ESIL
//! This module is used to parse ESIL strings and convert them into the IR.

use std::collections::HashMap;

#[derive(Debug, Copy, Clone)]
pub enum Arity {
    Zero,
    Unary,
    Binary,
    Ternary,
}

impl Arity {
    pub fn n(self) -> u8 {
        match self {
            Arity::Zero    => 0,
            Arity::Unary   => 1,
            Arity::Binary  => 2,
            Arity::Ternary => 3,
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub struct Operator<'a> {
    op: &'a str,
    arity: Arity,
}

impl<'a> Operator<'a> {
    pub fn new(op: &str, n: Arity) -> Operator {
        Operator { op: op, arity: n }
    }

    pub fn nop() -> Operator<'a> {
        Operator { op: "nop", arity: Arity::Zero }
    }
}

#[derive(Debug, Copy, Clone)]
pub enum Location {
    Memory,
    Register,
    Constant,
    Temporary,
    Unknown,
    Null,
}

#[derive(Debug, Clone)]
/// Value is used to represent operands to an operator in a statement.
pub struct Value {
    /// Name
    name: String,
    /// Size in bits.
    size: u8,
    /// Memory, Register, Constant or Temporary.
    location: Location,
    /// Value if evaluable.
    value: i64,
    // TODO: Convert from u32 to TypeSet.
    // Every value can be considered in terms of typesets rather than fixed
    // types which can then be narrowed down based on the analysis.
    // TypeSet can be implemented simply as a bit-vector.
    typeset: u32,
}

impl<'a> Value {
    pub fn new(name: &'a str, size: u8, location: Location, value: i64, typeset: u32) -> Value {
        Value {
            name: name.to_string().clone(),
            size: size,
            location: location,
            value: value,
            typeset: typeset,
        }
    }

    pub fn null() -> Value {
        Value {
            name: String::from(""),
            size: 0,
            location: Location::Null,
            value: 0,
            typeset: 0,
        }
    }

    pub fn new_tmp(i: u64) -> Value {
        Value {
            name: format!("tmp_{:x}", i),
            size: 0,
            location: Location::Temporary,
            value: 0,
            typeset: 0,
        }
    }
}

#[derive(Debug, Clone)]
pub struct Instruction<'a> {
    opcode: Operator<'a>,
    dest: Value,
    operand_1: Value,
    operand_2: Value,
}

impl<'a> Instruction<'a> {
    pub fn new(opcode: Operator<'a>, dest: Value, op1: Value, op2: Value) -> Instruction<'a> {
        Instruction {
            opcode: opcode,
            dest: dest,
            operand_1: op1,
            operand_2: op2,
        }
    }
    pub fn to_string(&self) -> String {
        if self.opcode.op == "=" {
            format!("{} {} {}", self.operand_1.name, self.opcode.op, self.operand_2.name)
        } else {
            format!("{} {} {} {} {}", self.dest.name, "=", self.operand_1.name, self.opcode.op, self.operand_2.name)
        }
    }
}

macro_rules! hash {
    ( $( ($x:expr, $y:expr) ),* ) => {
        {
            let mut temp_hash = HashMap::new();
            $(
                temp_hash.insert($x, $y);
             )*
            temp_hash
        }
    };
}

fn init_opset() -> HashMap<&'static str, Operator<'static>> {
    // Make a map from esil string to struct Operator.
    // (operator: &str, arity: Arity).
    // Possible Operatortimization:  Move to compile-time generation ?
    hash![
        ("==" , Operator::new("==", Arity::Binary)),
        ("<"  , Operator::new("<" , Arity::Binary)),
        (">"  , Operator::new(">" , Arity::Binary)),
        ("<=" , Operator::new("<=", Arity::Binary)),
        (">=" , Operator::new(">=", Arity::Binary)),
        ("<<" , Operator::new("<<", Arity::Binary)),
        (">>" , Operator::new(">>", Arity::Binary)),
        ("&"  , Operator::new("&" , Arity::Binary)),
        ("|"  , Operator::new("|" , Arity::Binary)),
        ("="  , Operator::new("=" , Arity::Binary)),
        ("*"  , Operator::new("*" , Arity::Binary)),
        ("^"  , Operator::new("^" , Arity::Binary)),
        ("+"  , Operator::new("+" , Arity::Binary)),
        ("-"  , Operator::new("-" , Arity::Binary)),
        ("/"  , Operator::new("/" , Arity::Binary)),
        ("%"  , Operator::new("%" , Arity::Binary)),
        ("?{" , Operator::new("?{", Arity::Unary)),
        ("!"  , Operator::new("!" , Arity::Unary)),
        ("--" , Operator::new("--", Arity::Unary)),
        ("++" , Operator::new("++", Arity::Unary)),
        ("}"  , Operator::new("}" , Arity::Zero))
    ]
}

fn init_regset() -> HashMap<&'static str, u8> {
    // Use from sdb later, probably a better option.
    hash![
        ("rax", 64),
        ("rbx", 64),
        ("rcx", 64),
        ("rdx", 64),
        ("rsp", 64),
        ("rbp", 64),
        ("rsi", 64),
        ("rdi", 64),
        ("rip", 64)
    ]
}

pub struct Parser<'a> {
    stack: Vec<Value>,
    insts: Vec<Instruction<'a>>,
    opset: HashMap<&'a str, Operator<'a>>,
    regset: HashMap<&'a str, u8>,
    tmp_index: u64,
}

impl<'a> Parser<'a> {
    pub fn new() -> Parser<'a> {
        Parser {
            stack: Vec::new(),
            insts: Vec::new(),
            opset: init_opset(),
            regset: init_regset(),
            tmp_index: 0,
        }
    }

    fn get_tmp_register(&mut self) -> Value {
        self.tmp_index += 1;
        Value::new_tmp(self.tmp_index)
    }

    fn add_inst(&mut self, op: Operator<'a>) {
        let op2 = match self.stack.pop() {
            Some(ele) => ele,
            None => return,
        };
        let mut op1 = Value::null();
        if op.arity.n() == 2 {
            op1 = match self.stack.pop() {
                Some(ele) => ele,
                None => return,
            };
        }
        let dst: Value = match op.op {
            "=" => Value::null(),
            _ => self.get_tmp_register(),
        };
        self.insts.push(Instruction::new(op, dst.clone(), op2, op1));
        self.stack.push(dst);
    }

    pub fn parse(&mut self, esil: &'a str) {
        for token in esil.split(',') {
            let op = match self.opset.get(token) {
                Some(op) => op.clone(),
                None => Operator::nop(),
            };

            if op.op != "nop" {
                self.add_inst(op);
                continue;
            }

            if !token.contains('=') {
                // Treat it as a operand.
                let mut val_type: Location;
                let mut val: i64 = 0;
                // Change this default based on arch.
                let mut size: u8 = 64;

                if let Some(s) = self.regset.get(token) {
                    val_type = Location::Register;
                    // For now, reg is just a u8.
                    size = *s;
                } else if token.contains("temp") {
                    val_type = Location::Temporary;
                } else if let Ok(v) = token.parse::<i64>() {
                    val_type = Location::Constant;
                    val = v;
                } else {
                    val_type = Location::Unknown;
                }

                let v = Value::new(token, size, val_type, val, 0);
                self.stack.push(v);
                continue;
            }

            // Expand the 'composite' operators to 'basic' ones.
            for t in token.split_terminator('=') {
                let o = match self.opset.get(t) {
                    Some(op) => op.clone(),
                    None => Operator::nop(),
                };
                if o.op == "nop" {
                    // Return error here instead.
                    return;
                }
                let dst = self.stack.last().unwrap().clone();
                self.add_inst(o);
                self.stack.push(dst);
                self.add_inst(Operator::new("=", Arity::Binary));
            }
        }
    }

    // Should actually return an iter instead.
    pub fn emit_insts(&self) -> Vec<Instruction<'a>> {
        (self).insts.clone()
    }
}

#[test]
fn testing() {
	let mut p = Parser::new();
	p.parse("0,0x204db1,rip,+,[1],==,%z,zf,=,%b8,cf,=,%p,pf,=,%s,sf,=");
}
<commit_msg>Added some comments for documentation<commit_after>// TODO: Add License information.

//! Module to parse ESIL strings and convert them into the IR.
//!
//! ESIL (Evaluable Strings Intermediate Language) is the IL used by radare2.
//!
//! For a complete documentation of ESIL please check
//!  [wiki](https://github.com/radare/radare2/wiki/ESIL).
//!
//! # Details
//!
//! The `Parser` struct provides methods needed to convert a valid ESIL string
//! into the IR. `Parser::parse()` parses the ESIL string and returns an `Err`
//! if the ESIL string is Invalid.
//!
//! `Parser` also provides `Parser::emit_insts()` to extract the `Instructions`
//! it generates. Calling `Parser::parse()` several times will add more instructions.
//!
//! # Example
//!
//! ```
//! let esil = String::from("eax,ebx,^=");
//! let p = Parser::new();
//! p.parse(esil)
//! for inst in &p.emit_insts() {
//!     println!("{}", inst.to_string());
//! }
//! ```

use std::collections::HashMap;

#[derive(Debug, Copy, Clone)]
pub enum Arity {
    Zero,
    Unary,
    Binary,
    Ternary,
}

impl Arity {
    pub fn n(self) -> u8 {
        match self {
            Arity::Zero    => 0,
            Arity::Unary   => 1,
            Arity::Binary  => 2,
            Arity::Ternary => 3,
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub struct Operator<'a> {
    op: &'a str,
    arity: Arity,
}

impl<'a> Operator<'a> {
    pub fn new(op: &str, n: Arity) -> Operator {
        Operator { op: op, arity: n }
    }

    pub fn nop() -> Operator<'a> {
        Operator { op: "nop", arity: Arity::Zero }
    }
}

#[derive(Debug, Copy, Clone)]
pub enum Location {
    Memory,
    Register,
    Constant,
    Temporary,
    Unknown,
    Null,
}

#[derive(Debug, Clone)]
/// Value is used to represent operands to an operator in a statement.
pub struct Value {
    /// Name
    name: String,
    /// Size in bits.
    size: u8,
    /// Memory, Register, Constant or Temporary.
    location: Location,
    /// Value if evaluable.
    value: i64,
    // TODO: Convert from u32 to TypeSet.
    // Every value can be considered in terms of typesets rather than fixed
    // types which can then be narrowed down based on the analysis.
    // TypeSet can be implemented simply as a bit-vector.
    typeset: u32,
}

impl Value {
    pub fn new(name: String, size: u8, location: Location, value: i64, typeset: u32) -> Value {
        Value {
            name: name.clone(),
            size: size,
            location: location,
            value: value,
            typeset: typeset,
        }
    }

    pub fn null() -> Value {
        Value::new("".to_string(), 0, Location::Null, 0, 0)
    }

    pub fn tmp(i: u64) -> Value {
        Value::new(format!("tmp_{:x}", i), 0, Location::Temporary, 0, 0)
    }
}

#[derive(Debug, Clone)]
pub struct Instruction<'a> {
    opcode: Operator<'a>,
    dest: Value,
    operand_1: Value,
    operand_2: Value,
}

impl<'a> Instruction<'a> {
    pub fn new(opcode: Operator<'a>, dest: Value, op1: Value, op2: Value) -> Instruction<'a> {
        Instruction {
            opcode: opcode,
            dest: dest,
            operand_1: op1,
            operand_2: op2,
        }
    }
    pub fn to_string(&self) -> String {
        if self.opcode.op == "=" {
            format!("{} {} {}", self.operand_1.name, self.opcode.op, self.operand_2.name)
        } else {
            format!("{} {} {} {} {}", self.dest.name, "=", self.operand_1.name, self.opcode.op, self.operand_2.name)
        }
    }
}

macro_rules! hash {
    ( $( ($x:expr, $y:expr) ),* ) => {
        {
            let mut temp_hash = HashMap::new();
            $(
                temp_hash.insert($x, $y);
             )*
            temp_hash
        }
    };
}

fn init_opset() -> HashMap<&'static str, Operator<'static>> {
    // Make a map from esil string to struct Operator.
    // (operator: &str, arity: Arity).
    // Possible Optimization:  Move to compile-time generation ?
    hash![
        ("==" , Operator::new("==", Arity::Binary)),
        ("<"  , Operator::new("<" , Arity::Binary)),
        (">"  , Operator::new(">" , Arity::Binary)),
        ("<=" , Operator::new("<=", Arity::Binary)),
        (">=" , Operator::new(">=", Arity::Binary)),
        ("<<" , Operator::new("<<", Arity::Binary)),
        (">>" , Operator::new(">>", Arity::Binary)),
        ("&"  , Operator::new("&" , Arity::Binary)),
        ("|"  , Operator::new("|" , Arity::Binary)),
        ("="  , Operator::new("=" , Arity::Binary)),
        ("*"  , Operator::new("*" , Arity::Binary)),
        ("^"  , Operator::new("^" , Arity::Binary)),
        ("+"  , Operator::new("+" , Arity::Binary)),
        ("-"  , Operator::new("-" , Arity::Binary)),
        ("/"  , Operator::new("/" , Arity::Binary)),
        ("%"  , Operator::new("%" , Arity::Binary)),
        ("?{" , Operator::new("?{", Arity::Unary)),
        ("!"  , Operator::new("!" , Arity::Unary)),
        ("--" , Operator::new("--", Arity::Unary)),
        ("++" , Operator::new("++", Arity::Unary)),
        ("}"  , Operator::new("}" , Arity::Zero))
    ]
}

fn init_regset() -> HashMap<&'static str, u8> {
    // Use from sdb later, probably a better option.
    hash![
        ("rax", 64),
        ("rbx", 64),
        ("rcx", 64),
        ("rdx", 64),
        ("rsp", 64),
        ("rbp", 64),
        ("rsi", 64),
        ("rdi", 64),
        ("rip", 64)
    ]
}

pub struct Parser<'a> {
    stack: Vec<Value>,
    insts: Vec<Instruction<'a>>,
    opset: HashMap<&'a str, Operator<'a>>,
    regset: HashMap<&'a str, u8>,
    tmp_index: u64,
    default_size: u8,
}

impl<'a> Parser<'a> {
    pub fn new() -> Parser<'a> {
        Parser {
            stack: Vec::new(),
            insts: Vec::new(),
            opset: init_opset(),
            regset: init_regset(),
            tmp_index: 0,
            default_size: 64,
        }
    }

    fn get_tmp_register(&mut self) -> Value {
        self.tmp_index += 1;
        Value::tmp(self.tmp_index)
    }

    fn add_inst(&mut self, op: Operator<'a>) {
        let op2 = match self.stack.pop() {
            Some(ele) => ele,
            None => return,
        };
        let mut op1 = Value::null();
        if op.arity.n() == 2 {
            op1 = match self.stack.pop() {
                Some(ele) => ele,
                None => return,
            };
        }
        let dst: Value = match op.op {
            "=" => Value::null(),
            _ => self.get_tmp_register(),
        };
        self.insts.push(Instruction::new(op, dst.clone(), op2, op1));
        self.stack.push(dst);
    }

    pub fn parse(&mut self, esil: &'a str) {
        for token in esil.split(',') {
            let op = match self.opset.get(token) {
                Some(op) => op.clone(),
                None => Operator::nop(),
            };

            if op.op != "nop" {
                self.add_inst(op);
                continue;
            }

            if !token.contains('=') {
                // Treat it as a operand.
                let mut val_type = Location::Unknown;
                let mut val: i64 = 0;
                // Change this default based on arch.
                let mut size: u8 = self.default_size;
                if let Some(s) = self.regset.get(token) {
                    val_type = Location::Register;
                    // For now, reg is just a u8.
                    size = *s;
                } else if let Ok(v) = token.parse::<i64>() {
                    val_type = Location::Constant;
                    val = v;
                }

                let v = Value::new(String::from(token), size, val_type, val, 0);
                self.stack.push(v);
                continue;
            }

            // Expand the 'composite' operators to 'basic' ones.
            for t in token.split_terminator('=') {
                let o = match self.opset.get(t) {
                    Some(op) => op.clone(),
                    None => Operator::nop(),
                };
                if o.op == "nop" {
                    // Return error here instead.
                    return;
                }
                let dst = self.stack.last().unwrap().clone();
                self.add_inst(o);
                self.stack.push(dst);
                self.add_inst(Operator::new("=", Arity::Binary));
            }
        }
    }

    pub fn emit_insts(&self) -> Vec<Instruction<'a>> {
        (self).insts.clone()
    }
}

#[test]
fn testing() {
	let mut p = Parser::new();
	p.parse("0,0x204db1,rip,+,[1],==,%z,zf,=,%b8,cf,=,%p,pf,=,%s,sf,=");
}
<|endoftext|>


<commit_before>// Register allocator.
//
// Before this pass, it is assumed that we have infinite number of
// registers. This pass maps them to a finite number of registers.
// We actually have only 7 registers.
//
// We allocate registers only within a single expression. In other
// words, there are no registers that live beyond semicolons.
// This design choice simplifies the implementation a lot, since
// practically we don't have to think about the case in which
// registers are exhausted and need to be spilled to memory.

#![allow(non_upper_case_globals)]

use crate::gen_ir::*;
use crate::irdump::*;
use crate::*;

lazy_static! {
    static ref USED: Mutex<Vec<bool>> = Mutex::new(Vec::new());
    static ref REG_MAP: Mutex<Vec<i32>> = Mutex::new(Vec::new());
}

pub const regs: [&'static str; 8] = ["rbp", "r10", "r11", "rbx", "r12", "r13", "r14", "r15"];
pub const regs8: [&'static str; 8] = ["rbl", "r10b", "r11b", "bl", "r12b", "r13b", "r14b", "r15b"];
pub const regs32: [&'static str; 8] = ["ebp", "r10d", "r11d", "ebx", "r12d", "r13d", "r14d", "r15d"];

// Code generator

fn alloc(ir_reg: i32) -> i32 {
    let mut reg_map = REG_MAP.lock().unwrap();
    if reg_map[ir_reg as usize] != -1 {
        let r = reg_map[ir_reg as usize];
        assert!(USED.lock().unwrap()[r as usize], "");
        return r;
    }

    let mut used = USED.lock().unwrap();
    for i in 0..regs.len() {
        if used[i] {
            continue;
        }
        reg_map[ir_reg as usize] = i as i32;
        used[i] = true;
        return i as i32;
    }
    panic!("register exhausted");
}

fn visit(irv: &mut Vec<IR>) {
    // r0 is a reserved register that is always mapped to rbp.

    match (REG_MAP.lock(), USED.lock()) {
        (Ok(mut reg_map), Ok(mut used)) => {
            reg_map[0] = 0;
            used[0] = true;
        }
        _ => {
            panic!();
        }
    }

    for i in 0..irv.len() {
        let mut ir = &mut irv[i];

        match irinfo_get(&ir.op).unwrap().ty {
            IRInfoType::REG |
                IRInfoType::REG_IMM |
                IRInfoType::REG_LABEL |
                IRInfoType::LABEL_ADDR => {
                ir.lhs = alloc(ir.lhs);
            }
            IRInfoType::REG_REG => {
                ir.lhs = alloc(ir.lhs);
                ir.rhs = alloc(ir.rhs);
            }
            IRInfoType::CALL => {
                ir.lhs = alloc(ir.lhs);
                for i in 0..ir.nargs {
                    ir.args[i] = alloc(ir.args[i]);
                }
            }
            _ => {}
        }

        if ir.op == IRType::KILL {
            let mut used = USED.lock().unwrap();
            assert!(used[ir.lhs as usize]);
            used[ir.lhs as usize] = false;
            ir.op = IRType::NOP;
        }
    }
}

pub fn alloc_regs(fns: &mut Vec<IR>) {

    for i in 0..fns.len() {
        let mut fun = &mut fns[i];
        match (REG_MAP.lock(), USED.lock()) {
            (Ok(mut reg_map), Ok(mut used)) => {
                *reg_map = Vec::new();
                for _i in 0..fun.ir.len() {
                    reg_map.push(-1);
                }

                *used = Vec::new();
                for _i in 0..regs.len() {
                    used.push(false);
                }
            }
            _ => {
                panic!();
            }
        }

        visit(&mut fun.ir);
    }
}
<commit_msg>Since we no longer have IR_SUB_IMM, we can remove this code.<commit_after>// Register allocator.
//
// Before this pass, it is assumed that we have infinite number of
// registers. This pass maps them to a finite number of registers.
// We actually have only 7 registers.
//
// We allocate registers only within a single expression. In other
// words, there are no registers that live beyond semicolons.
// This design choice simplifies the implementation a lot, since
// practically we don't have to think about the case in which
// registers are exhausted and need to be spilled to memory.

#![allow(non_upper_case_globals)]

use crate::gen_ir::*;
use crate::irdump::*;
use crate::*;

lazy_static! {
    static ref USED: Mutex<Vec<bool>> = Mutex::new(Vec::new());
    static ref REG_MAP: Mutex<Vec<i32>> = Mutex::new(Vec::new());
}

pub const regs: [&'static str; 7] = ["r10", "r11", "rbx", "r12", "r13", "r14", "r15"];
pub const regs8: [&'static str; 7] = ["r10b", "r11b", "bl", "r12b", "r13b", "r14b", "r15b"];
pub const regs32: [&'static str; 7] = ["r10d", "r11d", "ebx", "r12d", "r13d", "r14d", "r15d"];

// Code generator

fn alloc(ir_reg: i32) -> i32 {
    let mut reg_map = REG_MAP.lock().unwrap();
    if reg_map[ir_reg as usize] != -1 {
        let r = reg_map[ir_reg as usize];
        assert!(USED.lock().unwrap()[r as usize], "");
        return r;
    }

    let mut used = USED.lock().unwrap();
    for i in 0..regs.len() {
        if used[i] {
            continue;
        }
        reg_map[ir_reg as usize] = i as i32;
        used[i] = true;
        return i as i32;
    }
    panic!("register exhausted");
}

fn visit(irv: &mut Vec<IR>) {
    for i in 0..irv.len() {
        let mut ir = &mut irv[i];

        match irinfo_get(&ir.op).unwrap().ty {
            IRInfoType::REG |
                IRInfoType::REG_IMM |
                IRInfoType::REG_LABEL |
                IRInfoType::LABEL_ADDR => {
                ir.lhs = alloc(ir.lhs);
            }
            IRInfoType::REG_REG => {
                ir.lhs = alloc(ir.lhs);
                ir.rhs = alloc(ir.rhs);
            }
            IRInfoType::CALL => {
                ir.lhs = alloc(ir.lhs);
                for i in 0..ir.nargs {
                    ir.args[i] = alloc(ir.args[i]);
                }
            }
            _ => {}
        }

        if ir.op == IRType::KILL {
            let mut used = USED.lock().unwrap();
            assert!(used[ir.lhs as usize]);
            used[ir.lhs as usize] = false;
            ir.op = IRType::NOP;
        }
    }
}

pub fn alloc_regs(fns: &mut Vec<IR>) {

    for i in 0..fns.len() {
        let mut fun = &mut fns[i];
        match (REG_MAP.lock(), USED.lock()) {
            (Ok(mut reg_map), Ok(mut used)) => {
                *reg_map = Vec::new();
                for _i in 0..fun.ir.len() {
                    reg_map.push(-1);
                }

                *used = Vec::new();
                for _i in 0..regs.len() {
                    used.push(false);
                }
            }
            _ => {
                panic!();
            }
        }

        visit(&mut fun.ir);
    }
}
<|endoftext|>


<commit_before>//! way-cooler registry.

use std::collections::HashMap;
use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};

pub type RegKey = String;
pub type RegVal = RegistryValue;
pub type RegMap = HashMap<RegKey, RegVal>;

lazy_static! {
    /// Registry variable for the registry
    static ref REGISTRY: RwLock<RegMap> =
        RwLock::new(HashMap::new());
}

/// A value in the registry.
#[derive(Debug, Clone)]
pub enum RegistryValue {
    /// An integer value
    Integer(i32),
    /// A Lua numeric value
    Double(f64),
    /// A boolean value
    Boolean(bool),
    /// A string value
    String(String),
    /// A list of values.
    List(Vec<RegistryValue>)
}

/// Acquires a read lock on the registry.
fn read_lock<'a>() -> RwLockReadGuard<'a, RegMap> {
    REGISTRY.read().unwrap()
}

/// Acquires a write lock on the registry.
fn write_lock<'a>() -> RwLockWriteGuard<'a, RegMap> {
    REGISTRY.write().unwrap()
}

/// Gets a value from the regsitry.
pub fn get(name: &RegKey) -> Option<RegVal> {
    trace!("get: key {}", name);
    let ref reg = *read_lock();
    // cloned() is a method on Option<T> where T: Clone
    reg.get(name).cloned()
}

/// Gets a value from the registry, or a default
/// if the value is not found.
pub fn get_or_default(name: &RegKey, value: RegVal) -> RegVal {
    match get(name) {
        Some(value) => value,
        None => value
    }
}

/// Set a value in the registry.
///
/// If the key already exists, returns the old value.
/// Returns `None` for new keys.
pub fn set(name: RegKey, value: RegVal) -> Option<RegVal> {
    trace!("set: {} = {:?}", &name, &value);
    let ref mut reg = *write_lock();
    reg.insert(name, value)
}

use hlua::{Push, AsMutLua, PushGuard, LuaContext};

impl<L: AsMutLua> Push<L> for RegistryValue {
    fn push_to_lua(self, lua: L) -> PushGuard<L> {
        unsafe {
            match self {
                RegistryValue::Boolean(val) => {
                    val.push_to_lua(lua)
                },
                RegistryValue::Double(val) => {
                    val.push_to_lua(lua)
                },
                RegistryValue::Integer(val) => {
                    val.push_to_lua(lua)
                },
                RegistryValue::String(text) => {
                    text.push_to_lua(lua)
                },
                RegistryValue::List(list) => {
                    list.push_to_lua(lua)
                }
            }
        }
    }
}
<commit_msg>Vain attempts to convert json to lua<commit_after>//! way-cooler registry.

use std::collections::HashMap;
use std::fmt::{Debug, Display, Formatter};
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};

use rustc_serialize::{Encodable, Decodable};
use rustc_serialize::json::{Json, ToJson};

pub type RegKey = String;
pub type RegVal = RegistryValue;
pub type RegMap = HashMap<RegKey, RegVal>;

lazy_static! {
    /// Registry variable for the registry
    static ref REGISTRY: RwLock<RegMap> =
        RwLock::new(HashMap::new());
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum RegistryAccess {
    Public,
    Lua,
    Private
}

/// Values stored in the registry
pub struct RegistryValue {
    access: RegistryAccess,
    object: Arc<ToJson + Sync + Send>
}

impl RegistryValue {
    /// What access the module has to it
    pub fn access(&self) -> RegistryAccess {
        self.access
    }

    /// Gets the json of a registry value
    pub fn get_json(&self) -> Json {
        self.object.to_json()
    }
}

/// Acquires a read lock on the registry.
fn read_lock<'a>() -> RwLockReadGuard<'a, RegMap> {
    REGISTRY.read().unwrap()
}

/// Acquires a write lock on the registry.
fn write_lock<'a>() -> RwLockWriteGuard<'a, RegMap> {
    REGISTRY.write().unwrap()
}

/// Gets a Json object from a registry key
pub fn get_json(name: &RegKey) -> Option<Json> {
    trace!("get_json: {}", name);
    let ref reg = *read_lock();
    if let Some(ref val) = reg.get(name) {
        Some(val.get_json())
    }
    else {
        None
    }
}

pub struct JsonWrap(Json);

use hlua::{Push, PushGuard, AsMutLua, LuaContext};
use hlua::LuaTable;
use hlua_ffi;

fn push_iter<L, V, I>(mut lua: L, iterator: I) -> PushGuard<L>
    where L: AsMutLua, V: for<'b> Push<&'b mut L>, I: Iterator<Item=V>
{
    // creating empty table
    unsafe { hlua_ffi::lua_newtable(lua.as_mut_lua().0) };

    for (elem, index) in iterator.zip((1 ..)) {
        let size = elem.push_to_lua(&mut lua).forget();

        match size {
            0 => continue,
            1 => {
                let index = index as u32;
                index.push_to_lua(&mut lua).forget();
                unsafe { hlua_ffi::lua_insert(lua.as_mut_lua().0, -2) }
                unsafe { hlua_ffi::lua_settable(lua.as_mut_lua().0, -3) }
            },
            2 => unsafe { hlua_ffi::lua_settable(lua.as_mut_lua().0, -3) },
            _ => unreachable!()
        }
    }

    PushGuard { lua: lua, size: 1 }
}

fn push_rec_iter<L, V, I>(mut lua: L, iterator: I) -> PushGuard<L>
    where L: AsMutLua, V: for<'a> Push<&'a mut L>, I: Iterator<Item=V>
{
    let (nrec, _) = iterator.size_hint();

    // creating empty table with pre-allocated non-array elements
    unsafe { hlua_ffi::lua_createtable(lua.as_mut_lua().0, 0, nrec as i32) };

    for elem in iterator {
        let size = elem.push_to_lua(&mut lua).forget();

        match size {
            0 => continue,
            2 => unsafe { hlua_ffi::lua_settable(lua.as_mut_lua().0, -3) },
            _ => unreachable!()
        }
    }

    PushGuard { lua: lua, size: 1 }
}

fn push_table_field<L, V>(mut lua: L, name: String, val: V)
    where L: AsMutLua, V: for<'a> Push<&'a mut L> {
    name.push_to_lua(&mut lua);
    val.push_to_lua(&mut lua);
    unsafe { hlua_ffi::lua_settable(lua.as_mut_lua().0, -3); }
}

use std::collections::BTreeMap;

impl<L: AsMutLua> Push<L> for JsonWrap {
    fn push_to_lua(self, lua: L) -> PushGuard<L> {
        unsafe {
            match self.0 {
                Json::I64(val) => {
                    (val as i32).push_to_lua(lua)
                },
                Json::U64(val) => {
                    (val as u32).push_to_lua(lua)
                },
                Json::F64(val) => {
                    val.push_to_lua(lua)
                },
                Json::String(val) => {
                    val.push_to_lua(lua)
                },
                Json::Null => {
                    ().push_to_lua(lua)
                },
                Json::Boolean(val) => {
                    val.push_to_lua(lua)
                },
                Json::Array(vals) => {
                    push_iter(lua, vals.into_iter()
                              .map(|j| JsonWrap(j)))
                },
                Json::Object(map_l) => {
                    let map: BTreeMap<String, Json> = map_l;

                    let len = map.len();

                    unsafe {
                        hlua_ffi::lua_newtable(lua.as_mut_lua().0);
                    }

                    for (key, value) in map.into_iter() {
                        push_table_field(&mut lua, key, JsonWrap(value));
                    }
                    PushGuard { lua: lua, size: 1 }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use common::{HTTPDate, HTTPVersion, Header, StatusCode};

use std::cmp::Ordering;
use std::sync::mpsc::Receiver;

use std::io::Result as IoResult;
use std::io::{self, Cursor, Read, Write};

use std::fs::File;

use std::str::FromStr;

/// Object representing an HTTP response whose purpose is to be given to a `Request`.
///
/// Some headers cannot be changed. Trying to define the value
/// of one of these will have no effect:
///
///  - `Accept-Ranges`
///  - `Connection`
///  - `Content-Range`
///  - `Trailer`
///  - `Transfer-Encoding`
///  - `Upgrade`
///
/// Some headers have special behaviors:
///
///  - `Content-Encoding`: If you define this header, the library
///     will assume that the data from the `Read` object has the specified encoding
///     and will just pass-through.
///
///  - `Content-Length`: The length of the data should be set manually
///     using the `Reponse` object's API. Attempting to set the value of this
///     header will be equivalent to modifying the size of the data but the header
///     itself may not be present in the final result.
///
pub struct Response<R>
where
    R: Read,
{
    reader: R,
    status_code: StatusCode,
    headers: Vec<Header>,
    data_length: Option<usize>,
    chunked_threshold: Option<usize>,
}

/// A `Response` without a template parameter.
pub type ResponseBox = Response<Box<dyn Read + Send>>;

/// Transfer encoding to use when sending the message.
/// Note that only *supported* encoding are listed here.
#[derive(Copy, Clone)]
enum TransferEncoding {
    Identity,
    Chunked,
}

impl FromStr for TransferEncoding {
    type Err = ();

    fn from_str(input: &str) -> Result<TransferEncoding, ()> {
        if input.eq_ignore_ascii_case("identity") {
            Ok(TransferEncoding::Identity)
        } else if input.eq_ignore_ascii_case("chunked") {
            Ok(TransferEncoding::Chunked)
        } else {
            Err(())
        }
    }
}

/// Builds a Date: header with the current date.
fn build_date_header() -> Header {
    let d = HTTPDate::new();
    Header::from_bytes(&b"Date"[..], &d.to_string().into_bytes()[..]).unwrap()
}

fn write_message_header<W>(
    mut writer: W,
    http_version: &HTTPVersion,
    status_code: &StatusCode,
    headers: &[Header],
) -> IoResult<()>
where
    W: Write,
{
    // writing status line
    write!(
        &mut writer,
        "HTTP/{}.{} {} {}\r\n",
        http_version.0,
        http_version.1,
        status_code.0,
        status_code.default_reason_phrase()
    )?;

    // writing headers
    for header in headers.iter() {
        writer.write_all(header.field.as_str().as_ref())?;
        write!(&mut writer, ": ")?;
        writer.write_all(header.value.as_str().as_ref())?;
        write!(&mut writer, "\r\n")?;
    }

    // separator between header and data
    write!(&mut writer, "\r\n")?;

    Ok(())
}

fn choose_transfer_encoding(
    request_headers: &[Header],
    http_version: &HTTPVersion,
    entity_length: &Option<usize>,
    has_additional_headers: bool,
    chunked_threshold: usize,
) -> TransferEncoding {
    use util;

    // HTTP 1.0 doesn't support other encoding
    if *http_version <= (1, 0) {
        return TransferEncoding::Identity;
    }

    // parsing the request's TE header
    let user_request = request_headers
        .iter()
        // finding TE
        .find(|h| h.field.equiv(&"TE"))
        // getting its value
        .map(|h| h.value.clone())
        // getting the corresponding TransferEncoding
        .and_then(|value| {
            // getting list of requested elements
            let mut parse = util::parse_header_value(value.as_str()); // TODO: remove conversion

            // sorting elements by most priority
            parse.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));

            // trying to parse each requested encoding
            for value in parse.iter() {
                // q=0 are ignored
                if value.1 <= 0.0 {
                    continue;
                }

                if let Ok(te) = TransferEncoding::from_str(value.0) {
                    return Some(te);
                }
            }

            // encoding not found
            None
        });

    //
    if let Some(user_request) = user_request {
        return user_request;
    }

    // if we have additional headers, using chunked
    if has_additional_headers {
        return TransferEncoding::Chunked;
    }

    // if we don't have a Content-Length, or if the Content-Length is too big, using chunks writer
    if entity_length
        .as_ref()
        .map_or(true, |val| *val >= chunked_threshold)
    {
        return TransferEncoding::Chunked;
    }

    // Identity by default
    TransferEncoding::Identity
}

impl<R> Response<R>
where
    R: Read,
{
    /// Creates a new Response object.
    ///
    /// The `additional_headers` argument is a receiver that
    ///  may provide headers even after the response has been sent.
    ///
    /// All the other arguments are straight-forward.
    pub fn new(
        status_code: StatusCode,
        headers: Vec<Header>,
        data: R,
        data_length: Option<usize>,
        additional_headers: Option<Receiver<Header>>,
    ) -> Response<R> {
        let mut response = Response {
            reader: data,
            status_code,
            headers: Vec::with_capacity(16),
            data_length,
            chunked_threshold: None,
        };

        for h in headers {
            response.add_header(h)
        }

        // dummy implementation
        if let Some(additional_headers) = additional_headers {
            for h in additional_headers.iter() {
                response.add_header(h)
            }
        }

        response
    }

    /// Set a threshold for `Content-Length` where we chose chunked
    /// transfer. Notice that chunked transfer might happen regardless of
    /// this threshold, for instance when the request headers indicate
    /// it is wanted or when there is no `Content-Length`.
    pub fn with_chunked_threshold(mut self, length: usize) -> Response<R> {
        self.chunked_threshold = Some(length);
        self
    }

    /// The current `Content-Length` threshold for switching over to
    /// chunked transfer. The default is 32768 bytes. Notice that
    /// chunked transfer is mutually exclusive with sending a
    /// `Content-Length` header as per the HTTP spec.
    pub fn chunked_threshold(&self) -> usize {
        self.chunked_threshold.unwrap_or(32768)
    }

    /// Adds a header to the list.
    /// Does all the checks.
    pub fn add_header<H>(&mut self, header: H)
    where
        H: Into<Header>,
    {
        let header = header.into();

        // ignoring forbidden headers
        if header.field.equiv(&"Accept-Ranges")
            || header.field.equiv(&"Connection")
            || header.field.equiv(&"Content-Range")
            || header.field.equiv(&"Trailer")
            || header.field.equiv(&"Transfer-Encoding")
            || header.field.equiv(&"Upgrade")
        {
            return;
        }

        // if the header is Content-Length, setting the data length
        if header.field.equiv(&"Content-Length") {
            if let Ok(val) = usize::from_str(header.value.as_str()) {
                self.data_length = Some(val)
            }

            return;
        }

        self.headers.push(header);
    }

    /// Returns the same request, but with an additional header.
    ///
    /// Some headers cannot be modified and some other have a
    ///  special behavior. See the documentation above.
    #[inline]
    pub fn with_header<H>(mut self, header: H) -> Response<R>
    where
        H: Into<Header>,
    {
        self.add_header(header.into());
        self
    }

    /// Returns the same request, but with a different status code.
    #[inline]
    pub fn with_status_code<S>(mut self, code: S) -> Response<R>
    where
        S: Into<StatusCode>,
    {
        self.status_code = code.into();
        self
    }

    /// Returns the same request, but with different data.
    pub fn with_data<S>(self, reader: S, data_length: Option<usize>) -> Response<S>
    where
        S: Read,
    {
        Response {
            reader,
            headers: self.headers,
            status_code: self.status_code,
            data_length,
            chunked_threshold: self.chunked_threshold,
        }
    }

    /// Prints the HTTP response to a writer.
    ///
    /// This function is the one used to send the response to the client's socket.
    /// Therefore you shouldn't expect anything pretty-printed or even readable.
    ///
    /// The HTTP version and headers passed as arguments are used to
    ///  decide which features (most notably, encoding) to use.
    ///
    /// Note: does not flush the writer.
    pub fn raw_print<W: Write>(
        mut self,
        mut writer: W,
        http_version: HTTPVersion,
        request_headers: &[Header],
        do_not_send_body: bool,
        upgrade: Option<&str>,
    ) -> IoResult<()> {
        let mut transfer_encoding = Some(choose_transfer_encoding(
            request_headers,
            &http_version,
            &self.data_length,
            false, /* TODO */
            self.chunked_threshold(),
        ));

        // add `Date` if not in the headers
        if self
            .headers
            .iter()
            .find(|h| h.field.equiv(&"Date"))
            .is_none()
        {
            self.headers.insert(0, build_date_header());
        }

        // add `Server` if not in the headers
        if self
            .headers
            .iter()
            .find(|h| h.field.equiv(&"Server"))
            .is_none()
        {
            self.headers.insert(
                0,
                Header::from_bytes(&b"Server"[..], &b"tiny-http (Rust)"[..]).unwrap(),
            );
        }

        // handling upgrade
        if let Some(upgrade) = upgrade {
            self.headers.insert(
                0,
                Header::from_bytes(&b"Upgrade"[..], upgrade.as_bytes()).unwrap(),
            );
            self.headers.insert(
                0,
                Header::from_bytes(&b"Connection"[..], &b"upgrade"[..]).unwrap(),
            );
            transfer_encoding = None;
        }

        // if the transfer encoding is identity, the content length must be known ; therefore if
        // we don't know it, we buffer the entire response first here
        // while this is an expensive operation, it is only ever needed for clients using HTTP 1.0
        let (mut reader, data_length) = match (self.data_length, transfer_encoding) {
            (Some(l), _) => (Box::new(self.reader) as Box<dyn Read>, Some(l)),
            (None, Some(TransferEncoding::Identity)) => {
                let mut buf = Vec::new();
                self.reader.read_to_end(&mut buf)?;
                let l = buf.len();
                (Box::new(Cursor::new(buf)) as Box<dyn Read>, Some(l))
            }
            _ => (Box::new(self.reader) as Box<dyn Read>, None),
        };

        // checking whether to ignore the body of the response
        let do_not_send_body = do_not_send_body
            || match self.status_code.0 {
                // sattus code 1xx, 204 and 304 MUST not include a body
                100..=199 | 204 | 304 => true,
                _ => false,
            };

        // preparing headers for transfer
        match transfer_encoding {
            Some(TransferEncoding::Chunked) => self
                .headers
                .push(Header::from_bytes(&b"Transfer-Encoding"[..], &b"chunked"[..]).unwrap()),

            Some(TransferEncoding::Identity) => {
                assert!(data_length.is_some());
                let data_length = data_length.unwrap();

                self.headers.push(
                    Header::from_bytes(
                        &b"Content-Length"[..],
                        format!("{}", data_length).as_bytes(),
                    )
                    .unwrap(),
                )
            }

            _ => (),
        };

        // sending headers
        write_message_header(
            writer.by_ref(),
            &http_version,
            &self.status_code,
            &self.headers,
        )?;

        // sending the body
        if !do_not_send_body {
            match transfer_encoding {
                Some(TransferEncoding::Chunked) => {
                    use chunked_transfer::Encoder;

                    let mut writer = Encoder::new(writer);
                    io::copy(&mut reader, &mut writer)?;
                }

                Some(TransferEncoding::Identity) => {
                    use util::EqualReader;

                    assert!(data_length.is_some());
                    let data_length = data_length.unwrap();

                    if data_length >= 1 {
                        let (mut equ_reader, _) = EqualReader::new(reader.by_ref(), data_length);
                        io::copy(&mut equ_reader, &mut writer)?;
                    }
                }

                _ => (),
            }
        }

        Ok(())
    }
}

impl<R> Response<R>
where
    R: Read + Send + 'static,
{
    /// Turns this response into a `Response<Box<Read + Send>>`.
    pub fn boxed(self) -> ResponseBox {
        Response {
            reader: Box::new(self.reader) as Box<dyn Read + Send>,
            status_code: self.status_code,
            headers: self.headers,
            data_length: self.data_length,
            chunked_threshold: self.chunked_threshold,
        }
    }
}

impl Response<File> {
    /// Builds a new `Response` from a `File`.
    ///
    /// The `Content-Type` will **not** be automatically detected,
    ///  you must set it yourself.
    pub fn from_file(file: File) -> Response<File> {
        let file_size = file.metadata().ok().map(|v| v.len() as usize);

        Response::new(
            StatusCode(200),
            Vec::with_capacity(0),
            file,
            file_size,
            None,
        )
    }
}

impl Response<Cursor<Vec<u8>>> {
    pub fn from_data<D>(data: D) -> Response<Cursor<Vec<u8>>>
    where
        D: Into<Vec<u8>>,
    {
        let data = data.into();
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            Vec::with_capacity(0),
            Cursor::new(data),
            Some(data_len),
            None,
        )
    }

    pub fn from_string<S>(data: S) -> Response<Cursor<Vec<u8>>>
    where
        S: Into<String>,
    {
        let data = data.into();
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            vec![
                Header::from_bytes(&b"Content-Type"[..], &b"text/plain; charset=UTF-8"[..])
                    .unwrap(),
            ],
            Cursor::new(data.into_bytes()),
            Some(data_len),
            None,
        )
    }
}

impl Response<io::Empty> {
    /// Builds an empty `Response` with the given status code.
    pub fn empty<S>(status_code: S) -> Response<io::Empty>
    where
        S: Into<StatusCode>,
    {
        Response::new(
            status_code.into(),
            Vec::with_capacity(0),
            io::empty(),
            Some(0),
            None,
        )
    }

    /// DEPRECATED. Use `empty` instead.
    pub fn new_empty(status_code: StatusCode) -> Response<io::Empty> {
        Response::empty(status_code)
    }
}

impl Clone for Response<io::Empty> {
    fn clone(&self) -> Response<io::Empty> {
        Response {
            reader: io::empty(),
            status_code: self.status_code.clone(),
            headers: self.headers.clone(),
            data_length: self.data_length,
            chunked_threshold: self.chunked_threshold,
        }
    }
}
<commit_msg>Response: Allow manual handling of Range requests<commit_after>use common::{HTTPDate, HTTPVersion, Header, StatusCode};

use std::cmp::Ordering;
use std::sync::mpsc::Receiver;

use std::io::Result as IoResult;
use std::io::{self, Cursor, Read, Write};

use std::fs::File;

use std::str::FromStr;

/// Object representing an HTTP response whose purpose is to be given to a `Request`.
///
/// Some headers cannot be changed. Trying to define the value
/// of one of these will have no effect:
///
///  - `Connection`
///  - `Trailer`
///  - `Transfer-Encoding`
///  - `Upgrade`
///
/// Some headers have special behaviors:
///
///  - `Content-Encoding`: If you define this header, the library
///     will assume that the data from the `Read` object has the specified encoding
///     and will just pass-through.
///
///  - `Content-Length`: The length of the data should be set manually
///     using the `Reponse` object's API. Attempting to set the value of this
///     header will be equivalent to modifying the size of the data but the header
///     itself may not be present in the final result.
///
pub struct Response<R>
where
    R: Read,
{
    reader: R,
    status_code: StatusCode,
    headers: Vec<Header>,
    data_length: Option<usize>,
    chunked_threshold: Option<usize>,
}

/// A `Response` without a template parameter.
pub type ResponseBox = Response<Box<dyn Read + Send>>;

/// Transfer encoding to use when sending the message.
/// Note that only *supported* encoding are listed here.
#[derive(Copy, Clone)]
enum TransferEncoding {
    Identity,
    Chunked,
}

impl FromStr for TransferEncoding {
    type Err = ();

    fn from_str(input: &str) -> Result<TransferEncoding, ()> {
        if input.eq_ignore_ascii_case("identity") {
            Ok(TransferEncoding::Identity)
        } else if input.eq_ignore_ascii_case("chunked") {
            Ok(TransferEncoding::Chunked)
        } else {
            Err(())
        }
    }
}

/// Builds a Date: header with the current date.
fn build_date_header() -> Header {
    let d = HTTPDate::new();
    Header::from_bytes(&b"Date"[..], &d.to_string().into_bytes()[..]).unwrap()
}

fn write_message_header<W>(
    mut writer: W,
    http_version: &HTTPVersion,
    status_code: &StatusCode,
    headers: &[Header],
) -> IoResult<()>
where
    W: Write,
{
    // writing status line
    write!(
        &mut writer,
        "HTTP/{}.{} {} {}\r\n",
        http_version.0,
        http_version.1,
        status_code.0,
        status_code.default_reason_phrase()
    )?;

    // writing headers
    for header in headers.iter() {
        writer.write_all(header.field.as_str().as_ref())?;
        write!(&mut writer, ": ")?;
        writer.write_all(header.value.as_str().as_ref())?;
        write!(&mut writer, "\r\n")?;
    }

    // separator between header and data
    write!(&mut writer, "\r\n")?;

    Ok(())
}

fn choose_transfer_encoding(
    request_headers: &[Header],
    http_version: &HTTPVersion,
    entity_length: &Option<usize>,
    has_additional_headers: bool,
    chunked_threshold: usize,
) -> TransferEncoding {
    use util;

    // HTTP 1.0 doesn't support other encoding
    if *http_version <= (1, 0) {
        return TransferEncoding::Identity;
    }

    // parsing the request's TE header
    let user_request = request_headers
        .iter()
        // finding TE
        .find(|h| h.field.equiv(&"TE"))
        // getting its value
        .map(|h| h.value.clone())
        // getting the corresponding TransferEncoding
        .and_then(|value| {
            // getting list of requested elements
            let mut parse = util::parse_header_value(value.as_str()); // TODO: remove conversion

            // sorting elements by most priority
            parse.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal));

            // trying to parse each requested encoding
            for value in parse.iter() {
                // q=0 are ignored
                if value.1 <= 0.0 {
                    continue;
                }

                if let Ok(te) = TransferEncoding::from_str(value.0) {
                    return Some(te);
                }
            }

            // encoding not found
            None
        });

    //
    if let Some(user_request) = user_request {
        return user_request;
    }

    // if we have additional headers, using chunked
    if has_additional_headers {
        return TransferEncoding::Chunked;
    }

    // if we don't have a Content-Length, or if the Content-Length is too big, using chunks writer
    if entity_length
        .as_ref()
        .map_or(true, |val| *val >= chunked_threshold)
    {
        return TransferEncoding::Chunked;
    }

    // Identity by default
    TransferEncoding::Identity
}

impl<R> Response<R>
where
    R: Read,
{
    /// Creates a new Response object.
    ///
    /// The `additional_headers` argument is a receiver that
    ///  may provide headers even after the response has been sent.
    ///
    /// All the other arguments are straight-forward.
    pub fn new(
        status_code: StatusCode,
        headers: Vec<Header>,
        data: R,
        data_length: Option<usize>,
        additional_headers: Option<Receiver<Header>>,
    ) -> Response<R> {
        let mut response = Response {
            reader: data,
            status_code,
            headers: Vec::with_capacity(16),
            data_length,
            chunked_threshold: None,
        };

        for h in headers {
            response.add_header(h)
        }

        // dummy implementation
        if let Some(additional_headers) = additional_headers {
            for h in additional_headers.iter() {
                response.add_header(h)
            }
        }

        response
    }

    /// Set a threshold for `Content-Length` where we chose chunked
    /// transfer. Notice that chunked transfer might happen regardless of
    /// this threshold, for instance when the request headers indicate
    /// it is wanted or when there is no `Content-Length`.
    pub fn with_chunked_threshold(mut self, length: usize) -> Response<R> {
        self.chunked_threshold = Some(length);
        self
    }

    /// The current `Content-Length` threshold for switching over to
    /// chunked transfer. The default is 32768 bytes. Notice that
    /// chunked transfer is mutually exclusive with sending a
    /// `Content-Length` header as per the HTTP spec.
    pub fn chunked_threshold(&self) -> usize {
        self.chunked_threshold.unwrap_or(32768)
    }

    /// Adds a header to the list.
    /// Does all the checks.
    pub fn add_header<H>(&mut self, header: H)
    where
        H: Into<Header>,
    {
        let header = header.into();

        // ignoring forbidden headers
        if header.field.equiv(&"Connection")
            || header.field.equiv(&"Trailer")
            || header.field.equiv(&"Transfer-Encoding")
            || header.field.equiv(&"Upgrade")
        {
            return;
        }

        // if the header is Content-Length, setting the data length
        if header.field.equiv(&"Content-Length") {
            if let Ok(val) = usize::from_str(header.value.as_str()) {
                self.data_length = Some(val)
            }

            return;
        }

        self.headers.push(header);
    }

    /// Returns the same request, but with an additional header.
    ///
    /// Some headers cannot be modified and some other have a
    ///  special behavior. See the documentation above.
    #[inline]
    pub fn with_header<H>(mut self, header: H) -> Response<R>
    where
        H: Into<Header>,
    {
        self.add_header(header.into());
        self
    }

    /// Returns the same request, but with a different status code.
    #[inline]
    pub fn with_status_code<S>(mut self, code: S) -> Response<R>
    where
        S: Into<StatusCode>,
    {
        self.status_code = code.into();
        self
    }

    /// Returns the same request, but with different data.
    pub fn with_data<S>(self, reader: S, data_length: Option<usize>) -> Response<S>
    where
        S: Read,
    {
        Response {
            reader,
            headers: self.headers,
            status_code: self.status_code,
            data_length,
            chunked_threshold: self.chunked_threshold,
        }
    }

    /// Prints the HTTP response to a writer.
    ///
    /// This function is the one used to send the response to the client's socket.
    /// Therefore you shouldn't expect anything pretty-printed or even readable.
    ///
    /// The HTTP version and headers passed as arguments are used to
    ///  decide which features (most notably, encoding) to use.
    ///
    /// Note: does not flush the writer.
    pub fn raw_print<W: Write>(
        mut self,
        mut writer: W,
        http_version: HTTPVersion,
        request_headers: &[Header],
        do_not_send_body: bool,
        upgrade: Option<&str>,
    ) -> IoResult<()> {
        let mut transfer_encoding = Some(choose_transfer_encoding(
            request_headers,
            &http_version,
            &self.data_length,
            false, /* TODO */
            self.chunked_threshold(),
        ));

        // add `Date` if not in the headers
        if self
            .headers
            .iter()
            .find(|h| h.field.equiv(&"Date"))
            .is_none()
        {
            self.headers.insert(0, build_date_header());
        }

        // add `Server` if not in the headers
        if self
            .headers
            .iter()
            .find(|h| h.field.equiv(&"Server"))
            .is_none()
        {
            self.headers.insert(
                0,
                Header::from_bytes(&b"Server"[..], &b"tiny-http (Rust)"[..]).unwrap(),
            );
        }

        // handling upgrade
        if let Some(upgrade) = upgrade {
            self.headers.insert(
                0,
                Header::from_bytes(&b"Upgrade"[..], upgrade.as_bytes()).unwrap(),
            );
            self.headers.insert(
                0,
                Header::from_bytes(&b"Connection"[..], &b"upgrade"[..]).unwrap(),
            );
            transfer_encoding = None;
        }

        // if the transfer encoding is identity, the content length must be known ; therefore if
        // we don't know it, we buffer the entire response first here
        // while this is an expensive operation, it is only ever needed for clients using HTTP 1.0
        let (mut reader, data_length) = match (self.data_length, transfer_encoding) {
            (Some(l), _) => (Box::new(self.reader) as Box<dyn Read>, Some(l)),
            (None, Some(TransferEncoding::Identity)) => {
                let mut buf = Vec::new();
                self.reader.read_to_end(&mut buf)?;
                let l = buf.len();
                (Box::new(Cursor::new(buf)) as Box<dyn Read>, Some(l))
            }
            _ => (Box::new(self.reader) as Box<dyn Read>, None),
        };

        // checking whether to ignore the body of the response
        let do_not_send_body = do_not_send_body
            || match self.status_code.0 {
                // sattus code 1xx, 204 and 304 MUST not include a body
                100..=199 | 204 | 304 => true,
                _ => false,
            };

        // preparing headers for transfer
        match transfer_encoding {
            Some(TransferEncoding::Chunked) => self
                .headers
                .push(Header::from_bytes(&b"Transfer-Encoding"[..], &b"chunked"[..]).unwrap()),

            Some(TransferEncoding::Identity) => {
                assert!(data_length.is_some());
                let data_length = data_length.unwrap();

                self.headers.push(
                    Header::from_bytes(
                        &b"Content-Length"[..],
                        format!("{}", data_length).as_bytes(),
                    )
                    .unwrap(),
                )
            }

            _ => (),
        };

        // sending headers
        write_message_header(
            writer.by_ref(),
            &http_version,
            &self.status_code,
            &self.headers,
        )?;

        // sending the body
        if !do_not_send_body {
            match transfer_encoding {
                Some(TransferEncoding::Chunked) => {
                    use chunked_transfer::Encoder;

                    let mut writer = Encoder::new(writer);
                    io::copy(&mut reader, &mut writer)?;
                }

                Some(TransferEncoding::Identity) => {
                    use util::EqualReader;

                    assert!(data_length.is_some());
                    let data_length = data_length.unwrap();

                    if data_length >= 1 {
                        let (mut equ_reader, _) = EqualReader::new(reader.by_ref(), data_length);
                        io::copy(&mut equ_reader, &mut writer)?;
                    }
                }

                _ => (),
            }
        }

        Ok(())
    }
}

impl<R> Response<R>
where
    R: Read + Send + 'static,
{
    /// Turns this response into a `Response<Box<Read + Send>>`.
    pub fn boxed(self) -> ResponseBox {
        Response {
            reader: Box::new(self.reader) as Box<dyn Read + Send>,
            status_code: self.status_code,
            headers: self.headers,
            data_length: self.data_length,
            chunked_threshold: self.chunked_threshold,
        }
    }
}

impl Response<File> {
    /// Builds a new `Response` from a `File`.
    ///
    /// The `Content-Type` will **not** be automatically detected,
    ///  you must set it yourself.
    pub fn from_file(file: File) -> Response<File> {
        let file_size = file.metadata().ok().map(|v| v.len() as usize);

        Response::new(
            StatusCode(200),
            Vec::with_capacity(0),
            file,
            file_size,
            None,
        )
    }
}

impl Response<Cursor<Vec<u8>>> {
    pub fn from_data<D>(data: D) -> Response<Cursor<Vec<u8>>>
    where
        D: Into<Vec<u8>>,
    {
        let data = data.into();
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            Vec::with_capacity(0),
            Cursor::new(data),
            Some(data_len),
            None,
        )
    }

    pub fn from_string<S>(data: S) -> Response<Cursor<Vec<u8>>>
    where
        S: Into<String>,
    {
        let data = data.into();
        let data_len = data.len();

        Response::new(
            StatusCode(200),
            vec![
                Header::from_bytes(&b"Content-Type"[..], &b"text/plain; charset=UTF-8"[..])
                    .unwrap(),
            ],
            Cursor::new(data.into_bytes()),
            Some(data_len),
            None,
        )
    }
}

impl Response<io::Empty> {
    /// Builds an empty `Response` with the given status code.
    pub fn empty<S>(status_code: S) -> Response<io::Empty>
    where
        S: Into<StatusCode>,
    {
        Response::new(
            status_code.into(),
            Vec::with_capacity(0),
            io::empty(),
            Some(0),
            None,
        )
    }

    /// DEPRECATED. Use `empty` instead.
    pub fn new_empty(status_code: StatusCode) -> Response<io::Empty> {
        Response::empty(status_code)
    }
}

impl Clone for Response<io::Empty> {
    fn clone(&self) -> Response<io::Empty> {
        Response {
            reader: io::empty(),
            status_code: self.status_code.clone(),
            headers: self.headers.clone(),
            data_length: self.data_length,
            chunked_threshold: self.chunked_threshold,
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Random number generators and adapters for common usage:
//!
//! - [`ThreadRng`], a fast, secure, auto-seeded thread-local generator
//! - [`StdRng`] and [`SmallRng`], algorithms to cover typical usage
//! - [`OsRng`] as an entropy source
//! - [`mock::StepRng`] as a simple counter for tests
//! - [`adapter::ReadRng`] to read from a file/stream
//! - [`adapter::ReseedingRng`] to reseed a PRNG on clone / process fork etc.
//!
//! # Background — Random number generators (RNGs)
//!
//! Computers are inherently deterministic, so to get *random* numbers one
//! either has to use a hardware generator or collect bits of *entropy* from
//! various sources (e.g. event timestamps, or jitter). This is a relatively
//! slow and complicated operation.
//!
//! Generally the operating system will collect some entropy, remove bias, and
//! use that to seed its own PRNG; [`OsRng`] provides an interface to this.
//! Some alternatives are available although not normally required; see the
//! [`rdrand`] and [`rand_jitter`] crates.
//!
//! ## Pseudo-random number generators
//!
//! What is commonly used instead of "true" random number renerators, are
//! *pseudo-random number generators* (PRNGs), deterministic algorithms that
//! produce an infinite stream of pseudo-random numbers from a small random
//! seed. PRNGs are faster, and have better provable properties. The numbers
//! produced can be statistically of very high quality and can be impossible to
//! predict. (They can also have obvious correlations and be trivial to predict;
//! quality varies.)
//!
//! There are two different types of PRNGs: those developed for simulations
//! and statistics, and those developed for use in cryptography; the latter are
//! called Cryptographically Secure PRNGs (CSPRNG or CPRNG). Both types can
//! have good statistical quality but the latter also have to be impossible to
//! predict, even after seeing many previous output values. Rand provides a good
//! default algorithm from each class:
//!
//! - [`SmallRng`] is a PRNG chosen for low memory usage, high performance and
//!   good statistical quality.
//! - [`StdRng`] is a CSPRNG chosen for good performance and trust of security
//!   (based on reviews, maturity and usage). The current algorithm is HC-128,
//!   which is one of the recommendations by ECRYPT's eSTREAM project.
//!
//! The above PRNGs do not cover all use-cases; more algorithms can be found in
//! the [`prng`][crate::prng] module, as well as in several other crates. For example, you
//! may wish a CSPRNG with significantly lower memory usage than [`StdRng`]
//! while being less concerned about performance, in which case [`ChaChaRng`]
//! (from the `rand_chacha` crate) is a good choice.
//!
//! One complexity is that the internal state of a PRNG must change with every
//! generated number. For APIs this generally means a mutable reference to the
//! state of the PRNG has to be passed around.
//!
//! A solution is [`ThreadRng`]. This is a thread-local implementation of
//! [`StdRng`] with automatic seeding on first use. It is the best choice if you
//! "just" want a convenient, secure, fast random number source. Use via the
//! [`thread_rng`] function, which gets a reference to the current thread's
//! local instance.
//!
//! ## Seeding
//!
//! As mentioned above, PRNGs require a random seed in order to produce random
//! output. This is especially important for CSPRNGs, which are still
//! deterministic algorithms, thus can only be secure if their seed value is
//! also secure. To seed a PRNG, use one of:
//!
//! - [`FromEntropy::from_entropy`]; this is the most convenient way to seed
//!   with fresh, secure random data.
//! - [`SeedableRng::from_rng`]; this allows seeding from another PRNG or
//!   from an entropy source such as [`OsRng`].
//! - [`SeedableRng::from_seed`]; this is mostly useful if you wish to be able
//!   to reproduce the output sequence by using a fixed seed. (Don't use
//!   [`StdRng`] or [`SmallRng`] in this case since different algorithms may be
//!   used by future versions of Rand; use an algorithm from the
//!   [`prng`] module.)
//!
//! ## Conclusion
//!
//! - [`thread_rng`] is what you often want to use.
//! - If you want more control, flexibility, or better performance, use
//!   [`StdRng`], [`SmallRng`] or an algorithm from the [`prng`] module.
//! - Use [`FromEntropy::from_entropy`] to seed new PRNGs.
//! - If you need reproducibility, use [`SeedableRng::from_seed`] combined with
//!   a named PRNG.
//!
//! More information and notes on cryptographic security can be found
//! in the [`prng`] module.
//!
//! ## Examples
//!
//! Examples of seeding PRNGs:
//!
//! ```
//! use rand::prelude::*;
//! # use rand::Error;
//!
//! // StdRng seeded securely by the OS or local entropy collector:
//! let mut rng = StdRng::from_entropy();
//! # let v: u32 = rng.gen();
//!
//! // SmallRng seeded from thread_rng:
//! # fn try_inner() -> Result<(), Error> {
//! let mut rng = SmallRng::from_rng(thread_rng())?;
//! # let v: u32 = rng.gen();
//! # Ok(())
//! # }
//! # try_inner().unwrap();
//!
//! // SmallRng seeded by a constant, for deterministic results:
//! let seed = [1,2,3,4, 5,6,7,8, 9,10,11,12, 13,14,15,16]; // byte array
//! let mut rng = SmallRng::from_seed(seed);
//! # let v: u32 = rng.gen();
//! ```
//!
//!
//! # Implementing custom RNGs
//!
//! If you want to implement custom RNG, see the [`rand_core`] crate. The RNG
//! will have to implement the [`RngCore`] trait, where the [`Rng`] trait is
//! build on top of.
//!
//! If the RNG needs seeding, also implement the [`SeedableRng`] trait.
//!
//! [`CryptoRng`] is a marker trait cryptographically secure PRNGs can
//! implement.
//!
//! [`SmallRng`]: rngs::SmallRng
//! [`StdRng`]: rngs::StdRng
//! [`ThreadRng`]: rngs::ThreadRng
//! [`mock::StepRng`]: rngs::mock::StepRng
//! [`adapter::ReadRng`]: rngs::adapter::ReadRng
//! [`adapter::ReseedingRng`]: rngs::adapter::ReseedingRng
//! [`ChaChaRng`]: ../../rand_chacha/struct.ChaChaRng.html
//! [`rdrand`]: https://crates.io/crates/rdrand
//! [`rand_jitter`]: https://crates.io/crates/rand_jitter

pub mod adapter;

#[cfg(feature="std")] mod entropy;
pub mod mock;   // Public so we don't export `StepRng` directly, making it a bit
                // more clear it is intended for testing.
mod small;
mod std;
#[cfg(feature="std")] pub(crate) mod thread;

#[allow(deprecated)]
#[cfg(feature="std")] pub use self::entropy::EntropyRng;

pub use self::small::SmallRng;
pub use self::std::StdRng;
#[cfg(feature="std")] pub use self::thread::ThreadRng;

#[cfg(feature="getrandom")] mod os;
#[cfg(feature="getrandom")] pub use self::os::OsRng;
<commit_msg>Re-write rngs module documentation<commit_after>// Copyright 2018 Developers of the Rand project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Random number generators and adapters
//!
//! ## Background: Random number generators (RNGs)
//!
//! Computers cannot produce random numbers from nowhere. We classify
//! random number generators as follows:
//!
//! -   "True" random number generators (TRNGs) use hard-to-predict data sources
//!     (e.g. the high-resolution parts of event timings and sensor jitter) to
//!     harvest random bit-sequences, apply algorithms to remove bias and
//!     estimate available entropy, then combine these bits into a byte-sequence
//!     or an entropy pool. This job is usually done by the operating system or
//!     a hardware generator (HRNG).
//! -   "Pseudo"-random number generators (PRNGs) use algorithms to transform a
//!     seed into a sequence of pseudo-random numbers. These generators can be
//!     fast and produce well-distributed unpredictable random numbers (or not).
//!     They are usually deterministic: given algorithm and seed, the output
//!     sequence can be reproduced. They have finite period and eventually loop;
//!     with many algorithms this period is fixed and can be proven sufficiently
//!     long, while others are chaotic and the period depends on the seed.
//! -   "Cryptographically secure" pseudo-random number generators (CSPRNGs)
//!     are the sub-set of PRNGs which are secure. Security of the generator
//!     relies both on hiding the internal state and using a strong algorithm.
//!
//! ## Traits and functionality
//!
//! All RNGs implement the [`RngCore`] trait, as a consequence of which the
//! [`Rng`] extension trait is automatically implemented. Secure RNGs may
//! additionally implement the [`CryptoRng`] trait.
//!
//! All PRNGs require a seed to produce their random number sequence. The
//! [`SeedableRng`] trait provides three ways of constructing PRNGs:
//!
//! -   `from_seed` accepts a type specific to the PRNG
//! -   `from_rng` allows a PRNG to be seeded from any other RNG
//! -   `seed_from_u64` allows any PRNG to be seeded from a `u64` insecurely
//!
//! Additionally, [`FromEntropy::from_entropy`] is a shortcut for seeding from
//! [`OsRng`].
//!
//! Use the [`rand_core`] crate when implementing your own RNGs.
//!
//! ## Our generators
//!
//! This crate provides several random number generators:
//!
//! -   [`OsRng`] is an interface to the operating system's random number
//!     source. Typically the operating system uses a CSPRNG with entropy
//!     provided by a TRNG and some type of on-going re-seeding.
//! -   [`ThreadRng`], provided by the [`thread_rng`] function, is a handle to a
//!     thread-local CSPRNG with periodic seeding from [`OsRng`]. Because this
//!     is local, it is typically much faster than [`OsRng`]. It should be
//!     secure, though the paranoid may prefer [`OsRng`].
//! -   [`StdRng`] is a CSPRNG chosen for good performance and trust of security
//!     (based on reviews, maturity and usage). The current algorithm is HC-128,
//!     which is one of the recommendations by ECRYPT's eSTREAM project.
//!     [`StdRng`] provides the algorithm used by [`ThreadRng`] but without
//!     periodic reseeding.
//! -   [`SmallRng`] is an **insecure** PRNG designed to be fast, simple, require
//!     little memory, and have good output quality.
//!
//! The algorithms selected for [`StdRng`] and [`SmallRng`] may change in any
//! release and may be platform-dependent, therefore they should be considered
//! **not reproducible**.
//!
//! ## Additional generators
//!
//! **TRNGs**: The [`rdrand`] crate provides an interface to the RDRAND and
//! RDSEED instructions available in modern Intel and AMD CPUs.
//! The [`rand_jitter`] crate provides a user-space implementation of
//! entropy harvesting from CPU timer jitter, but is very slow and has
//! [security issues](https://github.com/rust-random/rand/issues/699).
//!
//! **PRNGs**: Several companion crates are available, providing individual or
//! families of PRNG algorithms. These provide the implementations behind
//! [`StdRng`] and [`SmallRng`] but can also be used directly, indeed *should*
//! be used directly when **reproducibility** matters.
//! Some suggestions are: [`rand_chacha`], [`rand_pcg`], [`rand_xoshiro`].
//! A full list can be found by searching for crates with the [`rng` tag].
//!
//! [`SmallRng`]: rngs::SmallRng
//! [`StdRng`]: rngs::StdRng
//! [`OsRng`]: rngs::OsRng
//! [`ThreadRng`]: rngs::ThreadRng
//! [`mock::StepRng`]: rngs::mock::StepRng
//! [`adapter::ReadRng`]: rngs::adapter::ReadRng
//! [`adapter::ReseedingRng`]: rngs::adapter::ReseedingRng
//! [`ChaChaRng`]: ../../rand_chacha/struct.ChaChaRng.html
//! [`rdrand`]: https://crates.io/crates/rdrand
//! [`rand_jitter`]: https://crates.io/crates/rand_jitter
//! [`rand_chacha`]: https://crates.io/crates/rand_chacha
//! [`rand_pcg`]: https://crates.io/crates/rand_pcg
//! [`rand_xoshiro`]: https://crates.io/crates/rand_xoshiro
//! [`rng` tag]: https://crates.io/keywords/rng

pub mod adapter;

#[cfg(feature="std")] mod entropy;
pub mod mock;   // Public so we don't export `StepRng` directly, making it a bit
                // more clear it is intended for testing.
mod small;
mod std;
#[cfg(feature="std")] pub(crate) mod thread;

#[allow(deprecated)]
#[cfg(feature="std")] pub use self::entropy::EntropyRng;

pub use self::small::SmallRng;
pub use self::std::StdRng;
#[cfg(feature="std")] pub use self::thread::ThreadRng;

#[cfg(feature="getrandom")] mod os;
#[cfg(feature="getrandom")] pub use self::os::OsRng;
<|endoftext|>


<commit_before>use std::fmt;

use openssl::error::ErrorStack;
use reqwest::Error as ReqwestError;

pub enum Error {
    Openssl(ErrorStack),
    Reqwest(ReqwestError),
}

impl From<ErrorStack> for Error {
    fn from(stack: ErrorStack) -> Self {
        Self::Openssl(stack)
    }
}

impl From<ReqwestError> for Error {
    fn from(error: ReqwestError) -> Self {
        Self::Reqwest(error)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Openssl(stack) => stack.fmt(f),
            Self::Reqwest(error) => error.fmt(f),
        }
    }
}
<commit_msg>impl Debug and std::error::Error for s3::Error<commit_after>use std::fmt;

use openssl::error::ErrorStack;
use reqwest::Error as ReqwestError;

#[derive(Debug)]
pub enum Error {
    Openssl(ErrorStack),
    Reqwest(ReqwestError),
}

impl From<ErrorStack> for Error {
    fn from(stack: ErrorStack) -> Self {
        Self::Openssl(stack)
    }
}

impl From<ReqwestError> for Error {
    fn from(error: ReqwestError) -> Self {
        Self::Reqwest(error)
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Openssl(stack) => stack.fmt(f),
            Self::Reqwest(error) => error.fmt(f),
        }
    }
}

impl std::error::Error for Error {}
<|endoftext|>


<commit_before>use core::fmt;

#[lang="eh_personality"]
pub fn eh_personality() { }

#[lang="panic_fmt"]
extern fn panic_fmt(fmt: fmt::Arguments, file: &'static str, line: usize) -> ! {
    println!("panic: {} ({}:{})", fmt, file, line);
    loop {
        unsafe {
            asm!("hlt");
        }
    }
}

#[lang="stack_exhausted"]
fn stack_exhausted() -> ! { loop { } }
<commit_msg>Disable interrupts after panicking.<commit_after>use core::fmt;

#[lang="eh_personality"]
pub fn eh_personality() { }

#[lang="panic_fmt"]
extern fn panic_fmt(fmt: fmt::Arguments, file: &'static str, line: usize) -> ! {
    println!("panic: {} ({}:{})", fmt, file, line);
    unsafe {
        asm!("cli");
        loop {
            asm!("hlt");
        }
    }
}

#[lang="stack_exhausted"]
fn stack_exhausted() -> ! { loop { } }
<|endoftext|>