<commit_before>//! *lamport* implements one-time hash-based signatures using the Lamport signature scheme.

extern crate crypto;
extern crate rand;

use rand::OsRng;
use rand::Rng;
use crypto::digest::Digest;

pub type LamportSignatureData = Vec<Vec<u8>>;

/// A one-time signing public key
#[derive(Clone)]
pub struct PublicKey<T: Digest + Clone> {
    zero_values: Vec<Vec<u8>>,
    one_values: Vec<Vec<u8>>,
    digest: T,
}

/// A one-time signing private key
#[derive(Clone)]
pub struct PrivateKey<T: Digest + Clone> {
    // For a n bits hash function: (n * n/8 bytes) for zero_values and one_values
    zero_values: Vec<Vec<u8>>,
    one_values: Vec<Vec<u8>>,
    digest: T,
    used: bool,
}

impl<T: Digest + Clone> From<PublicKey<T>> for Vec<u8> {
    fn from(original: PublicKey<T>) -> Vec<u8> {
        original.to_bytes()
    }
}

impl<T: Digest + Clone> PublicKey<T> {
    pub fn values(&self) -> (Vec<Vec<u8>>, Vec<Vec<u8>>) {
        (self.zero_values.clone(), self.one_values.clone())
    }

    pub fn from_vec(vec: Vec<u8>, digest: T) -> Option<PublicKey<T>> {
        let size = vec.len();
        let hash_output_size = digest.output_bytes();

        let mut zero_values_merged = vec;
        let one_values_merged = zero_values_merged.split_off(size / 2);

        let mut zero_values = Vec::new();
        for i in (0..zero_values_merged.len()).filter(|x| x % hash_output_size == 0) {
            // indexes for heads
            let mut sub_vec = Vec::new();
            for j in 0..hash_output_size {
                sub_vec.push(zero_values_merged[i + j]);
            }

            zero_values.push(sub_vec);
        }

        let mut one_values = Vec::new();
        for i in (0..one_values_merged.len()).filter(|x| x % hash_output_size == 0) {
            // indexes for heads
            let mut sub_vec = Vec::new();
            for j in 0..hash_output_size {
                sub_vec.push(one_values_merged[i + j]);
            }

            one_values.push(sub_vec);
        }

        Some(PublicKey {
            zero_values: zero_values,
            one_values: one_values,
            digest: digest,
        })
    }

    pub fn to_bytes(&self) -> Vec<u8> {
        let mut bytes = self.zero_values.clone();
        bytes.extend(self.one_values.clone());
        bytes.iter().fold(Vec::new(), |mut acc, i| {
            acc.append(&mut i.clone());
            acc
        })
    }

    /// Verifies that the signature of the data is correctly signed with the given key
    pub fn verify_signature(&self, signature: &LamportSignatureData, data: &[u8]) -> bool {
        let mut digest = self.digest.clone();
        digest.input(data);
        let mut data_hash = vec![0 as u8; digest.output_bytes()];
        digest.result(data_hash.as_mut_slice());
        digest.reset();

        for (i, byte) in data_hash.iter().enumerate() {
            for j in 0..8 {
                let offset = i * 8 + j;
                if (byte & (1 << j)) > 0 {
                    digest.input(signature[offset].as_slice());
                    let mut hashed_value = vec![0 as u8; digest.output_bytes()];
                    digest.result(hashed_value.as_mut_slice());
                    digest.reset();
                    if hashed_value != self.one_values[offset] {
                        return false;
                    }
                } else {
                    digest.input(signature[offset].as_slice());
                    let mut hashed_value = vec![0 as u8; digest.output_bytes()];
                    digest.result(hashed_value.as_mut_slice());
                    digest.reset();
                    if hashed_value != self.zero_values[offset] {
                        return false;
                    }
                }
            }
        }

        true
    }
}

impl<T: Digest + Clone> PrivateKey<T> {
    /// Generates a new random one-time signing key. This method can panic if OS RNG fails
    pub fn new(digest: T) -> PrivateKey<T> {
        let generate_bit_hash_values = |hasher: &T| -> Vec<Vec<u8>> {
            let mut rng = match OsRng::new() {
                Ok(g) => g,
                Err(e) => panic!("Failed to obtain OS RNG: {}", e),
            };
            let buffer_byte = vec![0 as u8; hasher.output_bytes()];
            let mut buffer = vec![buffer_byte; hasher.output_bits()];

            for hash in &mut buffer {
                rng.fill_bytes(hash)
            }

            println!("{:?}", buffer);

            buffer
        };

        let zero_values = generate_bit_hash_values(&digest);
        let one_values = generate_bit_hash_values(&digest);

        PrivateKey {
            zero_values: zero_values,
            one_values: one_values,
            digest: digest,
            used: false,
        }
    }

    /// Returns the public key associated with this private key
    pub fn public_key(&self) -> PublicKey<T> {
        let mut digest = self.digest.clone();

        let hash_values = |x: &Vec<Vec<u8>>, hash_func: &mut Digest| -> Vec<Vec<u8>> {
            let buffer_byte = vec![0 as u8; hash_func.output_bytes()];
            let mut buffer = vec![buffer_byte; hash_func.output_bits()];

            for i in 0..hash_func.output_bits() {
                hash_func.input(x[i].as_slice());
                hash_func.result(buffer[i].as_mut_slice());
                hash_func.reset();
            }

            buffer
        };

        let hashed_zero_values = hash_values(&self.zero_values, &mut digest);
        let hashed_one_values = hash_values(&self.one_values, &mut digest);

        PublicKey {
            zero_values: hashed_zero_values,
            one_values: hashed_one_values,
            digest: digest,
        }
    }

    /// Signs the data with the private key and returns the result if successful.
    /// If unsuccesful, an explanation string is returned
    pub fn sign(&mut self, data: &[u8]) -> Result<LamportSignatureData, &'static str> {
        if self.used {
            return Err("Attempting to sign more than once.");
        }
        self.digest.input(data);
        let mut data_hash = vec![0 as u8; self.digest.output_bytes()];
        self.digest.result(data_hash.as_mut_slice());
        self.digest.reset();

        let signature_len = data_hash.len() * 8;
        let mut signature = Vec::with_capacity(signature_len);

        for (i, byte) in data_hash.iter().enumerate() {
            for j in 0..8 {
                let offset = i * 8 + j;
                if (byte & (1 << j)) > 0 {
                    // Bit is 1
                    signature.push(self.one_values[offset].clone());
                } else {
                    // Bit is 0
                    signature.push(self.zero_values[offset].clone());
                }
            }
        }
        self.used = true;
        Ok(signature)
    }
}

impl<T: Digest + Clone> Drop for PrivateKey<T> {
    fn drop(&mut self) {
        let zeroize_vector = |vector: &mut Vec<Vec<u8>>| {
            for v2 in vector.iter_mut() {
                for byte in v2.iter_mut() {
                    *byte = 0;
                }
            }
        };

        zeroize_vector(&mut self.zero_values);
        zeroize_vector(&mut self.one_values);
    }
}

impl<T: Digest + Clone> PartialEq for PrivateKey<T> {
    // ⚠ This is not a constant-time implementation
    fn eq(&self, other: &PrivateKey<T>) -> bool {
        if self.one_values.len() != other.one_values.len() {
            return false;
        }
        if self.zero_values.len() != other.zero_values.len() {
            return false;
        }

        for i in 0..self.zero_values.len() {
            if self.zero_values[i] != other.zero_values[i] ||
               self.one_values[i] != other.one_values[i] {
                return false;
            }
        }
        true
    }
}

#[cfg(test)]
pub mod tests;
<commit_msg>Chaining iterators, no copy.<commit_after>//! *lamport* implements one-time hash-based signatures using the Lamport signature scheme.

extern crate crypto;
extern crate rand;

use rand::OsRng;
use rand::Rng;
use crypto::digest::Digest;

pub type LamportSignatureData = Vec<Vec<u8>>;

/// A one-time signing public key
#[derive(Clone)]
pub struct PublicKey<T: Digest + Clone> {
    zero_values: Vec<Vec<u8>>,
    one_values: Vec<Vec<u8>>,
    digest: T,
}

/// A one-time signing private key
#[derive(Clone)]
pub struct PrivateKey<T: Digest + Clone> {
    // For a n bits hash function: (n * n/8 bytes) for zero_values and one_values
    zero_values: Vec<Vec<u8>>,
    one_values: Vec<Vec<u8>>,
    digest: T,
    used: bool,
}

impl<T: Digest + Clone> From<PublicKey<T>> for Vec<u8> {
    fn from(original: PublicKey<T>) -> Vec<u8> {
        original.to_bytes()
    }
}

impl<T: Digest + Clone> PublicKey<T> {
    pub fn values(&self) -> (Vec<Vec<u8>>, Vec<Vec<u8>>) {
        (self.zero_values.clone(), self.one_values.clone())
    }

    pub fn from_vec(vec: Vec<u8>, digest: T) -> Option<PublicKey<T>> {
        let size = vec.len();
        let hash_output_size = digest.output_bytes();

        let mut zero_values_merged = vec;
        let one_values_merged = zero_values_merged.split_off(size / 2);

        let mut zero_values = Vec::new();
        for i in (0..zero_values_merged.len()).filter(|x| x % hash_output_size == 0) {
            // indexes for heads
            let mut sub_vec = Vec::new();
            for j in 0..hash_output_size {
                sub_vec.push(zero_values_merged[i + j]);
            }

            zero_values.push(sub_vec);
        }

        let mut one_values = Vec::new();
        for i in (0..one_values_merged.len()).filter(|x| x % hash_output_size == 0) {
            // indexes for heads
            let mut sub_vec = Vec::new();
            for j in 0..hash_output_size {
                sub_vec.push(one_values_merged[i + j]);
            }

            one_values.push(sub_vec);
        }

        Some(PublicKey {
            zero_values: zero_values,
            one_values: one_values,
            digest: digest,
        })
    }

    pub fn to_bytes(&self) -> Vec<u8> {
        self.zero_values.iter().chain(self.one_values.iter()).fold(Vec::new(), |mut acc, i| {
            acc.append(&mut i.clone());
            acc
        })
    }

    /// Verifies that the signature of the data is correctly signed with the given key
    pub fn verify_signature(&self, signature: &LamportSignatureData, data: &[u8]) -> bool {
        let mut digest = self.digest.clone();
        digest.input(data);
        let mut data_hash = vec![0 as u8; digest.output_bytes()];
        digest.result(data_hash.as_mut_slice());
        digest.reset();

        for (i, byte) in data_hash.iter().enumerate() {
            for j in 0..8 {
                let offset = i * 8 + j;
                if (byte & (1 << j)) > 0 {
                    digest.input(signature[offset].as_slice());
                    let mut hashed_value = vec![0 as u8; digest.output_bytes()];
                    digest.result(hashed_value.as_mut_slice());
                    digest.reset();
                    if hashed_value != self.one_values[offset] {
                        return false;
                    }
                } else {
                    digest.input(signature[offset].as_slice());
                    let mut hashed_value = vec![0 as u8; digest.output_bytes()];
                    digest.result(hashed_value.as_mut_slice());
                    digest.reset();
                    if hashed_value != self.zero_values[offset] {
                        return false;
                    }
                }
            }
        }

        true
    }
}

impl<T: Digest + Clone> PrivateKey<T> {
    /// Generates a new random one-time signing key. This method can panic if OS RNG fails
    pub fn new(digest: T) -> PrivateKey<T> {
        let generate_bit_hash_values = |hasher: &T| -> Vec<Vec<u8>> {
            let mut rng = match OsRng::new() {
                Ok(g) => g,
                Err(e) => panic!("Failed to obtain OS RNG: {}", e),
            };
            let buffer_byte = vec![0 as u8; hasher.output_bytes()];
            let mut buffer = vec![buffer_byte; hasher.output_bits()];

            for hash in &mut buffer {
                rng.fill_bytes(hash)
            }

            println!("{:?}", buffer);

            buffer
        };

        let zero_values = generate_bit_hash_values(&digest);
        let one_values = generate_bit_hash_values(&digest);

        PrivateKey {
            zero_values: zero_values,
            one_values: one_values,
            digest: digest,
            used: false,
        }
    }

