<commit_before>//
// gash.rs
//
// Starting code for PS2
// Running on Rust 0.9
//
// University of Virginia - cs4414 Spring 2014
// Weilin Xu, David Evans
// Version 0.4
//
/*
Vikram Bhasin vb8nd and Justin Ingram jci5kb
*/

extern mod extra;

use std::{io, run, os};
use std::io::buffered::BufferedReader;
use std::io::stdin;
use std::os::getcwd;
use extra::getopts;
use std::Path;

struct Shell {
    cmd_prompt: ~str,
    cwd: Path,
}

impl Shell {
    fn new(prompt_str: &str) -> Shell {
        Shell {
            cmd_prompt: prompt_str.to_owned(),
            cwd: getcwd(),
        }
    }

    fn run(&mut self) {
        let mut stdin = BufferedReader::new(stdin());

        loop {
            print(self.cmd_prompt);
            io::stdio::flush();

            let line = stdin.read_line().unwrap();
            let cmd_line = line.trim().to_owned();
            let program = cmd_line.splitn(' ', 1).nth(0).expect("no program");

            match program {
                ""      =>  { continue; }
                "exit"  =>  { return; }
                "cd"	=>  { self.changeDir(cmd_line); }
                _       =>  { self.run_cmdline(cmd_line); }
            }
        }
    }

    fn run_cmdline(&mut self, cmd_line: &str) {
        let mut argv: ~[~str] =
            cmd_line.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();

        if argv.len() > 0 {
            let program: ~str = argv.remove(0);
            self.run_cmd(program, argv);
        }
    }

    fn run_cmd(&mut self, program: &str, argv: &[~str]) {
        if self.cmd_exists(program) {
            run::process_status(program, argv);
        } else {
            println!("{:s}: command not found", program);
        }
    }

    fn cmd_exists(&mut self, cmd_path: &str) -> bool {
        let ret = run::process_output("which", [cmd_path.to_owned()]);
        return ret.expect("exit code error.").status.success();
    }

    //Justin's function for cd
    fn changeDir(&mut self, cmd_line: &str) {
	let mut argv: ~[~str] =
            cmd_line.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();

        if argv.len() > 1 {
            let pstring: ~str = argv.remove(1);

            //println!("You want to change cwd to {:s}", pstring);

            let mut npath = Path::new(pstring);

            let mut cpath = self.cwd.clone();
            cpath.push(npath);

            if cpath.exists() {
		self.cwd = cpath;
            } else {
		println("Path does not exist!");
            }

        }
        match self.cwd.as_str() {
	    Some(path_str) => {println!("{:s}", path_str); }
	    None	=> {println("Path not representable as string!"); }
        }
    }
}

fn get_cmdline_from_args() -> Option<~str> {
    /* Begin processing program arguments and initiate the parameters. */
    let args = os::args();

    let opts = ~[
        getopts::optopt("c")
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m }
        Err(f) => { fail!(f.to_err_msg()) }
    };

    if matches.opt_present("c") {
        let cmd_str = match matches.opt_str("c") {
                                                Some(cmd_str) => {cmd_str.to_owned()},
                                                None => {~""}
                                              };
        return Some(cmd_str);
    } else {
        return None;
    }
}

fn main() {
    let opt_cmd_line = get_cmdline_from_args();

    match opt_cmd_line {
        Some(cmd_line) => Shell::new("").run_cmdline(cmd_line),
        None           => Shell::new("gash > ").run()
    }
}
<commit_msg>All warnings gone<commit_after>//
// gash.rs
//
// Starting code for PS2
// Running on Rust 0.9
//
// University of Virginia - cs4414 Spring 2014
// Weilin Xu, David Evans
// Version 0.4
//
/*
Vikram Bhasin vb8nd and Justin Ingram jci5kb
*/

extern mod extra;

use std::{io, run, os};
use std::io::buffered::BufferedReader;
use std::io::stdin;
use std::os::getcwd;
use extra::getopts;

struct Shell {
    cmd_prompt: ~str,
    cwd: Path,
}

impl Shell {
    fn new(prompt_str: &str) -> Shell {
        Shell {
            cmd_prompt: prompt_str.to_owned(),
            cwd: getcwd(),
        }
    }

    fn run(&mut self) {
        let mut stdin = BufferedReader::new(stdin());

        loop {
            print(self.cmd_prompt);
            io::stdio::flush();

            let line = stdin.read_line().unwrap();
            let cmd_line = line.trim().to_owned();
            let program = cmd_line.splitn(' ', 1).nth(0).expect("no program");

            match program {
                ""      =>  { continue; }
                "exit"  =>  { return; }
                "cd"	=>  { self.changeDir(cmd_line); }
                _       =>  { self.run_cmdline(cmd_line); }
            }
        }
    }

    fn run_cmdline(&mut self, cmd_line: &str) {
        let mut argv: ~[~str] =
            cmd_line.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();

        if argv.len() > 0 {
            let program: ~str = argv.remove(0);
            self.run_cmd(program, argv);
        }
    }

    fn run_cmd(&mut self, program: &str, argv: &[~str]) {
        if self.cmd_exists(program) {
            run::process_status(program, argv);
        } else {
            println!("{:s}: command not found", program);
        }
    }

    fn cmd_exists(&mut self, cmd_path: &str) -> bool {
        let ret = run::process_output("which", [cmd_path.to_owned()]);
        return ret.expect("exit code error.").status.success();
    }

    //Justin's function for cd
    fn changeDir(&mut self, cmd_line: &str) {
	let mut argv: ~[~str] =
            cmd_line.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();

        if argv.len() > 1 {
            let pstring: ~str = argv.remove(1);

            //println!("You want to change cwd to {:s}", pstring);

            let npath = Path::new(pstring);

            let mut cpath = self.cwd.clone();
            cpath.push(npath);

            if cpath.exists() {
		self.cwd = cpath;
            } else {
		println("Path does not exist!");
            }

        }
        match self.cwd.as_str() {
	    Some(path_str) => {println!("{:s}", path_str); }
	    None	=> {println("Path not representable as string!"); }
        }
    }
}

fn get_cmdline_from_args() -> Option<~str> {
    /* Begin processing program arguments and initiate the parameters. */
    let args = os::args();

    let opts = ~[
        getopts::optopt("c")
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m }
        Err(f) => { fail!(f.to_err_msg()) }
    };

    if matches.opt_present("c") {
        let cmd_str = match matches.opt_str("c") {
                                                Some(cmd_str) => {cmd_str.to_owned()},
                                                None => {~""}
                                              };
        return Some(cmd_str);
    } else {
        return None;
    }
}

fn main() {
    let opt_cmd_line = get_cmdline_from_args();

    match opt_cmd_line {
        Some(cmd_line) => Shell::new("").run_cmdline(cmd_line),
        None           => Shell::new("gash > ").run()
    }
}
<|endoftext|>


<commit_before>
use protos::{LogFileEntry, StudyQuestion, StudyQuestions, StudyQuestionAnswer};
use std::error::Error;
use std::path::PathBuf;
extern crate regex;
use self::regex::Regex;

pub struct UserStudyQuestions {
    pub replay_filename: String,
    pub answer_lines: Vec<String>,
}
impl UserStudyQuestions {
    #[allow(unused_assignments)]
    pub fn get_questions(&mut self) -> Result<Option<StudyQuestions>, Box<Error>> {
        use scaii_core;
        use std::fs::File;
        use std::io::prelude::*;
        use std::io::BufReader;
        println!("Looking for question file for replay {}", self.replay_filename);
        let mut questionfile_path = scaii_core::get_default_replay_dir()?;
        let question_filename_root = get_question_filename_root_for_replay(&self.replay_filename);
        println!("question_filename root is {}", question_filename_root);
        let question_filename = get_filename_for_root(question_filename_root)?;
        let mut user_id = "".to_string();
        let mut treatment_id = "".to_string();
        let mut answer_filename = "".to_string();
        match question_filename  {
            None => { return Ok(None); },
            Some(ref filename)  => {
                user_id = get_field_from_filename(filename, 2 as u32);
                treatment_id = get_field_from_filename(filename, 3 as u32);
                answer_filename = filename.replace("question", "answer");
                let result = str::replace("Hello World!", "!", "?");
println!("{}", result); // => "Hello World!"
                println!("user_id {} ", user_id);
                println!("treatment_id {}", treatment_id);
            },
        }
        let q_file = question_filename.clone();
        questionfile_path.push(&question_filename.unwrap().as_str());
        println!("path is {:?}", questionfile_path);

        if check_file(questionfile_path.clone()).len() != 0 {
            panic!();
        }

        let f = File::open(questionfile_path)?;
        let reader = BufReader::new(f);

        let mut questions_vec: Vec<StudyQuestion> = Vec::new();
        for line in reader.lines() {
            let question_string = line.unwrap();
            println!("STUDY QUESTION: {}", question_string);
            let question: StudyQuestion = self.get_study_question(&question_string)?;
            questions_vec.push(question);
        }
        let study_questions = StudyQuestions {
            user_id: user_id,
            treatment_id: treatment_id,
            study_questions: questions_vec,
            answer_filename: answer_filename,
        };
        Ok(Some(study_questions))
    }

    fn get_study_question(&mut self, line : &String) -> Result<StudyQuestion, Box<Error>> {
        use super::ReplayError;
        let line_parts_iterator = line.split(";");
        let vec: Vec<&str> = line_parts_iterator.collect();

        let step: String = vec[0].to_string();

        let question_index: String = vec[1].to_string();

        let question_type: String = vec[2].to_string();

        let question : String = vec[3].to_string();

        let mut answer_vec : Vec<String> = Vec::new();
        for x in 4..vec.len() {
            answer_vec.push(vec[x].to_string());
        }

        Ok(StudyQuestion {
            step: step,
            question: question,
            answers: answer_vec,
            question_id_for_step: question_index,
            question_type: question_type,
        })
    }

    pub fn persist_log_entry_incremental(&mut self, lfe: LogFileEntry) -> Result<(), Box<Error>> {
        use std::fs::{File, OpenOptions,remove_file};
        use std::io::Write;
        use scaii_core;
        println!("{:?}", &lfe);
        let output_line = lfe.entry;

        let mut answerfile_path = scaii_core::get_default_replay_dir()?;
        let filename = lfe.filename;
        answerfile_path.push(filename);

        if !answerfile_path.exists(){
            File::create(&answerfile_path)?;
        }
        let mut file = OpenOptions::new().append(true).open(answerfile_path).unwrap();
        file.write_all(output_line.as_bytes())?;
        file.write_all("\n".to_string().as_bytes())?;
        Ok(())
    }
    pub fn persist_log_entry(&mut self, lfe: LogFileEntry) -> Result<(), Box<Error>> {
        use std::fs::{File, OpenOptions,remove_file};
        use std::io::Write;
        use scaii_core;
        println!("{:?}", &lfe);
        let output_line = lfe.entry;
        self.answer_lines.push(output_line);

        if lfe.is_last_line {
            let mut answerfile_path = scaii_core::get_default_replay_dir()?;
            // append to file called replayX_answers_userID_treatmentID.txt
            // or ??? replayX_answers_userID_sessionID_treatmentID.txt
            let filename = lfe.filename;
            answerfile_path.push(filename);
            if answerfile_path.exists(){
                remove_file(&answerfile_path)?;
            }
            println!("answerfile_path is {:?}", answerfile_path);
            File::create(&answerfile_path)?;
            let mut file = OpenOptions::new().append(true).open(answerfile_path).unwrap();
            for line in &self.answer_lines {
                file.write_all(line.as_bytes())?;
                file.write_all("\n".to_string().as_bytes())?;
                //writeln!(file, format!("{}",line))?;
            }
        }

        Ok(())
    }

