<commit_before>/*
 * Copyright 2015-2016 Torrie Fischer <tdfischer@hackerbots.net>
 *
 * 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 std::str::FromStr;

#[derive(Debug,PartialEq,Eq)]
pub enum Command {
    Nick,
    User,
    Join,
    Part,
    Quit,
    Ping,
    Mode,
    Pong,
    Pass,
    Privmsg,
    Topic,
    Numeric(u32),
    Unknown(String)
}

impl Command {
    pub fn as_string(&self) -> String {
        match self {
            &Command::Nick => "NICK".to_string(),
            &Command::Join => "JOIN".to_string(),
            &Command::Part => "PART".to_string(),
            &Command::Pong => "PONG".to_string(),
            &Command::Privmsg => "PRIVMSG".to_string(),
            &Command::User => "USER".to_string(),
            &Command::Quit => "QUIT".to_string(),
            &Command::Ping => "PING".to_string(),
            &Command::Mode => "MODE".to_string(),
            &Command::Pass => "PASS".to_string(),
            &Command::Topic => "TOPIC".to_string(),
            &Command::Numeric(n)=> format!("{:0>3}", n),
            &Command::Unknown(ref s) => s.clone()
        }
    }
}

impl Message {
    pub fn to_string(&self) -> String {
        let mut ret = String::new();
        match self.prefix {
            Some(ref pfx) => {
                ret.push(':');
                ret.push_str(pfx.trim());
                ret.push(' ');
            },
            None => ()
        };
        ret.push_str(self.command.as_string().trim());
        for ref arg in self.args.iter() {
            ret.push(' ');
            ret.push_str(arg.trim());
        }

        match self.suffix {
            Some(ref sfx) => {
                ret.push_str(" :");
                ret.push_str(sfx.trim());
            },
            None => ()
        };

        return ret;
    }

    fn split_parts(line: &str) -> (Option<String>, &str, Option<String>) {
        let mut prefix_end = 0;
        if line.starts_with(":") {
            for c in line[1..].chars() {
                prefix_end += 1;
                if c == ' ' {
                    break
                }
            }
        }

        let mut args_end = prefix_end;
        let mut found_space = false;
        for c in line[prefix_end..].chars() {
            if c == ' ' {
                found_space = true;
            } else if c == ':' && found_space {
                args_end -= 1;
                break
            } else {
                found_space = false;
            }
            args_end += 1;
        }

        let prefix = match prefix_end {
            0 => None,
            _ => Some(line[1..prefix_end].to_string())
        };
        let args = &line[prefix_end..args_end].trim();
        let len = line.len();
        let suffix = if args_end == len {
            None
        } else {
            Some(line[args_end+2..].to_string())
        };

        (prefix, args, suffix)
    }

    pub fn from_str(line: &str) -> Self {
        let parts = Self::split_parts(line.trim());
        let split: Vec<&str> = parts.1.split(" ").collect();
        let mut args = Vec::new();
        for s in split[1..].iter() {
            args.push(s.to_string());
        }
        let parsed_command: Result<Command, Command> = split[0].parse();
        Message{
            prefix: parts.0,
            command: parsed_command.ok().unwrap(),
            args: args,
            suffix: parts.2
        }
    }
}

impl From<Command> for Message {
    fn from(c: Command) -> Message {
        Message {
            prefix: None,
            command: c,
            args: vec![],
            suffix: None
        }
    }
}

impl FromStr for Command {
    type Err = Command;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "NICK" => Ok(Command::Nick),
            "USER" => Ok(Command::User),
            "JOIN" => Ok(Command::Join),
            "PART" => Ok(Command::Part),
            "QUIT" => Ok(Command::Quit),
            "PING" => Ok(Command::Ping),
            "MODE" => Ok(Command::Mode),
            "PASS" => Ok(Command::Pass),
            "TOPIC" => Ok(Command::Topic),
            "PRIVMSG" => Ok(Command::Privmsg),
            _ => Ok(Command::Unknown(s.to_string()))
        }
    }
}

#[derive(Debug)]
pub struct Message {
    pub prefix: Option<String>,
    pub command: Command,
    pub args: Vec<String>,
    pub suffix: Option<String>
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn classic_irc_session() {
        let msg = Message::from_str("USER nick 0 * hostname");
        assert_eq!(msg.prefix, None);
        assert_eq!(msg.command, Command::User);
        assert_eq!(msg.args, &["nick", "0", "*", "hostname"]);
        assert_eq!(msg.suffix, None);

        let msg = Message::from_str("NICK nick");
        assert_eq!(msg.prefix, None);
        assert_eq!(msg.command, Command::Nick);
        assert_eq!(msg.args, &["nick"]);
        assert_eq!(msg.suffix, None);

        let msg = Message::from_str(":nick!nick@hostname JOIN #foo");
        assert_eq!(msg.prefix, Some("nick!nick@hostname".to_owned()));
        assert_eq!(msg.command, Command::Join);
        assert_eq!(msg.args, &["#foo"]);
        assert_eq!(msg.suffix, None);

        let msg = Message::from_str(":nick!nick@hostname PRIVMSG #foo :Hello World!");
        assert_eq!(msg.prefix, Some("nick!nick@hostname".to_owned()));
        assert_eq!(msg.command, Command::Privmsg);
        assert_eq!(msg.args, &["#foo"]);
        assert_eq!(msg.suffix, Some("Hello World!".to_owned()));

        let msg = Message::from_str(":nick!nick@hostname QUIT :Goodbye!");
        assert_eq!(msg.prefix, Some("nick!nick@hostname".to_owned()));
        assert_eq!(msg.command, Command::Quit);
        assert!(msg.args.len() == 0);
        assert_eq!(msg.suffix, Some("Goodbye!".to_owned()));
    }

    #[test]
    fn weird_cases() {
        let msg = Message::from_str("USER nick:name 0 * hostname");
        assert_eq!(msg.prefix, None);
        assert_eq!(msg.command, Command::User);
        assert_eq!(msg.args, &["nick:name", "0", "*", "hostname"]);
        assert_eq!(msg.suffix, None);

        let msg = Message::from_str("USER  nick  0  *  hostname");
        assert_eq!(msg.prefix, None);
        assert_eq!(msg.command, Command::User);
        assert_eq!(msg.args, &["", "nick", "", "0", "", "*", "", "hostname"]);
        assert_eq!(msg.suffix, None);
    }

    #[test]
    fn utf8_messages() {
        let msg = Message::from_str(":nick!nick@hostname PRIVMSG #foo :Some utf8 fun éèàåöþœðßä");
        assert_eq!(msg.prefix, Some("nick!nick@hostname".to_owned()));
        assert_eq!(msg.command, Command::Privmsg);
        assert_eq!(msg.args, &["#foo"]);
        assert_eq!(msg.suffix, Some("Some utf8 fun éèàåöþœðßä".to_owned()));

        let msg = Message::from_str(":nick!nick@hostname PRIVMSG #héhé :In a chan with utf8 in its name!");
        assert_eq!(msg.prefix, Some("nick!nick@hostname".to_owned()));
        assert_eq!(msg.command, Command::Privmsg);
        assert_eq!(msg.args, &["#héhé"]);
        assert_eq!(msg.suffix, Some("In a chan with utf8 in its name!".to_owned()));
    }
}<commit_msg>Correctly handle utf8 in irc parsing.<commit_after>/*
 * Copyright 2015-2016 Torrie Fischer <tdfischer@hackerbots.net>
 *
 * 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 std::str::FromStr;

#[derive(Debug,PartialEq,Eq)]
pub enum Command {
    Nick,
    User,
    Join,
    Part,
    Quit,
    Ping,
    Mode,
    Pong,
    Pass,
    Privmsg,
    Topic,
    Numeric(u32),
    Unknown(String)
}

impl Command {
    pub fn as_string(&self) -> String {
        match self {
            &Command::Nick => "NICK".to_string(),
            &Command::Join => "JOIN".to_string(),
            &Command::Part => "PART".to_string(),
            &Command::Pong => "PONG".to_string(),
            &Command::Privmsg => "PRIVMSG".to_string(),
            &Command::User => "USER".to_string(),
            &Command::Quit => "QUIT".to_string(),
            &Command::Ping => "PING".to_string(),
            &Command::Mode => "MODE".to_string(),
            &Command::Pass => "PASS".to_string(),
            &Command::Topic => "TOPIC".to_string(),
            &Command::Numeric(n)=> format!("{:0>3}", n),
            &Command::Unknown(ref s) => s.clone()
        }
    }
}

impl Message {
    pub fn to_string(&self) -> String {
        let mut ret = String::new();
        match self.prefix {
            Some(ref pfx) => {
                ret.push(':');
                ret.push_str(pfx.trim());
                ret.push(' ');
            },
            None => ()
        };
        ret.push_str(self.command.as_string().trim());
        for ref arg in self.args.iter() {
            ret.push(' ');
            ret.push_str(arg.trim());
        }

        match self.suffix {
            Some(ref sfx) => {
                ret.push_str(" :");
                ret.push_str(sfx.trim());
            },
            None => ()
        };

        return ret;
    }

    fn split_parts(line: &str) -> (Option<String>, &str, Option<String>) {
        let mut prefix_end = 0;
        if line.starts_with(":") {
            for c in line[1..].chars() {
                prefix_end += 1;
                if c == ' ' {
                    break
                }
            }
        }

        let mut args_end = line.len();
        let mut found_space = false;
        for (i, c) in line[prefix_end..].char_indices() {
            if c == ' ' {
                found_space = true;
            } else if c == ':' && found_space {
                args_end = prefix_end + i - 1;
                break
            } else {
                found_space = false;
            }
        }

        let prefix = match prefix_end {
            0 => None,
            _ => Some(line[1..prefix_end].to_string())
        };
        let args = &line[prefix_end..args_end].trim();
        let len = line.len();
        let suffix = if args_end == len {
            None
        } else {
            Some(line[args_end+2..].to_string())
        };

        (prefix, args, suffix)
    }

    pub fn from_str(line: &str) -> Self {
        let parts = Self::split_parts(line.trim());
        let split: Vec<&str> = parts.1.split(" ").collect();
        let mut args = Vec::new();
        for s in split[1..].iter() {
            args.push(s.to_string());
        }
        let parsed_command: Result<Command, Command> = split[0].parse();
        Message{
            prefix: parts.0,
            command: parsed_command.ok().unwrap(),
            args: args,
            suffix: parts.2
        }
    }
}

impl From<Command> for Message {
    fn from(c: Command) -> Message {
        Message {
            prefix: None,
            command: c,
            args: vec![],
            suffix: None
        }
    }
}

impl FromStr for Command {
    type Err = Command;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "NICK" => Ok(Command::Nick),
            "USER" => Ok(Command::User),
            "JOIN" => Ok(Command::Join),
            "PART" => Ok(Command::Part),
            "QUIT" => Ok(Command::Quit),
            "PING" => Ok(Command::Ping),
            "MODE" => Ok(Command::Mode),
            "PASS" => Ok(Command::Pass),
            "TOPIC" => Ok(Command::Topic),
            "PRIVMSG" => Ok(Command::Privmsg),
            _ => Ok(Command::Unknown(s.to_string()))
        }
    }
}

#[derive(Debug)]
pub struct Message {
    pub prefix: Option<String>,
    pub command: Command,
    pub args: Vec<String>,
    pub suffix: Option<String>
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn classic_irc_session() {
        let msg = Message::from_str("USER nick 0 * hostname");
        assert_eq!(msg.prefix, None);
        assert_eq!(msg.command, Command::User);
        assert_eq!(msg.args, &["nick", "0", "*", "hostname"]);
        assert_eq!(msg.suffix, None);

        let msg = Message::from_str("NICK nick");
        assert_eq!(msg.prefix, None);
        assert_eq!(msg.command, Command::Nick);
        assert_eq!(msg.args, &["nick"]);
        assert_eq!(msg.suffix, None);

        let msg = Message::from_str(":nick!nick@hostname JOIN #foo");
        assert_eq!(msg.prefix, Some("nick!nick@hostname".to_owned()));
        assert_eq!(msg.command, Command::Join);
        assert_eq!(msg.args, &["#foo"]);
        assert_eq!(msg.suffix, None);

        let msg = Message::from_str(":nick!nick@hostname PRIVMSG #foo :Hello World!");
        assert_eq!(msg.prefix, Some("nick!nick@hostname".to_owned()));
        assert_eq!(msg.command, Command::Privmsg);
        assert_eq!(msg.args, &["#foo"]);
        assert_eq!(msg.suffix, Some("Hello World!".to_owned()));

        let msg = Message::from_str(":nick!nick@hostname QUIT :Goodbye!");
        assert_eq!(msg.prefix, Some("nick!nick@hostname".to_owned()));
        assert_eq!(msg.command, Command::Quit);
        assert!(msg.args.len() == 0);
        assert_eq!(msg.suffix, Some("Goodbye!".to_owned()));
    }

    #[test]
    fn weird_cases() {
        let msg = Message::from_str("USER nick:name 0 * hostname");
        assert_eq!(msg.prefix, None);
        assert_eq!(msg.command, Command::User);
        assert_eq!(msg.args, &["nick:name", "0", "*", "hostname"]);
        assert_eq!(msg.suffix, None);

        let msg = Message::from_str("USER  nick  0  *  hostname");
        assert_eq!(msg.prefix, None);
        assert_eq!(msg.command, Command::User);
        assert_eq!(msg.args, &["", "nick", "", "0", "", "*", "", "hostname"]);
        assert_eq!(msg.suffix, None);
    }

    #[test]
    fn utf8_messages() {
        let msg = Message::from_str(":nick!nick@hostname PRIVMSG #foo :Some utf8 fun éèàåöþœðßä");
        assert_eq!(msg.prefix, Some("nick!nick@hostname".to_owned()));
        assert_eq!(msg.command, Command::Privmsg);
        assert_eq!(msg.args, &["#foo"]);
        assert_eq!(msg.suffix, Some("Some utf8 fun éèàåöþœðßä".to_owned()));

        let msg = Message::from_str(":nick!nick@hostname PRIVMSG #héhé :In a chan with utf8 in its name!");
        assert_eq!(msg.prefix, Some("nick!nick@hostname".to_owned()));
        assert_eq!(msg.command, Command::Privmsg);
        assert_eq!(msg.args, &["#héhé"]);
        assert_eq!(msg.suffix, Some("In a chan with utf8 in its name!".to_owned()));
    }
}<|endoftext|>


<commit_before>fn main() {
    // --- lifetime prevent dangling ref
    /*
    let r;

    {   // x not live long enough
        let x = 5;
        r = &x;
    }

    println!("r: {}", r);
    */

    // ---
    let string1 = String::from("abcd");
    let string2 = "xyz";

    let result = longest(string1.as_str(), string2);
    println!("Then longest string is {}", result);

    // --- how borrow checker helps after adding lifetime
    // annotations
    /*
    let s1 = String::from("long string is long");
    let result;

    {
        let s2 = String::from("xyz");
        result = longest(s1.as_str(), s2.as_str());
    }
    println!("Then longest string is {}", result);
    */
}