    /// Returns the public key associated with this private key
    pub fn public_key(&self) -> PublicKey<T> {
        let mut digest = self.digest.clone();

        let hash_values = |x: &Vec<Vec<u8>>, hash_func: &mut Digest| -> Vec<Vec<u8>> {
            let buffer_byte = vec![0 as u8; hash_func.output_bytes()];
            let mut buffer = vec![buffer_byte; hash_func.output_bits()];

            for i in 0..hash_func.output_bits() {
                hash_func.input(x[i].as_slice());
                hash_func.result(buffer[i].as_mut_slice());
                hash_func.reset();
            }

            buffer
        };

        let hashed_zero_values = hash_values(&self.zero_values, &mut digest);
        let hashed_one_values = hash_values(&self.one_values, &mut digest);

        PublicKey {
            zero_values: hashed_zero_values,
            one_values: hashed_one_values,
            digest: digest,
        }
    }

    /// Signs the data with the private key and returns the result if successful.
    /// If unsuccesful, an explanation string is returned
    pub fn sign(&mut self, data: &[u8]) -> Result<LamportSignatureData, &'static str> {
        if self.used {
            return Err("Attempting to sign more than once.");
        }
        self.digest.input(data);
        let mut data_hash = vec![0 as u8; self.digest.output_bytes()];
        self.digest.result(data_hash.as_mut_slice());
        self.digest.reset();

        let signature_len = data_hash.len() * 8;
        let mut signature = Vec::with_capacity(signature_len);

        for (i, byte) in data_hash.iter().enumerate() {
            for j in 0..8 {
                let offset = i * 8 + j;
                if (byte & (1 << j)) > 0 {
                    // Bit is 1
                    signature.push(self.one_values[offset].clone());
                } else {
                    // Bit is 0
                    signature.push(self.zero_values[offset].clone());
                }
            }
        }
        self.used = true;
        Ok(signature)
    }
}

impl<T: Digest + Clone> Drop for PrivateKey<T> {
    fn drop(&mut self) {
        let zeroize_vector = |vector: &mut Vec<Vec<u8>>| {
            for v2 in vector.iter_mut() {
                for byte in v2.iter_mut() {
                    *byte = 0;
                }
            }
        };

        zeroize_vector(&mut self.zero_values);
        zeroize_vector(&mut self.one_values);
    }
}

impl<T: Digest + Clone> PartialEq for PrivateKey<T> {
    // ⚠ This is not a constant-time implementation
    fn eq(&self, other: &PrivateKey<T>) -> bool {
        if self.one_values.len() != other.one_values.len() {
            return false;
        }
        if self.zero_values.len() != other.zero_values.len() {
            return false;
        }

        for i in 0..self.zero_values.len() {
            if self.zero_values[i] != other.zero_values[i] ||
               self.one_values[i] != other.one_values[i] {
                return false;
            }
        }
        true
    }
}

#[cfg(test)]
pub mod tests;
<|endoftext|>


<commit_before>//! A formatted and aligned table printer written in rust
use std::io::{stdout, Write, Error, ErrorKind};
use std::fmt;
use std::str;
use std::string::ToString;

pub mod cell;
pub mod row;

use row::Row;
use cell::Cell;

#[cfg(any(unix, macos))]
static LINEFEED: &'static [u8] = b"\n";
#[cfg(windows)]
static LINEFEED: &'static [u8] = b"\r\n";

/// A Struct representing a printable table
#[derive(Clone, Debug)]
pub struct Table {
	num_cols: usize,
	titles: Row,
	rows: Vec<Row>,
	col_sep: char,
	line_sep: char,
	sep_cross: char
}

impl Table {
	/// Create a new table with the number of columns equals to the length of `titles`
	pub fn new(titles: Row) -> Table {
		let n = titles.len();
		return Table {
			num_cols: n,
			titles: titles,
			rows: Vec::new(),
			col_sep: '|',
			line_sep: '-',
			sep_cross: '+'
		};
	}

	/// Change separators
	///
	/// `col` is the column separator
	/// `line` is the line separator
	/// `cross` is a special separator used when line and collumn separators meet
	/// Default separators used are '|', '-' and '+'
	pub fn separators(&mut self, col: char, line: char, cross: char) {
		self.col_sep = col;
		self.line_sep = line;
		self.sep_cross = cross;
	}

	/// Get the number of column
	pub fn get_column_num(&self) -> usize {
		return self.num_cols;
	}

	/// Get the number of rows
	pub fn len(&self) -> usize {
		return self.rows.len();
	}

	/// Get a mutable reference to a row
	pub fn get_mut_row(&mut self, row: usize) -> &mut Row {
		return &mut self.rows[row];
	}

	/// Get an immutable reference to a row
	pub fn get_row(&self, row: usize) -> &Row {
		return &self.rows[row];
	}

	/// Append a row in the table, transferring ownership of this row to the table
	/// and returning a mutable reference to the row
	pub fn add_row(&mut self, row: Row) -> Result<&mut Row, &str> {
		if row.len() != self.num_cols {
			return Err("Row does not have the proper number of column");
		}
		self.rows.push(row);
		let l = self.rows.len()-1;
		return Ok(self.get_mut_row(l));
	}

	/// Append an empty row in the table. Return a mutable reference to this new row.
	pub fn add_empty_row(&mut self) -> Result<&mut Row, &str> {
		let n = self.num_cols;
		return Ok(try!(self.add_row(Row::empty(n))));
	}

	/// Modify a single element in the table
	pub fn set_element(&mut self, element: &String, column: usize, row: usize) -> Result<(), &str> {
		if column >= self.num_cols {
			return Err("Column index is higher than expected");
		}
		if row > self.rows.len() {
			return Err("Row index is higher than contained number of rows");
		}
		let rowline = self.get_mut_row(row);
		rowline.set_cell(Cell::new(element), column);
		return Ok(());
	}

	/// Remove a row. Silently skip if row with index `row` does not exist
	pub fn remove_row(&mut self, row: usize) {
		if row < self.rows.len() {
			self.rows.remove(row);
		}
	}

	fn get_col_width(&self, col_idx: usize) -> Result<usize, &str> {
		if col_idx >= self.num_cols {
			return Err("Column index is too high");
		}
		let mut width = self.titles.get_cell_width(col_idx);
		for r in &self.rows {
			let l = r.get_cell_width(col_idx);
			if l > width {
				width = l;
			}
		}
		return Ok(width);
	}

	fn print_line_separator<T: Write>(&self, out: &mut T, col_width: &[usize]) -> Result<(), Error> {
		try!(out.write_all(self.sep_cross.to_string().as_bytes()));
		for i in 0..self.num_cols {
			for _ in 0..(col_width[i] + 2) {
				try!(out.write_all(self.line_sep.to_string().as_bytes()));
			}
			try!(out.write_all(self.sep_cross.to_string().as_bytes()));
		}
		return out.write_all(LINEFEED);
	}

	/// Print the table to `out`
	pub fn print<T: Write>(&self, out: &mut T) -> Result<(), Error> {
		// Compute columns width
		let mut col_width = vec![0usize; self.num_cols];
		for i in 0..self.num_cols {
			col_width[i] = self.get_col_width(i).unwrap();
		}
		// Print titles line
		try!(self.print_line_separator(out, &col_width));
		try!(self.titles.print(out, self.col_sep, &col_width));
		try!(self.print_line_separator(out, &col_width));
		// Print rows
		for r in &self.rows {
			try!(r.print(out, self.col_sep, &col_width));
			try!(self.print_line_separator(out, &col_width));
		}
		return out.flush();
	}

	/// Print the table to standard output
	/// # Panic
	/// Panic if writing to standard output fails
	pub fn printstd(&self) {
		self.print(&mut stdout())
			.ok()
			.expect("Cannot print table to standard output");
	}
}

impl fmt::Display for Table {
	fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
		let mut writer =  StringWriter::new();
		if let Err(_) = self.print(&mut writer) {
			return Err(fmt::Error)
		}
		return fmt.write_str(writer.as_string());
	}
}

/// Internal utility for writing data into a string
struct StringWriter {
	string: String
}

impl StringWriter {
	/// Create a new `StringWriter`
	fn new() -> StringWriter {
		return StringWriter{string: String::new()};
	}

	/// Return a reference to the internally written `String`
	fn as_string(&self) -> &String {
		return &self.string;
	}
}

impl Write for StringWriter {
	fn write(&mut self, data: &[u8]) -> Result<usize, Error> {
		let string = match str::from_utf8(data) {
			Ok(s) => s,
			Err(e) => return Err(Error::new(ErrorKind::Other, format!("Cannot decode utf8 string : {}", e)))
		};
		self.string.push_str(string);
		return Ok(data.len());
	}

	fn flush(&mut self) -> Result<(), Error> {
		// Nothing to do here
		return Ok(());
	}
}

/// Create a table with column's titles from arguments.
///
/// The table can also be initialized with some values.
/// All the arguments used for titles and elements must implement the `std::string::ToString` trait
/// # Syntax
/// table!([Title1, Title2, ...], [Element1_ row1, Element2_ row1, ...], [Element1_row2, ...], ...);
/// # Panic
/// May panic if some rows could not be inserted
/// # Example
/// ```
/// # #[macro_use] extern crate tabprint;
/// # fn main() {
/// // Create an empty table with titles :
/// let tab = table!(["Title1", "Title2", "Title3"]);
/// // Create a table initialized with some rows :
/// let tab = table!(["Title1", "Title2", "Title3"],
/// 				 ["Element1", "Element2", "Element3"],
/// 				 [1, 2, 3]
/// 				 );
/// # drop(tab);
/// # }
/// ```
#[macro_export]
macro_rules! table {
	([$($title: expr), *]) => ($crate::Table::new(row![$($title), *]));
	(
		[$($title: expr), *], $([$($key:expr), *]), *
	) => (
		{
			let mut tab = table!([$($title), *]);
			$(
				if let Err(e) = tab.add_row(row![$($key), *]) {
					panic!("Cannot create table from : {}", e);
				}
			)*
			tab
		}
	)
}

/// Create a table with `table!` macro, print it to standard output, then return this table for future usage.
///
/// The syntax is the same that the one for the `table!` macro
/// # Panic
/// May panic if some rows could not be inserted
#[macro_export]
macro_rules! ptable {
	([$($title: expr), *]) => (
		{
			let tab = table!([$($title), *]);
			tab.printstd();
			tab
		}
	);
	(
		[$($title: expr), *], $([$($key:expr), *]), *
	) => (
		{
			let tab = table!([$($title), *], $([$($key), *]), *);
			tab.printstd();
			tab
		}
	)
}<commit_msg>Fixed typo in a panic message<commit_after>//! A formatted and aligned table printer written in rust
use std::io::{stdout, Write, Error, ErrorKind};
use std::fmt;
use std::str;
use std::string::ToString;

pub mod cell;
pub mod row;

use row::Row;
use cell::Cell;

#[cfg(any(unix, macos))]
static LINEFEED: &'static [u8] = b"\n";
#[cfg(windows)]
static LINEFEED: &'static [u8] = b"\r\n";

/// A Struct representing a printable table
#[derive(Clone, Debug)]
pub struct Table {
	num_cols: usize,
	titles: Row,
	rows: Vec<Row>,
	col_sep: char,
	line_sep: char,
	sep_cross: char
}

impl Table {
	/// Create a new table with the number of columns equals to the length of `titles`
	pub fn new(titles: Row) -> Table {
		let n = titles.len();
		return Table {
			num_cols: n,
			titles: titles,
			rows: Vec::new(),
			col_sep: '|',
			line_sep: '-',
			sep_cross: '+'
		};
	}

