<commit_before>/*!
Provides a way to read lines from an unbuffered reader without needing to either push back or buffer bytes.
*/

use std::io::{IoError, IoResult, OtherIoError};

/**
Reads a line of input from the given `Reader`.  This does not require a push back buffer.  It returns the line *with* the line terminator.

Note that this function *does not* support old-school Mac OS newlines (i.e. a single carriage return).  If it encounters a carriage return which is *not* immediately followed by a line feed, the carriage return will be included as part of the line.
*/
pub fn read_line<R: Reader>(r: &mut R) -> IoResult<String> {
	let mut line = String::new();
	loop {
		match try!(read_utf8_char(r)) {
			'\n' => {
				line.push('\n');
				break;
			},
			c => {
				line.push(c);
			}
		}
	}
	Ok(line)
}

#[test]
fn test_read_line() {
	let s = "line one\nline two\r\nline three\n";
	let mut r = ::std::io::BufReader::new(s.as_bytes());
	let oks = |s:&str| Ok(s.into_string());

	assert_eq!(read_line(&mut r), oks("line one\n"));
	assert_eq!(read_line(&mut r), oks("line two\r\n"));
	assert_eq!(read_line(&mut r), oks("line three\n"));
}

/**
Reads a single UTF-8 encoded Unicode code point from a `Reader`.
*/
pub fn read_utf8_char<R: Reader>(r: &mut R) -> IoResult<char> {
	fn invalid_utf8<T>(b: u8, initial: bool) -> IoResult<T> {
		Err(IoError {
			kind: OtherIoError,
			desc: "invalid utf-8 sequence",
			detail: if initial {
				Some(format!("invalid initial code unit {:#02x}", b))
			} else {
				Some(format!("invalid continuation code unit {:#02x}", b))
			}
		})
	}

	fn invalid_cp<T>(cp: u32) -> IoResult<T> {
		Err(IoError {
			kind: OtherIoError,
			desc: "invalid Unicode code point",
			detail: Some(format!("invalid code point {:#08x}", cp))
		})
	}

	// Why not use std::str::utf8_char_width?  We need to know the encoding to mask away the size bits anyway.
	let (mut cp, n) = match try!(r.read_u8()) {
		b @ 0b0000_0000 ... 0b0111_1111 => (b as u32, 0),
		b @ 0b1100_0000 ... 0b1101_1111 => ((b & 0b0001_1111) as u32, 1),
		b @ 0b1110_0000 ... 0b1110_1111 => ((b & 0b0000_1111) as u32, 2),
		b @ 0b1111_0000 ... 0b1111_0111 => ((b & 0b0000_0111) as u32, 3),
		b @ 0b1111_1000 ... 0b1111_1011 => ((b & 0b0000_0011) as u32, 4),
		b @ 0b1111_1100 ... 0b1111_1101 => ((b & 0b0000_0001) as u32, 5),
		b => return invalid_utf8(b, true)
	};

	for _ in range(0u, n) {
		let b = match try!(r.read_u8()) {
			b @ 0b10_000000 ... 0b10_111111 => (b & 0b00_111111) as u32,
			b => return invalid_utf8(b, false)
		};
		cp = (cp << 6) | b;
	}

	::std::char::from_u32(cp)
		.map(|c| Ok(c))
		.unwrap_or_else(|| invalid_cp(cp))
}

#[test]
fn test_read_utf8_char() {
	fn test_str(s: &str) {
		let mut reader = ::std::io::BufReader::new(s.as_bytes());
		for c in s.chars() {
			assert_eq!(Ok(c), read_utf8_char(&mut reader))
		}
	}

	fn first(s: &[u8]) -> IoResult<char> {
		let mut reader = ::std::io::BufReader::new(s);
		read_utf8_char(&mut reader)
	}

	test_str("abcdef");
	test_str("私の日本語わ下手ですよ！");

	assert!(first([0b1000_0000u8]).is_err());
	assert!(first([0b1100_0000u8, 0b0000_0000]).is_err());
}

/**
Reads a single line from standard input.
*/
pub fn stdin_read_line() -> IoResult<String> {
	read_line(&mut ::std::io::stdio::stdin_raw())
}
<commit_msg>Changed io::read_line to *not* error on incomplete lines.<commit_after>/*!
Provides a way to read lines from an unbuffered reader without needing to either push back or buffer bytes.
*/

use std::io::{IoError, IoResult, OtherIoError};

/**
Reads a line of input from the given `Reader`.  This does not require a push back buffer.  It returns the line *with* the line terminator.

Note that this function *does not* support old-school Mac OS newlines (i.e. a single carriage return).  If it encounters a carriage return which is *not* immediately followed by a line feed, the carriage return will be included as part of the line.
*/
pub fn read_line<R: Reader>(r: &mut R) -> IoResult<String> {
	let mut line = String::new();
	loop {
		match read_utf8_char(r) {
			Ok('\n') => {
				line.push('\n');
				break;
			},
			Ok(c) => {
				line.push(c);
			}
			Err(err) => {
				if err.kind == ::std::io::EndOfFile && line.len() > 0 {
					break
				} else {
					return Err(err)
				}
			}
		}
	}
	Ok(line)
}

#[test]
fn test_read_line() {
	let s = "line one\nline two\r\nline three\n";
	let mut r = ::std::io::BufReader::new(s.as_bytes());
	let oks = |s:&str| Ok(s.into_string());

	assert_eq!(read_line(&mut r), oks("line one\n"));
	assert_eq!(read_line(&mut r), oks("line two\r\n"));
	assert_eq!(read_line(&mut r), oks("line three\n"));
}

/**
Reads a single UTF-8 encoded Unicode code point from a `Reader`.
*/
pub fn read_utf8_char<R: Reader>(r: &mut R) -> IoResult<char> {
	fn invalid_utf8<T>(b: u8, initial: bool) -> IoResult<T> {
		Err(IoError {
			kind: OtherIoError,
			desc: "invalid utf-8 sequence",
			detail: if initial {
				Some(format!("invalid initial code unit {:#02x}", b))
			} else {
				Some(format!("invalid continuation code unit {:#02x}", b))
			}
		})
	}

	fn invalid_cp<T>(cp: u32) -> IoResult<T> {
		Err(IoError {
			kind: OtherIoError,
			desc: "invalid Unicode code point",
			detail: Some(format!("invalid code point {:#08x}", cp))
		})
	}

	// Why not use std::str::utf8_char_width?  We need to know the encoding to mask away the size bits anyway.
	let (mut cp, n) = match try!(r.read_u8()) {
		b @ 0b0000_0000 ... 0b0111_1111 => (b as u32, 0),
		b @ 0b1100_0000 ... 0b1101_1111 => ((b & 0b0001_1111) as u32, 1),
		b @ 0b1110_0000 ... 0b1110_1111 => ((b & 0b0000_1111) as u32, 2),
		b @ 0b1111_0000 ... 0b1111_0111 => ((b & 0b0000_0111) as u32, 3),
		b @ 0b1111_1000 ... 0b1111_1011 => ((b & 0b0000_0011) as u32, 4),
		b @ 0b1111_1100 ... 0b1111_1101 => ((b & 0b0000_0001) as u32, 5),
		b => return invalid_utf8(b, true)
	};

	for _ in range(0u, n) {
		let b = match try!(r.read_u8()) {
			b @ 0b10_000000 ... 0b10_111111 => (b & 0b00_111111) as u32,
			b => return invalid_utf8(b, false)
		};
		cp = (cp << 6) | b;
	}

	::std::char::from_u32(cp)
		.map(|c| Ok(c))
		.unwrap_or_else(|| invalid_cp(cp))
}

#[test]
fn test_read_utf8_char() {
	fn test_str(s: &str) {
		let mut reader = ::std::io::BufReader::new(s.as_bytes());
		for c in s.chars() {
			assert_eq!(Ok(c), read_utf8_char(&mut reader))
		}
	}

	fn first(s: &[u8]) -> IoResult<char> {
		let mut reader = ::std::io::BufReader::new(s);
		read_utf8_char(&mut reader)
	}

	test_str("abcdef");
	test_str("私の日本語わ下手ですよ！");

	assert!(first([0b1000_0000u8]).is_err());
	assert!(first([0b1100_0000u8, 0b0000_0000]).is_err());
}

/**
Reads a single line from standard input.
*/
pub fn stdin_read_line() -> IoResult<String> {
	read_line(&mut ::std::io::stdio::stdin_raw())
}
<|endoftext|>


<commit_before>extern crate sodiumoxide;
// extern crate "rustc-serialize" as rustc_serialize;

use std::collections::HashMap;
use std::vec;
use sodiumoxide::crypto::hash::sha512;
use sodiumoxide::crypto::sign::ed25519::{
    self, PublicKey, SecretKey, PUBLICKEYBYTES, SECRETKEYBYTES};
use protobuf::Message;
use simples_pb;

fn slice_to_pk(bytes: &[u8]) -> Option<PublicKey> {
    if bytes.len() != PUBLICKEYBYTES { return None; }
    let mut key:[u8; PUBLICKEYBYTES] = [0; PUBLICKEYBYTES];
    for i in range(0, PUBLICKEYBYTES) {
        key[i] = bytes[i];
    }
    Some(PublicKey(key))
}

fn hash_message<M: Message>(obj: &M) -> Vec<u8> {
    let bytes = &obj.write_to_bytes().unwrap()[];
    let sha512::Digest(commit_hash) = sha512::hash(&bytes[]);
    vec::as_vec(&commit_hash[]).clone()
}

pub trait Transaction {
    fn check_signatures(&self) -> Result<(), &'static str>;
}

impl Transaction for simples_pb::Transaction {
    fn check_signatures(&self) -> Result<(), &'static str> {
        let commit_hash = hash_message(self.get_commit());
        if commit_hash[] != self.get_commit_hash()[] {
            return Err("Invalid hash.");
        }
        let mut sign_map = HashMap::<&[u8], &[u8]>::new();
        for sign in self.get_signatures().iter() {
            sign_map.insert(sign.get_public_key(), sign.get_payload());
        }
        for transfer in self.get_commit().get_transfers().iter() {
            match sign_map.get(transfer.get_source_pk()) {
                Some(signature) => {
                    let pk = slice_to_pk(transfer.get_source_pk()).unwrap();
                    match ed25519::verify(&signature[], &pk) {
                        Some(hash) => {
                            if commit_hash[] != hash[] {
                                return Err("The incorrect hash is signed.")
                            }
                        },
                        None => return Err("Invalid signature.")
                    }
                },
                None => return Err("Missing key.")
            }
        }
        Ok(())
    }
}

#[derive(Default)]
pub struct TransactionBuilder {
    transfer_secret_keys: Vec<SecretKey>,
    bounty_secret_key: Option<SecretKey>,
    commit: simples_pb::Commitment
}