// compiler don't know if the result is ref to str1 or str2,
// so the borrow checker don't know enough info to check
// anything here. We programmers need to express our intentions
// mannually, to let borrow checker help us
// fn longest(str1: &str, str2: &str) -> &str {
// since scope is always nested, will be the shortest of str1 and str2
// need to make the lifetime of result relative to at least
// one of the args
fn longest<'a>(str1: &'a str, str2: &'a str) -> &'a str {
    if str1.len() > str2.len() {
        str1
    } else {
        str2
    }
}
<commit_msg>On lifetime of struct.<commit_after>#[derive(Debug)]
struct ImportantExcerpt<'a> {
    part: &'a str   // if hold ref, need lifetime anno
}

fn main() {
    // --- lifetime prevent dangling ref
    /*
    let r;

    {   // x not live long enough
        let x = 5;
        r = &x;
    }

    println!("r: {}", r);
    */

    // ---
    let string1 = String::from("abcd");
    let string2 = "xyz";

    let result = longest(string1.as_str(), string2);
    println!("Then longest string is {}", result);

    // --- how borrow checker helps after adding lifetime
    // annotations
    /*
    let s1 = String::from("long string is long");
    let result;

    {
        let s2 = String::from("xyz");
        result = longest(s1.as_str(), s2.as_str());
    }
    println!("Then longest string is {}", result);
    */

    let novel = String::from("Call me Ishmael. Some years ago...");
    let first_sentence = novel.split('.')
        .next()
        .expect("Could not find a '.'");
    let i = ImportantExcerpt {part: first_sentence};
}

// compiler don't know if the result is ref to str1 or str2,
// so the borrow checker don't know enough info to check
// anything here. We programmers need to express our intentions
// mannually, to let borrow checker help us
// fn longest(str1: &str, str2: &str) -> &str {
// since scope is always nested, will be the shortest of str1 and str2
// need to make the lifetime of result relative to at least
// one of the args
fn longest<'a>(str1: &'a str, str2: &'a str) -> &'a str {
    if str1.len() > str2.len() {
        str1
    } else {
        str2
    }
}
<|endoftext|>


<commit_before>use std::iter::repeat;

fn generate_expression(size: &usize) -> Vec<char> {
    let opens = repeat('[').take(*size);
    let closes = repeat(']').take(*size);
    let expr = opens.chain(&closes).collect();

    expr
}

fn well_formed(expr: &Vec<char>) -> bool {
    let mut count = 0;

    for b in expr {
        if count < 0 {
            break;
        }

        match *b {
            '[' => count += 1,
            _ => count -= 1,
        }
    }

    count == 0
}

fn main() {
    for n in (1..10) {
        let expr = generate_expression(&n);
        let wf = well_formed(&expr);
        let wf_str: String = expr.into_iter().collect();

        println!("'{}' is {}",
                 wf_str,
                 if wf {"well formed"} else {"NOT well formed"});
    }
}
<commit_msg>Generating vecs of n [ and n ]<commit_after>use std::iter::repeat;

fn generate_expression(size: &usize) -> Vec<char> {
    let opens = repeat('[').take(*size);
    let closes = repeat(']').take(*size);
    let expr = opens.chain(closes);

    expr.collect()
}

fn well_formed(expr: &Vec<char>) -> bool {
    let mut count = 0;

    for b in expr {
        if count < 0 {
            break;
        }

        match *b {
            '[' => count += 1,
            _ => count -= 1,
        }
    }

    count == 0
}

fn main() {
    for n in (1..10) {
        let expr = generate_expression(&n);
        let wf = well_formed(&expr);

        println!("'{}' is {}",
                 expr.into_iter().collect::<String>(),
                 if wf {"well formed"} else {"NOT well formed"});
    }
}
<|endoftext|>


<commit_before>fn generate_expression(size: &i32) -> Vec<char> {
    vec!['[', ']', ']', '[']
}

fn well_formed(expr: &Vec<char>) -> bool {
    let mut count = 0;

    for b in expr {
        if count < 0 {
            break;
        }

        match *b {
            '[' => count += 1,
            _ => count -= 1,
        }
    }

    count == 0
}

fn main() {
    for n in (1..10) {
        let expr = generate_expression(&n);
        let wf = well_formed(&expr);

        println!("{:?} is {}", expr, if wf {"well formed"} else {"NOT well formed"});
    }
}
<commit_msg>Checking for well-formedness<commit_after>fn generate_expression(size: &i32) -> Vec<char> {
    vec!['[', ']', ']', '[']
}