	/// Change separators
	///
	/// `col` is the column separator
	/// `line` is the line separator
	/// `cross` is a special separator used when line and collumn separators meet
	/// Default separators used are '|', '-' and '+'
	pub fn separators(&mut self, col: char, line: char, cross: char) {
		self.col_sep = col;
		self.line_sep = line;
		self.sep_cross = cross;
	}

	/// Get the number of column
	pub fn get_column_num(&self) -> usize {
		return self.num_cols;
	}

	/// Get the number of rows
	pub fn len(&self) -> usize {
		return self.rows.len();
	}

	/// Get a mutable reference to a row
	pub fn get_mut_row(&mut self, row: usize) -> &mut Row {
		return &mut self.rows[row];
	}

	/// Get an immutable reference to a row
	pub fn get_row(&self, row: usize) -> &Row {
		return &self.rows[row];
	}

	/// Append a row in the table, transferring ownership of this row to the table
	/// and returning a mutable reference to the row
	pub fn add_row(&mut self, row: Row) -> Result<&mut Row, &str> {
		if row.len() != self.num_cols {
			return Err("Row does not have the proper number of column");
		}
		self.rows.push(row);
		let l = self.rows.len()-1;
		return Ok(self.get_mut_row(l));
	}

	/// Append an empty row in the table. Return a mutable reference to this new row.
	pub fn add_empty_row(&mut self) -> Result<&mut Row, &str> {
		let n = self.num_cols;
		return Ok(try!(self.add_row(Row::empty(n))));
	}

	/// Modify a single element in the table
	pub fn set_element(&mut self, element: &String, column: usize, row: usize) -> Result<(), &str> {
		if column >= self.num_cols {
			return Err("Column index is higher than expected");
		}
		if row > self.rows.len() {
			return Err("Row index is higher than contained number of rows");
		}
		let rowline = self.get_mut_row(row);
		rowline.set_cell(Cell::new(element), column);
		return Ok(());
	}

	/// Remove a row. Silently skip if row with index `row` does not exist
	pub fn remove_row(&mut self, row: usize) {
		if row < self.rows.len() {
			self.rows.remove(row);
		}
	}

	fn get_col_width(&self, col_idx: usize) -> Result<usize, &str> {
		if col_idx >= self.num_cols {
			return Err("Column index is too high");
		}
		let mut width = self.titles.get_cell_width(col_idx);
		for r in &self.rows {
			let l = r.get_cell_width(col_idx);
			if l > width {
				width = l;
			}
		}
		return Ok(width);
	}

	fn print_line_separator<T: Write>(&self, out: &mut T, col_width: &[usize]) -> Result<(), Error> {
		try!(out.write_all(self.sep_cross.to_string().as_bytes()));
		for i in 0..self.num_cols {
			for _ in 0..(col_width[i] + 2) {
				try!(out.write_all(self.line_sep.to_string().as_bytes()));
			}
			try!(out.write_all(self.sep_cross.to_string().as_bytes()));
		}
		return out.write_all(LINEFEED);
	}

	/// Print the table to `out`
	pub fn print<T: Write>(&self, out: &mut T) -> Result<(), Error> {
		// Compute columns width
		let mut col_width = vec![0usize; self.num_cols];
		for i in 0..self.num_cols {
			col_width[i] = self.get_col_width(i).unwrap();
		}
		// Print titles line
		try!(self.print_line_separator(out, &col_width));
		try!(self.titles.print(out, self.col_sep, &col_width));
		try!(self.print_line_separator(out, &col_width));
		// Print rows
		for r in &self.rows {
			try!(r.print(out, self.col_sep, &col_width));
			try!(self.print_line_separator(out, &col_width));
		}
		return out.flush();
	}

	/// Print the table to standard output
	/// # Panic
	/// Panic if writing to standard output fails
	pub fn printstd(&self) {
		self.print(&mut stdout())
			.ok()
			.expect("Cannot print table to standard output");
	}
}

impl fmt::Display for Table {
	fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
		let mut writer =  StringWriter::new();
		if let Err(_) = self.print(&mut writer) {
			return Err(fmt::Error)
		}
		return fmt.write_str(writer.as_string());
	}
}

/// Internal utility for writing data into a string
struct StringWriter {
	string: String
}

impl StringWriter {
	/// Create a new `StringWriter`
	fn new() -> StringWriter {
		return StringWriter{string: String::new()};
	}

	/// Return a reference to the internally written `String`
	fn as_string(&self) -> &String {
		return &self.string;
	}
}

impl Write for StringWriter {
	fn write(&mut self, data: &[u8]) -> Result<usize, Error> {
		let string = match str::from_utf8(data) {
			Ok(s) => s,
			Err(e) => return Err(Error::new(ErrorKind::Other, format!("Cannot decode utf8 string : {}", e)))
		};
		self.string.push_str(string);
		return Ok(data.len());
	}

	fn flush(&mut self) -> Result<(), Error> {
		// Nothing to do here
		return Ok(());
	}
}

/// Create a table with column's titles from arguments.
///
/// The table can also be initialized with some values.
/// All the arguments used for titles and elements must implement the `std::string::ToString` trait
/// # Syntax
/// table!([Title1, Title2, ...], [Element1_ row1, Element2_ row1, ...], [Element1_row2, ...], ...);
/// # Panic
/// May panic if some rows could not be inserted
/// # Example
/// ```
/// # #[macro_use] extern crate tabprint;
/// # fn main() {
/// // Create an empty table with titles :
/// let tab = table!(["Title1", "Title2", "Title3"]);
/// // Create a table initialized with some rows :
/// let tab = table!(["Title1", "Title2", "Title3"],
/// 				 ["Element1", "Element2", "Element3"],
/// 				 [1, 2, 3]
/// 				 );
/// # drop(tab);
/// # }
/// ```
#[macro_export]
macro_rules! table {
	([$($title: expr), *]) => ($crate::Table::new(row![$($title), *]));
	(
		[$($title: expr), *], $([$($key:expr), *]), *
	) => (
		{
			let mut tab = table!([$($title), *]);
			$(
				if let Err(e) = tab.add_row(row![$($key), *]) {
					panic!("Cannot create table : {}", e);
				}
			)*
			tab
		}
	)
}

/// Create a table with `table!` macro, print it to standard output, then return this table for future usage.
///
/// The syntax is the same that the one for the `table!` macro
/// # Panic
/// May panic if some rows could not be inserted
#[macro_export]
macro_rules! ptable {
	([$($title: expr), *]) => (
		{
			let tab = table!([$($title), *]);
			tab.printstd();
			tab
		}
	);
	(
		[$($title: expr), *], $([$($key:expr), *]), *
	) => (
		{
			let tab = table!([$($title), *], $([$($key), *]), *);
			tab.printstd();
			tab
		}
	)
}<|endoftext|>


<commit_before>// TODO: deny missing docs
#![allow(missing_docs)]

extern crate futures;
extern crate tokio_core;
extern crate trust_dns;

use futures::Future;
use futures::future::join_all;
use std::net::{Ipv4Addr, SocketAddr};
use std::str::FromStr;
use tokio_core::reactor::Handle;
use trust_dns::client::{ClientFuture, ClientHandle};
use trust_dns::rr::domain;
use trust_dns::rr::{DNSClass, RData, RecordType};
use trust_dns::udp::UdpClientStream;

pub fn lookup<T: Into<SocketAddr>>(
    loop_handle: &Handle,
    domain_name: &str,
    server: T,
) -> Box<Future<Item = Option<Ipv4Addr>, Error = ()>> {
    let socket_address = server.into();
    let (stream, sender) = UdpClientStream::new(socket_address, loop_handle);
    let mut client = ClientFuture::new(stream, sender, loop_handle, None);
    let domain_name = domain::Name::from_str(domain_name).unwrap();

    Box::new(
        client
            .query(domain_name, DNSClass::IN, RecordType::A)
            .map(move |response| {
                let record = &response.answers()[0];
                if let RData::A(address) = *record.rdata() {
                    Some(address)
                } else {
                    None
                }
            })
            .map_err(move |_| ()),
    )
}

pub fn multiple_lookup<T: Into<SocketAddr>>(
    loop_handle: &Handle,
    domain_name: &str,
    servers: Vec<T>,
) -> Box<Future<Item = Vec<Option<Ipv4Addr>>, Error = ()>> {
    let futures: Vec<_> = servers
        .into_iter()
        .map(|server| lookup(loop_handle, domain_name, server))
        .collect();

    Box::new(join_all(futures))
}

#[cfg(test)]
mod test {
    use super::*;
    use tokio_core::reactor::Core;

    #[test]
    fn lookup_with_google() {
        let mut io_loop = Core::new().unwrap();
        let host = "example.com";
        let server = (Ipv4Addr::from_str("8.8.8.8").unwrap(), 53);

        let lookup = lookup(&io_loop.handle(), host, server);
        let results = io_loop.run(lookup);

        assert_eq!(results.unwrap(), Some(Ipv4Addr::new(93, 184, 216, 34)));
    }

    #[test]
    fn multiple_lookup_with_google() {
        let mut io_loop = Core::new().unwrap();
        let host = "example.com";
        let servers = vec![
            (Ipv4Addr::from_str("8.8.8.8").unwrap(), 53),
            (Ipv4Addr::from_str("8.8.4.4").unwrap(), 53),
        ];

        let lookup = multiple_lookup(&io_loop.handle(), host, servers);
        let results = io_loop.run(lookup);

        assert!(results.is_ok());
        let addresses = results.unwrap();
        assert_eq!(addresses[0], Some(Ipv4Addr::new(93, 184, 216, 34)));
        assert_eq!(addresses[1], Some(Ipv4Addr::new(93, 184, 216, 34)));
    }
}
<commit_msg>Change lookup result from string to Record<commit_after>// TODO: deny missing docs
#![allow(missing_docs)]
// for mem::discriminant_value
#![feature(discriminant_value)]

extern crate futures;
extern crate tokio_core;
extern crate trust_dns;

use futures::Future;
use futures::future::join_all;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::str::FromStr;
use tokio_core::reactor::Handle;
use trust_dns::client::{ClientFuture, ClientHandle};
use trust_dns::rr::domain;
use trust_dns::rr::{DNSClass, RData, Record, RecordType};
use trust_dns::udp::UdpClientStream;

#[derive(Debug)]
pub struct DnsResponse {
    server: IpAddr,
    answers: Vec<Record>,
}

pub type Result = std::result::Result<DnsResponse, Box<std::error::Error>>;

pub fn lookup<T: Into<SocketAddr>>(
    loop_handle: &Handle,
    domain_name: &str,
    server: T,
) -> Box<Future<Item=Result, Error=()>> {
    let socket_address = server.into();
    let (stream, sender) = UdpClientStream::new(socket_address, loop_handle);
    let mut client = ClientFuture::new(stream, sender, loop_handle, None);
    let domain_name = domain::Name::from_str(domain_name).unwrap();

    Box::new(
        client
            .query(domain_name, DNSClass::IN, RecordType::A)
            .map(move |mut response| {
                Ok(DnsResponse { server: socket_address.ip(), answers: response.take_answers() })
            })
            .map_err(move |_| ()),
    )
}

pub fn multiple_lookup<T: Into<SocketAddr>>(
    loop_handle: &Handle,
    domain_name: &str,
    servers: Vec<T>,
) -> Box<Future<Item=Vec<Result>, Error=()>> {
    let futures: Vec<_> = servers
        .into_iter()
        .map(|server| lookup(loop_handle, domain_name, server))
        .collect();

    Box::new(join_all(futures))
}

#[cfg(test)]
mod test {
    use super::*;
    use tokio_core::reactor::Core;
    use std::mem;