impl TransactionBuilder {
    pub fn new() -> TransactionBuilder {
        TransactionBuilder {
            transfer_secret_keys: Vec::<SecretKey>::new(),
            bounty_secret_key: None,
            commit: simples_pb::Commitment::new()
        }
    }

    pub fn add_transfer(
        mut self, sk: &SecretKey, source: &PublicKey, destination: &PublicKey,
        tokens: u64, op_index:u32) -> TransactionBuilder {
        let mut transfer = simples_pb::Transfer::new();
        transfer.set_op_index(op_index);
        transfer.set_tokens(tokens);
        transfer.mut_source_pk().push_all(&source.0);
        transfer.mut_destination_pk().push_all(&destination.0);

        self.transfer_secret_keys.push(sk.clone());
        self.commit.mut_transfers().push(transfer);
        self
    }

    pub fn set_bounty(mut self, sk: &SecretKey, source: &PublicKey,
                      bounty: u64) -> TransactionBuilder {
        self.bounty_secret_key = Some(sk.clone());
        self.commit.mut_bounty_pk().push_all(&source.0);
        self.commit.set_bounty(bounty);
        self
    }

    pub fn build(self) -> Result<simples_pb::Transaction, &'static str> {
        let commit_hash = hash_message(&self.commit);
        let mut transaction = simples_pb::Transaction::new();
        {
            let mut signatures = transaction.mut_signatures();
            let mut transfers_and_skeys = self.commit.get_transfers().iter()
                .zip(self.transfer_secret_keys.iter());

            for (transfer, sk) in transfers_and_skeys {
                let signature: Vec<u8> = ed25519::sign(&commit_hash[], sk);
                let pk_bytes = vec::as_vec(transfer.get_source_pk()).clone();
                let pk = slice_to_pk(&pk_bytes[]).unwrap();
                match ed25519::verify(&signature[0..], &pk) {
                    Some(_) => {
                        let mut sign = simples_pb::DetachedSignature::new();
                        sign.set_public_key(pk_bytes);
                        sign.set_payload(signature);
                        signatures.push(sign);
                    },
                    None => return Err("Invalid key for source account.")
                }
            }
        }
        transaction.set_commit(self.commit);
        transaction.set_commit_hash(commit_hash.clone());
        match transaction.check_signatures() {
            Ok(_) => Ok(transaction),
            Err(msg) => Err(msg)
        }
    }
}
<commit_msg>Added slice_to_sk utility function and dep cleanup.<commit_after>extern crate sodiumoxide;

use std::collections::HashMap;
use std::vec;
use sodiumoxide::crypto::hash::sha512;
use sodiumoxide::crypto::sign::ed25519::{
    self, PublicKey, SecretKey, PUBLICKEYBYTES, SECRETKEYBYTES};
use protobuf::Message;
use simples_pb;

pub fn slice_to_pk(bytes: &[u8]) -> Option<PublicKey> {
    if bytes.len() != PUBLICKEYBYTES { return None; }
    let mut key:[u8; PUBLICKEYBYTES] = [0; PUBLICKEYBYTES];
    for i in range(0, PUBLICKEYBYTES) {
        key[i] = bytes[i];
    }
    Some(PublicKey(key))
}

pub fn slice_to_sk(bytes: &[u8]) -> Option<SecretKey> {
    if bytes.len() != SECRETKEYBYTES { return None; }
    let mut key:[u8; SECRETKEYBYTES] = [0; SECRETKEYBYTES];
    for i in range(0, SECRETKEYBYTES) {
        key[i] = bytes[i];
    }
    Some(SecretKey(key))
}

fn hash_message<M: Message>(obj: &M) -> Vec<u8> {
    let bytes = &obj.write_to_bytes().unwrap()[];
    let sha512::Digest(commit_hash) = sha512::hash(&bytes[]);
    vec::as_vec(&commit_hash[]).clone()
}

pub trait Transaction {
    fn check_signatures(&self) -> Result<(), &'static str>;
}

impl Transaction for simples_pb::Transaction {
    fn check_signatures(&self) -> Result<(), &'static str> {
        let commit_hash = hash_message(self.get_commit());
        if commit_hash[] != self.get_commit_hash()[] {
            return Err("Invalid hash.");
        }
        let mut sign_map = HashMap::<&[u8], &[u8]>::new();
        for sign in self.get_signatures().iter() {
            sign_map.insert(sign.get_public_key(), sign.get_payload());
        }
        for transfer in self.get_commit().get_transfers().iter() {
            match sign_map.get(transfer.get_source_pk()) {
                Some(signature) => {
                    let pk = slice_to_pk(transfer.get_source_pk()).unwrap();
                    match ed25519::verify(&signature[], &pk) {
                        Some(hash) => {
                            if commit_hash[] != hash[] {
                                return Err("The incorrect hash is signed.")
                            }
                        },
                        None => return Err("Invalid signature.")
                    }
                },
                None => return Err("Missing key.")
            }
        }
        Ok(())
    }
}

#[derive(Default)]
pub struct TransactionBuilder {
    transfer_secret_keys: Vec<SecretKey>,
    bounty_secret_key: Option<SecretKey>,
    commit: simples_pb::Commitment
}

impl TransactionBuilder {
    pub fn new() -> TransactionBuilder {
        TransactionBuilder {
            transfer_secret_keys: Vec::<SecretKey>::new(),
            bounty_secret_key: None,
            commit: simples_pb::Commitment::new()
        }
    }

    pub fn add_transfer(
        mut self, sk: &SecretKey, source: &PublicKey, destination: &PublicKey,
        tokens: u64, op_index:u32) -> TransactionBuilder {
        let mut transfer = simples_pb::Transfer::new();
        transfer.set_op_index(op_index);
        transfer.set_tokens(tokens);
        transfer.mut_source_pk().push_all(&source.0);
        transfer.mut_destination_pk().push_all(&destination.0);

        self.transfer_secret_keys.push(sk.clone());
        self.commit.mut_transfers().push(transfer);
        self
    }

    pub fn set_bounty(mut self, sk: &SecretKey, source: &PublicKey,
                      bounty: u64) -> TransactionBuilder {
        self.bounty_secret_key = Some(sk.clone());
        self.commit.mut_bounty_pk().push_all(&source.0);
        self.commit.set_bounty(bounty);
        self
    }

    pub fn build(self) -> Result<simples_pb::Transaction, &'static str> {
        let commit_hash = hash_message(&self.commit);
        let mut transaction = simples_pb::Transaction::new();
        {
            let mut signatures = transaction.mut_signatures();
            let mut transfers_and_skeys = self.commit.get_transfers().iter()
                .zip(self.transfer_secret_keys.iter());

            for (transfer, sk) in transfers_and_skeys {
                let signature: Vec<u8> = ed25519::sign(&commit_hash[], sk);
                let pk_bytes = vec::as_vec(transfer.get_source_pk()).clone();
                let pk = slice_to_pk(&pk_bytes[]).unwrap();
                match ed25519::verify(&signature[0..], &pk) {
                    Some(_) => {
                        let mut sign = simples_pb::DetachedSignature::new();
                        sign.set_public_key(pk_bytes);
                        sign.set_payload(signature);
                        signatures.push(sign);
                    },
                    None => return Err("Invalid key for source account.")
                }
            }
        }
        transaction.set_commit(self.commit);
        transaction.set_commit_hash(commit_hash.clone());
        match transaction.check_signatures() {
            Ok(_) => Ok(transaction),
            Err(msg) => Err(msg)
        }
    }
}
<|endoftext|>


<commit_before>//! Virtual machine.

use std::rc::Rc;
use std::{error, result};
use std::fmt;
use std::collections::HashMap;
use std::error::Error as StdError;
use num::integer::Integer;
use num::{zero, one};
use item::{Block, BlockItem, Stack, StackItem};

pub type Result<T> = result::Result<T, Error>;

#[derive(PartialEq, Eq, Clone, Debug)]
pub enum Error {
    TypeError,
    DivideByZero,
    StackUnderflow,
    UnknownMethod(String),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::UnknownMethod(ref s) => write!(f, "{}: {}", self.description(), s),
            _ => write!(f, "{}", self.description()),
        }
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::DivideByZero => "Divided by zero",
            Error::TypeError => "Type error",
            Error::StackUnderflow => "Stack underflow",
            Error::UnknownMethod(_) => "Unknown method",
        }
    }
}

pub enum Method<I> {
    Builtin(Box<Fn(&mut Vm<I>) -> Result<()>>),
    Block(Block<I>),
}

pub struct Vm<I> {
    pub stack: Stack<I>,
    pub methods: HashMap<String, Rc<Method<I>>>,
}

impl<I> Vm<I> where I: Integer + Clone {
    pub fn run(&mut self, item: &BlockItem<I>) -> Result<()> {
        match *item {
            BlockItem::Literal(ref stack_item) =>
                self.stack.push(stack_item.clone()),
            BlockItem::Call(ref name) => {
                let method = match self.methods.get(&*name) {
                    Some(m) => m.clone(),
                    None => return Err(Error::UnknownMethod(name.clone())),
                };
                try!(match *method {
                    Method::Builtin(ref f) => (**f)(self),
                    Method::Block(ref b) => self.run_block(b),
                })
            },
        }
        Ok(())
    }

    pub fn run_block(&mut self, block: &Block<I>) -> Result<()> {
        for item in block.iter() {
            try!(self.run(item));
        }
        Ok(())
    }

    pub fn builtin<S>(&mut self, name: S, method: Box<Fn(&mut Vm<I>)
                   -> Result<()>>) where S: Into<String> {
        self.methods.insert(name.into(), Rc::new(Method::Builtin(method)));
    }

    pub fn pop_integer(&mut self) -> Result<I> {
        match self.stack.pop() {
            Some(StackItem::Integer(i)) => Ok(i),
            Some(..) => Err(Error::TypeError),
            None => Err(Error::StackUnderflow),
        }
    }

    pub fn new() -> Vm<I> {
        Vm {
            stack: Vec::new(),
            methods: HashMap::new(),
        }
    }