fn well_formed(expr: &Vec<char>) -> bool {
    let mut count = 0;

    for b in expr {
        if count < 0 {
            break;
        }

        match *b {
            '[' => count += 1,
            _ => count -= 1,
        }
    }

    count == 0
}

fn main() {
    for n in (1..10) {
        let expr = generate_expression(&n);
        let wf = well_formed(&expr);

        println!("{} is {}", expr.iter().collect(), if wf {"well formed"} else {"NOT well formed"});
    }
}
<|endoftext|>


<commit_before>use traits::*;

#[derive(PartialEq, Clone, Debug, Copy)]
pub enum BFTokenType {
    IncrementPtr(usize),      // >
    DecrementPtr(usize),      // <
    IncrementData(usize),     // +
    DecrementData(usize),     // -
    Output,                   // .
    Input,                    // ,

    // These store their corresponding tokens
    //  aka LoopStart stores the position of
    //  the matching LoopEnd and vice versa.
    LoopStart(usize),         // [
    LoopEnd(usize),           // ]

    // Optimised instructions
    SetBlock(usize,i8),        // Sets blocks of sells to value
    AddCurrentUp(usize),      // Adds current cell to cell in [current + value]
    AddCurrentDown(usize)     // Adds current cell to cell in [current - value]
}

#[derive(Clone, Debug, Copy)]
pub struct BFToken {
    pub token_type : BFTokenType,

    // Used for matching brackets
    pub pos : usize
}

pub struct BFLexer {
}

impl Lexer<Vec<BFToken>> for BFLexer {

    fn parse(input_string:String) -> LexResult<Vec<BFToken>> {
        // Import enum -> Allows for using enum values without
        //  BFTokenType:: prefix
        use self::BFTokenType::*;

        // Create empty vector
        let mut tokens : Vec<BFToken> = vec![];

        // Store the index of the token in the vector
        let mut pos : usize = 0;

        // This stores the positions of open brackets
        // This is used to match up bracket pairs
        //  so that they can jump to each other in O(1)
        let mut loop_stack : Vec<usize> = vec![];

        // Store the previous 5 tokens, these are used
        //  for optimisations.
        let mut last_tokens = [BFToken { pos : 0, token_type : Input }; 5];

        // Loop through each character
        for character in input_string.chars() {
            // Match our character to a TokenType.
            let token_type = match character {
                '>' => {
                    // Batch up all the '>' tokens
                    let mut new_x = 1;
                    if last_tokens[0].pos > 0 {
                        if let IncrementPtr(x) = last_tokens[0].token_type {
                            tokens.pop();
                            pos-=1;
                            new_x += x;
                        }
                    }
                    IncrementPtr(new_x)
                },
                '<' => {
                    let mut new_x = 1;
                    if last_tokens[0].pos > 0 {
                        if let DecrementPtr(x) = last_tokens[0].token_type {
                            tokens.pop();
                            pos-=1;
                            new_x += x;
                        }
                    }
                    DecrementPtr(new_x)
                },
                '+' =>  {
                    let mut new_x = 1;
                    if last_tokens[0].pos > 0 {
                        if let IncrementData(x) = last_tokens[0].token_type {
                            tokens.pop();
                            pos-=1;
                            new_x += x;
                        }
                    }
                    IncrementData(new_x)
                },
                '-' => {
                    let mut new_x = 1;
                    if last_tokens[0].pos > 0 {
                        if let DecrementData(x) = last_tokens[0].token_type {
                            tokens.pop();
                            pos-=1;
                            new_x += x;
                        }
                    }
                    DecrementData(new_x)
                },
                '.' => Output,
                ',' => Input,
                '[' => {
                    // Push this pos onto stack
                    loop_stack.push(pos);

                    // Store temp value of 0 for now - will be updated
                    //  once we know matching ]
                    LoopStart(0)
                },
                ']' => {

                    // Get matching bracket.
                    let index = loop_stack.pop().expect("Bracket mismatch");
                    let mut ret_token : BFTokenType = LoopEnd(index);

                    // Update it's data position to this ] token
                    tokens[index].token_type = LoopStart(pos);

                    match last_tokens[0].token_type {
                        // This currently checks for [-] or [+] and replaces
                        //  those with a set current cell to 0 instruction.
                        // It also looks for a [-]>[-] type pattern where
                        //  multiple cells are set to 0
                        IncrementData(_)
                        | DecrementData(_) => {
                            if let LoopStart(_) = last_tokens[1].token_type {
                                let mut size = 1;
                                // Check for blocks
                                if let IncrementPtr(1) = last_tokens[2].token_type {
                                    if let SetBlock(x, 0) = last_tokens[3].token_type {
                                        size = x + 1;
                                        tokens.pop();
                                        tokens.pop();
                                        pos-=2;
                                    };
                                }

                                ret_token = SetBlock(size, 0);
                                tokens.pop();
                                tokens.pop();
                                pos-=2;
                            }
                        },

                        // Optimisation for [-<+>] or [->+<] pattern.
                        // These loops will add the current cell value
                        //  to the cell value offset by the amount of
                        //  ptr increments '>' and '<'.
                        // As it adds to the offset cell, it takes from
                        //  the current cell. This will loop until the
                        //  current cell is 0.
                        // So it can be equated to addition and setting
                        //  the current value to 0
                        //
                        // Examples:
                        //
                        //  [->+<]         adds mem[current] to mem[current+1]
                        //                 sets mem[current] to 0
                        //
                        //  [->>>>+<<<<]   adds mem[current] to mem[current+4]
                        //                 sets mem[current] to 0
                        //
                        // Note that due to the combination optimisation
                        //  this only needs to check for [->+<] or [-<+>]
                        DecrementPtr(x)
                        | IncrementPtr(x) => {
                            if let IncrementData(a) = last_tokens[1].token_type {
                                match last_tokens[2].token_type {
                                    IncrementPtr(y)
                                    | DecrementPtr(y) => {
                                        if x == y {
                                            if let DecrementData(b) = last_tokens[3].token_type  {
                                                if a == b {
                                                    if let LoopStart(_) = last_tokens[4].token_type {
                                                        if last_tokens[0].token_type == DecrementPtr(x) {
                                                            ret_token = AddCurrentUp(x);
                                                        } else {
                                                            ret_token = AddCurrentDown(x);
                                                        }
                                                        tokens.pop();
                                                        tokens.pop();
                                                        tokens.pop();
                                                        tokens.pop();
                                                        tokens.pop();
                                                        pos-=5;
                                                    }
                                                }
                                            }
                                        }
                                    },
                                    _ => { }
                                }
                            };
                        },

                        // Empty Loop
                        LoopStart(_) => {
                            tokens.pop();
                            pos-=1;
                            continue;
                        },

                        _ => { }
                    }

                    ret_token
                },

                // BF is a very simple language - if there is an
                //  unrecognised character, it is ignored.
                _ => continue
            };

            let mut i = 1;
            // Recopy last tokens instead of shifting as they
            //  may have chagned.
            while i < 5 && pos > i-1 {
                last_tokens[i] = tokens[pos-i].clone();
                i+=1;
            }

            pos+=1;

            // Add it to the list of tokens
            let token = BFToken { token_type : token_type, pos : pos};

            last_tokens[0] = token.clone();

            tokens.push(token);
        }

        // No way for this to currently fail.
        LexResult::Success(tokens)
    }
}
<commit_msg>Added more comments. Fixed a typo.<commit_after>use traits::*;

#[derive(PartialEq, Clone, Debug, Copy)]
pub enum BFTokenType {
    IncrementPtr(usize),      // >
    DecrementPtr(usize),      // <
    IncrementData(usize),     // +
    DecrementData(usize),     // -
    Output,                   // .
    Input,                    // ,

    // These store their corresponding tokens
    //  aka LoopStart stores the position of
    //  the matching LoopEnd and vice versa.
    LoopStart(usize),         // [
    LoopEnd(usize),           // ]

    // Optimised instructions
    SetBlock(usize,i8),        // Sets blocks of sells to value
    AddCurrentUp(usize),      // Adds current cell to cell in [current + value]
    AddCurrentDown(usize)     // Adds current cell to cell in [current - value]
}

#[derive(Clone, Debug, Copy)]
pub struct BFToken {
    pub token_type : BFTokenType,

    // Used for matching brackets
    pub pos : usize
}

pub struct BFLexer {
}

impl Lexer<Vec<BFToken>> for BFLexer {