    #[test]
    fn lookup_with_google() {
        let mut io_loop = Core::new().unwrap();
        let host = "example.com";
        let server = (Ipv4Addr::from_str("8.8.8.8").unwrap(), 53);

        let lookup = lookup(&io_loop.handle(), host, server);
        let result = io_loop.run(lookup).unwrap();
        let response: DnsResponse = result.unwrap();

        assert_eq!(response.server, Ipv4Addr::from_str("8.8.8.8").unwrap());
        assert_eq!(response.answers.len(), 1);
        assert!(is_A_record(response.answers[0].rdata()));
        if let RData::A(ip) = *response.answers[0].rdata() {
            assert_eq!(ip, Ipv4Addr::new(93, 184, 216, 34));
        }
    }

    #[test]
    fn multiple_lookup_with_google() {
        let mut io_loop = Core::new().unwrap();
        let host = "example.com";
        let servers = vec![
            (Ipv4Addr::from_str("8.8.8.8").unwrap(), 53),
            (Ipv4Addr::from_str("8.8.4.4").unwrap(), 53),
        ];

        let lookup = multiple_lookup(&io_loop.handle(), host, servers);
        let mut responses: Vec<_> = io_loop.run(lookup).unwrap();
        assert_eq!(responses.len(), 2);

        let response = responses.pop().unwrap().unwrap();
        assert_eq!(response.server, Ipv4Addr::from_str("8.8.4.4").unwrap());
        assert_eq!(response.answers.len(), 1);
        assert!(is_A_record(response.answers[0].rdata()));
        if let RData::A(ip) = *response.answers[0].rdata() {
            assert_eq!(ip, Ipv4Addr::new(93, 184, 216, 34));
        }

        let response = responses.pop().unwrap().unwrap();
        assert_eq!(response.server, Ipv4Addr::from_str("8.8.8.8").unwrap());
        assert_eq!(response.answers.len(), 1);
        assert!(is_A_record(response.answers[0].rdata()));
        if let RData::A(ip) = *response.answers[0].rdata() {
            assert_eq!(ip, Ipv4Addr::new(93, 184, 216, 34));
        }
    }

    #[allow(non_snake_case)]
    fn is_A_record(rdata: &RData) -> bool {
        mem::discriminant(rdata) == mem::discriminant(&RData::A(Ipv4Addr::new(127, 0, 0, 1)))
    }
}
<|endoftext|>


<commit_before>#![no_std]
#![crate_name="cpu"]
#![crate_type="rlib"]
#![feature(asm)]

#[allow(unstable)] #[macro_use]
extern crate core;

pub use cpu::*;

#[allow(dead_code)]
#[static_assert]
static SUPPORTED : bool = cfg!(any(target_arch="x86", target_arch="x86_64"));

#[macro_use] mod bitflags;

#[cfg(target_arch="x86")]
#[path = "x86.rs"]
mod cpu;
#[cfg(target_arch="x86_64")]
#[path = "x86_64.rs"]
mod cpu;
#[cfg(not(any(target_arch="x86", target_arch="x86_64")))]
mod cpu { }

pub mod std {
	pub use core::fmt;
	pub use core::num;
	pub use core::option;
	pub use core::cmp;
	pub use core::clone;
	pub use core::marker;
	pub use core::ops;
}
<commit_msg>Either #![cfg(..)] works, or will error anyways<commit_after>#![cfg(any(target_arch="x86", target_arch="x86_64"))]

#![no_std]
#![crate_name="cpu"]
#![crate_type="rlib"]
#![feature(asm)]

#[allow(unstable)] #[macro_use]
extern crate core;

pub use cpu::*;

#[macro_use] mod bitflags;

#[cfg(target_arch="x86")]
#[path = "x86.rs"]
mod cpu;
#[cfg(target_arch="x86_64")]
#[path = "x86_64.rs"]
mod cpu;

pub mod std {
	pub use core::fmt;
	pub use core::num;
	pub use core::option;
	pub use core::cmp;
	pub use core::clone;
	pub use core::marker;
	pub use core::ops;
}
<|endoftext|>


<commit_before>//! A simple but effective sign library, written in Rust.
//!
//! Ported from [nobi](https://github.com/cyx/nobi),
//! which itself is a port of [itsdangerous](http://pythonhosted.org/itsdangerous/).
//!
//! # Possible use cases
//!
//! * Creating an activation link for users
//! * Creating a password reset link
//!
//! # Basic Example:
//!
//! ```rust
//! use nobsign::Signer;
//! let signer = Signer::new(b"my secret");
//!
//! // Let's say the user's ID is 101
//! let signed = signer.sign("101");
//!
//! // You can now email this url to your users!
//! let url = format!("http://yoursite.com/activate/?key={}", signed);
//!
//! // Later check the signature and get the value back
//! let unsigned = signer.unsign(&signed).unwrap();
//! ```
//!
//! # Example with timestamped signatures
//!
//! ```rust
//! use nobsign::TimestampSigner;
//! let signer = TimestampSigner::new(b"my secret");
//!
//! // Let's say the user's ID is 101
//! let signed = signer.sign("101");
//!
//! // You can now email this url to your users!
//! let url = format!("http://yoursite.com/activate/?key={}", signed);
//!
//! // In your code, you can verify the expiration:
//! signer.unsign(&signed, 86400).unwrap(); // 1 day expiration
//! ```
extern crate ring;
extern crate byteorder;
extern crate base64;

// Use same EPOCH as nobi, the Ruby implementation
const EPOCH : u64 = 1293840000;

use ring::{digest, hmac};
use byteorder::{ByteOrder, LittleEndian};
use base64::URL_SAFE_NO_PAD;

static ALGORITHM: &'static digest::Algorithm = &digest::SHA1;

#[derive(Debug,PartialEq)]
pub enum Error {
    BadData,
    BadSignature,
    BadTimeSignature,
    SignatureExpired,
}

pub struct Signer {
    separator: char,
    key: hmac::SigningKey,
}

pub struct TimestampSigner {
    signer: Signer,
}

fn int_to_bytes(n: i32) -> [u8; 4] {
    let mut out = [0; 4];

    LittleEndian::write_i32(&mut out, n);
    out
}
fn bytes_to_int(n: &[u8]) -> i32 {
    assert!(n.len() == 4);
    LittleEndian::read_i32(n)
}

impl Signer {
    pub fn new(secret: &[u8]) -> Signer {
        let initial_key = hmac::SigningKey::new(ALGORITHM, secret);
        let derived_key = hmac::sign(&initial_key, b"nobi.Signer");

        Signer {
            separator: '.',
            key: hmac::SigningKey::new(ALGORITHM, derived_key.as_ref()),
        }
    }

    pub fn sign(&self, value: &str) -> String {
        format!("{}{}{}", value, self.separator, &self.signature(value))
    }

    pub fn unsign(&self, value: &str) -> Result<String, Error> {
        let mut splitted = value.rsplitn(2, |c| c == self.separator);

        let sig = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };
        let value = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };

        let sig = try!(base64::decode_config(sig, URL_SAFE_NO_PAD).map_err(|_| Error::BadSignature));
        try!(hmac::verify_with_own_key(&self.key, value.as_bytes(), &sig)
                .map_err(|_| Error::BadSignature));

        Ok(value.into())
    }

    fn signature(&self, value: &str) -> String {
        let sig = hmac::sign(&self.key, value.as_bytes());
        base64::encode_config(sig.as_ref(), URL_SAFE_NO_PAD)
    }
}

impl TimestampSigner {
    pub fn new(secret: &[u8]) -> TimestampSigner {
        TimestampSigner { signer: Signer::new(secret) }
    }

    fn get_timestamp(&self) -> i32 {
        let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH);
        (match now {
            Ok(dur) => dur,
            Err(err) => err.duration(),
        }.as_secs() - EPOCH) as i32
    }

    pub fn sign(&self, value: &str) -> String {
        let timestamp = self.get_timestamp();
        let timestamp = int_to_bytes(timestamp);
        let timestamp = base64::encode_config(&timestamp, URL_SAFE_NO_PAD);
        let value = format!("{}{}{}", value, self.signer.separator, timestamp);

        format!("{}{}{}",
                value,
                self.signer.separator,
                &self.signature(&value))
    }

    pub fn unsign(&self, value: &str, max_age: u32) -> Result<String, Error> {
        let result = try!(self.signer.unsign(value));

        if !result.contains(self.signer.separator) {
            return Err(Error::BadTimeSignature);
        }

        let mut splitted = result.rsplitn(2, |c| c == self.signer.separator);

        let timestamp = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };
        let value = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };

        let timestamp = base64::decode_config(timestamp, URL_SAFE_NO_PAD)
            .map_err(|_| Error::BadTimeSignature)?;

        if timestamp.len() != 4 {
            return Err(Error::BadTimeSignature);
        }
        let timestamp = bytes_to_int(&timestamp);
        let age = self.get_timestamp() - timestamp;

        if age > max_age as i32 {
            return Err(Error::SignatureExpired);
        }

        Ok(value.into())
    }

    fn signature(&self, value: &str) -> String {
        self.signer.signature(value)
    }
}

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

    #[test]
    fn signs() {
        let signer = Signer::new(b"my-key");

        let signed = signer.sign("value");
        assert_eq!("value.EWkF3-80sipsPgLQ01NuTuPb0jQ".to_owned(), signed);

        assert_eq!("value".to_owned(), signer.unsign(&signed).unwrap());
    }

    #[test]
    fn bad_unsign() {
        let signer = Signer::new(b"my-key");

        let signed = "value.ABCDEF";
        assert_eq!(Err(Error::BadSignature), signer.unsign(&signed));
    }

    #[test]
    fn signs_with_timestamp() {
        let signer = TimestampSigner::new(b"my-key");

        let signed = signer.sign("value");

        assert_eq!("value".to_owned(), signer.unsign(&signed, 100).unwrap());
    }

    #[test]
    fn bad_unsign_with_timestamp() {
        let signer = TimestampSigner::new(b"my-key");

        let signed = "value.ABCDEF";
        assert_eq!(Err(Error::BadSignature), signer.unsign(&signed, 10));

        let signed = "value.EWkF3-80sipsPgLQ01NuTuPb0jQ";
        assert_eq!(Err(Error::BadTimeSignature), signer.unsign(&signed, 10));

        let signed = "value.AB.HHdWpuF7QVDZ_02wvECHvtV8vIc";
        assert_eq!(Err(Error::BadTimeSignature), signer.unsign(&signed, 10));
    }

    #[test]
    fn unsign_expired() {
        use std::thread;
        use std::time::Duration;

        let signer = TimestampSigner::new(b"my-key");
        let signed = signer.sign("value");
        thread::sleep(Duration::from_secs(1));
        assert_eq!(Err(Error::SignatureExpired), signer.unsign(&signed, 0));
    }

    #[test]
    fn with_secure_secret() {
        use ring::rand::{SystemRandom, SecureRandom};
        let sys_rand = SystemRandom::new();
        let mut key = vec![0u8; ALGORITHM.output_len];
        sys_rand.fill(&mut key).unwrap();

        let signer = Signer::new(&key);
        let signed = signer.sign("Hello world!");

        assert_eq!("Hello world!", signer.unsign(&signed).unwrap());
    }
}
<commit_msg>cargo-fmt the code<commit_after>//! A simple but effective sign library, written in Rust.
//!
//! Ported from [nobi](https://github.com/cyx/nobi),
//! which itself is a port of [itsdangerous](http://pythonhosted.org/itsdangerous/).
//!
//! # Possible use cases
//!
//! * Creating an activation link for users
//! * Creating a password reset link
//!
//! # Basic Example:
//!
//! ```rust
//! use nobsign::Signer;
//! let signer = Signer::new(b"my secret");
//!
//! // Let's say the user's ID is 101
//! let signed = signer.sign("101");
//!
//! // You can now email this url to your users!
//! let url = format!("http://yoursite.com/activate/?key={}", signed);
//!
//! // Later check the signature and get the value back
//! let unsigned = signer.unsign(&signed).unwrap();
//! ```
//!
//! # Example with timestamped signatures
//!
//! ```rust
//! use nobsign::TimestampSigner;
//! let signer = TimestampSigner::new(b"my secret");
//!
//! // Let's say the user's ID is 101
//! let signed = signer.sign("101");
//!
//! // You can now email this url to your users!
//! let url = format!("http://yoursite.com/activate/?key={}", signed);
//!
//! // In your code, you can verify the expiration:
//! signer.unsign(&signed, 86400).unwrap(); // 1 day expiration
//! ```
extern crate base64;
extern crate byteorder;
extern crate ring;