    pub fn new_with_builtins() -> Vm<I> {
        let mut vm = Vm::<I>::new();
        vm.builtin("+", Box::new(|vm| {
            let n2 = try!(vm.pop_integer());
            let n1 = try!(vm.pop_integer());
            vm.stack.push(StackItem::Integer(n1 + n2));
            Ok(())
        }));
        vm.builtin("-", Box::new(|vm| {
            let n2 = try!(vm.pop_integer());
            let n1 = try!(vm.pop_integer());
            vm.stack.push(StackItem::Integer(n1 - n2));
            Ok(())
        }));
        vm.builtin("*", Box::new(|vm| {
            let n2 = try!(vm.pop_integer());
            let n1 = try!(vm.pop_integer());
            vm.stack.push(StackItem::Integer(n1 * n2));
            Ok(())
        }));
        vm.builtin("/", Box::new(|vm| {
            let n2 = try!(vm.pop_integer());
            let n1 = try!(vm.pop_integer());
            if n2 == zero() {
                return Err(Error::DivideByZero);
            }
            vm.stack.push(StackItem::Integer(n1 / n2));
            Ok(())
        }));
        vm.builtin("fn", Box::new(|vm| {
            let block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let name = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            match (name, block) {
                (StackItem::Symbol(s), StackItem::Block(b)) =>
                    { vm.methods.insert(s, Rc::new(Method::Block(b))); },
                _ => return Err(Error::TypeError),
            }
            Ok(())
        }));
        vm.builtin("swap", Box::new(|vm| {
            let b = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            vm.stack.push(b);
            vm.stack.push(a);
            Ok(())
        }));
        vm.builtin("over", Box::new(|vm| {
            let b = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            vm.stack.push(a.clone());
            vm.stack.push(b);
            vm.stack.push(a);
            Ok(())
        }));
        vm.builtin("rot", Box::new(|vm| {
            let c = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let b = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            vm.stack.push(b);
            vm.stack.push(c);
            vm.stack.push(a);
            Ok(())
        }));
        vm.builtin("dup", Box::new(|vm| {
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            vm.stack.push(a.clone());
            vm.stack.push(a);
            Ok(())
        }));
        vm.builtin("pop", Box::new(|vm| {
            let _ = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            Ok(())
        }));
        vm.builtin("false", Box::new(|vm| {
            vm.stack.push(StackItem::Boolean(false));
            Ok(())
        }));
        vm.builtin("true", Box::new(|vm| {
            vm.stack.push(StackItem::Boolean(true));
            Ok(())
        }));
        vm.builtin("eq", Box::new(|vm| {
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let b = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            vm.stack.push(StackItem::Boolean(a == b));
            Ok(())
        }));
        vm.builtin("not", Box::new(|vm| {
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let StackItem::Boolean(boolean) = a {
                vm.stack.push(StackItem::Boolean(!boolean));
            } else {
                return Err(Error::TypeError)
            }
            Ok(())
        }));
        vm.builtin("if", Box::new(|vm| {
            let block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let condition = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let (StackItem::Block(block), StackItem::Boolean(condition)) =
                    (block, condition) {
                if condition {
                    try!(vm.run_block(&block));
                }
            } else {
                return Err(Error::TypeError);
            }
            Ok(())
        }));
        vm.builtin("ifelse", Box::new(|vm| {
            let else_block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let if_block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let condition = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let (StackItem::Block(else_block), StackItem::Block(if_block), StackItem::Boolean(condition)) =
                    (else_block, if_block, condition) {
                if condition {
                    try!(vm.run_block(&if_block));
                } else {
                    try!(vm.run_block(&else_block));
                }
            } else {
                return Err(Error::TypeError);
            }
            Ok(())
        }));
        vm.builtin("or", Box::new(|vm| {
            let b = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let (StackItem::Boolean(a), StackItem::Boolean(b)) =
                    (a, b) {
                vm.stack.push(StackItem::Boolean(a || b));
            } else {
                return Err(Error::TypeError);
            }
            Ok(())
        }));
        vm.builtin("while", Box::new(|vm| {
            let action_block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let condition_block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let (StackItem::Block(action_block), StackItem::Block(condition_block)) =
                    (action_block, condition_block) {
                loop {
                    try!(vm.run_block(&condition_block));
                    let condition = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
                    if let StackItem::Boolean(condition) = condition {
                        if condition {
                            try!(vm.run_block(&action_block));
                        } else {
                            break;
                        }
                    } else {
                        return Err(Error::TypeError);
                    }
                }
            } else {
                return Err(Error::TypeError);
            }
            Ok(())
        }));
        vm.builtin("times", Box::new(|vm| {
            let block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let times = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let (StackItem::Block(block), StackItem::Integer(mut times)) =
                    (block, times) {
                while times > zero() {
                    try!(vm.run_block(&block));
                    times = times - one::<I>();
                }
            } else {
                return Err(Error::TypeError);
            }
            Ok(())
        }));
        vm
    }
}

impl<I> fmt::Display for Vm<I> where I: fmt::Display {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for item in &self.stack {
            try!(write!(f, "{} ", item));
        }
        Ok(())
    }
}
<commit_msg>Add cat builtin<commit_after>//! Virtual machine.

use std::rc::Rc;
use std::{error, result};
use std::fmt;
use std::collections::HashMap;
use std::error::Error as StdError;
use num::integer::Integer;
use num::{zero, one};
use item::{Block, BlockItem, Stack, StackItem};

pub type Result<T> = result::Result<T, Error>;

#[derive(PartialEq, Eq, Clone, Debug)]
pub enum Error {
    TypeError,
    DivideByZero,
    StackUnderflow,
    UnknownMethod(String),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::UnknownMethod(ref s) => write!(f, "{}: {}", self.description(), s),
            _ => write!(f, "{}", self.description()),
        }
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::DivideByZero => "Divided by zero",
            Error::TypeError => "Type error",
            Error::StackUnderflow => "Stack underflow",
            Error::UnknownMethod(_) => "Unknown method",
        }
    }
}

pub enum Method<I> {
    Builtin(Box<Fn(&mut Vm<I>) -> Result<()>>),
    Block(Block<I>),
}

pub struct Vm<I> {
    pub stack: Stack<I>,
    pub methods: HashMap<String, Rc<Method<I>>>,
}

impl<I> Vm<I> where I: Integer + Clone {
    pub fn run(&mut self, item: &BlockItem<I>) -> Result<()> {
        match *item {
            BlockItem::Literal(ref stack_item) =>
                self.stack.push(stack_item.clone()),
            BlockItem::Call(ref name) => {
                let method = match self.methods.get(&*name) {
                    Some(m) => m.clone(),
                    None => return Err(Error::UnknownMethod(name.clone())),
                };
                try!(match *method {
                    Method::Builtin(ref f) => (**f)(self),
                    Method::Block(ref b) => self.run_block(b),
                })
            },
        }
        Ok(())
    }

    pub fn run_block(&mut self, block: &Block<I>) -> Result<()> {
        for item in block.iter() {
            try!(self.run(item));
        }
        Ok(())
    }

    pub fn builtin<S>(&mut self, name: S, method: Box<Fn(&mut Vm<I>)
                   -> Result<()>>) where S: Into<String> {
        self.methods.insert(name.into(), Rc::new(Method::Builtin(method)));
    }

    pub fn pop_integer(&mut self) -> Result<I> {
        match self.stack.pop() {
            Some(StackItem::Integer(i)) => Ok(i),
            Some(..) => Err(Error::TypeError),
            None => Err(Error::StackUnderflow),
        }
    }

    pub fn new() -> Vm<I> {
        Vm {
            stack: Vec::new(),
            methods: HashMap::new(),
        }
    }

    pub fn new_with_builtins() -> Vm<I> {
        let mut vm = Vm::<I>::new();
        vm.builtin("+", Box::new(|vm| {
            let n2 = try!(vm.pop_integer());
            let n1 = try!(vm.pop_integer());
            vm.stack.push(StackItem::Integer(n1 + n2));
            Ok(())
        }));
        vm.builtin("-", Box::new(|vm| {
            let n2 = try!(vm.pop_integer());
            let n1 = try!(vm.pop_integer());
            vm.stack.push(StackItem::Integer(n1 - n2));
            Ok(())
        }));
        vm.builtin("*", Box::new(|vm| {
            let n2 = try!(vm.pop_integer());
            let n1 = try!(vm.pop_integer());
            vm.stack.push(StackItem::Integer(n1 * n2));
            Ok(())
        }));
        vm.builtin("/", Box::new(|vm| {
            let n2 = try!(vm.pop_integer());
            let n1 = try!(vm.pop_integer());
            if n2 == zero() {
                return Err(Error::DivideByZero);
            }
            vm.stack.push(StackItem::Integer(n1 / n2));
            Ok(())
        }));
        vm.builtin("fn", Box::new(|vm| {
            let block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let name = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            match (name, block) {
                (StackItem::Symbol(s), StackItem::Block(b)) =>
                    { vm.methods.insert(s, Rc::new(Method::Block(b))); },
                _ => return Err(Error::TypeError),
            }
            Ok(())
        }));
        vm.builtin("swap", Box::new(|vm| {
            let b = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            vm.stack.push(b);
            vm.stack.push(a);
            Ok(())
        }));
        vm.builtin("over", Box::new(|vm| {
            let b = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            vm.stack.push(a.clone());
            vm.stack.push(b);
            vm.stack.push(a);
            Ok(())
        }));
        vm.builtin("rot", Box::new(|vm| {
            let c = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let b = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            vm.stack.push(b);
            vm.stack.push(c);
            vm.stack.push(a);
            Ok(())
        }));
        vm.builtin("dup", Box::new(|vm| {
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            vm.stack.push(a.clone());
            vm.stack.push(a);
            Ok(())
        }));
        vm.builtin("pop", Box::new(|vm| {
            let _ = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            Ok(())
        }));
        vm.builtin("false", Box::new(|vm| {
            vm.stack.push(StackItem::Boolean(false));
            Ok(())
        }));
        vm.builtin("true", Box::new(|vm| {
            vm.stack.push(StackItem::Boolean(true));
            Ok(())
        }));
        vm.builtin("eq", Box::new(|vm| {
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let b = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            vm.stack.push(StackItem::Boolean(a == b));
            Ok(())
        }));
        vm.builtin("not", Box::new(|vm| {
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let StackItem::Boolean(boolean) = a {
                vm.stack.push(StackItem::Boolean(!boolean));
            } else {
                return Err(Error::TypeError)
            }
            Ok(())
        }));
        vm.builtin("if", Box::new(|vm| {
            let block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let condition = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let (StackItem::Block(block), StackItem::Boolean(condition)) =
                    (block, condition) {
                if condition {
                    try!(vm.run_block(&block));
                }
            } else {
                return Err(Error::TypeError);
            }
            Ok(())
        }));
        vm.builtin("ifelse", Box::new(|vm| {
            let else_block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let if_block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let condition = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let (StackItem::Block(else_block), StackItem::Block(if_block), StackItem::Boolean(condition)) =
                    (else_block, if_block, condition) {
                if condition {
                    try!(vm.run_block(&if_block));
                } else {
                    try!(vm.run_block(&else_block));
                }
            } else {
                return Err(Error::TypeError);
            }
            Ok(())
        }));
        vm.builtin("or", Box::new(|vm| {
            let b = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let (StackItem::Boolean(a), StackItem::Boolean(b)) =
                    (a, b) {
                vm.stack.push(StackItem::Boolean(a || b));
            } else {
                return Err(Error::TypeError);
            }
            Ok(())
        }));
        vm.builtin("while", Box::new(|vm| {
            let action_block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let condition_block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let (StackItem::Block(action_block), StackItem::Block(condition_block)) =
                    (action_block, condition_block) {
                loop {
                    try!(vm.run_block(&condition_block));
                    let condition = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
                    if let StackItem::Boolean(condition) = condition {
                        if condition {
                            try!(vm.run_block(&action_block));
                        } else {
                            break;
                        }
                    } else {
                        return Err(Error::TypeError);
                    }
                }
            } else {
                return Err(Error::TypeError);
            }
            Ok(())
        }));
        vm.builtin("times", Box::new(|vm| {
            let block = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let times = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let (StackItem::Block(block), StackItem::Integer(mut times)) =
                    (block, times) {
                while times > zero() {
                    try!(vm.run_block(&block));
                    times = times - one::<I>();
                }
            } else {
                return Err(Error::TypeError);
            }
            Ok(())
        }));
        vm.builtin("cat", Box::new(|vm| {
            let b = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            let a = try!(vm.stack.pop().ok_or(Error::StackUnderflow));
            if let (StackItem::String(b), StackItem::String(mut a)) =
                    (b, a) {
                a.push_str(&b);
                vm.stack.push(StackItem::String(a));
            }
            Ok(())
        }));
        vm
    }
}