    fn parse(input_string:String) -> LexResult<Vec<BFToken>> {
        // Import enum -> Allows for using enum values without
        //  BFTokenType:: prefix
        use self::BFTokenType::*;

        // Create empty vector
        let mut tokens : Vec<BFToken> = vec![];

        // Store the index of the token in the vector
        let mut pos : usize = 0;

        // This stores the positions of open brackets
        // This is used to match up bracket pairs
        //  so that they can jump to each other in O(1)
        let mut loop_stack : Vec<usize> = vec![];

        // Store the previous 5 tokens, these are used
        //  for optimisations.
        let mut last_tokens = [BFToken { pos : 0, token_type : Input }; 5];

        // Loop through each character
        for character in input_string.chars() {
            // Match our character to a TokenType.
            let token_type = match character {
                '>' => {
                    // Batch up all the '>' tokens
                    let mut new_x = 1;
                    if last_tokens[0].pos > 0 {
                        if let IncrementPtr(x) = last_tokens[0].token_type {
                            tokens.pop();
                            pos-=1;
                            new_x += x;
                        }
                    }
                    IncrementPtr(new_x)
                },
                '<' => {
                    let mut new_x = 1;
                    if last_tokens[0].pos > 0 {
                        if let DecrementPtr(x) = last_tokens[0].token_type {
                            tokens.pop();
                            pos-=1;
                            new_x += x;
                        }
                    }
                    DecrementPtr(new_x)
                },
                '+' =>  {
                    let mut new_x = 1;
                    if last_tokens[0].pos > 0 {
                        if let IncrementData(x) = last_tokens[0].token_type {
                            tokens.pop();
                            pos-=1;
                            new_x += x;
                        }
                    }
                    IncrementData(new_x)
                },
                '-' => {
                    let mut new_x = 1;
                    if last_tokens[0].pos > 0 {
                        if let DecrementData(x) = last_tokens[0].token_type {
                            tokens.pop();
                            pos-=1;
                            new_x += x;
                        }
                    }
                    DecrementData(new_x)
                },
                '.' => Output,
                ',' => Input,
                '[' => {
                    // Push this pos onto stack
                    loop_stack.push(pos);

                    // Store temp value of 0 for now - will be updated
                    //  once we know matching ]
                    LoopStart(0)
                },
                ']' => {

                    // Get matching bracket.
                    let index = loop_stack.pop().expect("Bracket mismatch");
                    let mut ret_token : BFTokenType = LoopEnd(index);

                    // Update it's data position to this ] token
                    tokens[index].token_type = LoopStart(pos);

                    // Lots of optimisations can be done when you have a loop.
                    //  At this point the lexer knows all the components of the
                    //  loop and can optimise out common loop patterns into single
                    //  instruction calls.
                    match last_tokens[0].token_type {
                        // This currently checks for [-] or [+] and replaces
                        //  those with a set current cell to 0 instruction.
                        // It also looks for a [-]>[-] type pattern where
                        //  multiple cells are set to 0
                        IncrementData(_)
                        | DecrementData(_) => {
                            if let LoopStart(_) = last_tokens[1].token_type {
                                let mut size = 1;
                                // Check for blocks
                                if let IncrementPtr(1) = last_tokens[2].token_type {
                                    if let SetBlock(x, 0) = last_tokens[3].token_type {
                                        size = x + 1;
                                        tokens.pop();
                                        tokens.pop();
                                        pos-=2;
                                    };
                                }

                                ret_token = SetBlock(size, 0);
                                tokens.pop();
                                tokens.pop();
                                pos-=2;
                            }
                        },

                        // TODO: Multiplication detection & optimisation?
                        // Optimisation for [-<+>] or [->+<] pattern.
                        // These loops will add the current cell value
                        //  to the cell value offset by the amount of
                        //  ptr increments '>' and '<'.
                        // As it adds to the offset cell, it takes from
                        //  the current cell. This will loop until the
                        //  current cell is 0.
                        // So it can be equated to addition and setting
                        //  the current value to 0
                        //
                        // Examples:
                        //
                        //  [->+<]         adds mem[current] to mem[current+1]
                        //                 sets mem[current] to 0
                        //
                        //  [->>>>+<<<<]   adds mem[current] to mem[current+4]
                        //                 sets mem[current] to 0
                        //
                        // Note that due to the combination optimisation
                        //  this only needs to check for [->+<] or [-<+>]
                        DecrementPtr(x)
                        | IncrementPtr(x) => {
                            if let IncrementData(a) = last_tokens[1].token_type {
                                match last_tokens[2].token_type {
                                    IncrementPtr(y)
                                    | DecrementPtr(y) => {
                                        if x == y {
                                            if let DecrementData(b) = last_tokens[3].token_type  {
                                                if a == b {
                                                    if let LoopStart(_) = last_tokens[4].token_type {
                                                        if last_tokens[0].token_type == DecrementPtr(x) {
                                                            ret_token = AddCurrentUp(x);
                                                        } else {
                                                            ret_token = AddCurrentDown(x);
                                                        }
                                                        tokens.pop();
                                                        tokens.pop();
                                                        tokens.pop();
                                                        tokens.pop();
                                                        tokens.pop();
                                                        pos-=5;
                                                    }
                                                }
                                            }
                                        }
                                    },
                                    _ => { }
                                }
                            };
                        },

                        // Empty Loop
                        LoopStart(_) => {
                            tokens.pop();
                            pos-=1;
                            continue;
                        },

                        _ => { }
                    }

                    ret_token
                },

                // BF is a very simple language - if there is an
                //  unrecognised character, it is ignored.
                _ => continue
            };

            let mut i = 1;
            // Recopy last tokens instead of shifting as they
            //  may have changed due to optimisations.
            while i < 5 && pos > i-1 {
                last_tokens[i] = tokens[pos-i].clone();
                i+=1;
            }

            pos+=1;

            // Add it to the list of tokens
            let token = BFToken { token_type : token_type, pos : pos};

            last_tokens[0] = token.clone();

            tokens.push(token);
        }

        // No way for this to currently fail.
        LexResult::Success(tokens)
    }
}
<|endoftext|>


<commit_before>//! Loads "sysroot" crate.
//!
//! One confusing point here is that normally sysroot is a bunch of `.rlib`s,
//! but we can't process `.rlib` and need source code instead. The source code
//! is typically installed with `rustup component add rust-src` command.

use std::{convert::TryFrom, env, ops, path::PathBuf, process::Command};

use anyhow::{format_err, Result};
use arena::{Arena, Idx};
use paths::{AbsPath, AbsPathBuf};

use crate::utf8_stdout;

#[derive(Default, Debug, Clone, Eq, PartialEq)]
pub struct Sysroot {
    crates: Arena<SysrootCrateData>,
}

pub type SysrootCrate = Idx<SysrootCrateData>;

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct SysrootCrateData {
    pub name: String,
    pub root: AbsPathBuf,
    pub deps: Vec<SysrootCrate>,
}

impl ops::Index<SysrootCrate> for Sysroot {
    type Output = SysrootCrateData;
    fn index(&self, index: SysrootCrate) -> &SysrootCrateData {
        &self.crates[index]
    }
}

impl Sysroot {
    pub fn public_deps(&self) -> impl Iterator<Item = (&'static str, SysrootCrate)> + '_ {
        // core is added as a dependency before std in order to
        // mimic rustcs dependency order
        vec!["core", "alloc", "std"].into_iter().filter_map(move |it| Some((it, self.by_name(it)?)))
    }

    pub fn proc_macro(&self) -> Option<SysrootCrate> {
        self.by_name("proc_macro")
    }

    pub fn crates<'a>(&'a self) -> impl Iterator<Item = SysrootCrate> + ExactSizeIterator + 'a {
        self.crates.iter().map(|(id, _data)| id)
    }

    pub fn discover(cargo_toml: &AbsPath) -> Result<Sysroot> {
        let current_dir = cargo_toml.parent().unwrap();
        let sysroot_src_dir = discover_sysroot_src_dir(current_dir)?;
        let res = Sysroot::load(&sysroot_src_dir)?;
        Ok(res)
    }

    pub fn load(sysroot_src_dir: &AbsPath) -> Result<Sysroot> {
        let mut sysroot = Sysroot { crates: Arena::default() };

        for name in SYSROOT_CRATES.trim().lines() {
            // FIXME: first path when 1.47 comes out
            // https://github.com/rust-lang/rust/pull/73265
            let root = [format!("lib{}/lib.rs", name), format!("{}/src/lib.rs", name)]
                .iter()
                .map(|it| sysroot_src_dir.join(it))
                .find(|it| it.exists());

            if let Some(root) = root {
                sysroot.crates.alloc(SysrootCrateData {
                    name: name.into(),
                    root,
                    deps: Vec::new(),
                });
            }
        }

        if let Some(std) = sysroot.by_name("std") {
            for dep in STD_DEPS.trim().lines() {
                if let Some(dep) = sysroot.by_name(dep) {
                    sysroot.crates[std].deps.push(dep)
                }
            }
        }

        if let Some(alloc) = sysroot.by_name("alloc") {
            if let Some(core) = sysroot.by_name("core") {
                sysroot.crates[alloc].deps.push(core);
            }
        }

        if sysroot.by_name("core").is_none() {
            anyhow::bail!(
                "could not find libcore in sysroot path `{}`",
                sysroot_src_dir.as_ref().display()
            );
        }

        Ok(sysroot)
    }

    fn by_name(&self, name: &str) -> Option<SysrootCrate> {
        let (id, _data) = self.crates.iter().find(|(_id, data)| data.name == name)?;
        Some(id)
    }
}

fn discover_sysroot_src_dir(current_dir: &AbsPath) -> Result<AbsPathBuf> {
    if let Ok(path) = env::var("RUST_SRC_PATH") {
        let path = AbsPathBuf::try_from(path.as_str())
            .map_err(|path| format_err!("RUST_SRC_PATH must be absolute: {}", path.display()))?;
        return Ok(path);
    }

    let sysroot_path = {
        let mut rustc = Command::new(toolchain::rustc());
        rustc.current_dir(current_dir).args(&["--print", "sysroot"]);
        let stdout = utf8_stdout(rustc)?;
        AbsPathBuf::assert(PathBuf::from(stdout))
    };

    get_rust_src(&sysroot_path)
        .or_else(|| {
            let mut rustup = Command::new(toolchain::rustup());
            rustup.current_dir(current_dir).args(&["component", "add", "rust-src"]);
            utf8_stdout(rustup).ok()?;
            get_rust_src(&sysroot_path)
        })
        .ok_or_else(|| {
            format_err!(
                "\
can't load standard library from sysroot
{}
(discovered via `rustc --print sysroot`)
try running `rustup component add rust-src` or set `RUST_SRC_PATH`",
                sysroot_path.display(),
            )
        })
}

fn get_rust_src(sysroot_path: &AbsPath) -> Option<AbsPathBuf> {
    // Try the new path first since the old one still exists.
    //
    // FIXME: remove `src` when 1.47 comes out
    // https://github.com/rust-lang/rust/pull/73265
    let rust_src = sysroot_path.join("lib/rustlib/src/rust");
    ["library", "src"].iter().map(|it| rust_src.join(it)).find(|it| it.exists())
}