    pub fn persist_study_question_answer(&mut self, sqa: StudyQuestionAnswer) -> Result<(), Box<Error>> {
        use std::fs::{File, OpenOptions,remove_file};
        use std::io::Write;
        use scaii_core;
        println!("{:?}", &sqa);
        let user_id = sqa.user_id;
        let treatment_id = sqa.treatment_id;
        let answer = sqa.answer;
        let question = sqa.question;
        let question_number = sqa.question_number;
        let step = sqa.step;
        let output_line = format!("{};{};{};{}", question_number, step, question, answer);
        self.answer_lines.push(output_line);

        if step == "summary".to_string() {
            let mut answerfile_path = scaii_core::get_default_replay_dir()?;
            // append to file called replayX_answers_userID_treatmentID.txt
            // or ??? replayX_answers_userID_sessionID_treatmentID.txt
            let filename = get_answer_filename(&self.replay_filename, &user_id, &treatment_id);
            answerfile_path.push(filename);
            if answerfile_path.exists(){
                remove_file(&answerfile_path)?;
            }
            println!("answerfile_path is {:?}", answerfile_path);
            File::create(&answerfile_path)?;
            let mut file = OpenOptions::new().append(true).open(answerfile_path).unwrap();
            for line in &self.answer_lines {
                file.write_all(line.as_bytes())?;
                file.write_all("\n".to_string().as_bytes())?;
                //writeln!(file, format!("{}",line))?;
            }
        }

        Ok(())
    }
}

fn get_field_from_filename(filename : &String, index : u32) -> String {
    let filename_parts_iterator = filename.split(".");
    let vec: Vec<&str> = filename_parts_iterator.collect();
    let root = vec[0];
    let root_parts_iterator = root.split("_");
    let vec:Vec<&str> = root_parts_iterator.collect();
    vec[index as usize].to_string()
}

fn get_filename_for_root(root: String) -> Result<Option<String>, Box<Error>>  {
    use scaii_core;
    use std::fs;
    let questionfile_path = scaii_core::get_default_replay_dir()?;
    for entry in fs::read_dir(questionfile_path)? {
        let entry = entry?;
        let path = entry.path();
        if !path.is_dir() {
            let filename = path.file_name().unwrap().to_str().unwrap().to_string();
            if filename.starts_with(root.as_str()) {
                println!("filename found was {}", filename);
                return Ok(Some(filename));
            }
        }
    }
    Ok(None)
}

fn get_question_filename_root_for_replay(filename : &String) -> String {
    let filename_parts_iterator = filename.split(".");
    let vec: Vec<&str> = filename_parts_iterator.collect();
    format!("{}_questions", vec[0])
}

// replayX_answers_userID_treatmentID.txt
fn get_answer_filename(filename : &String, user_id: &String, treatment_id: &String) -> String {
    let filename_parts_iterator = filename.split(".");
    let vec: Vec<&str> = filename_parts_iterator.collect();
    format!("{}_answers_{}_{}.txt", vec[0], user_id, treatment_id)
}

fn check_file(file: PathBuf) -> Vec<u8> {
    use std::io::{BufRead, BufReader};
    use std::fs::File;
    let f = File::open(&file).expect("Unable to open file");
    let f = BufReader::new(f);

    let mut errors: Vec<u8> = Vec::new();
    let mut line_num: u8 = 0;

    for line in f.lines() {
        line_num = line_num + 1;
        let q_string = line.unwrap();

        if check_comments_or_nl(&q_string, &line_num, &file.to_str().unwrap()) {
            let temp = check_question(&q_string, &line_num, &file.to_str().unwrap());
            if temp != 0 {
                errors.push(temp);
            }
        }
    }
    errors
}

#[allow(dead_code)]
fn check_comments_or_nl(line: &String, line_num: &u8, fname: &str) -> bool {
    let line_arr: Vec<char> = line.chars().collect();
    if line_arr.len() == 0 {      // Check for newline
        // println!("LINE\t{:?}\tWARNING: Newline/Empty line", line_num);
        return false;
    }

    if line_arr[0] == '#' {         // Check for comments
        //println!("LINE\t{:?}\t{:?}", line_num, line);
        return false;
    }
    true
}

fn check_question(line: &String, line_num: &u8, fname: &str) ->u8{
    let line_copy = line.clone();
    let line_parts_iterator = line_copy.split(";");
    let vec: Vec<&str> = line_parts_iterator.collect();

    let field0: String = vec[0].to_string();

    let mut temp:u8 = 0;

    if line.chars().filter(|&c| c == ';').count() < 3 {
        println!("\nERROR (1): Missing \';\' delimiter.\n--> {:?}:{:?}\n\t|\n Line {:?} | {:?}\n\t|\n", fname, line_num, line_num, line);
        return 1;
    }

    temp = check_regex(02, "^[0-9]*|^summary$".to_string(), "Invalid entry for step number in Field 0.".to_string(), &field0, &line, &line_num, &fname);
    if temp != 0 { return temp; }

    let field1: String = vec[1].to_string();
    temp = check_regex(03, "^[0-9]*".to_string(), "Invalid entry for question index Field 1.".to_string(), &field1, &line, &line_num, &fname);

    let field2: String = vec[2].to_string();
    let f2vec: Vec<&str> = field2.split(":").collect();
    if temp != 0 { return temp; }

    if f2vec.len() > 0 {
        if f2vec[0] == "plain" {
            temp = check_regex(04, "plain:NA:NA".to_string(), "Invalid entry for question type in Field 2. When question type is specified as plain, subfields 1 and 2 must be 'NA'. ".to_string(), &field2, &line, &line_num, &fname);
            if temp != 0 { return temp; }
        }else{
            temp = check_regex(05, "^waitForClick".to_string(), "Invalid entry for question type in Field 2. Question type must be [ plain | waitForClick ]".to_string(), &f2vec[0].to_string(), &line, &line_num, &fname);
            if temp != 0 { return temp; }
            if f2vec.len() > 2  && !f2vec.contains(&""){
                temp = check_regex(06, "^(gameboard_?|rewardBar_?|saliencyMap_?){0,3}$".to_string(), "Invalid entry for question type in Field 2, subfield 1. Valid entries are {gameboard, rewardBar, saliencyMap}, deliminated by _.".to_string(), &f2vec[1].to_string(), &line, &line_num, &fname);
                if temp != 0 { return temp; }
            }else{
                temp = check_regex(07, "$^".to_string(), "Invalid entry for question type in Field 2. Question type waitForClick cannot have blank subfields.".to_string(), &field2, &line, &line_num, &fname);
                if temp != 0 { return temp; }
            }
        }
    }else {
        temp = check_regex(08, "$^".to_string(), "Invalid entry for question type in Field 2.".to_string(), &field2, &line, &line_num, &fname);
        if temp != 0 { return temp; }
    }

    return 0;
}

/*
check_regex args:
1. error code (0 means no error)
2. regex expression, if string does not pass it then this is an error
3. message explaining what the error is
4. reference to the substring to check
5. the full parent string
6. line number the substring is from
7. the filename the substring is from
*/
fn check_regex(error_code: u8, regex: String, msg: String, str_to_check: &String, full_line: &String, line_num: &u8, fname: &str) -> u8 {
    let step_reg = Regex::new(&regex).unwrap();
    if !step_reg.is_match(&str_to_check) {
        println!("\nERROR ({}): {}\n\t--> {}:{}\n\t|\n Line {} | {}\n\t|\n\tInvalid substring: {:?}\n", error_code, msg, fname, line_num, line_num, full_line, str_to_check);
        return error_code;
    }
    0
}
<commit_msg>fixed issue with comments in question file<commit_after>
use protos::{LogFileEntry, StudyQuestion, StudyQuestions, StudyQuestionAnswer};
use std::error::Error;
use std::path::PathBuf;
extern crate regex;
use self::regex::Regex;

pub struct UserStudyQuestions {
    pub replay_filename: String,
    pub answer_lines: Vec<String>,
}
impl UserStudyQuestions {
    #[allow(unused_assignments)]
    pub fn get_questions(&mut self) -> Result<Option<StudyQuestions>, Box<Error>> {
        use scaii_core;
        use std::fs::File;
        use std::io::prelude::*;
        use std::io::BufReader;
        println!("Looking for question file for replay {}", self.replay_filename);
        let mut questionfile_path = scaii_core::get_default_replay_dir()?;
        let question_filename_root = get_question_filename_root_for_replay(&self.replay_filename);
        println!("question_filename root is {}", question_filename_root);
        let question_filename = get_filename_for_root(question_filename_root)?;
        let mut user_id = "".to_string();
        let mut treatment_id = "".to_string();
        let mut answer_filename = "".to_string();
        match question_filename  {
            None => { return Ok(None); },
            Some(ref filename)  => {
                user_id = get_field_from_filename(filename, 2 as u32);
                treatment_id = get_field_from_filename(filename, 3 as u32);
                answer_filename = filename.replace("question", "answer");
                let result = str::replace("Hello World!", "!", "?");
println!("{}", result); // => "Hello World!"
                println!("user_id {} ", user_id);
                println!("treatment_id {}", treatment_id);
            },
        }
        let q_file = question_filename.clone();
        questionfile_path.push(&question_filename.unwrap().as_str());
        println!("path is {:?}", questionfile_path);

        if check_file(questionfile_path.clone()).len() != 0 {
            panic!();
        }

        let f = File::open(questionfile_path)?;
        let reader = BufReader::new(f);

        let mut questions_vec: Vec<StudyQuestion> = Vec::new();
        for line in reader.lines() {
            let question_string = line.unwrap();
            if check_comments_or_nl(&question_string) {
                println!("STUDY QUESTION: {}", question_string);
                let question: StudyQuestion = self.get_study_question(&question_string)?;
                questions_vec.push(question);
            }
        }
        let study_questions = StudyQuestions {
            user_id: user_id,
            treatment_id: treatment_id,
            study_questions: questions_vec,
            answer_filename: answer_filename,
        };
        Ok(Some(study_questions))
    }

    fn get_study_question(&mut self, line : &String) -> Result<StudyQuestion, Box<Error>> {
        use super::ReplayError;
        let line_parts_iterator = line.split(";");
        let vec: Vec<&str> = line_parts_iterator.collect();

        let step: String = vec[0].to_string();

        let question_index: String = vec[1].to_string();

        let question_type: String = vec[2].to_string();

        let question : String = vec[3].to_string();

        let mut answer_vec : Vec<String> = Vec::new();
        for x in 4..vec.len() {
            answer_vec.push(vec[x].to_string());
        }

        Ok(StudyQuestion {
            step: step,
            question: question,
            answers: answer_vec,
            question_id_for_step: question_index,
            question_type: question_type,
        })
    }