impl<I> fmt::Display for Vm<I> where I: fmt::Display {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for item in &self.stack {
            try!(write!(f, "{} ", item));
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::ui::{UIWriter,
                UI};
use habitat_api_client as api_client;
use habitat_core as hcore;
use lazy_static::lazy_static;
use std::{collections::HashMap,
          env,
          ffi::OsStr,
          iter::FromIterator};

extern crate json;
#[macro_use]
extern crate log;
#[macro_use]
extern crate lazy_static;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[cfg_attr(test, macro_use)]
extern crate serde_json;
#[cfg(windows)]
extern crate winapi;

pub use self::error::{Error,
                      Result};

pub mod cli;
pub mod command;
pub mod error;
pub mod locked_env_var;
pub mod output;
pub mod package_graph;
pub mod templating;
pub mod types;
pub mod ui;
pub mod util;

lazy_static::lazy_static! {
    pub static ref PROGRAM_NAME: String = {
        match env::current_exe() {
            Ok(path) => path.file_stem().and_then(OsStr::to_str).unwrap().to_string(),
            Err(e) => {
                error!("Error getting path of current_exe: {}", e);
                String::from("hab-?")
            }
        }
    };
}

// TODO (CM): It would be nice to come up with a way to more
// programmatically manage these flags. It's a bit of a pain to define
// the flag, and then define the environment variables
// separately. Nothing statically guarantees that you've specified an
// variable for a flag.

// TODO (CM): It'd be great to have a built-in way to document them,
// too.

// TODO (CM): Part of that documentation might be *when* a flag was
// added. In general, long-lived flags are a code-smell.

// TODO (CM): It may also be useful to break out features by area of
// concern. We can have any number of bitflags-generated structs.

bitflags::bitflags! {
    /// All the feature flags that are recogized by Habitat.
    ///
    /// In general, feature flags are enabled by setting the corresponding
    /// environment variable.
    ///
    /// Your binary should call `FeatureFlag::from_env` to get a set
    /// of flags to use.
    ///
    /// To add a new feature flag, you will need to add the bit mask
    /// constant here, as well as a mapping from the feature to the
    /// environment variable to which it corresponds in the `ENV_VARS`
    /// map below.
    pub struct FeatureFlag: u32 {
        const LIST               = 0b0000_0000_0001;
        const TEST_EXIT          = 0b0000_0000_0010;
        const TEST_BOOT_FAIL     = 0b0000_0000_0100;
        const REDACT_HTTP        = 0b0000_0000_1000;
        const IGNORE_SIGNALS     = 0b0000_0001_0000;
        const OFFLINE_INSTALL    = 0b0000_0100_0000;
        const IGNORE_LOCAL       = 0b0000_1000_0000;
        const EVENT_STREAM       = 0b0001_0000_0000;
        const TRIGGER_ELECTION   = 0b0010_0000_0000;
        const CONFIGURE_SHUTDOWN = 0b0100_0000_0000;
    }
}

lazy_static! {
    static ref ENV_VARS: HashMap<FeatureFlag, &'static str> = {
        let mapping = vec![(FeatureFlag::LIST, "HAB_FEAT_LIST"),
                           (FeatureFlag::TEST_EXIT, "HAB_FEAT_TEST_EXIT"),
                           (FeatureFlag::TEST_BOOT_FAIL, "HAB_FEAT_BOOT_FAIL"),
                           (FeatureFlag::REDACT_HTTP, "HAB_FEAT_REDACT_HTTP"),
                           (FeatureFlag::IGNORE_SIGNALS, "HAB_FEAT_IGNORE_SIGNALS"),
                           (FeatureFlag::OFFLINE_INSTALL, "HAB_FEAT_OFFLINE_INSTALL"),
                           (FeatureFlag::IGNORE_LOCAL, "HAB_FEAT_IGNORE_LOCAL"),
                           (FeatureFlag::EVENT_STREAM, "HAB_FEAT_EVENT_STREAM"),
                           (FeatureFlag::TRIGGER_ELECTION, "HAB_FEAT_TRIGGER_ELECTION"),
                           (FeatureFlag::CONFIGURE_SHUTDOWN, "HAB_FEAT_CONFIGURE_SHUTDOWN")];
        HashMap::from_iter(mapping)
    };
}

impl FeatureFlag {
    /// If the environment variable for a flag is set to _anything_ but
    /// the empty string, it is activated.
    pub fn from_env(ui: &mut UI) -> Self {
        let mut flags = FeatureFlag::empty();

        for (feature, env_var) in ENV_VARS.iter() {
            if let Some(val) = env::var_os(env_var) {
                if !val.is_empty() {
                    flags.insert(*feature);
                    ui.warn(&format!("Enabling feature: {:?}", feature))
                      .unwrap();
                }
            }
        }

        // TODO (CM): Once the other TODOs above are done (especially the
        // documentation bits), it would be nice to extract this logic
        // into an actual discoverable CLI subcommand; it's a little weird
        // that you have to know how to enable a feature flag before you
        // can even find out that there *are* feature flags to enable.
        //
        // There's no reason why "list feature flags" should itself be a
        // feature-flag.
        if flags.contains(FeatureFlag::LIST) {
            ui.warn("Listing feature flags environment variables:")
              .unwrap();
            for (feature, env_var) in ENV_VARS.iter() {
                ui.warn(&format!("  * {:?}: {}={:?}",
                                 feature,
                                 env_var,
                                 env::var_os(env_var).unwrap_or_default()))
                  .unwrap();
            }
        }

        flags
    }
}

pub mod sync {

    #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
    type InnerLock<T> = parking_lot::RwLock<T>;
    #[cfg(feature = "lock_as_mutex")]
    type InnerLock<T> = parking_lot::Mutex<T>;

    #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
    pub type ReadGuard<'a, T> = parking_lot::RwLockReadGuard<'a, T>;
    #[cfg(feature = "lock_as_mutex")]
    pub type ReadGuard<'a, T> = parking_lot::MutexGuard<'a, T>;

    #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
    pub type WriteGuard<'a, T> = parking_lot::RwLockWriteGuard<'a, T>;
    #[cfg(feature = "lock_as_mutex")]
    pub type WriteGuard<'a, T> = parking_lot::MutexGuard<'a, T>;

    /// A lock which provides the interface of a read/write lock, but which has the option to
    /// internally use either a RwLock or a Mutex in order to make it easier to expose erroneous
    /// recursive locking in tests while still using an RwLock in production to avoid deadlocking
    /// as much as possible.
    #[derive(Debug)]
    pub struct Lock<T> {
        inner: InnerLock<T>,
    }

    impl<T> Lock<T> {
        pub fn new(val: T) -> Self {
            #[cfg(feature = "lock_as_mutex")]
            println!("Lock::new is using Mutex to help find recursive locking");

            Self { inner: InnerLock::new(val), }
        }

        /// This acquires a read lock and will not deadlock if the same thread tries do acquire
        /// the lock recursively. However, it may result in writer starvation. Once we are confident
        /// that all recursive locking has been eliminated, we may replace this implementation
        /// and try_read_for (or add an additional methods) to provide fair locking for readers
        /// and writers.
        ///
        /// See https://github.com/habitat-sh/habitat/issues/6435
        #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
        pub fn read(&self) -> ReadGuard<T> { self.inner.read_recursive() }

        #[cfg(feature = "lock_as_mutex")]
        pub fn read(&self) -> ReadGuard<T> { self.inner.lock() }

        #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
        pub fn try_read_for(&self, timeout: std::time::Duration) -> Option<ReadGuard<T>> {
            self.inner.try_read_recursive_for(timeout)
        }

        #[cfg(feature = "lock_as_mutex")]
        pub fn try_read_for(&self, timeout: std::time::Duration) -> Option<ReadGuard<T>> {
            self.inner.try_lock_for(timeout)
        }

        #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
        pub fn write(&self) -> WriteGuard<T> { self.inner.write() }

        #[cfg(feature = "lock_as_mutex")]
        pub fn write(&self) -> ReadGuard<T> { self.inner.lock() }
    }
}
<commit_msg>Spawn deadlock detector thread if try_lock_for doesn't succeed<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::ui::{UIWriter,
                UI};
use habitat_api_client as api_client;
use habitat_core as hcore;
use lazy_static::lazy_static;
use std::{collections::HashMap,
          env,
          ffi::OsStr,
          iter::FromIterator};

extern crate json;
#[macro_use]
extern crate log;
#[macro_use]
extern crate lazy_static;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[cfg_attr(test, macro_use)]
extern crate serde_json;
#[cfg(windows)]
extern crate winapi;

pub use self::error::{Error,
                      Result};

pub mod cli;
pub mod command;
pub mod error;
pub mod locked_env_var;
pub mod output;
pub mod package_graph;
pub mod templating;
pub mod types;
pub mod ui;
pub mod util;

lazy_static::lazy_static! {
    pub static ref PROGRAM_NAME: String = {
        match env::current_exe() {
            Ok(path) => path.file_stem().and_then(OsStr::to_str).unwrap().to_string(),
            Err(e) => {
                error!("Error getting path of current_exe: {}", e);
                String::from("hab-?")
            }
        }
    };
}

// TODO (CM): It would be nice to come up with a way to more
// programmatically manage these flags. It's a bit of a pain to define
// the flag, and then define the environment variables
// separately. Nothing statically guarantees that you've specified an
// variable for a flag.

// TODO (CM): It'd be great to have a built-in way to document them,
// too.

// TODO (CM): Part of that documentation might be *when* a flag was
// added. In general, long-lived flags are a code-smell.

// TODO (CM): It may also be useful to break out features by area of
// concern. We can have any number of bitflags-generated structs.