impl SysrootCrateData {
    pub fn root_dir(&self) -> &AbsPath {
        self.root.parent().unwrap()
    }
}

const SYSROOT_CRATES: &str = "
alloc
core
panic_abort
panic_unwind
proc_macro
profiler_builtins
rtstartup
std
stdarch
term
test
unwind";

const STD_DEPS: &str = "
alloc
core
panic_abort
panic_unwind
profiler_builtins
rtstartup
proc_macro
stdarch
term
test
unwind";
<commit_msg>Add note if RUST_SRC_PATH is likely to be wrong<commit_after>//! Loads "sysroot" crate.
//!
//! One confusing point here is that normally sysroot is a bunch of `.rlib`s,
//! but we can't process `.rlib` and need source code instead. The source code
//! is typically installed with `rustup component add rust-src` command.

use std::{convert::TryFrom, env, ops, path::PathBuf, process::Command};

use anyhow::{format_err, Result};
use arena::{Arena, Idx};
use paths::{AbsPath, AbsPathBuf};

use crate::utf8_stdout;

#[derive(Default, Debug, Clone, Eq, PartialEq)]
pub struct Sysroot {
    crates: Arena<SysrootCrateData>,
}

pub type SysrootCrate = Idx<SysrootCrateData>;

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct SysrootCrateData {
    pub name: String,
    pub root: AbsPathBuf,
    pub deps: Vec<SysrootCrate>,
}

impl ops::Index<SysrootCrate> for Sysroot {
    type Output = SysrootCrateData;
    fn index(&self, index: SysrootCrate) -> &SysrootCrateData {
        &self.crates[index]
    }
}

impl Sysroot {
    pub fn public_deps(&self) -> impl Iterator<Item = (&'static str, SysrootCrate)> + '_ {
        // core is added as a dependency before std in order to
        // mimic rustcs dependency order
        vec!["core", "alloc", "std"].into_iter().filter_map(move |it| Some((it, self.by_name(it)?)))
    }

    pub fn proc_macro(&self) -> Option<SysrootCrate> {
        self.by_name("proc_macro")
    }

    pub fn crates<'a>(&'a self) -> impl Iterator<Item = SysrootCrate> + ExactSizeIterator + 'a {
        self.crates.iter().map(|(id, _data)| id)
    }

    pub fn discover(cargo_toml: &AbsPath) -> Result<Sysroot> {
        let current_dir = cargo_toml.parent().unwrap();
        let sysroot_src_dir = discover_sysroot_src_dir(current_dir)?;
        let res = Sysroot::load(&sysroot_src_dir)?;
        Ok(res)
    }

    pub fn load(sysroot_src_dir: &AbsPath) -> Result<Sysroot> {
        let mut sysroot = Sysroot { crates: Arena::default() };

        for name in SYSROOT_CRATES.trim().lines() {
            // FIXME: first path when 1.47 comes out
            // https://github.com/rust-lang/rust/pull/73265
            let root = [format!("lib{}/lib.rs", name), format!("{}/src/lib.rs", name)]
                .iter()
                .map(|it| sysroot_src_dir.join(it))
                .find(|it| it.exists());

            if let Some(root) = root {
                sysroot.crates.alloc(SysrootCrateData {
                    name: name.into(),
                    root,
                    deps: Vec::new(),
                });
            }
        }

        if let Some(std) = sysroot.by_name("std") {
            for dep in STD_DEPS.trim().lines() {
                if let Some(dep) = sysroot.by_name(dep) {
                    sysroot.crates[std].deps.push(dep)
                }
            }
        }

        if let Some(alloc) = sysroot.by_name("alloc") {
            if let Some(core) = sysroot.by_name("core") {
                sysroot.crates[alloc].deps.push(core);
            }
        }

        if sysroot.by_name("core").is_none() {
            let var_note = if env::var_os("RUST_SRC_PATH").is_some() {
                " (`RUST_SRC_PATH` might be set incorrectly)"
            } else {
                ""
            };
            anyhow::bail!(
                "could not find libcore in sysroot path `{}`{}",
                sysroot_src_dir.as_ref().display(),
                var_note,
            );
        }

        Ok(sysroot)
    }

    fn by_name(&self, name: &str) -> Option<SysrootCrate> {
        let (id, _data) = self.crates.iter().find(|(_id, data)| data.name == name)?;
        Some(id)
    }
}

fn discover_sysroot_src_dir(current_dir: &AbsPath) -> Result<AbsPathBuf> {
    if let Ok(path) = env::var("RUST_SRC_PATH") {
        let path = AbsPathBuf::try_from(path.as_str())
            .map_err(|path| format_err!("RUST_SRC_PATH must be absolute: {}", path.display()))?;
        return Ok(path);
    }

    let sysroot_path = {
        let mut rustc = Command::new(toolchain::rustc());
        rustc.current_dir(current_dir).args(&["--print", "sysroot"]);
        let stdout = utf8_stdout(rustc)?;
        AbsPathBuf::assert(PathBuf::from(stdout))
    };

    get_rust_src(&sysroot_path)
        .or_else(|| {
            let mut rustup = Command::new(toolchain::rustup());
            rustup.current_dir(current_dir).args(&["component", "add", "rust-src"]);
            utf8_stdout(rustup).ok()?;
            get_rust_src(&sysroot_path)
        })
        .ok_or_else(|| {
            format_err!(
                "\
can't load standard library from sysroot
{}
(discovered via `rustc --print sysroot`)
try running `rustup component add rust-src` or set `RUST_SRC_PATH`",
                sysroot_path.display(),
            )
        })
}

fn get_rust_src(sysroot_path: &AbsPath) -> Option<AbsPathBuf> {
    // Try the new path first since the old one still exists.
    //
    // FIXME: remove `src` when 1.47 comes out
    // https://github.com/rust-lang/rust/pull/73265
    let rust_src = sysroot_path.join("lib/rustlib/src/rust");
    ["library", "src"].iter().map(|it| rust_src.join(it)).find(|it| it.exists())
}

impl SysrootCrateData {
    pub fn root_dir(&self) -> &AbsPath {
        self.root.parent().unwrap()
    }
}

const SYSROOT_CRATES: &str = "
alloc
core
panic_abort
panic_unwind
proc_macro
profiler_builtins
rtstartup
std
stdarch
term
test
unwind";

const STD_DEPS: &str = "
alloc
core
panic_abort
panic_unwind
profiler_builtins
rtstartup
proc_macro
stdarch
term
test
unwind";
<|endoftext|>


<commit_before>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::env;
use std::fmt::{self, Display};
use std::fs;
use std::ops::BitOrAssign;
use std::os::unix::io::RawFd;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;

use sys_util::{EventFd, GuestMemory, GuestMemoryError, PollContext, PollToken};
use tpm2;

use super::{DescriptorChain, Queue, VirtioDevice, INTERRUPT_STATUS_USED_RING, TYPE_TPM};

const QUEUE_SIZE: u16 = 1;
const QUEUE_SIZES: &[u16] = &[QUEUE_SIZE];

// Simply store TPM state in /tmp/tpm-simulator. Before shipping this feature,
// will need to move state under /run/vm instead. https://crbug.com/921841
const SIMULATOR_DIR: &str = "/tmp/tpm-simulator";

struct Worker {
    queue: Queue,
    mem: GuestMemory,
    interrupt_status: Arc<AtomicUsize>,
    queue_evt: EventFd,
    kill_evt: EventFd,
    interrupt_evt: EventFd,
    interrupt_resample_evt: EventFd,
    device: Device,
}

struct Device {
    simulator: tpm2::Simulator,
    response: Option<Vec<u8>>,
}

impl Device {
    fn perform_work(&mut self, mem: &GuestMemory, desc: DescriptorChain) -> Result<u32> {
        if desc.is_read_only() {
            self.perform_send(&mem, desc)
        } else if desc.is_write_only() {
            self.perform_recv(&mem, desc)
        } else {
            Err(Error::ExpectedReadOrWrite)
        }
    }

    fn perform_send(&mut self, mem: &GuestMemory, desc: DescriptorChain) -> Result<u32> {
        if self.response.is_some() {
            return Err(Error::UnexpectedSend);
        }

        let mut buf = vec![0u8; desc.len as usize];
        mem.read_exact_at_addr(&mut buf, desc.addr)
            .map_err(Error::Read)?;

        let response = self.simulator.execute_command(&buf);
        self.response = Some(response.to_owned());
        Ok(desc.len)
    }

    fn perform_recv(&mut self, mem: &GuestMemory, desc: DescriptorChain) -> Result<u32> {
        let response = self.response.take().ok_or(Error::UnexpectedRecv)?;

        if response.len() > desc.len as usize {
            return Err(Error::SmallBuffer {
                size: desc.len as usize,
                required: response.len(),
            });
        }

        mem.write_all_at_addr(&response, desc.addr)
            .map_err(Error::Write)?;

        Ok(response.len() as u32)
    }
}

impl Worker {
    fn process_queue(&mut self) -> NeedsInterrupt {
        let avail_desc = match self.queue.iter(&self.mem).next() {
            Some(avail_desc) => avail_desc,
            None => return NeedsInterrupt::No,
        };

        let index = avail_desc.index;

        let len = match self.device.perform_work(&self.mem, avail_desc) {
            Ok(len) => len,
            Err(err) => {
                error!("{}", err);
                0
            }
        };

        self.queue.add_used(&self.mem, index, len);
        NeedsInterrupt::Yes
    }

    fn signal_used_queue(&self) {
        self.interrupt_status
            .fetch_or(INTERRUPT_STATUS_USED_RING as usize, Ordering::SeqCst);
        let _ = self.interrupt_evt.write(1);
    }