    pub fn persist_log_entry_incremental(&mut self, lfe: LogFileEntry) -> Result<(), Box<Error>> {
        use std::fs::{File, OpenOptions,remove_file};
        use std::io::Write;
        use scaii_core;
        println!("{:?}", &lfe);
        let output_line = lfe.entry;

        let mut answerfile_path = scaii_core::get_default_replay_dir()?;
        let filename = lfe.filename;
        answerfile_path.push(filename);

        if !answerfile_path.exists(){
            File::create(&answerfile_path)?;
        }
        let mut file = OpenOptions::new().append(true).open(answerfile_path).unwrap();
        file.write_all(output_line.as_bytes())?;
        file.write_all("\n".to_string().as_bytes())?;
        Ok(())
    }
    pub fn persist_log_entry(&mut self, lfe: LogFileEntry) -> Result<(), Box<Error>> {
        use std::fs::{File, OpenOptions,remove_file};
        use std::io::Write;
        use scaii_core;
        println!("{:?}", &lfe);
        let output_line = lfe.entry;
        self.answer_lines.push(output_line);

        if lfe.is_last_line {
            let mut answerfile_path = scaii_core::get_default_replay_dir()?;
            // append to file called replayX_answers_userID_treatmentID.txt
            // or ??? replayX_answers_userID_sessionID_treatmentID.txt
            let filename = lfe.filename;
            answerfile_path.push(filename);
            if answerfile_path.exists(){
                remove_file(&answerfile_path)?;
            }
            println!("answerfile_path is {:?}", answerfile_path);
            File::create(&answerfile_path)?;
            let mut file = OpenOptions::new().append(true).open(answerfile_path).unwrap();
            for line in &self.answer_lines {
                file.write_all(line.as_bytes())?;
                file.write_all("\n".to_string().as_bytes())?;
                //writeln!(file, format!("{}",line))?;
            }
        }

        Ok(())
    }

    pub fn persist_study_question_answer(&mut self, sqa: StudyQuestionAnswer) -> Result<(), Box<Error>> {
        use std::fs::{File, OpenOptions,remove_file};
        use std::io::Write;
        use scaii_core;
        println!("{:?}", &sqa);
        let user_id = sqa.user_id;
        let treatment_id = sqa.treatment_id;
        let answer = sqa.answer;
        let question = sqa.question;
        let question_number = sqa.question_number;
        let step = sqa.step;
        let output_line = format!("{};{};{};{}", question_number, step, question, answer);
        self.answer_lines.push(output_line);

        if step == "summary".to_string() {
            let mut answerfile_path = scaii_core::get_default_replay_dir()?;
            // append to file called replayX_answers_userID_treatmentID.txt
            // or ??? replayX_answers_userID_sessionID_treatmentID.txt
            let filename = get_answer_filename(&self.replay_filename, &user_id, &treatment_id);
            answerfile_path.push(filename);
            if answerfile_path.exists(){
                remove_file(&answerfile_path)?;
            }
            println!("answerfile_path is {:?}", answerfile_path);
            File::create(&answerfile_path)?;
            let mut file = OpenOptions::new().append(true).open(answerfile_path).unwrap();
            for line in &self.answer_lines {
                file.write_all(line.as_bytes())?;
                file.write_all("\n".to_string().as_bytes())?;
                //writeln!(file, format!("{}",line))?;
            }
        }

        Ok(())
    }
}

fn get_field_from_filename(filename : &String, index : u32) -> String {
    let filename_parts_iterator = filename.split(".");
    let vec: Vec<&str> = filename_parts_iterator.collect();
    let root = vec[0];
    let root_parts_iterator = root.split("_");
    let vec:Vec<&str> = root_parts_iterator.collect();
    vec[index as usize].to_string()
}

fn get_filename_for_root(root: String) -> Result<Option<String>, Box<Error>>  {
    use scaii_core;
    use std::fs;
    let questionfile_path = scaii_core::get_default_replay_dir()?;
    for entry in fs::read_dir(questionfile_path)? {
        let entry = entry?;
        let path = entry.path();
        if !path.is_dir() {
            let filename = path.file_name().unwrap().to_str().unwrap().to_string();
            if filename.starts_with(root.as_str()) {
                println!("filename found was {}", filename);
                return Ok(Some(filename));
            }
        }
    }
    Ok(None)
}

fn get_question_filename_root_for_replay(filename : &String) -> String {
    let filename_parts_iterator = filename.split(".");
    let vec: Vec<&str> = filename_parts_iterator.collect();
    format!("{}_questions", vec[0])
}

// replayX_answers_userID_treatmentID.txt
fn get_answer_filename(filename : &String, user_id: &String, treatment_id: &String) -> String {
    let filename_parts_iterator = filename.split(".");
    let vec: Vec<&str> = filename_parts_iterator.collect();
    format!("{}_answers_{}_{}.txt", vec[0], user_id, treatment_id)
}

fn check_file(file: PathBuf) -> Vec<u8> {
    use std::io::{BufRead, BufReader};
    use std::fs::File;
    let f = File::open(&file).expect("Unable to open file");
    let f = BufReader::new(f);

    let mut errors: Vec<u8> = Vec::new();
    let mut line_num: u8 = 0;

    for line in f.lines() {
        line_num = line_num + 1;
        let q_string = line.unwrap();

        if check_comments_or_nl(&q_string) {
            let temp = check_question(&q_string, &line_num, &file.to_str().unwrap());
            if temp != 0 {
                errors.push(temp);
            }
        }
    }
    errors
}

#[allow(dead_code)]
fn check_comments_or_nl(line: &String) -> bool {
    let line_arr: Vec<char> = line.chars().collect();
    if line_arr.len() == 0 {      // Check for newline
        return false;
    }

    if line_arr[0] == '#' {         // Check for comments
        return false;
    }
    true
}

fn check_question(line: &String, line_num: &u8, fname: &str) ->u8{
    let line_copy = line.clone();
    let line_parts_iterator = line_copy.split(";");
    let vec: Vec<&str> = line_parts_iterator.collect();

    let field0: String = vec[0].to_string();

    let mut temp:u8 = 0;

    if line.chars().filter(|&c| c == ';').count() < 3 {
        println!("\nERROR [ID: 1]: Missing \';\' delimiter.\n\t--> {:?}:{:?}\n\t|\n Line {:?} | {:?}\n\t|\n", fname, line_num, line_num, line);
        return 1;
    }

    temp = check_regex(02, "^[0-9]*|^summary$".to_string(), "Invalid entry for step number in Field 0.".to_string(), &field0, &line, &line_num, &fname);
    if temp != 0 { return temp; }

    let field1: String = vec[1].to_string();
    temp = check_regex(03, "^[0-9]*".to_string(), "Invalid entry for question index Field 1.".to_string(), &field1, &line, &line_num, &fname);
    if temp != 0 { return temp; }

    let field2: String = vec[2].to_string();
    let f2vec: Vec<&str> = field2.split(":").collect();

    if f2vec.len() > 0 {
        if f2vec[0] == "plain" {
            temp = check_regex(04, "plain:NA:NA".to_string(), "Invalid entry for question type in Field 2. When question type is specified as [plain], subfields 1 and 2 must be 'NA'. ".to_string(), &field2, &line, &line_num, &fname);
            if temp != 0 { return temp; }
        }else{
            temp = check_regex(05, "^waitForClick".to_string(), "Invalid entry for question type in Field 2. Expected [plain] or [waitForClick]".to_string(), &f2vec[0].to_string(), &line, &line_num, &fname);
            if temp != 0 { return temp; }
            if f2vec.len() > 2  && !f2vec.contains(&""){
                temp = check_regex(06, "^(gameboard_?|rewardBar_?|saliencyMap_?){0,3}$".to_string(), "Invalid entry for question type in Field 2, subfield 1. Valid entries are {gameboard, rewardBar, saliencyMap}, deliminated by _.".to_string(), &f2vec[1].to_string(), &line, &line_num, &fname);
                if temp != 0 { return temp; }
            }else{
                temp = check_regex(07, "$^".to_string(), "Invalid entry for question type in Field 2. Question type [waitForClick] cannot have blank subfields.".to_string(), &field2, &line, &line_num, &fname);
                if temp != 0 { return temp; }
            }
        }
    }else {
        temp = check_regex(08, "$^".to_string(), "Invalid entry for question type in Field 2. Expected [plain] or [waitForClick]".to_string(), &field2, &line, &line_num, &fname);
        if temp != 0 { return temp; }
    }

    return 0;
}

/*
check_regex args:
1. error code (0 means no error)
2. regex expression, if string does not pass it then this is an error
3. message explaining what the error is
4. reference to the substring to check
5. the full parent string
6. line number the substring is from
7. the filename the substring is from
*/
fn check_regex(error_code: u8, regex: String, msg: String, str_to_check: &String, full_line: &String, line_num: &u8, fname: &str) -> u8 {
    let step_reg = Regex::new(&regex).unwrap();
    if !step_reg.is_match(&str_to_check) {
        println!("\nERROR [ID: {}]: {}\n\t--> {}:{}\n\t|\n Line {} | {}\n\t|\n\tInvalid substring: {:?}\n", error_code, msg, fname, line_num, line_num, full_line, str_to_check);
        return error_code;
    }
    0
}
<|endoftext|>


<commit_before>//! FIXME: write short doc here

use std::{
    env,
    path::{Path, PathBuf},
    process::Command,
};

use ra_arena::{impl_arena_id, Arena, RawId};

use crate::Result;

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

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SysrootCrate(RawId);
impl_arena_id!(SysrootCrate);

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

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

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

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

    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: &Path) -> Result<Sysroot> {
        let src = try_find_src_path(cargo_toml)?;

        if !src.exists() {
            Err(format!(
                "can't load standard library from sysroot\n\
                 {:?}\n\
                 try running `rustup component add rust-src` or set `RUST_SRC_PATH`",
                src,
            ))?;
        }

        let mut sysroot = Sysroot { crates: Arena::default() };
        for name in SYSROOT_CRATES.trim().lines() {
            let root = src.join(format!("lib{}", name)).join("lib.rs");
            if root.exists() {
                sysroot.crates.alloc(SysrootCrateData {
                    name: name.into(),
                    root,
                    deps: Vec::new(),
                });
            }
        }
        if let Some(std) = sysroot.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.core() {
                sysroot.crates[alloc].deps.push(core);
            }
        }
        Ok(sysroot)
    }

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

fn try_find_src_path(cargo_toml: &Path) -> Result<PathBuf> {
    if let Ok(path) = env::var("RUST_SRC_PATH") {
        return Ok(path.into());
    }

    let rustc_output = Command::new("rustc")
        .current_dir(cargo_toml.parent().unwrap())
        .args(&["--print", "sysroot"])
        .output()?;
    if !rustc_output.status.success() {
        Err("failed to locate sysroot")?;
    }
    let stdout = String::from_utf8(rustc_output.stdout)?;
    let sysroot_path = Path::new(stdout.trim());
    Ok(sysroot_path.join("lib/rustlib/src/rust/src"))
}

impl SysrootCrate {
    pub fn name(self, sysroot: &Sysroot) -> &str {
        &sysroot.crates[self].name
    }
    pub fn root(self, sysroot: &Sysroot) -> &Path {
        sysroot.crates[self].root.as_path()
    }
    pub fn root_dir(self, sysroot: &Sysroot) -> &Path {
        self.root(sysroot).parent().unwrap()
    }
    pub fn deps<'a>(self, sysroot: &'a Sysroot) -> impl Iterator<Item = SysrootCrate> + 'a {
        sysroot.crates[self].deps.iter().copied()
    }
}

const SYSROOT_CRATES: &str = "
std
core
alloc
collections
libc
panic_unwind
proc_macro
rustc_unicode
std_unicode
test
alloc_jemalloc
alloc_system
compiler_builtins
getopts
panic_unwind
panic_abort
rand
term
unwind
build_helper
rustc_asan
rustc_lsan
rustc_msan
rustc_tsan
syntax";

const STD_DEPS: &str = "
alloc
alloc_jemalloc
alloc_system
core
panic_abort
rand
compiler_builtins
unwind
rustc_asan
rustc_lsan
rustc_msan
rustc_tsan
build_helper";
<commit_msg>Better error message<commit_after>//! FIXME: write short doc here