bitflags::bitflags! {
    /// All the feature flags that are recogized by Habitat.
    ///
    /// In general, feature flags are enabled by setting the corresponding
    /// environment variable.
    ///
    /// Your binary should call `FeatureFlag::from_env` to get a set
    /// of flags to use.
    ///
    /// To add a new feature flag, you will need to add the bit mask
    /// constant here, as well as a mapping from the feature to the
    /// environment variable to which it corresponds in the `ENV_VARS`
    /// map below.
    pub struct FeatureFlag: u32 {
        const LIST               = 0b0000_0000_0001;
        const TEST_EXIT          = 0b0000_0000_0010;
        const TEST_BOOT_FAIL     = 0b0000_0000_0100;
        const REDACT_HTTP        = 0b0000_0000_1000;
        const IGNORE_SIGNALS     = 0b0000_0001_0000;
        const OFFLINE_INSTALL    = 0b0000_0100_0000;
        const IGNORE_LOCAL       = 0b0000_1000_0000;
        const EVENT_STREAM       = 0b0001_0000_0000;
        const TRIGGER_ELECTION   = 0b0010_0000_0000;
        const CONFIGURE_SHUTDOWN = 0b0100_0000_0000;
    }
}

lazy_static! {
    static ref ENV_VARS: HashMap<FeatureFlag, &'static str> = {
        let mapping = vec![(FeatureFlag::LIST, "HAB_FEAT_LIST"),
                           (FeatureFlag::TEST_EXIT, "HAB_FEAT_TEST_EXIT"),
                           (FeatureFlag::TEST_BOOT_FAIL, "HAB_FEAT_BOOT_FAIL"),
                           (FeatureFlag::REDACT_HTTP, "HAB_FEAT_REDACT_HTTP"),
                           (FeatureFlag::IGNORE_SIGNALS, "HAB_FEAT_IGNORE_SIGNALS"),
                           (FeatureFlag::OFFLINE_INSTALL, "HAB_FEAT_OFFLINE_INSTALL"),
                           (FeatureFlag::IGNORE_LOCAL, "HAB_FEAT_IGNORE_LOCAL"),
                           (FeatureFlag::EVENT_STREAM, "HAB_FEAT_EVENT_STREAM"),
                           (FeatureFlag::TRIGGER_ELECTION, "HAB_FEAT_TRIGGER_ELECTION"),
                           (FeatureFlag::CONFIGURE_SHUTDOWN, "HAB_FEAT_CONFIGURE_SHUTDOWN")];
        HashMap::from_iter(mapping)
    };
}

impl FeatureFlag {
    /// If the environment variable for a flag is set to _anything_ but
    /// the empty string, it is activated.
    pub fn from_env(ui: &mut UI) -> Self {
        let mut flags = FeatureFlag::empty();

        for (feature, env_var) in ENV_VARS.iter() {
            if let Some(val) = env::var_os(env_var) {
                if !val.is_empty() {
                    flags.insert(*feature);
                    ui.warn(&format!("Enabling feature: {:?}", feature))
                      .unwrap();
                }
            }
        }

        // TODO (CM): Once the other TODOs above are done (especially the
        // documentation bits), it would be nice to extract this logic
        // into an actual discoverable CLI subcommand; it's a little weird
        // that you have to know how to enable a feature flag before you
        // can even find out that there *are* feature flags to enable.
        //
        // There's no reason why "list feature flags" should itself be a
        // feature-flag.
        if flags.contains(FeatureFlag::LIST) {
            ui.warn("Listing feature flags environment variables:")
              .unwrap();
            for (feature, env_var) in ENV_VARS.iter() {
                ui.warn(&format!("  * {:?}: {}={:?}",
                                 feature,
                                 env_var,
                                 env::var_os(env_var).unwrap_or_default()))
                  .unwrap();
            }
        }

        flags
    }
}

pub mod sync {
    use std::time::Duration;

    #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
    type InnerLock<T> = parking_lot::RwLock<T>;
    #[cfg(feature = "lock_as_mutex")]
    type InnerLock<T> = parking_lot::Mutex<T>;

    #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
    pub type ReadGuard<'a, T> = parking_lot::RwLockReadGuard<'a, T>;
    #[cfg(feature = "lock_as_mutex")]
    pub type ReadGuard<'a, T> = parking_lot::MutexGuard<'a, T>;

    #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
    pub type WriteGuard<'a, T> = parking_lot::RwLockWriteGuard<'a, T>;
    #[cfg(feature = "lock_as_mutex")]
    pub type WriteGuard<'a, T> = parking_lot::MutexGuard<'a, T>;

    /// A lock which provides the interface of a read/write lock, but which has the option to
    /// internally use either a RwLock or a Mutex in order to make it easier to expose erroneous
    /// recursive locking in tests while still using an RwLock in production to avoid deadlocking
    /// as much as possible.
    #[derive(Debug)]
    pub struct Lock<T> {
        inner: InnerLock<T>,
    }

    impl<T> Lock<T> {
        pub fn new(val: T) -> Self {
            #[cfg(feature = "lock_as_mutex")]
            println!("Lock::new is using Mutex to help find recursive locking");

            Self { inner: InnerLock::new(val), }
        }

        /// This acquires a read lock and will not deadlock if the same thread tries do acquire
        /// the lock recursively. However, it may result in writer starvation. Once we are confident
        /// that all recursive locking has been eliminated, we may replace this implementation
        /// and try_read_for (or add an additional methods) to provide fair locking for readers
        /// and writers.
        ///
        /// See https://github.com/habitat-sh/habitat/issues/6435
        #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
        pub fn read(&self) -> ReadGuard<T> { self.inner.read_recursive() }

        #[cfg(feature = "lock_as_mutex")]
        pub fn read(&self) -> ReadGuard<T> {
            use parking_lot::deadlock;
            use std::thread;

            if let Some(guard) = self.inner.try_lock_for(Duration::from_secs(5)) {
                guard
            } else {
                let detector = thread::spawn(move || {
                    loop {
                        thread::sleep(Duration::from_secs(1));
                        let deadlocks = deadlock::check_deadlock();
                        if deadlocks.is_empty() {
                            continue;
                        }

                        println!("{} deadlocks detected", deadlocks.len());
                        for (i, threads) in deadlocks.iter().enumerate() {
                            println!("Deadlock #{}", i);
                            for t in threads {
                                println!("Thread Id {:#?}", t.thread_id());
                                println!("{:#?}", t.backtrace());
                            }
                        }
                    }
                });
                let guard = self.inner
                                .try_lock_for(Duration::from_secs(5))
                                .expect("timed out waiting for lock");
                detector.join().expect("join deadlock detector thread");
                guard
            }
        }

        #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
        pub fn try_read_for(&self, timeout: Duration) -> Option<ReadGuard<T>> {
            self.inner.try_read_recursive_for(timeout)
        }

        #[cfg(feature = "lock_as_mutex")]
        pub fn try_read_for(&self, timeout: Duration) -> Option<ReadGuard<T>> {
            self.inner.try_lock_for(timeout)
        }

        #[cfg(any(feature = "lock_as_rwlock", not(feature = "lock_as_mutex")))]
        pub fn write(&self) -> WriteGuard<T> { self.inner.write() }

        #[cfg(feature = "lock_as_mutex")]
        pub fn write(&self) -> ReadGuard<T> { self.inner.lock() }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use anyhow::Result;
use diem_config::config::{NodeConfig, OnDiskStorageConfig};
use diem_crypto::ed25519::Ed25519PrivateKey;
use diem_global_constants::{DIEM_ROOT_KEY, TREASURY_COMPLIANCE_KEY};
use diem_secure_storage::{CryptoStorage, OnDiskStorage};
use diem_types::waypoint::Waypoint;
use std::{
    fs::File,
    io::Write,
    path::{Path, PathBuf},
};

pub trait BuildSwarm {
    /// Generate the configs for a swarm
    fn build_swarm(&self) -> Result<(Vec<NodeConfig>, Ed25519PrivateKey)>;
}

pub struct SwarmConfig {
    pub config_files: Vec<PathBuf>,
    pub diem_root_key_path: PathBuf,
    pub root_storage: OnDiskStorageConfig,
    pub waypoint: Waypoint,
}

impl SwarmConfig {
    pub fn build<T: BuildSwarm>(config_builder: &T, output_dir: &Path) -> Result<Self> {
        let (mut configs, diem_root_key) = config_builder.build_swarm()?;
        let mut config_files = vec![];

        for (index, config) in configs.iter_mut().enumerate() {
            let node_dir = output_dir.join(index.to_string());
            std::fs::create_dir_all(&node_dir)?;

            let node_path = node_dir.join("node.yaml");
            config.set_data_dir(node_dir);
            config.save(&node_path)?;
            config_files.push(node_path);
        }

        let diem_root_key_path = output_dir.join("mint.key");
        let serialized_keys = bcs::to_bytes(&diem_root_key)?;
        let mut key_file = File::create(&diem_root_key_path)?;
        key_file.write_all(&serialized_keys)?;

        let mut root_storage_config = OnDiskStorageConfig::default();
        root_storage_config.path = output_dir.join("root-storage.json");
        let mut root_storage = OnDiskStorage::new(root_storage_config.path());
        root_storage.import_private_key(DIEM_ROOT_KEY, diem_root_key)?;
        let diem_root_key = root_storage.export_private_key(DIEM_ROOT_KEY).unwrap();
        root_storage.import_private_key(TREASURY_COMPLIANCE_KEY, diem_root_key)?;

        Ok(SwarmConfig {
            config_files,
            diem_root_key_path,
            root_storage: root_storage_config,
            waypoint: configs[0].base.waypoint.waypoint(),
        })
    }
}
<commit_msg>swarm-config: impl BuildSwarm for ValidatorBuilder<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::validator_builder::ValidatorBuilder;
use anyhow::Result;
use diem_config::config::{NodeConfig, OnDiskStorageConfig};
use diem_crypto::ed25519::Ed25519PrivateKey;
use diem_global_constants::{DIEM_ROOT_KEY, TREASURY_COMPLIANCE_KEY};
use diem_secure_storage::{CryptoStorage, OnDiskStorage};
use diem_types::waypoint::Waypoint;
use std::{
    fs::File,
    io::Write,
    path::{Path, PathBuf},
};

pub trait BuildSwarm {
    /// Generate the configs for a swarm
    fn build_swarm(&self) -> Result<(Vec<NodeConfig>, Ed25519PrivateKey)>;
}

pub struct SwarmConfig {
    pub config_files: Vec<PathBuf>,
    pub diem_root_key_path: PathBuf,
    pub root_storage: OnDiskStorageConfig,
    pub waypoint: Waypoint,
}

impl SwarmConfig {
    pub fn build<T: BuildSwarm>(config_builder: &T, output_dir: &Path) -> Result<Self> {
        let (mut configs, diem_root_key) = config_builder.build_swarm()?;
        let mut config_files = vec![];

        for (index, config) in configs.iter_mut().enumerate() {
            let node_dir = output_dir.join(index.to_string());
            std::fs::create_dir_all(&node_dir)?;

            let node_path = node_dir.join("node.yaml");
            config.set_data_dir(node_dir);
            config.save(&node_path)?;
            config_files.push(node_path);
        }

        let diem_root_key_path = output_dir.join("mint.key");
        let serialized_keys = bcs::to_bytes(&diem_root_key)?;
        let mut key_file = File::create(&diem_root_key_path)?;
        key_file.write_all(&serialized_keys)?;

        let mut root_storage_config = OnDiskStorageConfig::default();
        root_storage_config.path = output_dir.join("root-storage.json");
        let mut root_storage = OnDiskStorage::new(root_storage_config.path());
        root_storage.import_private_key(DIEM_ROOT_KEY, diem_root_key)?;
        let diem_root_key = root_storage.export_private_key(DIEM_ROOT_KEY).unwrap();
        root_storage.import_private_key(TREASURY_COMPLIANCE_KEY, diem_root_key)?;

        Ok(SwarmConfig {
            config_files,
            diem_root_key_path,
            root_storage: root_storage_config,
            waypoint: configs[0].base.waypoint.waypoint(),
        })
    }
}

impl BuildSwarm for ValidatorBuilder {
    fn build_swarm(&self) -> Result<(Vec<NodeConfig>, Ed25519PrivateKey)> {
        let (root_keys, validators) = self.clone().build()?;

        Ok((
            validators.into_iter().map(|v| v.config).collect(),
            root_keys.root_key,
        ))
    }
}
<|endoftext|>


<commit_before>#![warn(rust_2018_idioms)]
#![cfg(feature = "full")]
#![cfg(unix)]

use futures::future::poll_fn;
use tokio::io::ReadBuf;
use tokio::net::UnixDatagram;
use tokio::try_join;

use std::io;
use std::sync::Arc;

async fn echo_server(socket: UnixDatagram) -> io::Result<()> {
    let mut recv_buf = vec![0u8; 1024];
    loop {
        let (len, peer_addr) = socket.recv_from(&mut recv_buf[..]).await?;
        if let Some(path) = peer_addr.as_pathname() {
            socket.send_to(&recv_buf[..len], path).await?;
        }
    }
}

#[tokio::test]
async fn echo() -> io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let server_path = dir.path().join("server.sock");
    let client_path = dir.path().join("client.sock");