// Use same EPOCH as nobi, the Ruby implementation
const EPOCH: u64 = 1293840000;

use base64::URL_SAFE_NO_PAD;
use byteorder::{ByteOrder, LittleEndian};
use ring::{digest, hmac};

static ALGORITHM: &'static digest::Algorithm = &digest::SHA1;

#[derive(Debug, PartialEq)]
pub enum Error {
    BadData,
    BadSignature,
    BadTimeSignature,
    SignatureExpired,
}

pub struct Signer {
    separator: char,
    key: hmac::SigningKey,
}

pub struct TimestampSigner {
    signer: Signer,
}

fn int_to_bytes(n: i32) -> [u8; 4] {
    let mut out = [0; 4];

    LittleEndian::write_i32(&mut out, n);
    out
}
fn bytes_to_int(n: &[u8]) -> i32 {
    assert!(n.len() == 4);
    LittleEndian::read_i32(n)
}

impl Signer {
    pub fn new(secret: &[u8]) -> Signer {
        let initial_key = hmac::SigningKey::new(ALGORITHM, secret);
        let derived_key = hmac::sign(&initial_key, b"nobi.Signer");

        Signer {
            separator: '.',
            key: hmac::SigningKey::new(ALGORITHM, derived_key.as_ref()),
        }
    }

    pub fn sign(&self, value: &str) -> String {
        format!("{}{}{}", value, self.separator, &self.signature(value))
    }

    pub fn unsign(&self, value: &str) -> Result<String, Error> {
        let mut splitted = value.rsplitn(2, |c| c == self.separator);

        let sig = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };
        let value = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };

        let sig = base64::decode_config(sig, URL_SAFE_NO_PAD).map_err(|_| Error::BadSignature)?;
        hmac::verify_with_own_key(&self.key, value.as_bytes(), &sig)
            .map_err(|_| Error::BadSignature)?;

        Ok(value.into())
    }

    fn signature(&self, value: &str) -> String {
        let sig = hmac::sign(&self.key, value.as_bytes());
        base64::encode_config(sig.as_ref(), URL_SAFE_NO_PAD)
    }
}

impl TimestampSigner {
    pub fn new(secret: &[u8]) -> TimestampSigner {
        TimestampSigner {
            signer: Signer::new(secret),
        }
    }

    fn get_timestamp(&self) -> i32 {
        let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH);
        (match now {
            Ok(dur) => dur,
            Err(err) => err.duration(),
        }
        .as_secs()
            - EPOCH) as i32
    }

    pub fn sign(&self, value: &str) -> String {
        let timestamp = self.get_timestamp();
        let timestamp = int_to_bytes(timestamp);
        let timestamp = base64::encode_config(&timestamp, URL_SAFE_NO_PAD);
        let value = format!("{}{}{}", value, self.signer.separator, timestamp);

        format!(
            "{}{}{}",
            value,
            self.signer.separator,
            &self.signature(&value)
        )
    }

    pub fn unsign(&self, value: &str, max_age: u32) -> Result<String, Error> {
        let result = self.signer.unsign(value)?;

        if !result.contains(self.signer.separator) {
            return Err(Error::BadTimeSignature);
        }

        let mut splitted = result.rsplitn(2, |c| c == self.signer.separator);

        let timestamp = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };
        let value = match splitted.next() {
            Some(val) => val,
            None => return Err(Error::BadSignature),
        };

        let timestamp = base64::decode_config(timestamp, URL_SAFE_NO_PAD)
            .map_err(|_| Error::BadTimeSignature)?;

        if timestamp.len() != 4 {
            return Err(Error::BadTimeSignature);
        }
        let timestamp = bytes_to_int(&timestamp);
        let age = self.get_timestamp() - timestamp;

        if age > max_age as i32 {
            return Err(Error::SignatureExpired);
        }

        Ok(value.into())
    }

    fn signature(&self, value: &str) -> String {
        self.signer.signature(value)
    }
}

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

    #[test]
    fn signs() {
        let signer = Signer::new(b"my-key");

        let signed = signer.sign("value");
        assert_eq!("value.EWkF3-80sipsPgLQ01NuTuPb0jQ".to_owned(), signed);

        assert_eq!("value".to_owned(), signer.unsign(&signed).unwrap());
    }

    #[test]
    fn bad_unsign() {
        let signer = Signer::new(b"my-key");

        let signed = "value.ABCDEF";
        assert_eq!(Err(Error::BadSignature), signer.unsign(&signed));
    }

    #[test]
    fn signs_with_timestamp() {
        let signer = TimestampSigner::new(b"my-key");

        let signed = signer.sign("value");

        assert_eq!("value".to_owned(), signer.unsign(&signed, 100).unwrap());
    }

    #[test]
    fn bad_unsign_with_timestamp() {
        let signer = TimestampSigner::new(b"my-key");

        let signed = "value.ABCDEF";
        assert_eq!(Err(Error::BadSignature), signer.unsign(&signed, 10));

        let signed = "value.EWkF3-80sipsPgLQ01NuTuPb0jQ";
        assert_eq!(Err(Error::BadTimeSignature), signer.unsign(&signed, 10));

        let signed = "value.AB.HHdWpuF7QVDZ_02wvECHvtV8vIc";
        assert_eq!(Err(Error::BadTimeSignature), signer.unsign(&signed, 10));
    }

    #[test]
    fn unsign_expired() {
        use std::thread;
        use std::time::Duration;

        let signer = TimestampSigner::new(b"my-key");
        let signed = signer.sign("value");
        thread::sleep(Duration::from_secs(1));
        assert_eq!(Err(Error::SignatureExpired), signer.unsign(&signed, 0));
    }

    #[test]
    fn with_secure_secret() {
        use ring::rand::{SecureRandom, SystemRandom};
        let sys_rand = SystemRandom::new();
        let mut key = vec![0u8; ALGORITHM.output_len];
        sys_rand.fill(&mut key).unwrap();

        let signer = Signer::new(&key);
        let signed = signer.sign("Hello world!");

        assert_eq!("Hello world!", signer.unsign(&signed).unwrap());
    }
}
<|endoftext|>


<commit_before>#![feature(recover)]
#![allow(non_camel_case_types)]

//! [Github](https://github.com/hansihe/Rustler)
//! [Example](https://github.com/hansihe/Rustler_Example)
//!
//! Rustler is a library for writing Erlang NIFs in safe Rust code. That means there should be no
//! ways to crash the BEAM (Erlang VM). The library provides facilities for generating the
//! boilerplate for interacting with the BEAM, handles encoding and decoding of Erlang terms, and
//! catches rust panics before they unwind into C.
//!
//! The library provides functionality for both Erlang and Elixir, however Elixir is favored as of
//! now.
//!
//! This crate provides the entire runtime library for rustler. Code generators are located in the
//! rustler_codegen library.
//!
//! # Getting Started
//! When starting a project with Rustler, the first thing we need is Rust itself. Because this
//! library uses compiler plugins for code generation, we need the nightly version of the rust
//! compiler.  Rust nightly can be installed by following [the instructions on the install
//! page](https://www.rust-lang.org/downloads.html#nightly). When using rust for more than testing,
//! it is recommended to use [multirust](https://github.com/brson/multirust), which takes care of
//! downloading, installing, updating and switching between different versions of rust.
//!
//! ## Setting up the project
//! Creating and managing a rust project is done with the `cargo` command line utility. Creating a
//! new project is done with `cargo new <project_name>`. This will create a new folder with the
//! project name, which contains the initial scaffold for a rust project.
//!
//! Before we start writing rust code for our NIFs, we need to add the dependencies to the
//! `Cargo.toml` file, as well as specify that we want to generate a dynamically linked library for
//! the BEAM to load. Your `Cargo.toml` file should something like this when done:
//!
//! ```
//! [package]
//! name = "<project_name>"
//! version = "0.1.0"
//! authors = ["<you>"]
//!
//! [lib]
//! name = "<library name>"
//! path = "src/lib.rs"
//! crate_type = ["dylib"]
//!
//! [dependencies]
//! rustler = ">=0.5"
//! rustler_codegen = ">=0.5"
//! ```
//!
//! At this point we can try running `cargo test` to verify that our project compiles, and to run
//! the automatically generated test.
//!
//! ## Writing our NIF
//! As we specified in the `Cargo.toml` file, the entry point to our NIF library is located at
//! `src/lib.rs`. This file was automatically generated when we created the project. By default
//! this file contains one test. Feel free to delete the test for now.
//!
//! To import the rustler library, add these lines to the top of the `lib.rs` file:
//!
//! ```
//! #![feature(plugin)]
//! #![plugin(rustler_codegen)]
//!
//! extern crate rustler;
//! use rustler::{NifEnv, NifTerm, NifError, NifResult, NifDecoder, NifEncoder};
//! ```
//!
//! To create a basic NIF, we need to add two more things to our `lib.rs` file, a function to
//! execute, and the nif export.
//!
//! ```
//! rustler_export_nifs!(
//!     "Elixir.OurNifModule",
//!     [("square", 1, square)],
//!     None
//! );
//!
//! fn square<'a>(env: &'a NifEnv, args: &Vec<NifTerm>) -> NifResult<NifTerm<'a>> {
//!     Err(NifError::BadArg)
//! }
//! ```
//!
//! The first parameter in the `rustler_export_nifs!` macro is the name of the module we want to
//! register our NIF in. Since all Elixir modules are prefixed with `Elixir`, this will register
//! the nif library in the `OurNifModule` Elixir module.
//!
//! The second argument is a list of tuples representing the functions we want to export. Each
//! tuple consists of 3 elements, the name of the function as exposed to Erlang, the arity of our
//! function, and the identifier of the function we define below.
//!
//! The third argument is a function that will be executed when the NIF library is loaded. This is
//! optional, and we will pass None fow now.
//!
//! We also define the `square` function itself. The function takes two arguments, and returns a
//! result. The first argument is the nif environment, which represents the process the NIF is
//! called in. The second argument is a list of terms. This list will always be of the same length
//! as the arity we specified in the `rustler_export_nifs!` entry.
//!
//! The return type is either a NifTerm or a NifError. Right now we simply return a BadArg error.
//!
//! ## Running our NIF
//! At this point our project should successfully compile with the `cargo build` command, and
//! output a dynamically linked library that can be loaded by the BEAM.
//!
//! To test our NIF, we will write a simple `exs` script that loads and runs our NIF, we will name
//! this file `run.exs`, and we will place it in our project root directory.
//!
//! ```elixir
//! defmodule OurNifModule do
//!     @on_load {:init, 0}
//!
//!     def init do
//!         path = :filelib.wildcard('target/{debug,release}/lib<library_name>.*') |> hd
//!             |> :filename.rootname
//!         :ok = :erlang.load_nif(path, 0)
//!         :ok
//!     end
//!
//!     def square(_a), do: exit(:nif_not_loaded)
//! end
//!
//! IO.inspect OurNifModule.square(12)
//! ```
//!
//! The Elixir code for running the NIF is very simple. We specify that we want our init function
//! to run when the module is first loaded by using the @on_load attribute. In our init function we
//! first find the path of the dynalically loaded library we genrated earlier when running `cargo
//! build`. We then call `:erlang.load_nif/2 with our path and some optional data.
//!
//! If we try running it we will see that the OurNifModule.square function returns badarg. This is
//! correct, as that is the error we returned from our rust function.
//!
//! ## Making our NIF useful
//! Our nif now runs, but it doesn't really do anything useful yet. In order to do anything useful,
//! we need to be able to accept input and produce output. This is generally done through the
//! [NifEncoder](trait.NifEncoder.html) and [NifDecoder](trait.NifDecoder.html) traits.
//!
//! Let's try to extract number from the args. As rust is very close to the metal, and uses
//! native types, while erlangs type system is a bit more lenient, we need to inform rust what type
//! we want to extract from the term. If we type annotate the variable we put the decoded value
//! into, the rust compiler is smart enough to propagate the type information, and everything else
//! is handled for us.
//!
//! ```
//! fn square<'a>(env: &'a NifEnv, args: &Vec<NifTerm>) -> NifResult<NifTerm<'a>> {
//!     let number: i64 = NifDecoder::decode(args[0], env);
//!     Ok(NifEncoder::encode(number * number))
//! }
//! ```