use std::{
    env,
    path::{Path, PathBuf},
    process::Command,
};

use ra_arena::{impl_arena_id, Arena, RawId};

use crate::Result;

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

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SysrootCrate(RawId);
impl_arena_id!(SysrootCrate);

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

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

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

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

    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: &Path) -> Result<Sysroot> {
        let src = try_find_src_path(cargo_toml)?;

        if !src.exists() {
            Err(format!(
                "can't load standard library from sysroot\n\
                 {}\n\
                 (discovered via `rustc --print sysroot`)\n\
                 try running `rustup component add rust-src` or set `RUST_SRC_PATH`",
                src.display(),
            ))?;
        }

        let mut sysroot = Sysroot { crates: Arena::default() };
        for name in SYSROOT_CRATES.trim().lines() {
            let root = src.join(format!("lib{}", name)).join("lib.rs");
            if root.exists() {
                sysroot.crates.alloc(SysrootCrateData {
                    name: name.into(),
                    root,
                    deps: Vec::new(),
                });
            }
        }
        if let Some(std) = sysroot.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.core() {
                sysroot.crates[alloc].deps.push(core);
            }
        }
        Ok(sysroot)
    }

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

fn try_find_src_path(cargo_toml: &Path) -> Result<PathBuf> {
    if let Ok(path) = env::var("RUST_SRC_PATH") {
        return Ok(path.into());
    }

    let rustc_output = Command::new("rustc")
        .current_dir(cargo_toml.parent().unwrap())
        .args(&["--print", "sysroot"])
        .output()?;
    if !rustc_output.status.success() {
        Err("failed to locate sysroot")?;
    }
    let stdout = String::from_utf8(rustc_output.stdout)?;
    let sysroot_path = Path::new(stdout.trim());
    Ok(sysroot_path.join("lib/rustlib/src/rust/src"))
}

impl SysrootCrate {
    pub fn name(self, sysroot: &Sysroot) -> &str {
        &sysroot.crates[self].name
    }
    pub fn root(self, sysroot: &Sysroot) -> &Path {
        sysroot.crates[self].root.as_path()
    }
    pub fn root_dir(self, sysroot: &Sysroot) -> &Path {
        self.root(sysroot).parent().unwrap()
    }
    pub fn deps<'a>(self, sysroot: &'a Sysroot) -> impl Iterator<Item = SysrootCrate> + 'a {
        sysroot.crates[self].deps.iter().copied()
    }
}

const SYSROOT_CRATES: &str = "
std
core
alloc
collections
libc
panic_unwind
proc_macro
rustc_unicode
std_unicode
test
alloc_jemalloc
alloc_system
compiler_builtins
getopts
panic_unwind
panic_abort
rand
term
unwind
build_helper
rustc_asan
rustc_lsan
rustc_msan
rustc_tsan
syntax";

const STD_DEPS: &str = "
alloc
alloc_jemalloc
alloc_system
core
panic_abort
rand
compiler_builtins
unwind
rustc_asan
rustc_lsan
rustc_msan
rustc_tsan
build_helper";
<|endoftext|>


<commit_before>mod block_range;

use crate::ServerData;
use block_range::AsyncBlockRange;
use futures::{
    channel::mpsc::{self, Sender},
    prelude::*,
};
use godcoin::{get_epoch_time, net::*, prelude::*};
use log::{debug, error, info, warn};
use std::{
    io::Cursor,
    net::SocketAddr,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    time::Duration,
};
use tokio::{net::TcpStream, time};
use tokio_tungstenite::tungstenite::{protocol, Message as WsMessage};

pub struct WsClient {
    filter: Option<BlockFilter>,
    addr: SocketAddr,
    tx: Sender<WsMessage>,
    needs_pong: Arc<AtomicBool>,
}

impl WsClient {
    #[inline]
    pub fn new(addr: SocketAddr, tx: Sender<WsMessage>) -> Self {
        Self {
            filter: None,
            addr,
            tx,
            needs_pong: Arc::new(AtomicBool::new(false)),
        }
    }

    #[inline]
    pub fn needs_pong(&self) -> Arc<AtomicBool> {
        Arc::clone(&self.needs_pong)
    }

    #[inline]
    pub fn set_needs_pong(&self, flag: bool) {
        self.needs_pong.store(flag, Ordering::Release);
    }

    #[inline]
    pub fn addr(&self) -> SocketAddr {
        self.addr
    }

    #[inline]
    pub fn filter(&self) -> Option<&BlockFilter> {
        self.filter.as_ref()
    }

    #[inline]
    pub fn sender(&self) -> Sender<WsMessage> {
        self.tx.clone()
    }
}

pub fn handle_new_client(stream: TcpStream, peer_addr: SocketAddr, data: Arc<ServerData>) {
    let config = Some(protocol::WebSocketConfig {
        // # of protocol Message's
        max_send_queue: Some(16),
        // 64 MiB
        max_message_size: Some(64 << 20),
        // 16 MiB
        max_frame_size: Some(16 << 20),
    });
    tokio::spawn(async move {
        let ws_stream = match tokio_tungstenite::accept_async_with_config(stream, config).await {
            Ok(ws) => ws,
            Err(e) => {
                error!("[{}] WebSocket accept error: {:?}", peer_addr, e);
                return;
            }
        };
        info!("[{}] Connection opened", peer_addr);

        let (tx, rx) = mpsc::channel(32);
        let (sink, mut stream) = ws_stream.split();
        let mut state = WsClient::new(peer_addr, tx.clone());
        let needs_pong = state.needs_pong();

        let ws_reader = {
            let data = Arc::clone(&data);
            let mut tx = tx.clone();
            async move {
                while let Some(msg) = stream.next().await {
                    let res = process_ws_msg(&data, &mut state, msg.unwrap());
                    if let Some(res) = res {
                        if let Err(e) = tx.send(res).await {
                            warn!("[{}] Failed to send message: {:?}", peer_addr, e);
                        }
                    }
                }
            }
        };

        let ws_writer = rx.map(|msg| Ok(msg)).forward(sink.sink_map_err(move |e| {
            warn!("[{}] Sink send error: {:?}", peer_addr, e);
        }));

        let heartbeat_interval = async move {
            let dur = Duration::from_secs(20);
            let mut interval = time::interval_at(time::Instant::now() + dur, dur);
            loop {
                interval.tick().await;
                if needs_pong.swap(true, Ordering::AcqRel) {
                    debug!("[{}] Did not receive pong in time from peer", peer_addr);
                    break;
                }

                let nonce = get_epoch_time();
                let msg = Msg {
                    id: u32::max_value(),
                    body: Body::Ping(nonce),
                };
                debug!("[{}] Sending ping: {}", peer_addr, nonce);

                let mut buf = Vec::with_capacity(16);
                msg.serialize(&mut buf);

                if tx.clone().send(WsMessage::Binary(buf)).await.is_err() {
                    break;
                }
            }
        };

        tokio::select! {
            _ = ws_reader => {},
            _ = ws_writer => {},
            _ = heartbeat_interval => {},
        };

        info!("[{}] Connection closed", peer_addr);
        // Remove block subscriptions if there are any
        data.sub_pool.remove(peer_addr);
    });
}

pub fn process_ws_msg(
    data: &ServerData,
    state: &mut WsClient,
    msg: WsMessage,
) -> Option<WsMessage> {
    match msg {
        WsMessage::Binary(buf) => {
            state.set_needs_pong(false);

            let mut cur = Cursor::<&[u8]>::new(&buf);
            let msg = match Msg::deserialize(&mut cur) {
                Ok(msg) => {
                    let id = msg.id;
                    if cur.position() != buf.len() as u64 {
                        Msg {
                            id,
                            body: Body::Error(ErrorKind::BytesRemaining),
                        }
                    } else {
                        match handle_protocol_msg(data, state, msg) {
                            Some(body) => Msg { id, body },
                            None => return None,
                        }
                    }
                }
                Err(e) => {
                    error!(
                        "[{}] Error occurred during deserialization: {:?}",
                        state.addr(),
                        e
                    );
                    Msg {
                        id: u32::max_value(),
                        body: Body::Error(ErrorKind::Io),
                    }
                }
            };

            let mut buf = Vec::with_capacity(65536);
            msg.serialize(&mut buf);
            Some(WsMessage::Binary(buf))
        }
        WsMessage::Text(_) => Some(WsMessage::Close(Some(protocol::CloseFrame {
            code: protocol::frame::coding::CloseCode::Unsupported,
            reason: "text is not supported".into(),
        }))),
        _ => None,
    }
}

fn handle_protocol_msg(data: &ServerData, state: &mut WsClient, msg: Msg) -> Option<Body> {
    match msg.body {
        Body::Error(e) => {
            warn!(
                "[{}] Received error message from client: {:?}",
                state.addr(),
                e
            );
            None
        }
        Body::Request(req) => handle_rpc_request(data, state, msg.id, req),
        Body::Response(res) => {
            warn!(
                "[{}] Unexpected response from client: {:?}",
                state.addr(),
                res
            );
            None
        }
        Body::Ping(nonce) => {
            debug!("[{}] Received ping: {}", state.addr(), nonce);
            Some(Body::Pong(nonce))
        }
        Body::Pong(nonce) => {
            debug!("[{}] Received pong: {}", state.addr(), nonce);
            // We don't need to update the `needs_pong` state as it has already been updated when the message was
            // deserialized
            None
        }
    }
}

fn handle_rpc_request(
    data: &ServerData,
    state: &mut WsClient,
    id: u32,
    req: rpc::Request,
) -> Option<Body> {
    Some(match req {
        rpc::Request::Broadcast(tx) => {
            let res = data.minter.push_tx(tx);
            match res {
                Ok(_) => Body::Response(rpc::Response::Broadcast),
                Err(e) => Body::Error(ErrorKind::TxValidation(e)),
            }
        }
        rpc::Request::SetBlockFilter(filter) => {
            if filter.len() > 16 {
                return Some(Body::Error(ErrorKind::InvalidRequest));
            }
            state.filter = Some(filter);
            Body::Response(rpc::Response::SetBlockFilter)
        }
        rpc::Request::ClearBlockFilter => {
            state.filter = None;
            Body::Response(rpc::Response::ClearBlockFilter)
        }
        rpc::Request::Subscribe => {
            data.sub_pool.insert(state.addr(), state.sender());
            Body::Response(rpc::Response::Subscribe)
        }
        rpc::Request::Unsubscribe => {
            data.sub_pool.remove(state.addr());
            Body::Response(rpc::Response::Unsubscribe)
        }
        rpc::Request::GetProperties => {
            let props = data.chain.get_properties();
            Body::Response(rpc::Response::GetProperties(props))
        }
        rpc::Request::GetBlock(height) => match &state.filter {
            Some(filter) => match data.chain.get_filtered_block(height, filter) {
                Some(block) => Body::Response(rpc::Response::GetBlock(block)),
                None => Body::Error(ErrorKind::InvalidHeight),
            },
            None => match data.chain.get_block(height) {
                Some(block) => Body::Response(rpc::Response::GetBlock(FilteredBlock::Block(block))),
                None => Body::Error(ErrorKind::InvalidHeight),
            },
        },
        rpc::Request::GetFullBlock(height) => match data.chain.get_block(height) {
            Some(block) => Body::Response(rpc::Response::GetFullBlock(block)),
            None => Body::Error(ErrorKind::InvalidHeight),
        },
        rpc::Request::GetBlockRange(min_height, max_height) => {
            let range = AsyncBlockRange::try_new(Arc::clone(&data.chain), min_height, max_height);
            match range {
                Some(mut range) => {
                    if let Some(filter) = state.filter() {
                        range.set_filter(Some(filter.clone()));
                    }

                    let peer_addr = state.addr();
                    let mut tx = state.sender();
                    tokio::spawn(async move {
                        while let Some(block) = range.next().await {
                            let ws_msg = {
                                let msg = Msg {
                                    id,
                                    body: Body::Response(rpc::Response::GetBlock(block)),
                                };

                                let mut buf = Vec::with_capacity(65536);
                                msg.serialize(&mut buf);
                                WsMessage::Binary(buf)
                            };
                            if tx.send(ws_msg).await.is_err() {
                                warn!("[{}] Failed to send block range update", peer_addr);
                                return;
                            }
                        }

                        let ws_msg = {
                            let msg = Msg {
                                id,
                                body: Body::Response(rpc::Response::GetBlockRange),
                            };

                            let mut buf = Vec::with_capacity(32);
                            msg.serialize(&mut buf);
                            WsMessage::Binary(buf)
                        };
                        if tx.send(ws_msg).await.is_err() {
                            warn!("[{}] Failed to send block range finalizer", peer_addr);
                        }
                    });

                    return None;
                }
                None => Body::Error(ErrorKind::InvalidHeight),
            }
        }
        rpc::Request::GetAddressInfo(addr) => {
            let res = data.minter.get_addr_info(&addr);
            match res {
                Ok(info) => Body::Response(rpc::Response::GetAddressInfo(info)),
                Err(e) => Body::Error(ErrorKind::TxValidation(e)),
            }
        }
    })
}
<commit_msg>Remove unwrap when reading WS messages<commit_after>mod block_range;