    let server_socket = UnixDatagram::bind(server_path.clone())?;

    tokio::spawn(async move {
        if let Err(e) = echo_server(server_socket).await {
            eprintln!("Error in echo server: {}", e);
        }
    });

    {
        let socket = UnixDatagram::bind(&client_path).unwrap();
        socket.connect(server_path)?;
        socket.send(b"ECHO").await?;
        let mut recv_buf = [0u8; 16];
        let len = socket.recv(&mut recv_buf[..]).await?;
        assert_eq!(&recv_buf[..len], b"ECHO");
    }

    Ok(())
}

#[tokio::test]
async fn echo_from() -> io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let server_path = dir.path().join("server.sock");
    let client_path = dir.path().join("client.sock");

    let server_socket = UnixDatagram::bind(server_path.clone())?;

    tokio::spawn(async move {
        if let Err(e) = echo_server(server_socket).await {
            eprintln!("Error in echo server: {}", e);
        }
    });

    {
        let socket = UnixDatagram::bind(&client_path).unwrap();
        socket.connect(&server_path)?;
        socket.send(b"ECHO").await?;
        let mut recv_buf = [0u8; 16];
        let (len, addr) = socket.recv_from(&mut recv_buf[..]).await?;
        assert_eq!(&recv_buf[..len], b"ECHO");
        assert_eq!(addr.as_pathname(), Some(server_path.as_path()));
    }

    Ok(())
}

// Even though we use sync non-blocking io we still need a reactor.
#[tokio::test]
async fn try_send_recv_never_block() -> io::Result<()> {
    let mut recv_buf = [0u8; 16];
    let payload = b"PAYLOAD";
    let mut count = 0;

    let (dgram1, dgram2) = UnixDatagram::pair()?;

    // Send until we hit the OS `net.unix.max_dgram_qlen`.
    loop {
        dgram1.writable().await.unwrap();

        match dgram1.try_send(payload) {
            Err(err) => match err.kind() {
                io::ErrorKind::WouldBlock | io::ErrorKind::Other => break,
                _ => unreachable!("unexpected error {:?}", err),
            },
            Ok(len) => {
                assert_eq!(len, payload.len());
            }
        }
        count += 1;
    }

    // Read every dgram we sent.
    while count > 0 {
        dgram2.readable().await.unwrap();
        let len = dgram2.try_recv(&mut recv_buf[..])?;
        assert_eq!(len, payload.len());
        assert_eq!(payload, &recv_buf[..len]);
        count -= 1;
    }

    let err = dgram2.try_recv(&mut recv_buf[..]).unwrap_err();
    match err.kind() {
        io::ErrorKind::WouldBlock => (),
        _ => unreachable!("unexpected error {:?}", err),
    }

    Ok(())
}

#[tokio::test]
async fn split() -> std::io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("split.sock");
    let s = Arc::new(UnixDatagram::bind(path.clone())?);
    let r = s.clone();

    let msg = b"hello";
    let ((), ()) = try_join! {
        async {
            s.send_to(msg, path).await?;
            io::Result::Ok(())
        },
        async {
            let mut recv_buf = [0u8; 32];
            let (len, _) = r.recv_from(&mut recv_buf[..]).await?;
            assert_eq!(&recv_buf[..len], msg);
            Ok(())
        },
    }?;

    Ok(())
}

#[tokio::test]
async fn send_to_recv_from_poll() -> std::io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let sender_path = dir.path().join("sender.sock");
    let receiver_path = dir.path().join("receiver.sock");

    let sender = UnixDatagram::bind(&sender_path)?;
    let receiver = UnixDatagram::bind(&receiver_path)?;

    let msg = b"hello";
    poll_fn(|cx| sender.poll_send_to(cx, msg, &receiver_path)).await?;

    let mut recv_buf = [0u8; 32];
    let mut read = ReadBuf::new(&mut recv_buf);
    let addr = poll_fn(|cx| receiver.poll_recv_from(cx, &mut read)).await?;

    assert_eq!(read.filled(), msg);
    assert_eq!(addr.as_pathname(), Some(sender_path.as_ref()));
    Ok(())
}

#[tokio::test]
async fn send_recv_poll() -> std::io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let sender_path = dir.path().join("sender.sock");
    let receiver_path = dir.path().join("receiver.sock");

    let sender = UnixDatagram::bind(&sender_path)?;
    let receiver = UnixDatagram::bind(&receiver_path)?;

    sender.connect(&receiver_path)?;
    receiver.connect(&sender_path)?;

    let msg = b"hello";
    poll_fn(|cx| sender.poll_send(cx, msg)).await?;

    let mut recv_buf = [0u8; 32];
    let mut read = ReadBuf::new(&mut recv_buf);
    let _len = poll_fn(|cx| receiver.poll_recv(cx, &mut read)).await?;

    assert_eq!(read.filled(), msg);
    Ok(())
}

#[tokio::test]
async fn try_send_to_recv_from() -> std::io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let server_path = dir.path().join("server.sock");
    let client_path = dir.path().join("client.sock");

    // Create listener
    let server = UnixDatagram::bind(&server_path)?;

    // Create socket pair
    let client = UnixDatagram::bind(&client_path)?;

    for _ in 0..5 {
        loop {
            client.writable().await?;

            match client.try_send_to(b"hello world", &server_path) {
                Ok(n) => {
                    assert_eq!(n, 11);
                    break;
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
                Err(e) => panic!("{:?}", e),
            }
        }

        loop {
            server.readable().await?;

            let mut buf = [0; 512];

            match server.try_recv_from(&mut buf) {
                Ok((n, addr)) => {
                    assert_eq!(n, 11);
                    assert_eq!(addr.as_pathname(), Some(client_path.as_ref()));
                    assert_eq!(&buf[0..11], &b"hello world"[..]);
                    break;
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
                Err(e) => panic!("{:?}", e),
            }
        }
    }

    Ok(())
}

#[tokio::test]
async fn try_recv_buf_from() -> std::io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let server_path = dir.path().join("server.sock");
    let client_path = dir.path().join("client.sock");

    // Create listener
    let server = UnixDatagram::bind(&server_path)?;

    // Create socket pair
    let client = UnixDatagram::bind(&client_path)?;

    for _ in 0..5 {
        loop {
            client.writable().await?;

            match client.try_send_to(b"hello world", &server_path) {
                Ok(n) => {
                    assert_eq!(n, 11);
                    break;
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
                Err(e) => panic!("{:?}", e),
            }
        }

        loop {
            server.readable().await?;

            let mut buf = Vec::with_capacity(512);

            match server.try_recv_buf_from(&mut buf) {
                Ok((n, addr)) => {
                    assert_eq!(n, 11);
                    assert_eq!(addr.as_pathname(), Some(client_path.as_ref()));
                    assert_eq!(&buf[0..11], &b"hello world"[..]);
                    break;
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
                Err(e) => panic!("{:?}", e),
            }
        }
    }

    Ok(())
}

// Even though we use sync non-blocking io we still need a reactor.
#[tokio::test]
async fn try_recv_buf_never_block() -> io::Result<()> {
    let payload = b"PAYLOAD";
    let mut count = 0;

    let (dgram1, dgram2) = UnixDatagram::pair()?;

    // Send until we hit the OS `net.unix.max_dgram_qlen`.
    loop {
        dgram1.writable().await.unwrap();

        match dgram1.try_send(payload) {
            Err(err) => match err.kind() {
                io::ErrorKind::WouldBlock | io::ErrorKind::Other => break,
                _ => unreachable!("unexpected error {:?}", err),
            },
            Ok(len) => {
                assert_eq!(len, payload.len());
            }
        }
        count += 1;
    }

    // Read every dgram we sent.
    while count > 0 {
        let mut recv_buf = Vec::with_capacity(16);

        dgram2.readable().await.unwrap();
        let len = dgram2.try_recv_buf(&mut recv_buf)?;
        assert_eq!(len, payload.len());
        assert_eq!(payload, &recv_buf[..len]);
        count -= 1;
    }

    let mut recv_buf = vec![0; 16];
    let err = dgram2.try_recv_from(&mut recv_buf).unwrap_err();
    match err.kind() {
        io::ErrorKind::WouldBlock => (),
        _ => unreachable!("unexpected error {:?}", err),
    }

    Ok(())
}
<commit_msg>net: fix the uds_datagram tests with the latest nightly stdlib (#3952)<commit_after>#![warn(rust_2018_idioms)]
#![cfg(feature = "full")]
#![cfg(unix)]

use futures::future::poll_fn;
use tokio::io::ReadBuf;
use tokio::net::UnixDatagram;
use tokio::try_join;

use std::io;
use std::sync::Arc;

async fn echo_server(socket: UnixDatagram) -> io::Result<()> {
    let mut recv_buf = vec![0u8; 1024];
    loop {
        let (len, peer_addr) = socket.recv_from(&mut recv_buf[..]).await?;
        if let Some(path) = peer_addr.as_pathname() {
            socket.send_to(&recv_buf[..len], path).await?;
        }
    }
}

#[tokio::test]
async fn echo() -> io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let server_path = dir.path().join("server.sock");
    let client_path = dir.path().join("client.sock");

    let server_socket = UnixDatagram::bind(server_path.clone())?;

    tokio::spawn(async move {
        if let Err(e) = echo_server(server_socket).await {
            eprintln!("Error in echo server: {}", e);
        }
    });

    {
        let socket = UnixDatagram::bind(&client_path).unwrap();
        socket.connect(server_path)?;
        socket.send(b"ECHO").await?;
        let mut recv_buf = [0u8; 16];
        let len = socket.recv(&mut recv_buf[..]).await?;
        assert_eq!(&recv_buf[..len], b"ECHO");
    }

    Ok(())
}

#[tokio::test]
async fn echo_from() -> io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let server_path = dir.path().join("server.sock");
    let client_path = dir.path().join("client.sock");

    let server_socket = UnixDatagram::bind(server_path.clone())?;

    tokio::spawn(async move {
        if let Err(e) = echo_server(server_socket).await {
            eprintln!("Error in echo server: {}", e);
        }
    });

    {
        let socket = UnixDatagram::bind(&client_path).unwrap();
        socket.connect(&server_path)?;
        socket.send(b"ECHO").await?;
        let mut recv_buf = [0u8; 16];
        let (len, addr) = socket.recv_from(&mut recv_buf[..]).await?;
        assert_eq!(&recv_buf[..len], b"ECHO");
        assert_eq!(addr.as_pathname(), Some(server_path.as_path()));
    }

    Ok(())
}

// Even though we use sync non-blocking io we still need a reactor.
#[tokio::test]
async fn try_send_recv_never_block() -> io::Result<()> {
    let mut recv_buf = [0u8; 16];
    let payload = b"PAYLOAD";
    let mut count = 0;

    let (dgram1, dgram2) = UnixDatagram::pair()?;

    // Send until we hit the OS `net.unix.max_dgram_qlen`.
    loop {
        dgram1.writable().await.unwrap();

        match dgram1.try_send(payload) {
            Err(err) => match (err.kind(), err.raw_os_error()) {
                (io::ErrorKind::WouldBlock, _) => break,
                (_, Some(libc::ENOBUFS)) => break,
                _ => {
                    panic!("unexpected error {:?}", err);
                }
            },
            Ok(len) => {
                assert_eq!(len, payload.len());
            }
        }
        count += 1;
    }

    // Read every dgram we sent.
    while count > 0 {
        dgram2.readable().await.unwrap();
        let len = dgram2.try_recv(&mut recv_buf[..])?;
        assert_eq!(len, payload.len());
        assert_eq!(payload, &recv_buf[..len]);
        count -= 1;
    }

    let err = dgram2.try_recv(&mut recv_buf[..]).unwrap_err();
    match err.kind() {
        io::ErrorKind::WouldBlock => (),
        _ => unreachable!("unexpected error {:?}", err),
    }

    Ok(())
}

#[tokio::test]
async fn split() -> std::io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("split.sock");
    let s = Arc::new(UnixDatagram::bind(path.clone())?);
    let r = s.clone();

    let msg = b"hello";
    let ((), ()) = try_join! {
        async {
            s.send_to(msg, path).await?;
            io::Result::Ok(())
        },
        async {
            let mut recv_buf = [0u8; 32];
            let (len, _) = r.recv_from(&mut recv_buf[..]).await?;
            assert_eq!(&recv_buf[..len], msg);
            Ok(())
        },
    }?;

    Ok(())
}