    fn run(mut self) {
        #[derive(PollToken, Debug)]
        enum Token {
            // A request is ready on the queue.
            QueueAvailable,
            // Check if any interrupts need to be re-asserted.
            InterruptResample,
            // The parent thread requested an exit.
            Kill,
        }

        let poll_ctx = match PollContext::new()
            .and_then(|pc| pc.add(&self.queue_evt, Token::QueueAvailable).and(Ok(pc)))
            .and_then(|pc| {
                pc.add(&self.interrupt_resample_evt, Token::InterruptResample)
                    .and(Ok(pc))
            })
            .and_then(|pc| pc.add(&self.kill_evt, Token::Kill).and(Ok(pc)))
        {
            Ok(pc) => pc,
            Err(e) => {
                error!("vtpm failed creating PollContext: {}", e);
                return;
            }
        };

        'poll: loop {
            let events = match poll_ctx.wait() {
                Ok(v) => v,
                Err(e) => {
                    error!("vtpm failed polling for events: {}", e);
                    break;
                }
            };

            let mut needs_interrupt = NeedsInterrupt::No;
            for event in events.iter_readable() {
                match event.token() {
                    Token::QueueAvailable => {
                        if let Err(e) = self.queue_evt.read() {
                            error!("vtpm failed reading queue EventFd: {}", e);
                            break 'poll;
                        }
                        needs_interrupt |= self.process_queue();
                    }
                    Token::InterruptResample => {
                        let _ = self.interrupt_resample_evt.read();
                        if self.interrupt_status.load(Ordering::SeqCst) != 0 {
                            let _ = self.interrupt_evt.write(1);
                        }
                    }
                    Token::Kill => break 'poll,
                }
            }
            if needs_interrupt == NeedsInterrupt::Yes {
                self.signal_used_queue();
            }
        }
    }
}

/// Virtio vTPM device.
pub struct Tpm {
    kill_evt: Option<EventFd>,
}

impl Tpm {
    pub fn new() -> Tpm {
        Tpm { kill_evt: None }
    }
}

impl Drop for Tpm {
    fn drop(&mut self) {
        if let Some(kill_evt) = self.kill_evt.take() {
            let _ = kill_evt.write(1);
        }
    }
}

impl VirtioDevice for Tpm {
    fn keep_fds(&self) -> Vec<RawFd> {
        Vec::new()
    }

    fn device_type(&self) -> u32 {
        TYPE_TPM
    }

    fn queue_max_sizes(&self) -> &[u16] {
        QUEUE_SIZES
    }

    fn activate(
        &mut self,
        mem: GuestMemory,
        interrupt_evt: EventFd,
        interrupt_resample_evt: EventFd,
        interrupt_status: Arc<AtomicUsize>,
        mut queues: Vec<Queue>,
        mut queue_evts: Vec<EventFd>,
    ) {
        if queues.len() != 1 || queue_evts.len() != 1 {
            return;
        }
        let queue = queues.remove(0);
        let queue_evt = queue_evts.remove(0);

        if let Err(err) = fs::create_dir_all(SIMULATOR_DIR) {
            error!("vtpm failed to create directory for simulator: {}", err);
            return;
        }
        if let Err(err) = env::set_current_dir(SIMULATOR_DIR) {
            error!("vtpm failed to change into simulator directory: {}", err);
            return;
        }
        let simulator = tpm2::Simulator::singleton_in_current_directory();

        let (self_kill_evt, kill_evt) = match EventFd::new().and_then(|e| Ok((e.try_clone()?, e))) {
            Ok(v) => v,
            Err(err) => {
                error!("vtpm failed to create kill EventFd pair: {}", err);
                return;
            }
        };
        self.kill_evt = Some(self_kill_evt);

        let worker = Worker {
            queue,
            mem,
            interrupt_status,
            queue_evt,
            interrupt_evt,
            interrupt_resample_evt,
            kill_evt,
            device: Device {
                simulator,
                response: None,
            },
        };

        let worker_result = thread::Builder::new()
            .name("virtio_tpm".to_string())
            .spawn(|| worker.run());

        if let Err(e) = worker_result {
            error!("vtpm failed to spawn virtio_tpm worker: {}", e);
            return;
        }
    }
}

#[derive(PartialEq)]
enum NeedsInterrupt {
    Yes,
    No,
}

impl BitOrAssign for NeedsInterrupt {
    fn bitor_assign(&mut self, rhs: NeedsInterrupt) {
        if rhs == NeedsInterrupt::Yes {
            *self = NeedsInterrupt::Yes;
        }
    }
}

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

enum Error {
    ExpectedReadOrWrite,
    UnexpectedSend,
    Read(GuestMemoryError),
    UnexpectedRecv,
    SmallBuffer { size: usize, required: usize },
    Write(GuestMemoryError),
}

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        match self {
            ExpectedReadOrWrite => write!(f, "vtpm expected either read or write descriptor"),
            UnexpectedSend => write!(f, "vtpm encountered unexpected send"),
            Read(e) => write!(f, "vtpm failed to read from guest memory: {}", e),
            UnexpectedRecv => write!(f, "vtpm encountered unexpected recv"),
            SmallBuffer { size, required } => write!(
                f,
                "vtpm response buffer is too small: {} < {}",
                size, required
            ),
            Write(e) => write!(f, "vtpm failed to write to guest memory: {}", e),
        }
    }
}
<commit_msg>tpm: Handle send+recv as a single descriptor chain<commit_after>// Copyright 2018 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::env;
use std::fmt::{self, Display};
use std::fs;
use std::ops::BitOrAssign;
use std::os::unix::io::RawFd;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;

use sys_util::{EventFd, GuestMemory, GuestMemoryError, PollContext, PollToken};
use tpm2;

use super::{DescriptorChain, Queue, VirtioDevice, INTERRUPT_STATUS_USED_RING, TYPE_TPM};

// A single queue of size 2. The guest kernel driver will enqueue a single
// descriptor chain containing one command buffer and one response buffer at a
// time.
const QUEUE_SIZE: u16 = 2;
const QUEUE_SIZES: &[u16] = &[QUEUE_SIZE];

// Maximum command or response message size permitted by this device
// implementation. Named to match the equivalent constant in Linux's tpm.h.
// There is no hard requirement that the value is the same but it makes sense.
const TPM_BUFSIZE: usize = 4096;

// Simply store TPM state in /tmp/tpm-simulator. Before shipping this feature,
// will need to move state under /run/vm instead. https://crbug.com/921841
const SIMULATOR_DIR: &str = "/tmp/tpm-simulator";

struct Worker {
    queue: Queue,
    mem: GuestMemory,
    interrupt_status: Arc<AtomicUsize>,
    queue_evt: EventFd,
    kill_evt: EventFd,
    interrupt_evt: EventFd,
    interrupt_resample_evt: EventFd,
    device: Device,
}

struct Device {
    simulator: tpm2::Simulator,
}

// Checks that the input descriptor chain holds a read-only descriptor followed
// by a write-only descriptor, as required of the guest's virtio tpm driver.
//
// Returns those descriptors as a tuple: `(read_desc, write_desc)`.
fn two_input_descriptors(desc: DescriptorChain) -> Result<(DescriptorChain, DescriptorChain)> {
    let read_desc = desc;
    if !read_desc.is_read_only() {
        return Err(Error::ExpectedReadOnly);
    }

    let write_desc = match read_desc.next_descriptor() {
        Some(desc) => desc,
        None => return Err(Error::ExpectedSecondBuffer),
    };

    if !write_desc.is_write_only() {
        return Err(Error::ExpectedWriteOnly);
    }

    Ok((read_desc, write_desc))
}

impl Device {
    fn perform_work(&mut self, mem: &GuestMemory, desc: DescriptorChain) -> Result<u32> {
        let (read_desc, write_desc) = two_input_descriptors(desc)?;

        if read_desc.len > TPM_BUFSIZE as u32 {
            return Err(Error::CommandTooLong {
                size: read_desc.len as usize,
            });
        }

        let mut command = vec![0u8; read_desc.len as usize];
        mem.read_exact_at_addr(&mut command, read_desc.addr)
            .map_err(Error::Read)?;

        let response = self.simulator.execute_command(&command);

        if response.len() > TPM_BUFSIZE {
            return Err(Error::ResponseTooLong {
                size: response.len(),
            });
        }

        if response.len() > write_desc.len as usize {
            return Err(Error::BufferTooSmall {
                size: write_desc.len as usize,
                required: response.len(),
            });
        }

        mem.write_all_at_addr(&response, write_desc.addr)
            .map_err(Error::Write)?;

        Ok(response.len() as u32)
    }
}

impl Worker {
    fn process_queue(&mut self) -> NeedsInterrupt {
        let avail_desc = match self.queue.iter(&self.mem).next() {
            Some(avail_desc) => avail_desc,
            None => return NeedsInterrupt::No,
        };

        let index = avail_desc.index;

        let len = match self.device.perform_work(&self.mem, avail_desc) {
            Ok(len) => len,
            Err(err) => {
                error!("{}", err);
                0
            }
        };

        self.queue.add_used(&self.mem, index, len);
        NeedsInterrupt::Yes
    }

    fn signal_used_queue(&self) {
        self.interrupt_status
            .fetch_or(INTERRUPT_STATUS_USED_RING as usize, Ordering::SeqCst);
        let _ = self.interrupt_evt.write(1);
    }