use crate::ServerData;
use block_range::AsyncBlockRange;
use futures::{
    channel::mpsc::{self, Sender},
    prelude::*,
};
use godcoin::{get_epoch_time, net::*, prelude::*};
use log::{debug, error, info, warn};
use std::{
    io::Cursor,
    net::SocketAddr,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    time::Duration,
};
use tokio::{net::TcpStream, time};
use tokio_tungstenite::tungstenite::{protocol, Message as WsMessage};

pub struct WsClient {
    filter: Option<BlockFilter>,
    addr: SocketAddr,
    tx: Sender<WsMessage>,
    needs_pong: Arc<AtomicBool>,
}

impl WsClient {
    #[inline]
    pub fn new(addr: SocketAddr, tx: Sender<WsMessage>) -> Self {
        Self {
            filter: None,
            addr,
            tx,
            needs_pong: Arc::new(AtomicBool::new(false)),
        }
    }

    #[inline]
    pub fn needs_pong(&self) -> Arc<AtomicBool> {
        Arc::clone(&self.needs_pong)
    }

    #[inline]
    pub fn set_needs_pong(&self, flag: bool) {
        self.needs_pong.store(flag, Ordering::Release);
    }

    #[inline]
    pub fn addr(&self) -> SocketAddr {
        self.addr
    }

    #[inline]
    pub fn filter(&self) -> Option<&BlockFilter> {
        self.filter.as_ref()
    }

    #[inline]
    pub fn sender(&self) -> Sender<WsMessage> {
        self.tx.clone()
    }
}

pub fn handle_new_client(stream: TcpStream, peer_addr: SocketAddr, data: Arc<ServerData>) {
    let config = Some(protocol::WebSocketConfig {
        // # of protocol Message's
        max_send_queue: Some(16),
        // 64 MiB
        max_message_size: Some(64 << 20),
        // 16 MiB
        max_frame_size: Some(16 << 20),
    });
    tokio::spawn(async move {
        let ws_stream = match tokio_tungstenite::accept_async_with_config(stream, config).await {
            Ok(ws) => ws,
            Err(e) => {
                error!("[{}] WebSocket accept error: {:?}", peer_addr, e);
                return;
            }
        };
        info!("[{}] Connection opened", peer_addr);

        let (tx, rx) = mpsc::channel(32);
        let (sink, mut stream) = ws_stream.split();
        let mut state = WsClient::new(peer_addr, tx.clone());
        let needs_pong = state.needs_pong();

        let ws_reader = {
            let data = Arc::clone(&data);
            let mut tx = tx.clone();
            async move {
                while let Some(msg) = stream.next().await {
                    match msg {
                        Ok(msg) => {
                            let res = process_ws_msg(&data, &mut state, msg);
                            if let Some(res) = res {
                                if let Err(e) = tx.send(res).await {
                                    warn!("[{}] Failed to send message: {:?}", peer_addr, e);
                                }
                            }
                        }
                        Err(e) => {
                            warn!("[{}] Error reading WS message: {:?}", peer_addr, e);
                            break;
                        }
                    }
                }
            }
        };

        let ws_writer = rx.map(|msg| Ok(msg)).forward(sink.sink_map_err(move |e| {
            warn!("[{}] Sink send error: {:?}", peer_addr, e);
        }));

        let heartbeat_interval = async move {
            let dur = Duration::from_secs(20);
            let mut interval = time::interval_at(time::Instant::now() + dur, dur);
            loop {
                interval.tick().await;
                if needs_pong.swap(true, Ordering::AcqRel) {
                    debug!("[{}] Did not receive pong in time from peer", peer_addr);
                    break;
                }

                let nonce = get_epoch_time();
                let msg = Msg {
                    id: u32::max_value(),
                    body: Body::Ping(nonce),
                };
                debug!("[{}] Sending ping: {}", peer_addr, nonce);

                let mut buf = Vec::with_capacity(16);
                msg.serialize(&mut buf);

                if tx.clone().send(WsMessage::Binary(buf)).await.is_err() {
                    break;
                }
            }
        };

        tokio::select! {
            _ = ws_reader => {},
            _ = ws_writer => {},
            _ = heartbeat_interval => {},
        };

        info!("[{}] Connection closed", peer_addr);
        // Remove block subscriptions if there are any
        data.sub_pool.remove(peer_addr);
    });
}

pub fn process_ws_msg(
    data: &ServerData,
    state: &mut WsClient,
    msg: WsMessage,
) -> Option<WsMessage> {
    match msg {
        WsMessage::Binary(buf) => {
            state.set_needs_pong(false);

            let mut cur = Cursor::<&[u8]>::new(&buf);
            let msg = match Msg::deserialize(&mut cur) {
                Ok(msg) => {
                    let id = msg.id;
                    if cur.position() != buf.len() as u64 {
                        Msg {
                            id,
                            body: Body::Error(ErrorKind::BytesRemaining),
                        }
                    } else {
                        match handle_protocol_msg(data, state, msg) {
                            Some(body) => Msg { id, body },
                            None => return None,
                        }
                    }
                }
                Err(e) => {
                    error!(
                        "[{}] Error occurred during deserialization: {:?}",
                        state.addr(),
                        e
                    );
                    Msg {
                        id: u32::max_value(),
                        body: Body::Error(ErrorKind::Io),
                    }
                }
            };

            let mut buf = Vec::with_capacity(65536);
            msg.serialize(&mut buf);
            Some(WsMessage::Binary(buf))
        }
        WsMessage::Text(_) => Some(WsMessage::Close(Some(protocol::CloseFrame {
            code: protocol::frame::coding::CloseCode::Unsupported,
            reason: "text is not supported".into(),
        }))),
        _ => None,
    }
}

fn handle_protocol_msg(data: &ServerData, state: &mut WsClient, msg: Msg) -> Option<Body> {
    match msg.body {
        Body::Error(e) => {
            warn!(
                "[{}] Received error message from client: {:?}",
                state.addr(),
                e
            );
            None
        }
        Body::Request(req) => handle_rpc_request(data, state, msg.id, req),
        Body::Response(res) => {
            warn!(
                "[{}] Unexpected response from client: {:?}",
                state.addr(),
                res
            );
            None
        }
        Body::Ping(nonce) => {
            debug!("[{}] Received ping: {}", state.addr(), nonce);
            Some(Body::Pong(nonce))
        }
        Body::Pong(nonce) => {
            debug!("[{}] Received pong: {}", state.addr(), nonce);
            // We don't need to update the `needs_pong` state as it has already been updated when the message was
            // deserialized
            None
        }
    }
}

fn handle_rpc_request(
    data: &ServerData,
    state: &mut WsClient,
    id: u32,
    req: rpc::Request,
) -> Option<Body> {
    Some(match req {
        rpc::Request::Broadcast(tx) => {
            let res = data.minter.push_tx(tx);
            match res {
                Ok(_) => Body::Response(rpc::Response::Broadcast),
                Err(e) => Body::Error(ErrorKind::TxValidation(e)),
            }
        }
        rpc::Request::SetBlockFilter(filter) => {
            if filter.len() > 16 {
                return Some(Body::Error(ErrorKind::InvalidRequest));
            }
            state.filter = Some(filter);
            Body::Response(rpc::Response::SetBlockFilter)
        }
        rpc::Request::ClearBlockFilter => {
            state.filter = None;
            Body::Response(rpc::Response::ClearBlockFilter)
        }
        rpc::Request::Subscribe => {
            data.sub_pool.insert(state.addr(), state.sender());
            Body::Response(rpc::Response::Subscribe)
        }
        rpc::Request::Unsubscribe => {
            data.sub_pool.remove(state.addr());
            Body::Response(rpc::Response::Unsubscribe)
        }
        rpc::Request::GetProperties => {
            let props = data.chain.get_properties();
            Body::Response(rpc::Response::GetProperties(props))
        }
        rpc::Request::GetBlock(height) => match &state.filter {
            Some(filter) => match data.chain.get_filtered_block(height, filter) {
                Some(block) => Body::Response(rpc::Response::GetBlock(block)),
                None => Body::Error(ErrorKind::InvalidHeight),
            },
            None => match data.chain.get_block(height) {
                Some(block) => Body::Response(rpc::Response::GetBlock(FilteredBlock::Block(block))),
                None => Body::Error(ErrorKind::InvalidHeight),
            },
        },
        rpc::Request::GetFullBlock(height) => match data.chain.get_block(height) {
            Some(block) => Body::Response(rpc::Response::GetFullBlock(block)),
            None => Body::Error(ErrorKind::InvalidHeight),
        },
        rpc::Request::GetBlockRange(min_height, max_height) => {
            let range = AsyncBlockRange::try_new(Arc::clone(&data.chain), min_height, max_height);
            match range {
                Some(mut range) => {
                    if let Some(filter) = state.filter() {
                        range.set_filter(Some(filter.clone()));
                    }

                    let peer_addr = state.addr();
                    let mut tx = state.sender();
                    tokio::spawn(async move {
                        while let Some(block) = range.next().await {
                            let ws_msg = {
                                let msg = Msg {
                                    id,
                                    body: Body::Response(rpc::Response::GetBlock(block)),
                                };

                                let mut buf = Vec::with_capacity(65536);
                                msg.serialize(&mut buf);
                                WsMessage::Binary(buf)
                            };
                            if tx.send(ws_msg).await.is_err() {
                                warn!("[{}] Failed to send block range update", peer_addr);
                                return;
                            }
                        }

                        let ws_msg = {
                            let msg = Msg {
                                id,
                                body: Body::Response(rpc::Response::GetBlockRange),
                            };

                            let mut buf = Vec::with_capacity(32);
                            msg.serialize(&mut buf);
                            WsMessage::Binary(buf)
                        };
                        if tx.send(ws_msg).await.is_err() {
                            warn!("[{}] Failed to send block range finalizer", peer_addr);
                        }
                    });

                    return None;
                }
                None => Body::Error(ErrorKind::InvalidHeight),
            }
        }
        rpc::Request::GetAddressInfo(addr) => {
            let res = data.minter.get_addr_info(&addr);
            match res {
                Ok(info) => Body::Response(rpc::Response::GetAddressInfo(info)),
                Err(e) => Body::Error(ErrorKind::TxValidation(e)),
            }
        }
    })
}
<|endoftext|>


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