pub mod wrapper;
use wrapper::nif_interface::{NIF_ENV, NIF_TERM, enif_make_badarg, enif_make_atom_len};

#[macro_use]
extern crate lazy_static;

extern crate libc;

mod types;
pub use self::types::{ NifEncoder, NifDecoder };
pub mod resource;
pub mod binary;
pub mod tuple;
pub mod map;
pub mod atom;
pub mod codegen_runtime;
pub mod ex_struct;

pub type NifResult<T> = Result<T, NifError>;

/// On each NIF call, a NifEnv is passed in. The NifEnv is used for most operations that involve
/// communicating with the BEAM, like decoding and encoding terms.
///
/// There is no way to allocate a NifEnv at the moment, but this may be possible in the future.
#[derive(PartialEq)]
pub struct NifEnv {
    env: NIF_ENV,
}
impl NifEnv {
    pub fn as_c_arg(&self) -> NIF_ENV {
        self.env
    }
}

/// Represents usual errors that can happen in a nif. This enables you to return an error from
/// anywhere, even places where you don't have an NifEnv availible.
#[derive(Clone, Copy)]
pub enum NifError {
    /// Returned when the NIF has been called with the wrong number or type of arguments.
    BadArg,
    /// Returned when an allocation fails. Example: binary term, resource struct
    AllocFail,
    /// Encodes the string into an atom and returns it from the NIF.
    Atom(&'static str),
}
impl NifEncoder for NifError {
    fn encode<'a>(&self, env: &'a NifEnv) -> NifTerm<'a> {
        NifTerm::new(env, match *self {
            NifError::BadArg =>
                unsafe { enif_make_badarg(env.as_c_arg()) },
            NifError::AllocFail =>
                unsafe { enif_make_badarg(env.as_c_arg()) },
            NifError::Atom(name) =>
                unsafe { enif_make_atom_len(env.as_c_arg(),
                                            name.as_ptr() as *const u8,
                                            name.len() as libc::size_t) },
        })
    }
}

/// NifTerm is used to represent all erlang terms. Terms are always lifetime limited by a NifEnv.
///
/// NifTerm is cloneable and copyable, but it can not exist outside of the lifetime of the NifEnv
/// that owns it.
#[derive(Clone, Copy)]
pub struct NifTerm<'a> {
    term: NIF_TERM,
    env: &'a NifEnv,
}
impl<'a> NifTerm<'a> {
    pub fn new(env: &'a NifEnv, inner: NIF_TERM) -> Self {
        NifTerm {
            term: inner,
            env: env,
        }
    }
    /// This extracts the raw term pointer. It is usually used in order to obtain a type that can
    /// be passed to calls into the erlang vm.
    pub fn as_c_arg(&self) -> NIF_TERM {
        self.term
    }

    pub fn get_env(&self) -> &NifEnv {
        self.env
    }

    pub fn lifetime_cast<'b>(&self, env: &'b NifEnv) -> NifTerm<'b> {
        NifTerm::new(env, self.as_c_arg())
    }

    pub fn in_env<'b>(&self, env: &'b NifEnv) -> NifTerm<'b> {
        if self.get_env() == env {
            self.lifetime_cast(env)
        } else {
            NifTerm::new(env, wrapper::copy_term(env.as_c_arg(), self.as_c_arg()))
        }
    }
}

/// Exports a given list of functions to a Erlang module.
///
/// This should be called exactly once in every NIF library. It will wrap and export the given rust
/// functions into the Erlang module.
///
/// The first argument is a string specifying what Erlang/Elixir module you want the function
/// exported into. In Erlang this will simply be the atom you named your module. In Elixir, all
/// modules are prefixed with `Elixir.<module path>`
///
/// The second argument is a list of 3-tuples. Each tuple contains information on a single exported
/// NIF function. The first tuple item is the name you want to export the function into, the second
/// is the arity (number of arguments) of the exported function. The third argument is a
/// indentifier of a rust function. This is where your actual NIF will be implemented.
///
/// The third argument is an `Option<fn(env: &NifEnv, load_info: NifTerm) -> bool>`. If this is
/// `Some`, the function will execute when the NIF is first loaded by the BEAM.
///
/// WARNING: Only a stub, actual macro defined in compiler plugin.
#[macro_export]
macro_rules! rustler_export_nifs {
    ($module_name:expr, [$(($fun_export_name:expr, $arity:expr, $rust_fn:ident)),*], $init_function:expr) => { /* ... */ }
}
<commit_msg>update documentation for the new project generator<commit_after>#![feature(recover)]
#![allow(non_camel_case_types)]

//! [Github](https://github.com/hansihe/Rustler)
//! [Example](https://github.com/hansihe/Rustler_Example)
//!
//! Rustler is a library for writing Erlang NIFs in safe Rust code. That means there should be no
//! ways to crash the BEAM (Erlang VM). The library provides facilities for generating the
//! boilerplate for interacting with the BEAM, handles encoding and decoding of Erlang terms, and
//! catches rust panics before they unwind into C.
//!
//! The library provides functionality for both Erlang and Elixir, however Elixir is favored as of
//! now.
//!
//! This crate provides the entire runtime library for rustler. Code generators are located in the
//! rustler_codegen library.
//!
//! # Getting Started
//! Rustler has a Mix archive that provides several commands which makes working with Rustler
//! easier, including project generators, a tool for validating your Rust and Erlang environment,
//! as well as build utilities.
//!
//! As long as you have Elixir and Mix installed, you can install the archive by doing `mix
//! archive.install https://github.com/hansihe/rustler_archives/raw/master/rustler_installer.ez`.
//! When the archive is installed, you can generate a new Rustler project by doing `mix rustler.new
//! <path>`, and following the instructions.

pub mod wrapper;
use wrapper::nif_interface::{NIF_ENV, NIF_TERM, enif_make_badarg, enif_make_atom_len};

#[macro_use]
extern crate lazy_static;

extern crate libc;

mod types;
pub use self::types::{ NifEncoder, NifDecoder };
pub mod resource;
pub mod binary;
pub mod tuple;
pub mod map;
pub mod atom;
pub mod codegen_runtime;
pub mod ex_struct;

pub type NifResult<T> = Result<T, NifError>;

/// On each NIF call, a NifEnv is passed in. The NifEnv is used for most operations that involve
/// communicating with the BEAM, like decoding and encoding terms.
///
/// There is no way to allocate a NifEnv at the moment, but this may be possible in the future.
#[derive(PartialEq)]
pub struct NifEnv {
    env: NIF_ENV,
}
impl NifEnv {
    pub fn as_c_arg(&self) -> NIF_ENV {
        self.env
    }
}

/// Represents usual errors that can happen in a nif. This enables you to return an error from
/// anywhere, even places where you don't have an NifEnv availible.
#[derive(Clone, Copy)]
pub enum NifError {
    /// Returned when the NIF has been called with the wrong number or type of arguments.
    BadArg,
    /// Returned when an allocation fails. Example: binary term, resource struct
    AllocFail,
    /// Encodes the string into an atom and returns it from the NIF.
    Atom(&'static str),
}
impl NifEncoder for NifError {
    fn encode<'a>(&self, env: &'a NifEnv) -> NifTerm<'a> {
        NifTerm::new(env, match *self {
            NifError::BadArg =>
                unsafe { enif_make_badarg(env.as_c_arg()) },
            NifError::AllocFail =>
                unsafe { enif_make_badarg(env.as_c_arg()) },
            NifError::Atom(name) =>
                unsafe { enif_make_atom_len(env.as_c_arg(),
                                            name.as_ptr() as *const u8,
                                            name.len() as libc::size_t) },
        })
    }
}

/// NifTerm is used to represent all erlang terms. Terms are always lifetime limited by a NifEnv.
///
/// NifTerm is cloneable and copyable, but it can not exist outside of the lifetime of the NifEnv
/// that owns it.
#[derive(Clone, Copy)]
pub struct NifTerm<'a> {
    term: NIF_TERM,
    env: &'a NifEnv,
}
impl<'a> NifTerm<'a> {
    pub fn new(env: &'a NifEnv, inner: NIF_TERM) -> Self {
        NifTerm {
            term: inner,
            env: env,
        }
    }
    /// This extracts the raw term pointer. It is usually used in order to obtain a type that can
    /// be passed to calls into the erlang vm.
    pub fn as_c_arg(&self) -> NIF_TERM {
        self.term
    }

    pub fn get_env(&self) -> &NifEnv {
        self.env
    }

    pub fn lifetime_cast<'b>(&self, env: &'b NifEnv) -> NifTerm<'b> {
        NifTerm::new(env, self.as_c_arg())
    }

    pub fn in_env<'b>(&self, env: &'b NifEnv) -> NifTerm<'b> {
        if self.get_env() == env {
            self.lifetime_cast(env)
        } else {
            NifTerm::new(env, wrapper::copy_term(env.as_c_arg(), self.as_c_arg()))
        }
    }
}

/// Exports a given list of functions to a Erlang module.
///
/// This should be called exactly once in every NIF library. It will wrap and export the given rust
/// functions into the Erlang module.
///
/// The first argument is a string specifying what Erlang/Elixir module you want the function
/// exported into. In Erlang this will simply be the atom you named your module. In Elixir, all
/// modules are prefixed with `Elixir.<module path>`
///
/// The second argument is a list of 3-tuples. Each tuple contains information on a single exported
/// NIF function. The first tuple item is the name you want to export the function into, the second
/// is the arity (number of arguments) of the exported function. The third argument is a
/// indentifier of a rust function. This is where your actual NIF will be implemented.
///
/// The third argument is an `Option<fn(env: &NifEnv, load_info: NifTerm) -> bool>`. If this is
/// `Some`, the function will execute when the NIF is first loaded by the BEAM.
///
/// WARNING: Only a stub, actual macro defined in compiler plugin.
#[macro_export]
macro_rules! rustler_export_nifs {
    ($module_name:expr, [$(($fun_export_name:expr, $arity:expr, $rust_fn:ident)),*], $init_function:expr) => { /* ... */ }
}
<|endoftext|>


<commit_before>use std::fs;
use std::fs::File;
use std::io::prelude::*;
use std::io::{Cursor, SeekFrom, Take};
use std::marker::PhantomData;
use std::path::{Path, PathBuf};
use std::vec::Vec;

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use fs2::FileExt;
use time;
use regex::Regex;

use data::{Entry, Hint};
use util::{xxhash32, XxHash32, get_file_handle};

const DATA_FILE_EXTENSION: &'static str = "cask.data";
const HINT_FILE_EXTENSION: &'static str = "cask.hint";
const LOCK_FILE_NAME: &'static str = "cask.lock";

const DEFAULT_SIZE_THRESHOLD: usize = 100 * 1024 * 1024;

pub struct Log {
    path: PathBuf,
    sync: bool,
    size_threshold: usize,
    lock_file: File,
    current_file_id: u32,
    active_data_file_path: PathBuf,
    active_data_file: File,
    active_hint_file: File,
    active_hint_file_hasher: XxHash32,
}

pub struct Entries<'a> {
    data_file: Take<File>,
    data_file_pos: u64,
    phantom: PhantomData<&'a ()>,
}