    fn run(mut self) {
        #[derive(PollToken, Debug)]
        enum Token {
            // A request is ready on the queue.
            QueueAvailable,
            // Check if any interrupts need to be re-asserted.
            InterruptResample,
            // The parent thread requested an exit.
            Kill,
        }

        let poll_ctx = match PollContext::new()
            .and_then(|pc| pc.add(&self.queue_evt, Token::QueueAvailable).and(Ok(pc)))
            .and_then(|pc| {
                pc.add(&self.interrupt_resample_evt, Token::InterruptResample)
                    .and(Ok(pc))
            })
            .and_then(|pc| pc.add(&self.kill_evt, Token::Kill).and(Ok(pc)))
        {
            Ok(pc) => pc,
            Err(e) => {
                error!("vtpm failed creating PollContext: {}", e);
                return;
            }
        };

        'poll: loop {
            let events = match poll_ctx.wait() {
                Ok(v) => v,
                Err(e) => {
                    error!("vtpm failed polling for events: {}", e);
                    break;
                }
            };

            let mut needs_interrupt = NeedsInterrupt::No;
            for event in events.iter_readable() {
                match event.token() {
                    Token::QueueAvailable => {
                        if let Err(e) = self.queue_evt.read() {
                            error!("vtpm failed reading queue EventFd: {}", e);
                            break 'poll;
                        }
                        needs_interrupt |= self.process_queue();
                    }
                    Token::InterruptResample => {
                        let _ = self.interrupt_resample_evt.read();
                        if self.interrupt_status.load(Ordering::SeqCst) != 0 {
                            let _ = self.interrupt_evt.write(1);
                        }
                    }
                    Token::Kill => break 'poll,
                }
            }
            if needs_interrupt == NeedsInterrupt::Yes {
                self.signal_used_queue();
            }
        }
    }
}

/// Virtio vTPM device.
pub struct Tpm {
    kill_evt: Option<EventFd>,
}

impl Tpm {
    pub fn new() -> Tpm {
        Tpm { kill_evt: None }
    }
}

impl Drop for Tpm {
    fn drop(&mut self) {
        if let Some(kill_evt) = self.kill_evt.take() {
            let _ = kill_evt.write(1);
        }
    }
}

impl VirtioDevice for Tpm {
    fn keep_fds(&self) -> Vec<RawFd> {
        Vec::new()
    }

    fn device_type(&self) -> u32 {
        TYPE_TPM
    }

    fn queue_max_sizes(&self) -> &[u16] {
        QUEUE_SIZES
    }

    fn activate(
        &mut self,
        mem: GuestMemory,
        interrupt_evt: EventFd,
        interrupt_resample_evt: EventFd,
        interrupt_status: Arc<AtomicUsize>,
        mut queues: Vec<Queue>,
        mut queue_evts: Vec<EventFd>,
    ) {
        if queues.len() != 1 || queue_evts.len() != 1 {
            return;
        }
        let queue = queues.remove(0);
        let queue_evt = queue_evts.remove(0);

        if let Err(err) = fs::create_dir_all(SIMULATOR_DIR) {
            error!("vtpm failed to create directory for simulator: {}", err);
            return;
        }
        if let Err(err) = env::set_current_dir(SIMULATOR_DIR) {
            error!("vtpm failed to change into simulator directory: {}", err);
            return;
        }
        let simulator = tpm2::Simulator::singleton_in_current_directory();

        let (self_kill_evt, kill_evt) = match EventFd::new().and_then(|e| Ok((e.try_clone()?, e))) {
            Ok(v) => v,
            Err(err) => {
                error!("vtpm failed to create kill EventFd pair: {}", err);
                return;
            }
        };
        self.kill_evt = Some(self_kill_evt);

        let worker = Worker {
            queue,
            mem,
            interrupt_status,
            queue_evt,
            interrupt_evt,
            interrupt_resample_evt,
            kill_evt,
            device: Device { simulator },
        };

        let worker_result = thread::Builder::new()
            .name("virtio_tpm".to_string())
            .spawn(|| worker.run());

        if let Err(e) = worker_result {
            error!("vtpm failed to spawn virtio_tpm worker: {}", e);
            return;
        }
    }
}

#[derive(PartialEq)]
enum NeedsInterrupt {
    Yes,
    No,
}

impl BitOrAssign for NeedsInterrupt {
    fn bitor_assign(&mut self, rhs: NeedsInterrupt) {
        if rhs == NeedsInterrupt::Yes {
            *self = NeedsInterrupt::Yes;
        }
    }
}

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

enum Error {
    ExpectedReadOnly,
    ExpectedSecondBuffer,
    ExpectedWriteOnly,
    CommandTooLong { size: usize },
    Read(GuestMemoryError),
    ResponseTooLong { size: usize },
    BufferTooSmall { size: usize, required: usize },
    Write(GuestMemoryError),
}

impl Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        match self {
            ExpectedReadOnly => write!(f, "vtpm expected first descriptor to be read-only"),
            ExpectedSecondBuffer => write!(f, "vtpm expected a second descriptor"),
            ExpectedWriteOnly => write!(f, "vtpm expected second descriptor to be write-only"),
            CommandTooLong { size } => write!(
                f,
                "vtpm command is too long: {} > {} bytes",
                size, TPM_BUFSIZE
            ),
            Read(e) => write!(f, "vtpm failed to read from guest memory: {}", e),
            ResponseTooLong { size } => write!(
                f,
                "vtpm simulator generated a response that is unexpectedly long: {} > {} bytes",
                size, TPM_BUFSIZE
            ),
            BufferTooSmall { size, required } => write!(
                f,
                "vtpm response buffer is too small: {} < {} bytes",
                size, required
            ),
            Write(e) => write!(f, "vtpm failed to write to guest memory: {}", e),
        }
    }
}
<|endoftext|>


<commit_before>extern crate libc;
extern crate log;

use std::fmt;
use std::io;
use std::io::{Read, Write};
use std::slice;
use std::sync::mpsc::{channel, Sender, Receiver, RecvTimeoutError};
use std::time::Duration;

use core_foundation_sys::base::*;
use libc::c_void;

use consts::{CID_BROADCAST, HID_RPT_SIZE};
use u2ftypes::U2FDevice;

use super::iokit::*;

const READ_TIMEOUT: u64 = 15;

pub struct Device {
    device_ref: IOHIDDeviceRef,
    cid: [u8; 4],
    report_rx: Receiver<Vec<u8>>,
    report_send_void: *mut c_void,
    _scratch_buf: [u8; HID_RPT_SIZE],
}

impl Device {
    pub fn new(device_ref: IOHIDDeviceRef) -> Self {
        let _scratch_buf = [0; HID_RPT_SIZE];
        let (report_tx, report_rx) = channel();
        let boxed_report_tx = Box::new(report_tx);
        let report_send_void = Box::into_raw(boxed_report_tx) as *mut c_void;

        unsafe {
            IOHIDDeviceRegisterInputReportCallback(
                device_ref,
                _scratch_buf.as_ptr(),
                _scratch_buf.len() as CFIndex,
                read_new_data_cb,
                report_send_void,
            );
        }

        Self {
            device_ref,
            cid: CID_BROADCAST,
            report_rx,
            report_send_void,
            _scratch_buf,
        }
    }
}

impl Drop for Device {
    fn drop(&mut self) {
        debug!("Dropping U2F device {}", self);
        // Re-allocate this raw pointer for destruction
        let _ = unsafe { Box::from_raw(self.report_send_void) };
    }
}

impl fmt::Display for Device {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "InternalDevice(ref:{:?}, cid: {:02x}{:02x}{:02x}{:02x})",
            self.device_ref,
            self.cid[0],
            self.cid[1],
            self.cid[2],
            self.cid[3]
        )
    }
}

impl PartialEq for Device {
    fn eq(&self, other_device: &Device) -> bool {
        self.device_ref == other_device.device_ref
    }
}

impl Read for Device {
    fn read(&mut self, mut bytes: &mut [u8]) -> io::Result<usize> {
        let timeout = Duration::from_secs(READ_TIMEOUT);
        let data = match self.report_rx.recv_timeout(timeout) {
            Ok(v) => v,
            Err(e) if e == RecvTimeoutError::Timeout => {
                return Err(io::Error::new(io::ErrorKind::TimedOut, e));
            }
            Err(e) => {
                return Err(io::Error::new(io::ErrorKind::UnexpectedEof, e));
            }
        };
        bytes.write(&data)
    }
}

impl Write for Device {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        assert!(bytes.len() == HID_RPT_SIZE + 1);

        let report_id = bytes[0] as i64;
        // Skip report number when not using numbered reports.
        let start = if report_id == 0x0 { 1 } else { 0 };
        let data = &bytes[start..];

        let result = unsafe {
            IOHIDDeviceSetReport(
                self.device_ref,
                kIOHIDReportTypeOutput,
                report_id,
                data.as_ptr(),
                data.len() as CFIndex,
            )
        };
        if result != 0 {
            warn!("set_report sending failure = {0:X}", result);
            return Err(io::Error::from_raw_os_error(result));
        }
        trace!("set_report sending success = {0:X}", result);

        Ok(bytes.len())
    }

    // USB HID writes don't buffer, so this will be a nop.
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl U2FDevice for Device {
    fn get_cid<'a>(&'a self) -> &'a [u8; 4] {
        &self.cid
    }

    fn set_cid(&mut self, cid: [u8; 4]) {
        self.cid = cid;
    }
}

// This is called from the RunLoop thread
extern "C" fn read_new_data_cb(
    context: *mut c_void,
    _: IOReturn,
    _: *mut c_void,
    report_type: IOHIDReportType,
    report_id: u32,
    report: *mut u8,
    report_len: CFIndex,
) {
    unsafe {
        let tx = &mut *(context as *mut Sender<Vec<u8>>);

        trace!(
            "read_new_data_cb type={} id={} report={:?} len={}",
            report_type,
            report_id,
            report,
            report_len
        );

        let report_len = report_len as usize;
        if report_len > HID_RPT_SIZE {
            warn!(
                "read_new_data_cb got too much data! {} > {}",
                report_len,
                HID_RPT_SIZE
            );
            return;
        }

        let data = slice::from_raw_parts(report, report_len).to_vec();

        if let Err(e) = tx.send(data) {
            // TOOD: This happens when the channel closes before this thread
            // does. This is pretty common, but let's deal with stopping
            // properly later.
            warn!("Problem returning read_new_data_cb data for thread: {}", e);
        };
    }
}
<commit_msg>macOS: fix UAF for real by boxing Device.scratch_buf[]<commit_after>extern crate libc;
extern crate log;