use jsonway;
use url::Url;
use hyper::Client;
use hyper::method::Method;
use hyper::header::{UserAgent, Accept, ContentLength, ContentType, qitem};
use hyper::mime::{Mime, TopLevel, SubLevel, Attr, Value};
use rustc_serialize::{Encodable, Decodable};
use rustc_serialize::json::{self, Json, ToJson};

use errors::Error;
use response::{Dependency, NamedEntity, Tag};

/// 默认的 BosonNLP API 服务器地址
const DEFAULT_BOSONNLP_URL: &'static str = "http://api.bosonnlp.com";

/// BosonNLP API 鉴权 HTTP Header
header! { (XToken, "X-Token") => [String] }

/// BosonNLP API 返回结果类型
pub type Result<T> = ::std::result::Result<T, Error>;

/// [BosonNLP](http://bosonnlp.com) REST API 访问的封装
pub struct BosonNLP {
    /// 用于 API 鉴权的 API Token
    pub token: String,
    /// 是否压缩大于 10K 的请求体，默认为 true
    pub compress: bool,
    /// BosonNLP HTTP API 的 URL，默认为 `http://api.bosonnlp.com`
    bosonnlp_url: String,
    /// hyper http Client
    client: Client,
}

impl BosonNLP {
    /// 初始化一个新的 ``BosonNLP`` 实例
    pub fn new<T: Into<String>>(token: T) -> BosonNLP {
        BosonNLP {
            token: token.into(),
            compress: true,
            bosonnlp_url: DEFAULT_BOSONNLP_URL.to_owned(),
            client: Client::new(),
        }
    }

    /// 使用自定义参数初始化一个新的 ``BosonNLP`` 实例
    pub fn with_options<T: Into<String>>(token: T, bosonnlp_url: T, compress: bool) -> BosonNLP {
        BosonNLP {
            token: token.into(),
            compress: compress,
            bosonnlp_url: bosonnlp_url.into(),
            client: Client::new(),
        }
    }

    /// 使用自定义的 hyper Client 初始化一个新的 ``BosonNLP`` 实例
    pub fn with_client<T: Into<String>>(token: T, client: Client) -> BosonNLP {
        BosonNLP {
            token: token.into(),
            compress: true,
            bosonnlp_url: DEFAULT_BOSONNLP_URL.to_owned(),
            client: client,
        }
    }

    fn request<D, E>(&self, method: Method, endpoint: &str, params: Vec<(&str, &str)>, data: &E) -> Result<D>
        where D: Decodable,
              E: Encodable
    {
        let url_string = format!("{}{}", self.bosonnlp_url, endpoint);
        let mut url = Url::parse(&url_string).unwrap();
        url.set_query_from_pairs(params.into_iter());
        let body = match json::encode(data) {
            Ok(d) => d,
            Err(..) => "".to_owned(),
        };
        let req = self.client.request(method.clone(), url)
                             .header(UserAgent(format!("bosonnlp-rs/{}", env!("CARGO_PKG_VERSION"))))
                             .header(Accept(vec![qitem(Mime(TopLevel::Application, SubLevel::Json,
                                                            vec![(Attr::Charset, Value::Utf8)]))]))
                             .header(XToken(self.token.clone()));
        let mut res = if method == Method::Post {
            // TODO: support GZip compress
            let req = req.header(ContentType(Mime(TopLevel::Application, SubLevel::Json,
                                             vec![(Attr::Charset, Value::Utf8)])))
                         .body(&body);
            try!(req.send())
        } else {
            try!(req.send())
        };
        let mut body = match res.headers.clone().get::<ContentLength>() {
            Some(&ContentLength(len)) => String::with_capacity(len as usize),
            _ => String::new(),
        };
        try!(res.read_to_string(&mut body));
        debug!("rev response {:#?} {:#?}", res.status, body);
        if !res.status.is_success() {
            let result = match Json::from_str(&body) {
                Ok(obj) => obj,
                Err(..) => Json::Object(json::Object::new()),
            };
            let message = match result.find("message") {
                Some(msg) => msg.as_string().unwrap_or("").to_owned(),
                None => body,
            };
            return Err(Error::Api { code: res.status, reason: message });
        }
        Ok(try!(json::decode::<D>(&body)))
    }

    fn get<D: Decodable>(&self, endpoint: &str, params: Vec<(&str, &str)>) -> Result<D> {
        self.request(Method::Get, endpoint, params, &json::Object::new())
    }

    fn post<D>(&self, endpoint: &str, params: Vec<(&str, &str)>, data: &Json) -> Result<D>
        where D: Decodable,
    {
        self.request(Method::Post, endpoint, params, data)
    }

    /// [情感分析接口](http://docs.bosonnlp.com/sentiment.html)
    ///
    /// ``contents``: 需要做情感分析的文本序列
    ///
    /// ``model``: 使用不同的语料训练的模型
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.sentiment(&vec!["这家味道还不错".to_owned()], "food").unwrap();
    ///     assert_eq!(1, rs.len());
    /// }
    /// ```
    pub fn sentiment(&self, contents: &[String], model: &str) -> Result<Vec<(f32, f32)>> {
        let endpoint = format!("/sentiment/analysis?{}", model);
        let data = contents.to_json();
        self.post::<Vec<(f32, f32)>>(&endpoint, vec![], &data)
    }

    /// [新闻分类接口](http://docs.bosonnlp.com/classify.html)
    ///
    /// ``contents``: 需要做分类的新闻文本序列
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.classify(&vec!["俄否决安理会谴责叙军战机空袭阿勒颇平民".to_owned()]).unwrap();
    ///     assert_eq!(vec![5usize], rs);
    /// }
    /// ```
    pub fn classify(&self, contents: &[String]) -> Result<Vec<usize>> {
        let data = contents.to_json();
        self.post::<Vec<usize>>("/classify/analysis", vec![], &data)
    }

    /// [语义联想接口](http://docs.bosonnlp.com/suggest.html)
    ///
    /// ``word``: 需要做语义联想的词
    ///
    /// ``top_k``: 返回结果的条数，最大值可设定为 100
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.suggest("北京", 2).unwrap();
    ///     assert_eq!(2, rs.len());
    /// }
    /// ```
    pub fn suggest<T: AsRef<str>>(&self, word: T, top_k: usize) -> Result<Vec<(f32, String)>> {
        let data = word.as_ref().to_json();
        self.post::<Vec<(f32, String)>>("/suggest/analysis", vec![("top_k", &top_k.to_string())], &data)
    }

    /// [关键词提取接口](http://docs.bosonnlp.com/keywords.html)
    ///
    /// ``text``: 需要做关键词提取的文本
    ///
    /// ``top_k``: 返回结果的条数，最大值可设定为 100
    ///
    /// ``segmented``: `text` 是否已经进行了分词，若为 `true` 则不会再对内容进行分词处理
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.extract_keywords("病毒式媒体网站：让新闻迅速蔓延", 2, false).unwrap();
    ///     assert_eq!(2, rs.len());
    /// }
    /// ```
    pub fn extract_keywords<T: AsRef<str>>(&self, text: T, top_k: usize, segmented: bool) -> Result<Vec<(f32, String)>> {
        let data = text.as_ref().to_json();
        let top_k_str = top_k.to_string();
        let params = match segmented {
            true => vec![("top_k", top_k_str.as_ref()), ("segmented", "1")],
            false => vec![("top_k", top_k_str.as_ref())],
        };
        self.post::<Vec<(f32, String)>>("/keywords/analysis", params, &data)
    }

    /// [依存文法分析接口](http://docs.bosonnlp.com/depparser.html)
    ///
    /// ``contents``: 需要做依存文法分析的文本序列
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.depparser(&vec!["今天天气好".to_owned()]).unwrap();
    ///     assert_eq!(1, rs.len());
    ///     let dep0 = &rs[0];
    ///     assert_eq!(vec![2isize, 2isize, -1isize], dep0.head);
    ///     let rs = nlp.depparser(&vec!["今天天气好".to_owned(), "美好的世界".to_owned()]).unwrap();
    ///     assert_eq!(2, rs.len());
    /// }
    /// ```
    pub fn depparser(&self, contents: &[String]) -> Result<Vec<Dependency>> {
        let data = contents.to_json();
        self.post::<Vec<Dependency>>("/depparser/analysis", vec![], &data)
    }

    /// [命名实体识别接口](http://docs.bosonnlp.com/ner.html)
    ///
    /// ``contents``: 需要做命名实体识别的文本序列
    ///
    /// ``sensitivity``: 准确率与召回率之间的平衡。
    /// 设置成 1 能找到更多的实体，设置成 5 能以更高的精度寻找实体
    /// 一般设置为 3
    ///
    /// ``segmented``: 输入是否已经为分词结果
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.ner(&vec!["成都商报记者 姚永忠".to_owned()], 2, false).unwrap();
    ///     assert_eq!(1, rs.len());
    ///     let rs = nlp.ner(&vec!["成都商报记者 姚永忠".to_owned(), "微软XP操作系统今日正式退休".to_owned()], 2, false).unwrap();
    ///     assert_eq!(2, rs.len());
    /// }
    /// ```
    pub fn ner(&self, contents: &[String], sensitivity: usize, segmented: bool) -> Result<Vec<NamedEntity>> {
        let data = contents.to_json();
        let sensitivity_str = sensitivity.to_string();
        let params = match segmented {
            true => vec![("sensitivity", sensitivity_str.as_ref()), ("segmented", "1")],
            false => vec![("sensitivity", sensitivity_str.as_ref())],
        };
        self.post::<Vec<NamedEntity>>("/ner/analysis", params, &data)
    }

    /// [分词与词性标注接口](http://docs.bosonnlp.com/tag.html)
    pub fn tag(&self, contents: &[String], space_mode: usize, oov_level: usize, t2s: bool, special_char_conv: bool) -> () {
        unimplemented!();
    }

    /// [新闻摘要接口](http://docs.bosonnlp.com/summary.html)
    ///
    /// ``title``: 需要做摘要的新闻标题，如果没有则传入空字符串
    ///
    /// ``content``: 需要做摘要的新闻正文
    ///
    /// ``word_limit``: 摘要字数限制
    ///
    /// ``not_exceed``: 是否严格限制字数
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let title = "前优酷土豆技术副总裁黄冬加盟芒果TV任CTO";
    ///     let content = "腾讯科技讯（刘亚澜）10月22日消息，前优酷土豆技术副总裁黄冬已于日前正式加盟芒果TV，出任CTO一职。";
    ///     let rs = nlp.summary(title, content, 1.0, false);
    ///     assert!(rs.is_ok());
    /// }
    /// ```
    pub fn summary<T: Into<String>>(&self, title: T, content: T, word_limit: f32, not_exceed: bool) -> Result<String> {
        let data = jsonway::object(|obj| {
            obj.set("title", title.into());
            obj.set("content", content.into());
            obj.set("percentage", word_limit);
            if not_exceed {
                obj.set("not_exceed", 1);
            } else {
                obj.set("not_exceed", 0);
            }
        }).unwrap();
        self.post::<String>("/summary/analysis", vec![], &data)
    }