#[tokio::test]
async fn send_to_recv_from_poll() -> std::io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let sender_path = dir.path().join("sender.sock");
    let receiver_path = dir.path().join("receiver.sock");

    let sender = UnixDatagram::bind(&sender_path)?;
    let receiver = UnixDatagram::bind(&receiver_path)?;

    let msg = b"hello";
    poll_fn(|cx| sender.poll_send_to(cx, msg, &receiver_path)).await?;

    let mut recv_buf = [0u8; 32];
    let mut read = ReadBuf::new(&mut recv_buf);
    let addr = poll_fn(|cx| receiver.poll_recv_from(cx, &mut read)).await?;

    assert_eq!(read.filled(), msg);
    assert_eq!(addr.as_pathname(), Some(sender_path.as_ref()));
    Ok(())
}

#[tokio::test]
async fn send_recv_poll() -> std::io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let sender_path = dir.path().join("sender.sock");
    let receiver_path = dir.path().join("receiver.sock");

    let sender = UnixDatagram::bind(&sender_path)?;
    let receiver = UnixDatagram::bind(&receiver_path)?;

    sender.connect(&receiver_path)?;
    receiver.connect(&sender_path)?;

    let msg = b"hello";
    poll_fn(|cx| sender.poll_send(cx, msg)).await?;

    let mut recv_buf = [0u8; 32];
    let mut read = ReadBuf::new(&mut recv_buf);
    let _len = poll_fn(|cx| receiver.poll_recv(cx, &mut read)).await?;

    assert_eq!(read.filled(), msg);
    Ok(())
}

#[tokio::test]
async fn try_send_to_recv_from() -> std::io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let server_path = dir.path().join("server.sock");
    let client_path = dir.path().join("client.sock");

    // Create listener
    let server = UnixDatagram::bind(&server_path)?;

    // Create socket pair
    let client = UnixDatagram::bind(&client_path)?;

    for _ in 0..5 {
        loop {
            client.writable().await?;

            match client.try_send_to(b"hello world", &server_path) {
                Ok(n) => {
                    assert_eq!(n, 11);
                    break;
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
                Err(e) => panic!("{:?}", e),
            }
        }

        loop {
            server.readable().await?;

            let mut buf = [0; 512];

            match server.try_recv_from(&mut buf) {
                Ok((n, addr)) => {
                    assert_eq!(n, 11);
                    assert_eq!(addr.as_pathname(), Some(client_path.as_ref()));
                    assert_eq!(&buf[0..11], &b"hello world"[..]);
                    break;
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
                Err(e) => panic!("{:?}", e),
            }
        }
    }

    Ok(())
}

#[tokio::test]
async fn try_recv_buf_from() -> std::io::Result<()> {
    let dir = tempfile::tempdir().unwrap();
    let server_path = dir.path().join("server.sock");
    let client_path = dir.path().join("client.sock");

    // Create listener
    let server = UnixDatagram::bind(&server_path)?;

    // Create socket pair
    let client = UnixDatagram::bind(&client_path)?;

    for _ in 0..5 {
        loop {
            client.writable().await?;

            match client.try_send_to(b"hello world", &server_path) {
                Ok(n) => {
                    assert_eq!(n, 11);
                    break;
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
                Err(e) => panic!("{:?}", e),
            }
        }

        loop {
            server.readable().await?;

            let mut buf = Vec::with_capacity(512);

            match server.try_recv_buf_from(&mut buf) {
                Ok((n, addr)) => {
                    assert_eq!(n, 11);
                    assert_eq!(addr.as_pathname(), Some(client_path.as_ref()));
                    assert_eq!(&buf[0..11], &b"hello world"[..]);
                    break;
                }
                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => continue,
                Err(e) => panic!("{:?}", e),
            }
        }
    }

    Ok(())
}

// Even though we use sync non-blocking io we still need a reactor.
#[tokio::test]
async fn try_recv_buf_never_block() -> io::Result<()> {
    let payload = b"PAYLOAD";
    let mut count = 0;

    let (dgram1, dgram2) = UnixDatagram::pair()?;

    // Send until we hit the OS `net.unix.max_dgram_qlen`.
    loop {
        dgram1.writable().await.unwrap();

        match dgram1.try_send(payload) {
            Err(err) => match (err.kind(), err.raw_os_error()) {
                (io::ErrorKind::WouldBlock, _) => break,
                (_, Some(libc::ENOBUFS)) => break,
                _ => {
                    panic!("unexpected error {:?}", err);
                }
            },
            Ok(len) => {
                assert_eq!(len, payload.len());
            }
        }
        count += 1;
    }

    // Read every dgram we sent.
    while count > 0 {
        let mut recv_buf = Vec::with_capacity(16);

        dgram2.readable().await.unwrap();
        let len = dgram2.try_recv_buf(&mut recv_buf)?;
        assert_eq!(len, payload.len());
        assert_eq!(payload, &recv_buf[..len]);
        count -= 1;
    }

    let mut recv_buf = vec![0; 16];
    let err = dgram2.try_recv_from(&mut recv_buf).unwrap_err();
    match err.kind() {
        io::ErrorKind::WouldBlock => (),
        _ => unreachable!("unexpected error {:?}", err),
    }

    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::cell::RefCell;
use std::collections::BTreeMap;

use ast;
use ast::{Ident, Name, TokenTree};
use codemap::Span;
use ext::base::{ExtCtxt, MacEager, MacResult};
use ext::build::AstBuilder;
use parse::token;
use ptr::P;
use util::small_vector::SmallVector;

// Maximum width of any line in an extended error description (inclusive).
const MAX_DESCRIPTION_WIDTH: usize = 80;

thread_local! {
    static REGISTERED_DIAGNOSTICS: RefCell<ErrorMap> = {
        RefCell::new(BTreeMap::new())
    }
}

/// Error information type.
pub struct ErrorInfo {
    pub description: Option<Name>,
    pub use_site: Option
}

/// Mapping from error codes to metadata.
pub type ErrorMap = BTreeMap<Name, ErrorInfo>;

fn with_registered_diagnostics<T, F>(f: F) -> T where
    F: FnOnce(&mut ErrorMap) -> T,
{
    REGISTERED_DIAGNOSTICS.with(move |slot| {
        f(&mut *slot.borrow_mut())
    })
}

pub fn expand_diagnostic_used<'cx>(ecx: &'cx mut ExtCtxt,
                                   span: Span,
                                   token_tree: &[TokenTree])
                                   -> Box<MacResult+'cx> {
    let code = match (token_tree.len(), token_tree.get(0)) {
        (1, Some(&ast::TtToken(_, token::Ident(code, _)))) => code,
        _ => unreachable!()
    };

    with_registered_diagnostics(|diagnostics| {
        match diagnostics.get_mut(&code.name) {
            // Previously used errors.
            Some(&mut ErrorInfo { description: _, use_site: Some(previous_span) }) => {
                ecx.span_warn(span, &format!(
                    "diagnostic code {} already used", &token::get_ident(code)
                ));
                ecx.span_note(previous_span, "previous invocation");
            }
            // Newly used errors.
            Some(ref mut info) => {
                info.use_site = Some(span);
            }
            // Unregistered errors.
            None => {
                ecx.span_err(span, &format!(
                    "used diagnostic code {} not registered", &token::get_ident(code)
                ));
            }
        }
    });
    MacEager::expr(ecx.expr_tuple(span, Vec::new()))
}

pub fn expand_register_diagnostic<'cx>(ecx: &'cx mut ExtCtxt,
                                       span: Span,
                                       token_tree: &[TokenTree])
                                       -> Box<MacResult+'cx> {
    let (code, description) = match (
        token_tree.len(),
        token_tree.get(0),
        token_tree.get(1),
        token_tree.get(2)
    ) {
        (1, Some(&ast::TtToken(_, token::Ident(ref code, _))), None, None) => {
            (code, None)
        },
        (3, Some(&ast::TtToken(_, token::Ident(ref code, _))),
            Some(&ast::TtToken(_, token::Comma)),
            Some(&ast::TtToken(_, token::Literal(token::StrRaw(description, _), None)))) => {
            (code, Some(description))
        }
        _ => unreachable!()
    };
    // Check that the description starts and ends with a newline and doesn't
    // overflow the maximum line width.
    description.map(|raw_msg| {
        let msg = raw_msg.as_str();
        if !msg.starts_with("\n") || !msg.ends_with("\n") {
            ecx.span_err(span, &format!(
                "description for error code {} doesn't start and end with a newline",
                token::get_ident(*code)
            ));
        }
        if msg.lines().any(|line| line.len() > MAX_DESCRIPTION_WIDTH) {
            ecx.span_err(span, &format!(
                "description for error code {} contains a line longer than {} characters",
                token::get_ident(*code), MAX_DESCRIPTION_WIDTH
            ));
        }
    });
    // Add the error to the map.
    with_registered_diagnostics(|diagnostics| {
        let info = ErrorInfo {
            description: description,
            use_site: None
        };
        if diagnostics.insert(code.name, info).is_some() {
            ecx.span_err(span, &format!(
                "diagnostic code {} already registered", &token::get_ident(*code)
            ));
        }
    });
    let sym = Ident::new(token::gensym(&(
        "__register_diagnostic_".to_string() + &token::get_ident(*code)
    )));
    MacEager::items(SmallVector::many(vec![
        ecx.item_mod(
            span,
            span,
            sym,
            Vec::new(),
            Vec::new()
        )
    ]))
}