use std::fmt;
use std::io;
use std::io::{Read, Write};
use std::slice;
use std::sync::mpsc::{channel, Sender, Receiver, RecvTimeoutError};
use std::time::Duration;

use core_foundation_sys::base::*;
use libc::c_void;

use consts::{CID_BROADCAST, HID_RPT_SIZE};
use u2ftypes::U2FDevice;

use super::iokit::*;

const READ_TIMEOUT: u64 = 15;

pub struct Device {
    device_ref: IOHIDDeviceRef,
    cid: [u8; 4],
    report_rx: Receiver<Vec<u8>>,
    report_send_void: *mut c_void,
    scratch_buf_ptr: *mut u8,
}

impl Device {
    pub fn new(device_ref: IOHIDDeviceRef) -> Self {
        let (report_tx, report_rx) = channel();
        let report_send_void = Box::into_raw(Box::new(report_tx)) as *mut c_void;

        let scratch_buf = [0; HID_RPT_SIZE];
        let scratch_buf_ptr = Box::into_raw(Box::new(scratch_buf)) as *mut u8;

        unsafe {
            IOHIDDeviceRegisterInputReportCallback(
                device_ref,
                scratch_buf_ptr,
                HID_RPT_SIZE as CFIndex,
                read_new_data_cb,
                report_send_void,
            );
        }

        Self {
            device_ref,
            cid: CID_BROADCAST,
            report_rx,
            report_send_void,
            scratch_buf_ptr,
        }
    }
}

impl Drop for Device {
    fn drop(&mut self) {
        debug!("Dropping U2F device {}", self);

        unsafe {
            // Re-allocate raw pointers for destruction.
            let _ = Box::from_raw(self.report_send_void as *mut Sender<Vec<u8>>);
            let _ = Box::from_raw(self.scratch_buf_ptr as *mut [u8; HID_RPT_SIZE]);
        }
    }
}

impl fmt::Display for Device {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "InternalDevice(ref:{:?}, cid: {:02x}{:02x}{:02x}{:02x})",
            self.device_ref,
            self.cid[0],
            self.cid[1],
            self.cid[2],
            self.cid[3]
        )
    }
}

impl PartialEq for Device {
    fn eq(&self, other_device: &Device) -> bool {
        self.device_ref == other_device.device_ref
    }
}

impl Read for Device {
    fn read(&mut self, mut bytes: &mut [u8]) -> io::Result<usize> {
        let timeout = Duration::from_secs(READ_TIMEOUT);
        let data = match self.report_rx.recv_timeout(timeout) {
            Ok(v) => v,
            Err(e) if e == RecvTimeoutError::Timeout => {
                return Err(io::Error::new(io::ErrorKind::TimedOut, e));
            }
            Err(e) => {
                return Err(io::Error::new(io::ErrorKind::UnexpectedEof, e));
            }
        };
        bytes.write(&data)
    }
}

impl Write for Device {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        assert!(bytes.len() == HID_RPT_SIZE + 1);

        let report_id = bytes[0] as i64;
        // Skip report number when not using numbered reports.
        let start = if report_id == 0x0 { 1 } else { 0 };
        let data = &bytes[start..];

        let result = unsafe {
            IOHIDDeviceSetReport(
                self.device_ref,
                kIOHIDReportTypeOutput,
                report_id,
                data.as_ptr(),
                data.len() as CFIndex,
            )
        };
        if result != 0 {
            warn!("set_report sending failure = {0:X}", result);
            return Err(io::Error::from_raw_os_error(result));
        }
        trace!("set_report sending success = {0:X}", result);

        Ok(bytes.len())
    }

    // USB HID writes don't buffer, so this will be a nop.
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

impl U2FDevice for Device {
    fn get_cid<'a>(&'a self) -> &'a [u8; 4] {
        &self.cid
    }

    fn set_cid(&mut self, cid: [u8; 4]) {
        self.cid = cid;
    }
}

// This is called from the RunLoop thread
extern "C" fn read_new_data_cb(
    context: *mut c_void,
    _: IOReturn,
    _: *mut c_void,
    report_type: IOHIDReportType,
    report_id: u32,
    report: *mut u8,
    report_len: CFIndex,
) {
    unsafe {
        let tx = &mut *(context as *mut Sender<Vec<u8>>);

        trace!(
            "read_new_data_cb type={} id={} report={:?} len={}",
            report_type,
            report_id,
            report,
            report_len
        );

        let report_len = report_len as usize;
        if report_len > HID_RPT_SIZE {
            warn!(
                "read_new_data_cb got too much data! {} > {}",
                report_len,
                HID_RPT_SIZE
            );
            return;
        }

        let data = slice::from_raw_parts(report, report_len).to_vec();

        if let Err(e) = tx.send(data) {
            // TOOD: This happens when the channel closes before this thread
            // does. This is pretty common, but let's deal with stopping
            // properly later.
            warn!("Problem returning read_new_data_cb data for thread: {}", e);
        };
    }
}
<|endoftext|>


<commit_before>#![feature(duration)]
#![deny(unused_mut)]
extern crate zookeeper;
extern crate env_logger;

use zookeeper::{Acl, CreateMode, Watcher, WatchedEvent, ZooKeeper};
use zookeeper::perms;

use std::io::{BufRead, BufReader, Write};
use std::process::{Child, Command, Stdio};
use std::time::Duration;

struct LoggingWatcher;
impl Watcher for LoggingWatcher {
    fn handle(&self, e: &WatchedEvent) {
        println!("{:?}", e)
    }
}

fn start_zk() -> Child {
    match Command::new("java")
            .arg("-jar")
            .arg("zk-test-cluster/target/main.jar")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .spawn() {
        Ok(p) => p,
        Err(e) => panic!("failed to execute process: {}", e),
    }
}

fn get_connect_string(server: &mut Child) -> String {
    let mut reader = BufReader::new(server.stdout.as_mut().unwrap());

    let mut connect_string = String::new();
    if reader.read_line(&mut connect_string).is_err() {
        panic!("Couldn't read ZK connect_string")
    }
    connect_string.pop(); // remove '\n'
    connect_string
}

fn shutdown(server: &mut Child) {
    server.stdin.as_mut().unwrap().write(b"q").unwrap();
    assert!(server.wait().unwrap().success());
}

#[test]
fn simple_integration_test() {
    env_logger::init().unwrap();

    // Create a test cluster and obtain its connection string
    let mut server = start_zk();
    let connect_string = get_connect_string(&mut server);

    // Connect to the test cluster
    let client = ZooKeeper::connect(connect_string.as_ref(), Duration::from_secs(5), LoggingWatcher).unwrap();

    // Do the tests
    let acl1 = vec![Acl{perms: perms::ALL, scheme: "world".to_string(), id: "anyone".to_string()}];
    let create = client.create("/test", vec![8,8], acl1, CreateMode::Ephemeral);

    assert_eq!(create.ok(), Some("/test".to_string()));

    let exists = client.exists("/test", true);

    assert!(exists.is_ok());

    let children = client.get_children("/", true);

    assert!(children.is_ok());

    let mut sorted_children = children.unwrap();
    sorted_children.sort();

    assert_eq!(sorted_children, vec!["test".to_string(), "zookeeper".to_string()]);

    // After closing the client all operations return Err
    client.close();

    let exists = client.exists("/test", true);

    assert!(exists.is_err());

    // Close the server
    shutdown(&mut server);
}
<commit_msg>Use predefined Acl in integration test<commit_after>#![feature(duration)]
#![deny(unused_mut)]
extern crate zookeeper;
extern crate env_logger;

use zookeeper::{CreateMode, Watcher, WatchedEvent, ZooKeeper};
use zookeeper::acls;

use std::io::{BufRead, BufReader, Write};
use std::process::{Child, Command, Stdio};
use std::time::Duration;

struct LoggingWatcher;
impl Watcher for LoggingWatcher {
    fn handle(&self, e: &WatchedEvent) {
        println!("{:?}", e)
    }
}

fn start_zk() -> Child {
    match Command::new("java")
            .arg("-jar")
            .arg("zk-test-cluster/target/main.jar")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .spawn() {
        Ok(p) => p,
        Err(e) => panic!("failed to execute process: {}", e),
    }
}

fn get_connect_string(server: &mut Child) -> String {
    let mut reader = BufReader::new(server.stdout.as_mut().unwrap());

    let mut connect_string = String::new();
    if reader.read_line(&mut connect_string).is_err() {
        panic!("Couldn't read ZK connect_string")
    }
    connect_string.pop(); // remove '\n'
    connect_string
}

fn shutdown(server: &mut Child) {
    server.stdin.as_mut().unwrap().write(b"q").unwrap();
    assert!(server.wait().unwrap().success());
}

#[test]
fn simple_integration_test() {
    env_logger::init().unwrap();

    // Create a test cluster and obtain its connection string
    let mut server = start_zk();
    let connect_string = get_connect_string(&mut server);

    // Connect to the test cluster
    let client = ZooKeeper::connect(connect_string.as_ref(), Duration::from_secs(5), LoggingWatcher).unwrap();

    // Do the tests
    let create = client.create("/test", vec![8,8], acls::OPEN_ACL_UNSAFE.clone(), CreateMode::Ephemeral);

    assert_eq!(create.ok(), Some("/test".to_string()));

    let exists = client.exists("/test", true);

    assert!(exists.is_ok());

    let children = client.get_children("/", true);

    assert!(children.is_ok());

    let mut sorted_children = children.unwrap();
    sorted_children.sort();

    assert_eq!(sorted_children, vec!["test".to_string(), "zookeeper".to_string()]);

    // After closing the client all operations return Err
    client.close();

    let exists = client.exists("/test", true);

    assert!(exists.is_err());

    // Close the server
    shutdown(&mut server);
}
<|endoftext|>