    /// [文本聚类接口](http://docs.bosonnlp.com/cluster.html)
    pub fn cluster(&self, contents: &[String], task_id: &str, alpha: f32, beta: f32, timeout: usize) -> () {
        unimplemented!();
    }

    /// [典型意见接口](http://docs.bosonnlp.com/comments.html)
    pub fn comments(&self, contents: &[String], task_id: &str, alpha: f32, beta: f32, timeout: usize) -> () {
        unimplemented!();
    }
}
<commit_msg>Finish tag API<commit_after>use std::io::Read;

use jsonway;
use url::Url;
use hyper::Client;
use hyper::method::Method;
use hyper::header::{UserAgent, Accept, ContentLength, ContentType, qitem};
use hyper::mime::{Mime, TopLevel, SubLevel, Attr, Value};
use rustc_serialize::{Encodable, Decodable};
use rustc_serialize::json::{self, Json, ToJson};

use errors::Error;
use response::{Dependency, NamedEntity, Tag};

/// 默认的 BosonNLP API 服务器地址
const DEFAULT_BOSONNLP_URL: &'static str = "http://api.bosonnlp.com";

/// BosonNLP API 鉴权 HTTP Header
header! { (XToken, "X-Token") => [String] }

/// BosonNLP API 返回结果类型
pub type Result<T> = ::std::result::Result<T, Error>;

/// [BosonNLP](http://bosonnlp.com) REST API 访问的封装
pub struct BosonNLP {
    /// 用于 API 鉴权的 API Token
    pub token: String,
    /// 是否压缩大于 10K 的请求体，默认为 true
    pub compress: bool,
    /// BosonNLP HTTP API 的 URL，默认为 `http://api.bosonnlp.com`
    bosonnlp_url: String,
    /// hyper http Client
    client: Client,
}

impl BosonNLP {
    /// 初始化一个新的 ``BosonNLP`` 实例
    pub fn new<T: Into<String>>(token: T) -> BosonNLP {
        BosonNLP {
            token: token.into(),
            compress: true,
            bosonnlp_url: DEFAULT_BOSONNLP_URL.to_owned(),
            client: Client::new(),
        }
    }

    /// 使用自定义参数初始化一个新的 ``BosonNLP`` 实例
    pub fn with_options<T: Into<String>>(token: T, bosonnlp_url: T, compress: bool) -> BosonNLP {
        BosonNLP {
            token: token.into(),
            compress: compress,
            bosonnlp_url: bosonnlp_url.into(),
            client: Client::new(),
        }
    }

    /// 使用自定义的 hyper Client 初始化一个新的 ``BosonNLP`` 实例
    pub fn with_client<T: Into<String>>(token: T, client: Client) -> BosonNLP {
        BosonNLP {
            token: token.into(),
            compress: true,
            bosonnlp_url: DEFAULT_BOSONNLP_URL.to_owned(),
            client: client,
        }
    }

    fn request<D, E>(&self, method: Method, endpoint: &str, params: Vec<(&str, &str)>, data: &E) -> Result<D>
        where D: Decodable,
              E: Encodable
    {
        let url_string = format!("{}{}", self.bosonnlp_url, endpoint);
        let mut url = Url::parse(&url_string).unwrap();
        url.set_query_from_pairs(params.into_iter());
        let body = match json::encode(data) {
            Ok(d) => d,
            Err(..) => "".to_owned(),
        };
        let req = self.client.request(method.clone(), url)
                             .header(UserAgent(format!("bosonnlp-rs/{}", env!("CARGO_PKG_VERSION"))))
                             .header(Accept(vec![qitem(Mime(TopLevel::Application, SubLevel::Json,
                                                            vec![(Attr::Charset, Value::Utf8)]))]))
                             .header(XToken(self.token.clone()));
        let mut res = if method == Method::Post {
            // TODO: support GZip compress
            let req = req.header(ContentType(Mime(TopLevel::Application, SubLevel::Json,
                                             vec![(Attr::Charset, Value::Utf8)])))
                         .body(&body);
            try!(req.send())
        } else {
            try!(req.send())
        };
        let mut body = match res.headers.clone().get::<ContentLength>() {
            Some(&ContentLength(len)) => String::with_capacity(len as usize),
            _ => String::new(),
        };
        try!(res.read_to_string(&mut body));
        debug!("rev response {:#?} {:#?}", res.status, body);
        if !res.status.is_success() {
            let result = match Json::from_str(&body) {
                Ok(obj) => obj,
                Err(..) => Json::Object(json::Object::new()),
            };
            let message = match result.find("message") {
                Some(msg) => msg.as_string().unwrap_or("").to_owned(),
                None => body,
            };
            return Err(Error::Api { code: res.status, reason: message });
        }
        Ok(try!(json::decode::<D>(&body)))
    }

    fn get<D: Decodable>(&self, endpoint: &str, params: Vec<(&str, &str)>) -> Result<D> {
        self.request(Method::Get, endpoint, params, &json::Object::new())
    }

    fn post<D>(&self, endpoint: &str, params: Vec<(&str, &str)>, data: &Json) -> Result<D>
        where D: Decodable,
    {
        self.request(Method::Post, endpoint, params, data)
    }

    /// [情感分析接口](http://docs.bosonnlp.com/sentiment.html)
    ///
    /// ``contents``: 需要做情感分析的文本序列
    ///
    /// ``model``: 使用不同的语料训练的模型
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.sentiment(&vec!["这家味道还不错".to_owned()], "food").unwrap();
    ///     assert_eq!(1, rs.len());
    /// }
    /// ```
    pub fn sentiment(&self, contents: &[String], model: &str) -> Result<Vec<(f32, f32)>> {
        let endpoint = format!("/sentiment/analysis?{}", model);
        let data = contents.to_json();
        self.post::<Vec<(f32, f32)>>(&endpoint, vec![], &data)
    }

    /// [新闻分类接口](http://docs.bosonnlp.com/classify.html)
    ///
    /// ``contents``: 需要做分类的新闻文本序列
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.classify(&vec!["俄否决安理会谴责叙军战机空袭阿勒颇平民".to_owned()]).unwrap();
    ///     assert_eq!(vec![5usize], rs);
    /// }
    /// ```
    pub fn classify(&self, contents: &[String]) -> Result<Vec<usize>> {
        let data = contents.to_json();
        self.post::<Vec<usize>>("/classify/analysis", vec![], &data)
    }

    /// [语义联想接口](http://docs.bosonnlp.com/suggest.html)
    ///
    /// ``word``: 需要做语义联想的词
    ///
    /// ``top_k``: 返回结果的条数，最大值可设定为 100
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.suggest("北京", 2).unwrap();
    ///     assert_eq!(2, rs.len());
    /// }
    /// ```
    pub fn suggest<T: AsRef<str>>(&self, word: T, top_k: usize) -> Result<Vec<(f32, String)>> {
        let data = word.as_ref().to_json();
        self.post::<Vec<(f32, String)>>("/suggest/analysis", vec![("top_k", &top_k.to_string())], &data)
    }

    /// [关键词提取接口](http://docs.bosonnlp.com/keywords.html)
    ///
    /// ``text``: 需要做关键词提取的文本
    ///
    /// ``top_k``: 返回结果的条数，最大值可设定为 100
    ///
    /// ``segmented``: `text` 是否已经进行了分词，若为 `true` 则不会再对内容进行分词处理
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.extract_keywords("病毒式媒体网站：让新闻迅速蔓延", 2, false).unwrap();
    ///     assert_eq!(2, rs.len());
    /// }
    /// ```
    pub fn extract_keywords<T: AsRef<str>>(&self, text: T, top_k: usize, segmented: bool) -> Result<Vec<(f32, String)>> {
        let data = text.as_ref().to_json();
        let top_k_str = top_k.to_string();
        let params = match segmented {
            true => vec![("top_k", top_k_str.as_ref()), ("segmented", "1")],
            false => vec![("top_k", top_k_str.as_ref())],
        };
        self.post::<Vec<(f32, String)>>("/keywords/analysis", params, &data)
    }

    /// [依存文法分析接口](http://docs.bosonnlp.com/depparser.html)
    ///
    /// ``contents``: 需要做依存文法分析的文本序列
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.depparser(&vec!["今天天气好".to_owned()]).unwrap();
    ///     assert_eq!(1, rs.len());
    ///     let dep0 = &rs[0];
    ///     assert_eq!(vec![2isize, 2isize, -1isize], dep0.head);
    ///     let rs = nlp.depparser(&vec!["今天天气好".to_owned(), "美好的世界".to_owned()]).unwrap();
    ///     assert_eq!(2, rs.len());
    /// }
    /// ```
    pub fn depparser(&self, contents: &[String]) -> Result<Vec<Dependency>> {
        let data = contents.to_json();
        self.post::<Vec<Dependency>>("/depparser/analysis", vec![], &data)
    }

    /// [命名实体识别接口](http://docs.bosonnlp.com/ner.html)
    ///
    /// ``contents``: 需要做命名实体识别的文本序列
    ///
    /// ``sensitivity``: 准确率与召回率之间的平衡。
    /// 设置成 1 能找到更多的实体，设置成 5 能以更高的精度寻找实体
    /// 一般设置为 3
    ///
    /// ``segmented``: 输入是否已经为分词结果
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.ner(&vec!["成都商报记者 姚永忠".to_owned()], 2, false).unwrap();
    ///     assert_eq!(1, rs.len());
    ///     let rs = nlp.ner(&vec!["成都商报记者 姚永忠".to_owned(), "微软XP操作系统今日正式退休".to_owned()], 2, false).unwrap();
    ///     assert_eq!(2, rs.len());
    /// }
    /// ```
    pub fn ner(&self, contents: &[String], sensitivity: usize, segmented: bool) -> Result<Vec<NamedEntity>> {
        let data = contents.to_json();
        let sensitivity_str = sensitivity.to_string();
        let params = match segmented {
            true => vec![("sensitivity", sensitivity_str.as_ref()), ("segmented", "1")],
            false => vec![("sensitivity", sensitivity_str.as_ref())],
        };
        self.post::<Vec<NamedEntity>>("/ner/analysis", params, &data)
    }

    /// [分词与词性标注接口](http://docs.bosonnlp.com/tag.html)
    ///
    /// ``contents``: 需要做分词与词性标注的文本序列
    ///
    /// ``space_mode``: 空格保留选项，0-3 有效
    ///
    /// ``oov_level``: 枚举强度选项，0-4 有效
    ///
    /// ``t2s``: 是否开启繁体转简体
    ///
    /// ``special_char_conv``: 是否转化特殊字符，针对回车、Tab 等特殊字符。
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let rs = nlp.tag(&vec!["成都商报记者 姚永忠".to_owned()], 0, 3, false, false).unwrap();
    ///     assert_eq!(1, rs.len());
    /// }
    /// ```
    pub fn tag(&self, contents: &[String], space_mode: usize, oov_level: usize, t2s: bool, special_char_conv: bool) -> Result<Vec<Tag>> {
        let data = contents.to_json();
        let t2s_str = match t2s {
            true => "1",
            false => "0",
        };
        let special_char_conv_str = match special_char_conv {
            true => "1",
            false => "0",
        };
        let space_mode_str = space_mode.to_string();
        let oov_level_str = oov_level.to_string();
        let params = vec![
            ("space_mode", space_mode_str.as_ref()),
            ("oov_level", oov_level_str.as_ref()),
            ("t2s", t2s_str),
            ("special_char_conv", special_char_conv_str),
        ];
        self.post("/tag/analysis", params, &data)
    }