impl<'a> Iterator for Entries<'a> {
    type Item = (u64, Entry<'a>);

    fn next(&mut self) -> Option<(u64, Entry<'a>)> {
        if self.data_file.limit() == 0 {
            None
        } else {
            let entry = Entry::from_read(&mut self.data_file);
            let entry_pos = self.data_file_pos;

            self.data_file_pos += entry.size();

            Some((entry_pos, entry))
        }
    }
}

pub struct Hints<'a> {
    hint_file: Take<File>,
    phantom: PhantomData<&'a ()>,
}

impl<'a> Iterator for Hints<'a> {
    type Item = Hint<'a>;

    fn next(&mut self) -> Option<Hint<'a>> {
        if self.hint_file.limit() == 0 {
            None
        } else {
            Some(Hint::from_read(&mut self.hint_file))
        }
    }
}

pub struct RecreateHints<'a> {
    hint_file: File,
    hint_file_hasher: XxHash32,
    entries: Entries<'a>,
}

impl<'a> Iterator for RecreateHints<'a> {
    type Item = Hint<'a>;

    fn next(&mut self) -> Option<Hint<'a>> {
        self.entries.next().map(|e| {
            let (entry_pos, entry) = e;
            let hint = Hint::from(entry, entry_pos);
            hint.write_bytes(&mut self.hint_file);
            hint.write_bytes(&mut self.hint_file_hasher);
            hint
        })
    }
}

impl<'a> Drop for RecreateHints<'a> {
    fn drop(&mut self) {
        while self.next().is_some() {}
        self.hint_file
            .write_u32::<LittleEndian>(self.hint_file_hasher.get())
            .unwrap();
    }
}

impl Log {
    pub fn open(path: &str, sync: bool) -> Log {
        let path = PathBuf::from(path);

        if path.exists() {
            assert!(path.is_dir());
        } else {
            fs::create_dir(&path).unwrap();
        }

        let lock_file = File::create(path.join(LOCK_FILE_NAME)).unwrap();
        lock_file.try_lock_exclusive().unwrap();

        let current_file_id = time::now().to_timespec().sec as u32;
        let active_data_file_path = get_data_file_path(&path, current_file_id);
        let active_data_file = get_file_handle(&active_data_file_path, true);
        let active_hint_file = get_file_handle(&get_hint_file_path(&path, current_file_id), true);
        let active_hint_file_hasher = XxHash32::new();

        info!("Created new active data file {:?}", active_data_file_path);

        Log {
            path: path,
            sync: sync,
            size_threshold: DEFAULT_SIZE_THRESHOLD,
            lock_file: lock_file,
            current_file_id: current_file_id,
            active_data_file: active_data_file,
            active_data_file_path: active_data_file_path,
            active_hint_file: active_hint_file,
            active_hint_file_hasher: active_hint_file_hasher,
        }
    }

    pub fn find_files(&self) -> Vec<u32> {
        let files = fs::read_dir(&self.path).unwrap();

        lazy_static! {
            static ref RE: Regex =
                Regex::new(&format!("(\\d+).{}$", DATA_FILE_EXTENSION)).unwrap();
        }

        let mut files: Vec<u32> = files.flat_map(|f| {
                let file = f.unwrap();
                let file_metadata = file.metadata().unwrap();

                if file_metadata.is_file() {
                    let file_name = file.file_name();
                    let captures = RE.captures(file_name.to_str().unwrap());
                    captures.and_then(|c| c.at(1).and_then(|n| n.parse::<u32>().ok()))
                        .and_then(|id| if id != self.current_file_id {
                            Some(id)
                        } else {
                            None
                        })
                } else {
                    None
                }
            })
            .collect();

        files.sort();

        files
    }

    pub fn entries<'a>(&self, file_id: u32) -> Entries<'a> {
        let data_file_path = get_data_file_path(&self.path, file_id);
        info!("Loading data file: {:?}", data_file_path);
        let data_file = get_file_handle(&data_file_path, false);
        let data_file_size = data_file.metadata().unwrap().len();

        Entries {
            data_file: data_file.take(data_file_size),
            data_file_pos: 0,
            phantom: PhantomData,
        }
    }

    pub fn hints<'a>(&self, file_id: u32) -> Option<Hints<'a>> {
        let hint_file_path = get_hint_file_path(&self.path, file_id);
        if is_valid_hint_file(&hint_file_path) {
            info!("Loading hint file: {:?}", hint_file_path);
            let hint_file = get_file_handle(&hint_file_path, false);
            let hint_file_size = hint_file.metadata().unwrap().len();

            Some(Hints {
                hint_file: hint_file.take(hint_file_size - 4),
                phantom: PhantomData,
            })
        } else {
            None
        }
    }

    pub fn recreate_hints<'a>(&mut self, file_id: u32) -> RecreateHints<'a> {
        let hint_file_path = get_hint_file_path(&self.path, file_id);
        info!("Re-creating hint file: {:?}", hint_file_path);
        let hint_file = get_file_handle(&hint_file_path, true);
        let entries = self.entries(file_id);

        RecreateHints {
            hint_file: hint_file,
            hint_file_hasher: XxHash32::new(),
            entries: entries,
        }
    }

    pub fn read_entry<'a>(&self, file_id: u32, entry_pos: u64) -> Entry<'a> {
        let mut data_file = get_file_handle(&get_data_file_path(&self.path, file_id), false);
        data_file.seek(SeekFrom::Start(entry_pos)).unwrap();
        Entry::from_read(&mut data_file)
    }

    pub fn write_entry<'a>(&mut self, entry: &Entry<'a>) -> (u32, u64) {
        let mut active_data_file_pos = self.active_data_file.seek(SeekFrom::Current(0)).unwrap();

        if active_data_file_pos + entry.size() > self.size_threshold as u64 {
            info!("Active data file {:?} reached file limit",
                  self.active_data_file_path);
            self.new_active_file();
            info!("Created new active data file {:?}",
                  self.active_data_file_path);
            active_data_file_pos = 0;
        }

        let hint = Hint::new(&entry, active_data_file_pos);

        entry.write_bytes(&mut self.active_data_file);
        hint.write_bytes(&mut self.active_hint_file);
        hint.write_bytes(&mut self.active_hint_file_hasher);

        if self.sync {
            self.active_data_file.sync_data().unwrap();
        }

        (self.current_file_id, active_data_file_pos)
    }

    fn new_active_file(&mut self) {
        if self.sync {
            self.active_data_file.sync_data().unwrap();
        }

        self.active_hint_file
            .write_u32::<LittleEndian>(self.active_hint_file_hasher.get())
            .unwrap();

        self.current_file_id = time::now().to_timespec().sec as u32;

        self.active_data_file_path = get_data_file_path(&self.path, self.current_file_id);
        self.active_data_file = get_file_handle(&self.active_data_file_path, true);

        self.active_hint_file =
            get_file_handle(&get_hint_file_path(&self.path, self.current_file_id), true);

        self.active_hint_file_hasher = XxHash32::new();
    }
}

impl Drop for Log {
    fn drop(&mut self) {
        self.active_hint_file
            .write_u32::<LittleEndian>(self.active_hint_file_hasher.get())
            .unwrap();

        self.lock_file.unlock().unwrap();
    }
}

fn get_data_file_path(path: &Path, file_id: u32) -> PathBuf {
    path.join(file_id.to_string()).with_extension(DATA_FILE_EXTENSION)
}

fn get_hint_file_path(path: &Path, file_id: u32) -> PathBuf {
    path.join(file_id.to_string()).with_extension(HINT_FILE_EXTENSION)
}

fn is_valid_hint_file(path: &Path) -> bool {
    path.is_file() &&
    {
        let mut hint_file = get_file_handle(path, false);

        // FIXME: avoid reading the whole hint file into memory;
        let mut buf = Vec::new();
        hint_file.read_to_end(&mut buf).unwrap();

        buf.len() >= 4 &&
        {
            let hash = xxhash32(&buf[..buf.len() - 4]);

            let mut cursor = Cursor::new(&buf[buf.len() - 4..]);
            let checksum = cursor.read_u32::<LittleEndian>().unwrap();

            let valid = hash == checksum;

            if !valid {
                warn!("Found corrupt hint file: {:?}", &path);
            }

            valid
        }
    }
}
<commit_msg>extract function to close active file. add more logging<commit_after>use std::fs;
use std::fs::File;
use std::io::prelude::*;
use std::io::{Cursor, SeekFrom, Take};
use std::marker::PhantomData;
use std::path::{Path, PathBuf};
use std::vec::Vec;

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use fs2::FileExt;
use time;
use regex::Regex;

use data::{Entry, Hint};
use util::{xxhash32, XxHash32, get_file_handle};

const DATA_FILE_EXTENSION: &'static str = "cask.data";
const HINT_FILE_EXTENSION: &'static str = "cask.hint";
const LOCK_FILE_NAME: &'static str = "cask.lock";

const DEFAULT_SIZE_THRESHOLD: usize = 100 * 1024 * 1024;

pub struct Log {
    path: PathBuf,
    sync: bool,
    size_threshold: usize,
    lock_file: File,
    current_file_id: u32,
    active_data_file_path: PathBuf,
    active_data_file: File,
    active_hint_file: File,
    active_hint_file_hasher: XxHash32,
}

pub struct Entries<'a> {
    data_file: Take<File>,
    data_file_pos: u64,
    phantom: PhantomData<&'a ()>,
}

impl<'a> Iterator for Entries<'a> {
    type Item = (u64, Entry<'a>);

    fn next(&mut self) -> Option<(u64, Entry<'a>)> {
        if self.data_file.limit() == 0 {
            None
        } else {
            let entry = Entry::from_read(&mut self.data_file);
            let entry_pos = self.data_file_pos;

            self.data_file_pos += entry.size();

            Some((entry_pos, entry))
        }
    }
}

pub struct Hints<'a> {
    hint_file: Take<File>,
    phantom: PhantomData<&'a ()>,
}

impl<'a> Iterator for Hints<'a> {
    type Item = Hint<'a>;

    fn next(&mut self) -> Option<Hint<'a>> {
        if self.hint_file.limit() == 0 {
            None
        } else {
            Some(Hint::from_read(&mut self.hint_file))
        }
    }
}

pub struct RecreateHints<'a> {
    hint_file: File,
    hint_file_hasher: XxHash32,
    entries: Entries<'a>,
}

impl<'a> Iterator for RecreateHints<'a> {
    type Item = Hint<'a>;

    fn next(&mut self) -> Option<Hint<'a>> {
        self.entries.next().map(|e| {
            let (entry_pos, entry) = e;
            let hint = Hint::from(entry, entry_pos);
            hint.write_bytes(&mut self.hint_file);
            hint.write_bytes(&mut self.hint_file_hasher);
            hint
        })
    }
}

impl<'a> Drop for RecreateHints<'a> {
    fn drop(&mut self) {
        while self.next().is_some() {}
        self.hint_file
            .write_u32::<LittleEndian>(self.hint_file_hasher.get())
            .unwrap();
    }
}

impl Log {
    pub fn open(path: &str, sync: bool) -> Log {
        let path = PathBuf::from(path);

        if path.exists() {
            assert!(path.is_dir());
        } else {
            fs::create_dir(&path).unwrap();
        }

        let lock_file = File::create(path.join(LOCK_FILE_NAME)).unwrap();
        lock_file.try_lock_exclusive().unwrap();

        let current_file_id = time::now().to_timespec().sec as u32;
        let active_data_file_path = get_data_file_path(&path, current_file_id);
        let active_data_file = get_file_handle(&active_data_file_path, true);
        let active_hint_file = get_file_handle(&get_hint_file_path(&path, current_file_id), true);
        let active_hint_file_hasher = XxHash32::new();

        info!("Created new active data file {:?}", active_data_file_path);

        Log {
            path: path,
            sync: sync,
            size_threshold: DEFAULT_SIZE_THRESHOLD,
            lock_file: lock_file,
            current_file_id: current_file_id,
            active_data_file: active_data_file,
            active_data_file_path: active_data_file_path,
            active_hint_file: active_hint_file,
            active_hint_file_hasher: active_hint_file_hasher,
        }
    }