pub fn expand_build_diagnostic_array<'cx>(ecx: &'cx mut ExtCtxt,
                                          span: Span,
                                          token_tree: &[TokenTree])
                                          -> Box<MacResult+'cx> {
    assert_eq!(token_tree.len(), 3);
    let (_crate_name, name) = match (&token_tree[0], &token_tree[2]) {
        (
            // Crate name.
            &ast::TtToken(_, token::Ident(ref crate_name, _)),
            // DIAGNOSTICS ident.
            &ast::TtToken(_, token::Ident(ref name, _))
        ) => (crate_name.as_str(), name),
        _ => unreachable!()
    };

    // FIXME (#25705): we used to ensure error code uniqueness and
    // output error description JSON metadata here, but the approach
    // employed was too brittle.

    // Construct the output expression.
    let (count, expr) =
        with_registered_diagnostics(|diagnostics| {
            let descriptions: Vec<P<ast::Expr>> =
                diagnostics.iter().filter_map(|(code, info)| {
                    info.description.map(|description| {
                        ecx.expr_tuple(span, vec![
                            ecx.expr_str(span, token::get_name(*code)),
                            ecx.expr_str(span, token::get_name(description))
                        ])
                    })
                }).collect();
            (descriptions.len(), ecx.expr_vec(span, descriptions))
        });

    let static_ = ecx.lifetime(span, ecx.name_of("'static"));
    let ty_str = ecx.ty_rptr(
        span,
        ecx.ty_ident(span, ecx.ident_of("str")),
        Some(static_),
        ast::MutImmutable,
    );

    let ty = ecx.ty(
        span,
        ast::TyFixedLengthVec(
            ecx.ty(
                span,
                ast::TyTup(vec![ty_str.clone(), ty_str])
            ),
            ecx.expr_usize(span, count),
        ),
    );

    MacEager::items(SmallVector::many(vec![
        P(ast::Item {
            ident: name.clone(),
            attrs: Vec::new(),
            id: ast::DUMMY_NODE_ID,
            node: ast::ItemConst(
                ty,
                expr,
            ),
            vis: ast::Public,
            span: span,
        })
    ]))
}
<commit_msg>diagnostics: Allow long URLs in error explanations.<commit_after>// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::cell::RefCell;
use std::collections::BTreeMap;

use ast;
use ast::{Ident, Name, TokenTree};
use codemap::Span;
use ext::base::{ExtCtxt, MacEager, MacResult};
use ext::build::AstBuilder;
use parse::token;
use ptr::P;
use util::small_vector::SmallVector;

// Maximum width of any line in an extended error description (inclusive).
const MAX_DESCRIPTION_WIDTH: usize = 80;

thread_local! {
    static REGISTERED_DIAGNOSTICS: RefCell<ErrorMap> = {
        RefCell::new(BTreeMap::new())
    }
}

/// Error information type.
pub struct ErrorInfo {
    pub description: Option<Name>,
    pub use_site: Option
}

/// Mapping from error codes to metadata.
pub type ErrorMap = BTreeMap<Name, ErrorInfo>;

fn with_registered_diagnostics<T, F>(f: F) -> T where
    F: FnOnce(&mut ErrorMap) -> T,
{
    REGISTERED_DIAGNOSTICS.with(move |slot| {
        f(&mut *slot.borrow_mut())
    })
}

pub fn expand_diagnostic_used<'cx>(ecx: &'cx mut ExtCtxt,
                                   span: Span,
                                   token_tree: &[TokenTree])
                                   -> Box<MacResult+'cx> {
    let code = match (token_tree.len(), token_tree.get(0)) {
        (1, Some(&ast::TtToken(_, token::Ident(code, _)))) => code,
        _ => unreachable!()
    };

    with_registered_diagnostics(|diagnostics| {
        match diagnostics.get_mut(&code.name) {
            // Previously used errors.
            Some(&mut ErrorInfo { description: _, use_site: Some(previous_span) }) => {
                ecx.span_warn(span, &format!(
                    "diagnostic code {} already used", &token::get_ident(code)
                ));
                ecx.span_note(previous_span, "previous invocation");
            }
            // Newly used errors.
            Some(ref mut info) => {
                info.use_site = Some(span);
            }
            // Unregistered errors.
            None => {
                ecx.span_err(span, &format!(
                    "used diagnostic code {} not registered", &token::get_ident(code)
                ));
            }
        }
    });
    MacEager::expr(ecx.expr_tuple(span, Vec::new()))
}

pub fn expand_register_diagnostic<'cx>(ecx: &'cx mut ExtCtxt,
                                       span: Span,
                                       token_tree: &[TokenTree])
                                       -> Box<MacResult+'cx> {
    let (code, description) = match (
        token_tree.len(),
        token_tree.get(0),
        token_tree.get(1),
        token_tree.get(2)
    ) {
        (1, Some(&ast::TtToken(_, token::Ident(ref code, _))), None, None) => {
            (code, None)
        },
        (3, Some(&ast::TtToken(_, token::Ident(ref code, _))),
            Some(&ast::TtToken(_, token::Comma)),
            Some(&ast::TtToken(_, token::Literal(token::StrRaw(description, _), None)))) => {
            (code, Some(description))
        }
        _ => unreachable!()
    };

    // Check that the description starts and ends with a newline and doesn't
    // overflow the maximum line width.
    description.map(|raw_msg| {
        let msg = raw_msg.as_str();
        if !msg.starts_with("\n") || !msg.ends_with("\n") {
            ecx.span_err(span, &format!(
                "description for error code {} doesn't start and end with a newline",
                token::get_ident(*code)
            ));
        }

        // URLs can be unavoidably longer than the line limit, so we allow them.
        // Allowed format is: `[name]: http://rust-lang.org/`
        let is_url = |l: &str| l.starts_with('[') && l.contains("]:") && l.contains("http");

        if msg.lines().any(|line| line.len() > MAX_DESCRIPTION_WIDTH && !is_url(line)) {
            ecx.span_err(span, &format!(
                "description for error code {} contains a line longer than {} characters.\n\
                 if you're inserting a long URL use the footnote style to bypass this check.",
                token::get_ident(*code), MAX_DESCRIPTION_WIDTH
            ));
        }
    });
    // Add the error to the map.
    with_registered_diagnostics(|diagnostics| {
        let info = ErrorInfo {
            description: description,
            use_site: None
        };
        if diagnostics.insert(code.name, info).is_some() {
            ecx.span_err(span, &format!(
                "diagnostic code {} already registered", &token::get_ident(*code)
            ));
        }
    });
    let sym = Ident::new(token::gensym(&(
        "__register_diagnostic_".to_string() + &token::get_ident(*code)
    )));
    MacEager::items(SmallVector::many(vec![
        ecx.item_mod(
            span,
            span,
            sym,
            Vec::new(),
            Vec::new()
        )
    ]))
}

pub fn expand_build_diagnostic_array<'cx>(ecx: &'cx mut ExtCtxt,
                                          span: Span,
                                          token_tree: &[TokenTree])
                                          -> Box<MacResult+'cx> {
    assert_eq!(token_tree.len(), 3);
    let (_crate_name, name) = match (&token_tree[0], &token_tree[2]) {
        (
            // Crate name.
            &ast::TtToken(_, token::Ident(ref crate_name, _)),
            // DIAGNOSTICS ident.
            &ast::TtToken(_, token::Ident(ref name, _))
        ) => (crate_name.as_str(), name),
        _ => unreachable!()
    };

    // FIXME (#25705): we used to ensure error code uniqueness and
    // output error description JSON metadata here, but the approach
    // employed was too brittle.

    // Construct the output expression.
    let (count, expr) =
        with_registered_diagnostics(|diagnostics| {
            let descriptions: Vec<P<ast::Expr>> =
                diagnostics.iter().filter_map(|(code, info)| {
                    info.description.map(|description| {
                        ecx.expr_tuple(span, vec![
                            ecx.expr_str(span, token::get_name(*code)),
                            ecx.expr_str(span, token::get_name(description))
                        ])
                    })
                }).collect();
            (descriptions.len(), ecx.expr_vec(span, descriptions))
        });

    let static_ = ecx.lifetime(span, ecx.name_of("'static"));
    let ty_str = ecx.ty_rptr(
        span,
        ecx.ty_ident(span, ecx.ident_of("str")),
        Some(static_),
        ast::MutImmutable,
    );

    let ty = ecx.ty(
        span,
        ast::TyFixedLengthVec(
            ecx.ty(
                span,
                ast::TyTup(vec![ty_str.clone(), ty_str])
            ),
            ecx.expr_usize(span, count),
        ),
    );

    MacEager::items(SmallVector::many(vec![
        P(ast::Item {
            ident: name.clone(),
            attrs: Vec::new(),
            id: ast::DUMMY_NODE_ID,
            node: ast::ItemConst(
                ty,
                expr,
            ),
            vis: ast::Public,
            span: span,
        })
    ]))
}
<|endoftext|>