    /// [新闻摘要接口](http://docs.bosonnlp.com/summary.html)
    ///
    /// ``title``: 需要做摘要的新闻标题，如果没有则传入空字符串
    ///
    /// ``content``: 需要做摘要的新闻正文
    ///
    /// ``word_limit``: 摘要字数限制
    ///
    /// ``not_exceed``: 是否严格限制字数
    ///
    /// # 使用示例
    ///
    /// ```
    /// extern crate bosonnlp;
    ///
    /// use bosonnlp::BosonNLP;
    ///
    /// fn main() {
    ///     let nlp = BosonNLP::new(env!("BOSON_API_TOKEN"));
    ///     let title = "前优酷土豆技术副总裁黄冬加盟芒果TV任CTO";
    ///     let content = "腾讯科技讯（刘亚澜）10月22日消息，前优酷土豆技术副总裁黄冬已于日前正式加盟芒果TV，出任CTO一职。";
    ///     let rs = nlp.summary(title, content, 1.0, false);
    ///     assert!(rs.is_ok());
    /// }
    /// ```
    pub fn summary<T: Into<String>>(&self, title: T, content: T, word_limit: f32, not_exceed: bool) -> Result<String> {
        let data = jsonway::object(|obj| {
            obj.set("title", title.into());
            obj.set("content", content.into());
            obj.set("percentage", word_limit);
            if not_exceed {
                obj.set("not_exceed", 1);
            } else {
                obj.set("not_exceed", 0);
            }
        }).unwrap();
        self.post::<String>("/summary/analysis", vec![], &data)
    }

    /// [文本聚类接口](http://docs.bosonnlp.com/cluster.html)
    pub fn cluster(&self, contents: &[String], task_id: &str, alpha: f32, beta: f32, timeout: usize) -> () {
        unimplemented!();
    }

    /// [典型意见接口](http://docs.bosonnlp.com/comments.html)
    pub fn comments(&self, contents: &[String], task_id: &str, alpha: f32, beta: f32, timeout: usize) -> () {
        unimplemented!();
    }
}
<|endoftext|>


<commit_before>use providers::{self, parse, request, ProviderError};
use reqwest::{self, Client};
use std::time::Duration;

/// Url shortener: the way to retrieve a short url.
#[derive(Debug)]
pub struct UrlShortener {
    client: Client,
}

impl UrlShortener {
    /// Creates new `UrlShortener` with default (3 seconds) timeout.
    pub fn new() -> Result<UrlShortener, reqwest::Error> {
        UrlShortener::with_timeout(3)
    }

    /// Creates new `UrlShortener` with custom read timeout.
    pub fn with_timeout(seconds: u64) -> Result<UrlShortener, reqwest::Error> {
        let client = reqwest::ClientBuilder::new()
            .timeout(Duration::from_secs(seconds))
            .build()?;

        Ok(UrlShortener { client })
    }

    /// Try to generate a short URL from each provider, iterating over each
    /// provider until a short URL is successfully generated.
    /// If you wish to override the list or providers or their priority,
    /// provide your own list of providers as second argument.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use urlshortener::client::UrlShortener;
    ///
    /// let us = UrlShortener::new().unwrap();
    /// let long_url = "https://rust-lang.org";
    /// let _short_url = us.try_generate(long_url, None);
    /// ```
    ///
    /// ```rust,no_run
    /// use urlshortener::{client::UrlShortener, providers::Provider};
    ///
    /// let us = UrlShortener::new().unwrap();
    /// let providers = [
    ///     Provider::GooGl { api_key: "MY_API_KEY".to_owned() },
    ///     Provider::IsGd,
    /// ];
    /// let long_url = "https://rust-lang.org";
    /// let _short_url = us.try_generate(long_url, Some(&providers));
    /// ```
    ///
    /// # Errors
    ///
    /// Returns an `Error<ProviderError>` if there is an error generating a
    /// short URL from all providers.
    ///
    /// # Notes
    ///
    /// This function has been deprecated since it does not bring any UX improvements.
    /// The body could be easily re-written as:
    ///
    /// ```rust,no_run
    /// use urlshortener::{client::UrlShortener, providers::Provider};
    ///
    /// let us = UrlShortener::new().unwrap();
    /// let providers = [
    ///     Provider::GooGl { api_key: "MY_API_KEY".to_owned() },
    ///     Provider::IsGd,
    /// ];
    /// let long_url = "https://rust-lang.org";
    /// let mut short_url = None;
    /// for provider in &providers {
    ///     if let Ok(short_url_res) = us.generate(long_url, provider) {
    ///         short_url = Some(short_url_res);
    ///         break;
    ///     }
    /// }
    ///
    /// ```
    #[deprecated(since = "1.0.0", note = "Please use `generate` directly instead.")]
    pub fn try_generate(
        &self,
        url: &str,
        use_providers: Option<&[providers::Provider]>,
    ) -> Result<String, ProviderError> {
        let providers = use_providers.unwrap_or(providers::PROVIDERS);
        for provider in providers {
            let res = self.generate(url, provider);

            if res.is_ok() {
                return res;
            }
        }

        Err(ProviderError::Connection)
    }

    /// Attempts to get a short URL using the specified provider.
    ///
    /// # Examples
    ///
    /// ```rust, no_run
    /// use urlshortener::{providers::Provider, client::UrlShortener};
    ///
    /// let us = UrlShortener::new().unwrap();
    /// let long_url = "http://rust-lang.org";
    /// let _short_url = us.generate(long_url, &Provider::IsGd);
    /// ```
    ///
    /// ```rust,no_run
    /// use urlshortener::{providers::Provider, client::UrlShortener};
    ///
    /// let us = UrlShortener::new().unwrap();
    /// let api_key = "MY_API_KEY".to_owned();
    /// let long_url = "http://rust-lang.org";
    /// let _short_url = us.generate(long_url, &Provider::GooGl { api_key: api_key });
    /// ```
    pub fn generate<S: Into<String>>(
        &self,
        url: S,
        provider: &providers::Provider,
    ) -> Result<String, ProviderError> {
        let req = request(&url.into(), provider);

        if let Ok(mut response) = req.execute(&self.client) {
            response
                .text()
                .map_err(|_| ProviderError::Connection)
                .and_then(|t| parse(&t, provider))
        } else {
            Err(ProviderError::Connection)
        }
    }
}

#[cfg(test)]
mod tests {
    use client;
    use providers;

    /// This test does not cover services which require authentication for obvious reasons.
    #[test]
    fn providers() {
        let us = client::UrlShortener::with_timeout(5).unwrap();
        let url = "http://yandex.com";

        for provider in providers::PROVIDERS {
            println!("Request shortening via provider: {}", provider.to_name());
            if let Some(err) = us.generate(url, provider).err() {
                println!("Error: {:?}", err);
                assert_eq!(err, providers::ProviderError::Connection);
            }
        }
    }
}
<commit_msg>Rewrite the providers test.<commit_after>use providers::{self, parse, request, ProviderError};
use reqwest::{self, Client};
use std::time::Duration;

/// Url shortener: the way to retrieve a short url.
#[derive(Debug)]
pub struct UrlShortener {
    client: Client,
}

impl UrlShortener {
    /// Creates new `UrlShortener` with default (3 seconds) timeout.
    pub fn new() -> Result<UrlShortener, reqwest::Error> {
        UrlShortener::with_timeout(3)
    }

    /// Creates new `UrlShortener` with custom read timeout.
    pub fn with_timeout(seconds: u64) -> Result<UrlShortener, reqwest::Error> {
        let client = reqwest::ClientBuilder::new()
            .timeout(Duration::from_secs(seconds))
            .build()?;

        Ok(UrlShortener { client })
    }

    /// Try to generate a short URL from each provider, iterating over each
    /// provider until a short URL is successfully generated.
    /// If you wish to override the list or providers or their priority,
    /// provide your own list of providers as second argument.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use urlshortener::client::UrlShortener;
    ///
    /// let us = UrlShortener::new().unwrap();
    /// let long_url = "https://rust-lang.org";
    /// let _short_url = us.try_generate(long_url, None);
    /// ```
    ///
    /// ```rust,no_run
    /// use urlshortener::{client::UrlShortener, providers::Provider};
    ///
    /// let us = UrlShortener::new().unwrap();
    /// let providers = [
    ///     Provider::GooGl { api_key: "MY_API_KEY".to_owned() },
    ///     Provider::IsGd,
    /// ];
    /// let long_url = "https://rust-lang.org";
    /// let _short_url = us.try_generate(long_url, Some(&providers));
    /// ```
    ///
    /// # Errors
    ///
    /// Returns an `Error<ProviderError>` if there is an error generating a
    /// short URL from all providers.
    ///
    /// # Notes
    ///
    /// This function has been deprecated since it does not bring any UX improvements.
    /// The body could be easily re-written as:
    ///
    /// ```rust,no_run
    /// use urlshortener::{client::UrlShortener, providers::Provider};
    ///
    /// let us = UrlShortener::new().unwrap();
    /// let providers = [
    ///     Provider::GooGl { api_key: "MY_API_KEY".to_owned() },
    ///     Provider::IsGd,
    /// ];
    /// let long_url = "https://rust-lang.org";
    /// let mut short_url = None;
    /// for provider in &providers {
    ///     if let Ok(short_url_res) = us.generate(long_url, provider) {
    ///         short_url = Some(short_url_res);
    ///         break;
    ///     }
    /// }
    ///
    /// ```
    #[deprecated(since = "1.0.0", note = "Please use `generate` directly instead.")]
    pub fn try_generate(
        &self,
        url: &str,
        use_providers: Option<&[providers::Provider]>,
    ) -> Result<String, ProviderError> {
        let providers = use_providers.unwrap_or(providers::PROVIDERS);
        for provider in providers {
            let res = self.generate(url, provider);

            if res.is_ok() {
                return res;
            }
        }

        Err(ProviderError::Connection)
    }

    /// Attempts to get a short URL using the specified provider.
    ///
    /// # Examples
    ///
    /// ```rust, no_run
    /// use urlshortener::{providers::Provider, client::UrlShortener};
    ///
    /// let us = UrlShortener::new().unwrap();
    /// let long_url = "http://rust-lang.org";
    /// let _short_url = us.generate(long_url, &Provider::IsGd);
    /// ```
    ///
    /// ```rust,no_run
    /// use urlshortener::{providers::Provider, client::UrlShortener};
    ///
    /// let us = UrlShortener::new().unwrap();
    /// let api_key = "MY_API_KEY".to_owned();
    /// let long_url = "http://rust-lang.org";
    /// let _short_url = us.generate(long_url, &Provider::GooGl { api_key: api_key });
    /// ```
    pub fn generate<S: Into<String>>(
        &self,
        url: S,
        provider: &providers::Provider,
    ) -> Result<String, ProviderError> {
        let req = request(&url.into(), provider);

        if let Ok(mut response) = req.execute(&self.client) {
            response
                .text()
                .map_err(|_| ProviderError::Connection)
                .and_then(|t| parse(&t, provider))
        } else {
            Err(ProviderError::Connection)
        }
    }
}

#[cfg(test)]
mod tests {
    use client;
    use providers;

    /// This test does not cover services which require authentication for obvious reasons.
    #[test]
    fn providers() {
        let us = client::UrlShortener::with_timeout(5).unwrap();
        let url = "http://yandex.com";

        for provider in providers::PROVIDERS {
            if let Err(e) = us.generate(url, provider) {
                assert!(false, "{:?} -> {:?}", provider, e);
            }
        }
    }
}
<|endoftext|>