    pub fn find_files(&self) -> Vec<u32> {
        let files = fs::read_dir(&self.path).unwrap();

        lazy_static! {
            static ref RE: Regex =
                Regex::new(&format!("(\\d+).{}$", DATA_FILE_EXTENSION)).unwrap();
        }

        let mut files: Vec<u32> = files.flat_map(|f| {
                let file = f.unwrap();
                let file_metadata = file.metadata().unwrap();

                if file_metadata.is_file() {
                    let file_name = file.file_name();
                    let captures = RE.captures(file_name.to_str().unwrap());
                    captures.and_then(|c| c.at(1).and_then(|n| n.parse::<u32>().ok()))
                        .and_then(|id| if id != self.current_file_id {
                            Some(id)
                        } else {
                            None
                        })
                } else {
                    None
                }
            })
            .collect();

        files.sort();

        files
    }

    pub fn entries<'a>(&self, file_id: u32) -> Entries<'a> {
        let data_file_path = get_data_file_path(&self.path, file_id);
        info!("Loading data file: {:?}", data_file_path);
        let data_file = get_file_handle(&data_file_path, false);
        let data_file_size = data_file.metadata().unwrap().len();

        Entries {
            data_file: data_file.take(data_file_size),
            data_file_pos: 0,
            phantom: PhantomData,
        }
    }

    pub fn hints<'a>(&self, file_id: u32) -> Option<Hints<'a>> {
        let hint_file_path = get_hint_file_path(&self.path, file_id);
        if is_valid_hint_file(&hint_file_path) {
            info!("Loading hint file: {:?}", hint_file_path);
            let hint_file = get_file_handle(&hint_file_path, false);
            let hint_file_size = hint_file.metadata().unwrap().len();

            Some(Hints {
                hint_file: hint_file.take(hint_file_size - 4),
                phantom: PhantomData,
            })
        } else {
            None
        }
    }

    pub fn recreate_hints<'a>(&mut self, file_id: u32) -> RecreateHints<'a> {
        let hint_file_path = get_hint_file_path(&self.path, file_id);
        info!("Re-creating hint file: {:?}", hint_file_path);
        let hint_file = get_file_handle(&hint_file_path, true);
        let entries = self.entries(file_id);

        RecreateHints {
            hint_file: hint_file,
            hint_file_hasher: XxHash32::new(),
            entries: entries,
        }
    }

    pub fn read_entry<'a>(&self, file_id: u32, entry_pos: u64) -> Entry<'a> {
        let mut data_file = get_file_handle(&get_data_file_path(&self.path, file_id), false);
        data_file.seek(SeekFrom::Start(entry_pos)).unwrap();
        Entry::from_read(&mut data_file)
    }

    pub fn write_entry<'a>(&mut self, entry: &Entry<'a>) -> (u32, u64) {
        let mut active_data_file_pos = self.active_data_file.seek(SeekFrom::Current(0)).unwrap();

        if active_data_file_pos + entry.size() > self.size_threshold as u64 {
            info!("Active data file {:?} reached file limit",
                  self.active_data_file_path);

            self.close_active_file();
            self.new_active_file();

            active_data_file_pos = 0;
        }

        let hint = Hint::new(&entry, active_data_file_pos);

        entry.write_bytes(&mut self.active_data_file);
        hint.write_bytes(&mut self.active_hint_file);
        hint.write_bytes(&mut self.active_hint_file_hasher);

        if self.sync {
            self.active_data_file.sync_data().unwrap();
        }

        (self.current_file_id, active_data_file_pos)
    }

    fn close_active_file(&mut self) {
        if self.sync {
            self.active_data_file.sync_data().unwrap();
        }

        self.active_hint_file
            .write_u32::<LittleEndian>(self.active_hint_file_hasher.get())
            .unwrap();

        info!("Closed active data file {:?}", self.active_data_file_path);
    }

    fn new_active_file(&mut self) {
        self.current_file_id = time::now().to_timespec().sec as u32;

        self.active_data_file_path = get_data_file_path(&self.path, self.current_file_id);
        self.active_data_file = get_file_handle(&self.active_data_file_path, true);

        self.active_hint_file =
            get_file_handle(&get_hint_file_path(&self.path, self.current_file_id), true);

        self.active_hint_file_hasher = XxHash32::new();

        info!("Created new active data file {:?}",
              self.active_data_file_path);
    }
}

impl Drop for Log {
    fn drop(&mut self) {
        self.active_hint_file
            .write_u32::<LittleEndian>(self.active_hint_file_hasher.get())
            .unwrap();

        self.lock_file.unlock().unwrap();
    }
}

fn get_data_file_path(path: &Path, file_id: u32) -> PathBuf {
    path.join(file_id.to_string()).with_extension(DATA_FILE_EXTENSION)
}

fn get_hint_file_path(path: &Path, file_id: u32) -> PathBuf {
    path.join(file_id.to_string()).with_extension(HINT_FILE_EXTENSION)
}

fn is_valid_hint_file(path: &Path) -> bool {
    path.is_file() &&
    {
        let mut hint_file = get_file_handle(path, false);

        // FIXME: avoid reading the whole hint file into memory;
        let mut buf = Vec::new();
        hint_file.read_to_end(&mut buf).unwrap();

        buf.len() >= 4 &&
        {
            let hash = xxhash32(&buf[..buf.len() - 4]);

            let mut cursor = Cursor::new(&buf[buf.len() - 4..]);
            let checksum = cursor.read_u32::<LittleEndian>().unwrap();

            let valid = hash == checksum;

            if !valid {
                warn!("Found corrupt hint file: {:?}", &path);
            }

            valid
        }
    }
}
<|endoftext|>


<commit_before>use std::io::Write;

use crate::args::Args;

struct MultiLog {
    loggers: Vec<Box<dyn log::Log>>,
}

impl MultiLog {

    fn new() -> MultiLog {
        MultiLog {
            loggers: Vec::new(),
        }
    }

    fn push<L: log::Log + 'static>(&mut self, logger: L) {
        self.loggers.push(Box::new(logger));
    }
}

impl log::Log for MultiLog {

    fn enabled(&self, metadata: &log::Metadata) -> bool {
        self.loggers.iter().any(|logger| logger.enabled(metadata))
    }

    fn log(&self, record: &log::Record) {
        for log in &self.loggers {
            log.log(record);
        }
    }

    fn flush(&self) {
        for logger in &self.loggers {
            logger.flush();
        }
    }
}

struct WriterLog<W: Write> {
    writer: std::sync::Mutex<W>,
}

impl<W: Write + Send + Sync> WriterLog<W> {

    fn new(writer: W) -> WriterLog<W> {
        WriterLog { writer: std::sync::Mutex::new(writer) }
    }
}

impl<W: Write + Send + Sync> log::Log for WriterLog<W> {

    // TODO: Add support for log filtering. For now we just output everything as
    // there is no way to change the log verbosity.

    fn enabled(&self, _metadata: &log::Metadata) -> bool {
        true
    }

    fn log(&self, record: &log::Record) {
        let now = std::time::SystemTime::now();

        // We consider failures to write to the log stream critical. Otherwise,
        // if there is some other issue with the system it is not possible to
        // properly communicate it. Thus, we panic on all write errors.
        let mut writer = self.writer.lock()
            .expect("failed to acquire log output stream lock");

        || -> Result<(), std::io::Error> {
            write! {
                writer,
                "[{level} {timestamp} ",
                level = record.level(),
                timestamp = humantime::format_rfc3339_nanos(now)
            }?;
            match record.file() {
                Some(file) => write!(writer, "{file}")?,
                None => write!(writer, "<unknown>")?,
            }
            match record.line() {
                Some(line) => write!(writer, ":{line}]")?,
                None => write!(writer, ":<unknown>]")?,
            }

            write!(writer, " {}\n", record.args())?;

            Ok(())
        }().expect("failed to write to the log output stream")
    }

    fn flush(&self) {
        // Similarly to write errors, we also panic on flush errors as there is
        // nothing we can do otherwise.
        self.writer
            .lock()
            .expect("failed to acquire log output stream lock")
            .flush()
            .expect("failed to flush the log output stream");
    }
}

pub fn init(args: &Args) {
    let mut logger = MultiLog::new();
    if args.log_to_stdout {
        logger.push(WriterLog::new(std::io::stdout()));
    }
    if let Some(ref path) = args.log_to_file {
        let file = std::fs::OpenOptions::new()
            .append(true)
            .open(path)
            .expect("failed to open the log file");

        logger.push(WriterLog::new(file));
    }

    log::set_boxed_logger(Box::new(logger))
        .expect("failed to initialize logger");

    log::set_max_level(args.verbosity);
}
<commit_msg>Add basic support for log filtering.<commit_after>use std::io::Write;

use crate::args::Args;

struct MultiLog {
    loggers: Vec<Box<dyn log::Log>>,
}

impl MultiLog {

    fn new() -> MultiLog {
        MultiLog {
            loggers: Vec::new(),
        }
    }

    fn push<L: log::Log + 'static>(&mut self, logger: L) {
        self.loggers.push(Box::new(logger));
    }
}

impl log::Log for MultiLog {

    fn enabled(&self, metadata: &log::Metadata) -> bool {
        self.loggers.iter().any(|logger| logger.enabled(metadata))
    }

    fn log(&self, record: &log::Record) {
        for log in &self.loggers {
            log.log(record);
        }
    }

    fn flush(&self) {
        for logger in &self.loggers {
            logger.flush();
        }
    }
}

struct WriterLog<W: Write> {
    writer: std::sync::Mutex<W>,
}

impl<W: Write + Send + Sync> WriterLog<W> {

    fn new(writer: W) -> WriterLog<W> {
        WriterLog { writer: std::sync::Mutex::new(writer) }
    }
}

impl<W: Write + Send + Sync> log::Log for WriterLog<W> {

    fn enabled(&self, _metadata: &log::Metadata) -> bool {
        true
    }

    fn log(&self, record: &log::Record) {
        if !self.enabled(record.metadata()) {
            return;
        }

        let now = std::time::SystemTime::now();

        // We consider failures to write to the log stream critical. Otherwise,
        // if there is some other issue with the system it is not possible to
        // properly communicate it. Thus, we panic on all write errors.
        let mut writer = self.writer.lock()
            .expect("failed to acquire log output stream lock");

        || -> Result<(), std::io::Error> {
            write! {
                writer,
                "[{level} {timestamp} ",
                level = record.level(),
                timestamp = humantime::format_rfc3339_nanos(now)
            }?;
            match record.file() {
                Some(file) => write!(writer, "{file}")?,
                None => write!(writer, "<unknown>")?,
            }
            match record.line() {
                Some(line) => write!(writer, ":{line}]")?,
                None => write!(writer, ":<unknown>]")?,
            }

            write!(writer, " {}\n", record.args())?;

            Ok(())
        }().expect("failed to write to the log output stream")
    }

    fn flush(&self) {
        // Similarly to write errors, we also panic on flush errors as there is
        // nothing we can do otherwise.
        self.writer
            .lock()
            .expect("failed to acquire log output stream lock")
            .flush()
            .expect("failed to flush the log output stream");
    }
}

pub fn init(args: &Args) {
    let mut logger = MultiLog::new();
    if args.log_to_stdout {
        logger.push(WriterLog::new(std::io::stdout()));
    }
    if let Some(ref path) = args.log_to_file {
        let file = std::fs::OpenOptions::new()
            .append(true)
            .open(path)
            .expect("failed to open the log file");

        logger.push(WriterLog::new(file));
    }

    log::set_boxed_logger(Box::new(logger))
        .expect("failed to initialize logger");

    log::set_max_level(args.verbosity);
}
<|endoftext|>