<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Martin Kysel <code@martinkysel.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (methods) isnt

#[macro_use]
extern crate uucore;

use std::{
    fs::File,
    io::{stdin, stdout, BufReader, Read, Stdout, Write},
    path::Path,
    time::Duration,
};

#[cfg(all(unix, not(target_os = "fuchsia")))]
extern crate nix;

use clap::{crate_version, App, Arg};
use crossterm::{
    event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
    execute, queue,
    style::Attribute,
    terminal,
};

use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;

const BELL: &str = "\x07";

pub mod options {
    pub const SILENT: &str = "silent";
    pub const LOGICAL: &str = "logical";
    pub const NO_PAUSE: &str = "no-pause";
    pub const PRINT_OVER: &str = "print-over";
    pub const CLEAN_PRINT: &str = "clean-print";
    pub const SQUEEZE: &str = "squeeze";
    pub const PLAIN: &str = "plain";
    pub const LINES: &str = "lines";
    pub const NUMBER: &str = "number";
    pub const PATTERN: &str = "pattern";
    pub const FROM_LINE: &str = "from-line";
    pub const FILES: &str = "files";
}

const MULTI_FILE_TOP_PROMPT: &str = "::::::::::::::\n{}\n::::::::::::::\n";

pub fn uumain(args: impl uucore::Args) -> i32 {
    let matches = App::new(executable!())
        .about("A file perusal filter for CRT viewing.")
        .version(crate_version!())
        .arg(
            Arg::with_name(options::SILENT)
                .short("d")
                .long(options::SILENT)
                .help("Display help instead of ringing bell"),
        )
        // The commented arguments below are unimplemented:
        /*
        .arg(
            Arg::with_name(options::LOGICAL)
                .short("f")
                .long(options::LOGICAL)
                .help("Count logical rather than screen lines"),
        )
        .arg(
            Arg::with_name(options::NO_PAUSE)
                .short("l")
                .long(options::NO_PAUSE)
                .help("Suppress pause after form feed"),
        )
        .arg(
            Arg::with_name(options::PRINT_OVER)
                .short("c")
                .long(options::PRINT_OVER)
                .help("Do not scroll, display text and clean line ends"),
        )
        .arg(
            Arg::with_name(options::CLEAN_PRINT)
                .short("p")
                .long(options::CLEAN_PRINT)
                .help("Do not scroll, clean screen and display text"),
        )
        .arg(
            Arg::with_name(options::SQUEEZE)
                .short("s")
                .long(options::SQUEEZE)
                .help("Squeeze multiple blank lines into one"),
        )
        .arg(
            Arg::with_name(options::PLAIN)
                .short("u")
                .long(options::PLAIN)
                .help("Suppress underlining and bold"),
        )
        .arg(
            Arg::with_name(options::LINES)
                .short("n")
                .long(options::LINES)
                .value_name("number")
                .takes_value(true)
                .help("The number of lines per screen full"),
        )
        .arg(
            Arg::with_name(options::NUMBER)
                .allow_hyphen_values(true)
                .long(options::NUMBER)
                .required(false)
                .takes_value(true)
                .help("Same as --lines"),
        )
        .arg(
            Arg::with_name(options::FROM_LINE)
                .short("F")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .value_name("number")
                .help("Display file beginning from line number"),
        )
        .arg(
            Arg::with_name(options::PATTERN)
                .short("P")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .help("Display file beginning from pattern match"),
        )
        */
        .arg(
            Arg::with_name(options::FILES)
                .required(false)
                .multiple(true)
                .help("Path to the files to be read"),
        )
        .get_matches_from(args);

    let mut buff = String::new();
    let silent = matches.is_present(options::SILENT);
    if let Some(files) = matches.values_of(options::FILES) {
        let mut stdout = setup_term();
        let length = files.len();

        let mut files_iter = files.peekable();
        while let (Some(file), next_file) = (files_iter.next(), files_iter.peek()) {
            let file = Path::new(file);
            if file.is_dir() {
                terminal::disable_raw_mode().unwrap();
                show_usage_error!("'{}' is a directory.", file.display());
                return 1;
            }
            if !file.exists() {
                terminal::disable_raw_mode().unwrap();
                show_error!("cannot open {}: No such file or directory", file.display());
                return 1;
            }
            if length > 1 {
                buff.push_str(&MULTI_FILE_TOP_PROMPT.replace("{}", file.to_str().unwrap()));
            }
            let mut reader = BufReader::new(File::open(file).unwrap());
            reader.read_to_string(&mut buff).unwrap();
            more(&buff, &mut stdout, next_file.copied(), silent);
            buff.clear();
        }
        reset_term(&mut stdout);
    } else if atty::isnt(atty::Stream::Stdin) {
        stdin().read_to_string(&mut buff).unwrap();
        let mut stdout = setup_term();
        more(&buff, &mut stdout, None, silent);
        reset_term(&mut stdout);
    } else {
        show_usage_error!("bad usage");
    }
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn setup_term() -> std::io::Stdout {
    let stdout = stdout();
    terminal::enable_raw_mode().unwrap();
    stdout
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn setup_term() -> usize {
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn reset_term(stdout: &mut std::io::Stdout) {
    terminal::disable_raw_mode().unwrap();
    // Clear the prompt
    queue!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
    // Move cursor to the beginning without printing new line
    print!("\r");
    stdout.flush().unwrap();
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn reset_term(_: &mut usize) {}

fn more(buff: &str, mut stdout: &mut Stdout, next_file: Option<&str>, silent: bool) {
    let (cols, rows) = terminal::size().unwrap();
    let lines = break_buff(buff, usize::from(cols));

    let mut pager = Pager::new(rows, lines, next_file, silent);
    pager.draw(stdout, false);
    if pager.should_close() {
        return;
    }

    loop {
        let mut wrong_key = false;
        if event::poll(Duration::from_millis(10)).unwrap() {
            match event::read().unwrap() {
                Event::Key(KeyEvent {
                    code: KeyCode::Char('q'),
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char('c'),
                    modifiers: KeyModifiers::CONTROL,
                }) => {
                    reset_term(&mut stdout);
                    std::process::exit(0);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Down,
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char(' '),
                    modifiers: KeyModifiers::NONE,
                }) => {
                    if pager.should_close() {
                        return;
                    } else {
                        pager.page_down();
                    }
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Up,
                    modifiers: KeyModifiers::NONE,
                }) => {
                    pager.page_up();
                }
                _ => {
                    wrong_key = true;
                }
            }

            pager.draw(stdout, wrong_key);
        }
    }
}

struct Pager<'a> {
    // The current line at the top of the screen
    upper_mark: usize,
    // The number of rows that fit on the screen
    content_rows: u16,
    lines: Vec<String>,
    next_file: Option<&'a str>,
    line_count: usize,
    silent: bool,
}

impl<'a> Pager<'a> {
    fn new(rows: u16, lines: Vec<String>, next_file: Option<&'a str>, silent: bool) -> Self {
        let line_count = lines.len();
        Self {
            upper_mark: 0,
            content_rows: rows.saturating_sub(1),
            lines,
            next_file,
            line_count,
            silent,
        }
    }

    fn should_close(&mut self) -> bool {
        self.upper_mark
            .saturating_add(self.content_rows.into())
            .eq(&self.line_count)
    }

    fn page_down(&mut self) {
        self.upper_mark = self
            .upper_mark
            .saturating_add(self.content_rows.into())
            .min(self.line_count.saturating_sub(self.content_rows.into()));
    }

    fn page_up(&mut self) {
        self.upper_mark = self.upper_mark.saturating_sub(self.content_rows.into());
    }

    fn draw(&self, stdout: &mut std::io::Stdout, wrong_key: bool) {
        let lower_mark = self
            .line_count
            .min(self.upper_mark.saturating_add(self.content_rows.into()));
        self.draw_lines(stdout);
        self.draw_prompt(stdout, lower_mark, wrong_key);
        stdout.flush().unwrap();
    }

    fn draw_lines(&self, stdout: &mut std::io::Stdout) {
        execute!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
        let displayed_lines = self
            .lines
            .iter()
            .skip(self.upper_mark)
            .take(self.content_rows.into());

        for line in displayed_lines {
            stdout
                .write_all(format!("\r{}\n", line).as_bytes())
                .unwrap();
        }
    }

    fn draw_prompt(&self, stdout: &mut Stdout, lower_mark: usize, wrong_key: bool) {
        let status_inner = if lower_mark == self.line_count {
            format!("Next file: {}", self.next_file.unwrap_or_default())
        } else {
            format!(
                "{}%",
                (lower_mark as f64 / self.line_count as f64 * 100.0).round() as u16
            )
        };

        let status = format!("--More--({})", status_inner);

        let banner = match (self.silent, wrong_key) {
            (true, true) => "[Press 'h' for instructions. (unimplemented)]".to_string(),
            (true, false) => format!("{}[Press space to continue, 'q' to quit.]", status),
            (false, true) => format!("{}{}", status, BELL),
            (false, false) => status,
        };

        write!(
            stdout,
            "\r{}{}{}",
            Attribute::Reverse,
            banner,
            Attribute::Reset
        )
        .unwrap();
    }
}

// Break the lines on the cols of the terminal
fn break_buff(buff: &str, cols: usize) -> Vec<String> {
    let mut lines = Vec::with_capacity(buff.lines().count());

    for l in buff.lines() {
        lines.append(&mut break_line(l, cols));
    }
    lines
}

fn break_line(line: &str, cols: usize) -> Vec<String> {
    let width = UnicodeWidthStr::width(line);
    let mut lines = Vec::new();
    if width < cols {
        lines.push(line.to_string());
        return lines;
    }

    let gr_idx = UnicodeSegmentation::grapheme_indices(line, true);
    let mut last_index = 0;
    let mut total_width = 0;
    for (index, grapheme) in gr_idx {
        let width = UnicodeWidthStr::width(grapheme);
        total_width += width;

        if total_width > cols {
            lines.push(line[last_index..index].to_string());
            last_index = index;
            total_width = width;
        }
    }

    if last_index != line.len() {
        lines.push(line[last_index..].to_string());
    }
    lines
}

#[cfg(test)]
mod tests {
    use super::break_line;
    use unicode_width::UnicodeWidthStr;

    #[test]
    fn test_break_lines_long() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..200 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);
        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        assert_eq!((80, 80, 40), (widths[0], widths[1], widths[2]));
    }

    #[test]
    fn test_break_lines_short() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..20 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);

        assert_eq!(20, lines[0].len());
    }

    #[test]
    fn test_break_line_zwj() {
        let mut test_string = String::with_capacity(1100);
        for _ in 0..20 {
            test_string.push_str("👩🏻🔬");
        }

        let lines = break_line(&test_string, 80);

        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        // Each 👩🏻🔬 is 6 character width it break line to the closest number to 80 => 6 * 13 = 78
        assert_eq!((78, 42), (widths[0], widths[1]));
    }
}
<commit_msg>Fixing display when resizing terminal<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Martin Kysel <code@martinkysel.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (methods) isnt

#[macro_use]
extern crate uucore;

use std::{
    fs::File,
    io::{stdin, stdout, BufReader, Read, Stdout, Write},
    path::Path,
    time::Duration,
};

#[cfg(all(unix, not(target_os = "fuchsia")))]
extern crate nix;

use clap::{crate_version, App, Arg};
use crossterm::{
    event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
    execute, queue,
    style::Attribute,
    terminal,
};

use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;

const BELL: &str = "\x07";

pub mod options {
    pub const SILENT: &str = "silent";
    pub const LOGICAL: &str = "logical";
    pub const NO_PAUSE: &str = "no-pause";
    pub const PRINT_OVER: &str = "print-over";
    pub const CLEAN_PRINT: &str = "clean-print";
    pub const SQUEEZE: &str = "squeeze";
    pub const PLAIN: &str = "plain";
    pub const LINES: &str = "lines";
    pub const NUMBER: &str = "number";
    pub const PATTERN: &str = "pattern";
    pub const FROM_LINE: &str = "from-line";
    pub const FILES: &str = "files";
}

const MULTI_FILE_TOP_PROMPT: &str = "::::::::::::::\n{}\n::::::::::::::\n";

pub fn uumain(args: impl uucore::Args) -> i32 {
    let matches = App::new(executable!())
        .about("A file perusal filter for CRT viewing.")
        .version(crate_version!())
        .arg(
            Arg::with_name(options::SILENT)
                .short("d")
                .long(options::SILENT)
                .help("Display help instead of ringing bell"),
        )
        // The commented arguments below are unimplemented:
        /*
        .arg(
            Arg::with_name(options::LOGICAL)
                .short("f")
                .long(options::LOGICAL)
                .help("Count logical rather than screen lines"),
        )
        .arg(
            Arg::with_name(options::NO_PAUSE)
                .short("l")
                .long(options::NO_PAUSE)
                .help("Suppress pause after form feed"),
        )
        .arg(
            Arg::with_name(options::PRINT_OVER)
                .short("c")
                .long(options::PRINT_OVER)
                .help("Do not scroll, display text and clean line ends"),
        )
        .arg(
            Arg::with_name(options::CLEAN_PRINT)
                .short("p")
                .long(options::CLEAN_PRINT)
                .help("Do not scroll, clean screen and display text"),
        )
        .arg(
            Arg::with_name(options::SQUEEZE)
                .short("s")
                .long(options::SQUEEZE)
                .help("Squeeze multiple blank lines into one"),
        )
        .arg(
            Arg::with_name(options::PLAIN)
                .short("u")
                .long(options::PLAIN)
                .help("Suppress underlining and bold"),
        )
        .arg(
            Arg::with_name(options::LINES)
                .short("n")
                .long(options::LINES)
                .value_name("number")
                .takes_value(true)
                .help("The number of lines per screen full"),
        )
        .arg(
            Arg::with_name(options::NUMBER)
                .allow_hyphen_values(true)
                .long(options::NUMBER)
                .required(false)
                .takes_value(true)
                .help("Same as --lines"),
        )
        .arg(
            Arg::with_name(options::FROM_LINE)
                .short("F")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .value_name("number")
                .help("Display file beginning from line number"),
        )
        .arg(
            Arg::with_name(options::PATTERN)
                .short("P")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .help("Display file beginning from pattern match"),
        )
        */
        .arg(
            Arg::with_name(options::FILES)
                .required(false)
                .multiple(true)
                .help("Path to the files to be read"),
        )
        .get_matches_from(args);

    let mut buff = String::new();
    let silent = matches.is_present(options::SILENT);
    if let Some(files) = matches.values_of(options::FILES) {
        let mut stdout = setup_term();
        let length = files.len();

        let mut files_iter = files.peekable();
        while let (Some(file), next_file) = (files_iter.next(), files_iter.peek()) {
            let file = Path::new(file);
            if file.is_dir() {
                terminal::disable_raw_mode().unwrap();
                show_usage_error!("'{}' is a directory.", file.display());
                return 1;
            }
            if !file.exists() {
                terminal::disable_raw_mode().unwrap();
                show_error!("cannot open {}: No such file or directory", file.display());
                return 1;
            }
            if length > 1 {
                buff.push_str(&MULTI_FILE_TOP_PROMPT.replace("{}", file.to_str().unwrap()));
            }
            let mut reader = BufReader::new(File::open(file).unwrap());
            reader.read_to_string(&mut buff).unwrap();
            more(&buff, &mut stdout, next_file.copied(), silent);
            buff.clear();
        }
        reset_term(&mut stdout);
    } else if atty::isnt(atty::Stream::Stdin) {
        stdin().read_to_string(&mut buff).unwrap();
        let mut stdout = setup_term();
        more(&buff, &mut stdout, None, silent);
        reset_term(&mut stdout);
    } else {
        show_usage_error!("bad usage");
    }
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn setup_term() -> std::io::Stdout {
    let stdout = stdout();
    terminal::enable_raw_mode().unwrap();
    stdout
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn setup_term() -> usize {
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn reset_term(stdout: &mut std::io::Stdout) {
    terminal::disable_raw_mode().unwrap();
    // Clear the prompt
    queue!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
    // Move cursor to the beginning without printing new line
    print!("\r");
    stdout.flush().unwrap();
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn reset_term(_: &mut usize) {}

fn more(buff: &str, mut stdout: &mut Stdout, next_file: Option<&str>, silent: bool) {
    let (cols, rows) = terminal::size().unwrap();
    let lines = break_buff(buff, usize::from(cols));

    let mut pager = Pager::new(rows, lines, next_file, silent);
    pager.draw(stdout, false);
    if pager.should_close() {
        return;
    }

    loop {
        let mut wrong_key = false;
        if event::poll(Duration::from_millis(10)).unwrap() {
            match event::read().unwrap() {
                Event::Key(KeyEvent {
                    code: KeyCode::Char('q'),
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char('c'),
                    modifiers: KeyModifiers::CONTROL,
                }) => {
                    reset_term(&mut stdout);
                    std::process::exit(0);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Down,
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char(' '),
                    modifiers: KeyModifiers::NONE,
                }) => {
                    if pager.should_close() {
                        return;
                    } else {
                        pager.page_down();
                    }
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Up,
                    modifiers: KeyModifiers::NONE,
                }) => {
                    pager.page_up();
                }
                Event::Resize(col, row) => {
                    pager.page_resize(col, row);
                }
                // FIXME: Need to fix, there are more than just unknown keys.
                _ => {
                    wrong_key = true;
                }
            }

            pager.draw(stdout, wrong_key);
        }
    }
}

struct Pager<'a> {
    // The current line at the top of the screen
    upper_mark: usize,
    // The number of rows that fit on the screen
    content_rows: u16,
    lines: Vec<String>,
    next_file: Option<&'a str>,
    line_count: usize,
    silent: bool,
}

impl<'a> Pager<'a> {
    fn new(rows: u16, lines: Vec<String>, next_file: Option<&'a str>, silent: bool) -> Self {
        let line_count = lines.len();
        Self {
            upper_mark: 0,
            content_rows: rows.saturating_sub(1),
            lines,
            next_file,
            line_count,
            silent,
        }
    }

    fn should_close(&mut self) -> bool {
        self.upper_mark
            .saturating_add(self.content_rows.into())
            .eq(&self.line_count)
    }

    fn page_down(&mut self) {
        self.upper_mark = self
            .upper_mark
            .saturating_add(self.content_rows.into())
            .min(self.line_count.saturating_sub(self.content_rows.into()));
    }

    fn page_up(&mut self) {
        self.upper_mark = self.upper_mark.saturating_sub(self.content_rows.into());
    }

    // TODO: Deal with column size changes.
    fn page_resize(&mut self, _: u16, row: u16) {
        self.content_rows = row.saturating_sub(1);
    }

    fn draw(&self, stdout: &mut std::io::Stdout, wrong_key: bool) {
        let lower_mark = self
            .line_count
            .min(self.upper_mark.saturating_add(self.content_rows.into()));
        self.draw_lines(stdout);
        self.draw_prompt(stdout, lower_mark, wrong_key);
        stdout.flush().unwrap();
    }

    fn draw_lines(&self, stdout: &mut std::io::Stdout) {
        execute!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
        let displayed_lines = self
            .lines
            .iter()
            .skip(self.upper_mark)
            .take(self.content_rows.into());

        for line in displayed_lines {
            stdout
                .write_all(format!("\r{}\n", line).as_bytes())
                .unwrap();
        }
    }

    fn draw_prompt(&self, stdout: &mut Stdout, lower_mark: usize, wrong_key: bool) {
        let status_inner = if lower_mark == self.line_count {
            format!("Next file: {}", self.next_file.unwrap_or_default())
        } else {
            format!(
                "{}%",
                (lower_mark as f64 / self.line_count as f64 * 100.0).round() as u16
            )
        };

        let status = format!("--More--({})", status_inner);

        let banner = match (self.silent, wrong_key) {
            (true, true) => "[Press 'h' for instructions. (unimplemented)]".to_string(),
            (true, false) => format!("{}[Press space to continue, 'q' to quit.]", status),
            (false, true) => format!("{}{}", status, BELL),
            (false, false) => status,
        };

        write!(
            stdout,
            "\r{}{}{}",
            Attribute::Reverse,
            banner,
            Attribute::Reset
        )
        .unwrap();
    }
}

// Break the lines on the cols of the terminal
fn break_buff(buff: &str, cols: usize) -> Vec<String> {
    let mut lines = Vec::with_capacity(buff.lines().count());

    for l in buff.lines() {
        lines.append(&mut break_line(l, cols));
    }
    lines
}

fn break_line(line: &str, cols: usize) -> Vec<String> {
    let width = UnicodeWidthStr::width(line);
    let mut lines = Vec::new();
    if width < cols {
        lines.push(line.to_string());
        return lines;
    }

    let gr_idx = UnicodeSegmentation::grapheme_indices(line, true);
    let mut last_index = 0;
    let mut total_width = 0;
    for (index, grapheme) in gr_idx {
        let width = UnicodeWidthStr::width(grapheme);
        total_width += width;

        if total_width > cols {
            lines.push(line[last_index..index].to_string());
            last_index = index;
            total_width = width;
        }
    }

    if last_index != line.len() {
        lines.push(line[last_index..].to_string());
    }
    lines
}

#[cfg(test)]
mod tests {
    use super::break_line;
    use unicode_width::UnicodeWidthStr;

    #[test]
    fn test_break_lines_long() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..200 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);
        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        assert_eq!((80, 80, 40), (widths[0], widths[1], widths[2]));
    }

    #[test]
    fn test_break_lines_short() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..20 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);

        assert_eq!(20, lines[0].len());
    }

    #[test]
    fn test_break_line_zwj() {
        let mut test_string = String::with_capacity(1100);
        for _ in 0..20 {
            test_string.push_str("👩🏻🔬");
        }

        let lines = break_line(&test_string, 80);

        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        // Each 👩🏻🔬 is 6 character width it break line to the closest number to 80 => 6 * 13 = 78
        assert_eq!((78, 42), (widths[0], widths[1]));
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use imgui_sys::{self, ImVec4};

use model::{Color, FieldCoord, HexCoord};
use vec2::Vec2;

const SQRT_3: f32 = 1.732_050_807_568_877_f32;

// #f3e4cf
const FIELD_WHITE: [f32; 4] = [0.9529, 0.8941, 0.8118, 1.0];
// #998578
const FIELD_BLACK: [f32; 4] = [0.6, 0.5216, 0.4706, 1.0];
// #ffffff
const PIECE_WHITE: [f32; 4] = [1.0, 1.0, 1.0, 1.0];
// #000000
const PIECE_OUTLINE: [f32; 4] = [0.0, 0.0, 0.0, 1.0];
// #757575
const PIECE_BLACK: [f32; 4] = [0.4588, 0.4588, 0.4588, 1.0];

macro_rules! im_color {
    ($v:expr) => (imgui_sys::igColorConvertFloat4ToU32(ImVec4::from($v)))
}

pub fn draw_hex(coord: &HexCoord, origin: Vec2, size: f32) {
    for i in 0..6 {
        draw_field(&coord.to_field(i), origin, size);
    }
}

fn draw_field(coord: &FieldCoord, origin: Vec2, size: f32) {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    unsafe {
        let color = match coord.color() {
            Color::White => im_color!(FIELD_WHITE),
            Color::Black => im_color!(FIELD_BLACK),
        };

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddTriangleFilled(draw_list, v1.into(), v2.into(), v3.into(), color);
    }
}

pub fn highlight_field(color: [f32; 4], coord: &FieldCoord, origin: Vec2, size: f32) {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    unsafe {
        let highlight = im_color!(color);

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddTriangleFilled(
            draw_list,
            v1.into(),
            v2.into(),
            v3.into(),
            highlight,
        );
    }
}

pub fn highlight_field_dot(color: [f32; 4], coord: &FieldCoord, origin: Vec2, size: f32) {
    let center = field_center(coord, origin, size);
    unsafe {
        let highlight = im_color!(color);

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddCircleFilled(
            draw_list,
            center.into(),
            size / (4.0 * SQRT_3),
            highlight,
            15,
        );
    }
}

pub fn draw_piece(coord: &FieldCoord, origin: Vec2, size: f32) {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    let center = field_center(coord, origin, size);

    const SCALE: f32 = 0.7;

    let v1 = center + (v1 - center) * SCALE;
    let v2 = center + (v2 - center) * SCALE;
    let v3 = center + (v3 - center) * SCALE;

    // Linear equation derived by human testing and regression
    let outline_size = 0.032 * size - 0.335;

    unsafe {
        let color = match coord.color() {
            Color::White => im_color!(PIECE_WHITE),
            Color::Black => im_color!(PIECE_BLACK),
        };

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddTriangleFilled(draw_list, v1.into(), v2.into(), v3.into(), color);
        imgui_sys::ImDrawList_AddTriangle(
            draw_list,
            v1.into(),
            v2.into(),
            v3.into(),
            im_color!(PIECE_OUTLINE),
            outline_size,
        );
    }
}

fn field_center(coord: &FieldCoord, origin: Vec2, size: f32) -> Vec2 {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    let center_x = (v1.x + v2.x + v3.x) / 3.0;
    let min_y = (v1.y).min(v2.y).min(v3.y);
    let center_y = match coord.color() {
        Color::White => min_y + size / SQRT_3,
        Color::Black => min_y + size / (2.0 * SQRT_3),
    };

    Vec2::new(center_x, center_y)
}

fn field_vertexes(coord: &FieldCoord, origin: Vec2, size: f32) -> (Vec2, Vec2, Vec2) {
    let center = hex_to_pixel(&coord.to_hex(), origin, size);
    let height = size * SQRT_3 / 2.0;

    let west = Vec2::new(-size, 0.0);
    let east = Vec2::new(size, 0.0);

    let northwest = Vec2::new(-size / 2.0, -height);
    let southwest = Vec2::new(-size / 2.0, height);
    let northeast = Vec2::new(size / 2.0, -height);
    let southeast = Vec2::new(size / 2.0, height);

    let v1;
    let v2;
    match coord.f() {
        0 => {
            v1 = center + northwest;
            v2 = center + northeast;
        }
        1 => {
            v1 = center + northeast;
            v2 = center + east;
        }
        2 => {
            v1 = center + east;
            v2 = center + southeast;
        }
        3 => {
            v1 = center + southeast;
            v2 = center + southwest;
        }
        4 => {
            v1 = center + southwest;
            v2 = center + west;
        }
        5 => {
            v1 = center + west;
            v2 = center + northwest;
        }
        _ => panic!("You made the impossible possible."),
    };

    (center, v1, v2)
}

// Algorithm based on http://www.redblobgames.com/grids/hexagons/#hex-to-pixel
fn hex_to_pixel(coord: &HexCoord, origin: Vec2, size: f32) -> Vec2 {
    let x = coord.x() as f32;
    let y = coord.y() as f32;

    let p = Vec2::new(size * (3.0 / 2.0) * x, size * -SQRT_3 * (x / 2.0 + y));

    origin + p
}

// Algorithm based on http://www.redblobgames.com/grids/hexagons/#pixel-to-hex
pub fn pixel_to_field(p: Vec2, origin: Vec2, size: f32) -> Option<FieldCoord> {
    let v = p - origin;

    let q = v.x * (2.0 / 3.0) / size;
    let r = (-v.x - SQRT_3 * v.y) / (size * 3.0);

    if let Some(hex) = round_hex_coord(q, r) {
        /*
           To find the field, we subtract the converted hex coordinates (q, r) from the rounded hex
           coordinates (hex) to get the fractional part of the coordinates. Here is a diagram of a
           single hex, with the fractional coordinates of each of its vertexes in the hex coordinate
           system:

                (-1/3, 2/3) _______ (1/3, 1/3)
                           /\     /\
                          /  \   /  \
                         /    \ /    \
            (-2/3, 1/3) (----(0,0)----) (2/3, -1/3)
                         \    / \    /
                          \  /   \  /
              (-1/3, -1/3) \/_____\/ (1/3, -2/3)

           Suppose our converted coordinates were (q, r) = (1.333, 0.583). Our rounded coordinates
           are hex = (1, 1), and so our fractional coordinates are (0.333, -0.417).

           We now define three linear equations that will split this hexagon into its six fields:

                            _______ y = x (/)
                           /\     /\
                          /  \   /  \
                         /    \ /    \
                        (------X------) y = -x/2 (-)
                         \    / \    /
                          \  /   \  /
                           \/_____\/ y = -2x (\)

           Using these three equations as linear inequalities, we can check any pair of fractional
           coordinates and find which field it is in. (As a reminder, fields are numbered clockwise
           starting from 0 at the top.)

           For our example (0.333, -0.417), we see that:
               y + x/2 >= 0, so our field is above (-), and is either 5, 0, or 1.
               y +  2x >= 0, so our field is to the right of (\), and is either 0 or 1.
               y -   x <  0, so our field is to the right of (/), and is 1.
        */
        let x_diff = q - hex.x() as f32;
        let y_diff = r - hex.y() as f32;
        let mut i = 0;

        if y_diff + x_diff / 2.0 >= 0.0 {
            i |= 0b100;
        }
        if y_diff + x_diff * 2.0 >= 0.0 {
            i |= 0b010;
        }
        if y_diff - x_diff >= 0.0 {
            i |= 0b001;
        }

        // Using a lookup table because nested ifs are too confusing
        const INVALID: u32 = 6;
        let field_lookup = [3, 4, 2, INVALID, INVALID, 5, 1, 0];

        Some(hex.to_field(field_lookup[i]))
    } else {
        None
    }
}

// Algorithm from http://www.redblobgames.com/grids/hexagons/#rounding
fn round_hex_coord(x: f32, y: f32) -> Option<HexCoord> {
    let z = -x - y;

    let mut rx = x.round();
    let mut ry = y.round();
    let rz = z.round();

    let x_diff = (rx - x).abs();
    let y_diff = (ry - y).abs();
    let z_diff = (rz - z).abs();

    if x_diff > y_diff && x_diff > z_diff {
        rx = -ry - rz;
    } else if y_diff > z_diff {
        ry = -rx - rz;
    }

    let rx = rx as i32;
    let ry = ry as i32;

    HexCoord::try_new(rx, ry)
}
<commit_msg>Replace panic! with unreachable!<commit_after>/*
 * Copyright (C) 2017-2018 Ryan Huang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use imgui_sys::{self, ImVec4};

use model::{Color, FieldCoord, HexCoord};
use vec2::Vec2;

const SQRT_3: f32 = 1.732_050_807_568_877_f32;

// #f3e4cf
const FIELD_WHITE: [f32; 4] = [0.9529, 0.8941, 0.8118, 1.0];
// #998578
const FIELD_BLACK: [f32; 4] = [0.6, 0.5216, 0.4706, 1.0];
// #ffffff
const PIECE_WHITE: [f32; 4] = [1.0, 1.0, 1.0, 1.0];
// #000000
const PIECE_OUTLINE: [f32; 4] = [0.0, 0.0, 0.0, 1.0];
// #757575
const PIECE_BLACK: [f32; 4] = [0.4588, 0.4588, 0.4588, 1.0];

macro_rules! im_color {
    ($v:expr) => (imgui_sys::igColorConvertFloat4ToU32(ImVec4::from($v)))
}

pub fn draw_hex(coord: &HexCoord, origin: Vec2, size: f32) {
    for i in 0..6 {
        draw_field(&coord.to_field(i), origin, size);
    }
}

fn draw_field(coord: &FieldCoord, origin: Vec2, size: f32) {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    unsafe {
        let color = match coord.color() {
            Color::White => im_color!(FIELD_WHITE),
            Color::Black => im_color!(FIELD_BLACK),
        };

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddTriangleFilled(draw_list, v1.into(), v2.into(), v3.into(), color);
    }
}

pub fn highlight_field(color: [f32; 4], coord: &FieldCoord, origin: Vec2, size: f32) {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    unsafe {
        let highlight = im_color!(color);

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddTriangleFilled(
            draw_list,
            v1.into(),
            v2.into(),
            v3.into(),
            highlight,
        );
    }
}

pub fn highlight_field_dot(color: [f32; 4], coord: &FieldCoord, origin: Vec2, size: f32) {
    let center = field_center(coord, origin, size);
    unsafe {
        let highlight = im_color!(color);

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddCircleFilled(
            draw_list,
            center.into(),
            size / (4.0 * SQRT_3),
            highlight,
            15,
        );
    }
}

pub fn draw_piece(coord: &FieldCoord, origin: Vec2, size: f32) {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    let center = field_center(coord, origin, size);

    const SCALE: f32 = 0.7;

    let v1 = center + (v1 - center) * SCALE;
    let v2 = center + (v2 - center) * SCALE;
    let v3 = center + (v3 - center) * SCALE;

    // Linear equation derived by human testing and regression
    let outline_size = 0.032 * size - 0.335;

    unsafe {
        let color = match coord.color() {
            Color::White => im_color!(PIECE_WHITE),
            Color::Black => im_color!(PIECE_BLACK),
        };

        let draw_list = imgui_sys::igGetWindowDrawList();
        imgui_sys::ImDrawList_AddTriangleFilled(draw_list, v1.into(), v2.into(), v3.into(), color);
        imgui_sys::ImDrawList_AddTriangle(
            draw_list,
            v1.into(),
            v2.into(),
            v3.into(),
            im_color!(PIECE_OUTLINE),
            outline_size,
        );
    }
}

fn field_center(coord: &FieldCoord, origin: Vec2, size: f32) -> Vec2 {
    let (v1, v2, v3) = field_vertexes(coord, origin, size);
    let center_x = (v1.x + v2.x + v3.x) / 3.0;
    let min_y = (v1.y).min(v2.y).min(v3.y);
    let center_y = match coord.color() {
        Color::White => min_y + size / SQRT_3,
        Color::Black => min_y + size / (2.0 * SQRT_3),
    };

    Vec2::new(center_x, center_y)
}

fn field_vertexes(coord: &FieldCoord, origin: Vec2, size: f32) -> (Vec2, Vec2, Vec2) {
    let center = hex_to_pixel(&coord.to_hex(), origin, size);
    let height = size * SQRT_3 / 2.0;

    let west = Vec2::new(-size, 0.0);
    let east = Vec2::new(size, 0.0);

    let northwest = Vec2::new(-size / 2.0, -height);
    let southwest = Vec2::new(-size / 2.0, height);
    let northeast = Vec2::new(size / 2.0, -height);
    let southeast = Vec2::new(size / 2.0, height);

    let v1;
    let v2;
    match coord.f() {
        0 => {
            v1 = center + northwest;
            v2 = center + northeast;
        }
        1 => {
            v1 = center + northeast;
            v2 = center + east;
        }
        2 => {
            v1 = center + east;
            v2 = center + southeast;
        }
        3 => {
            v1 = center + southeast;
            v2 = center + southwest;
        }
        4 => {
            v1 = center + southwest;
            v2 = center + west;
        }
        5 => {
            v1 = center + west;
            v2 = center + northwest;
        }
        _ => unreachable!(),
    };

    (center, v1, v2)
}

// Algorithm based on http://www.redblobgames.com/grids/hexagons/#hex-to-pixel
fn hex_to_pixel(coord: &HexCoord, origin: Vec2, size: f32) -> Vec2 {
    let x = coord.x() as f32;
    let y = coord.y() as f32;

    let p = Vec2::new(size * (3.0 / 2.0) * x, size * -SQRT_3 * (x / 2.0 + y));

    origin + p
}

// Algorithm based on http://www.redblobgames.com/grids/hexagons/#pixel-to-hex
pub fn pixel_to_field(p: Vec2, origin: Vec2, size: f32) -> Option<FieldCoord> {
    let v = p - origin;

    let q = v.x * (2.0 / 3.0) / size;
    let r = (-v.x - SQRT_3 * v.y) / (size * 3.0);

    if let Some(hex) = round_hex_coord(q, r) {
        /*
           To find the field, we subtract the converted hex coordinates (q, r) from the rounded hex
           coordinates (hex) to get the fractional part of the coordinates. Here is a diagram of a
           single hex, with the fractional coordinates of each of its vertexes in the hex coordinate
           system:

                (-1/3, 2/3) _______ (1/3, 1/3)
                           /\     /\
                          /  \   /  \
                         /    \ /    \
            (-2/3, 1/3) (----(0,0)----) (2/3, -1/3)
                         \    / \    /
                          \  /   \  /
              (-1/3, -1/3) \/_____\/ (1/3, -2/3)

           Suppose our converted coordinates were (q, r) = (1.333, 0.583). Our rounded coordinates
           are hex = (1, 1), and so our fractional coordinates are (0.333, -0.417).

           We now define three linear equations that will split this hexagon into its six fields:

                            _______ y = x (/)
                           /\     /\
                          /  \   /  \
                         /    \ /    \
                        (------X------) y = -x/2 (-)
                         \    / \    /
                          \  /   \  /
                           \/_____\/ y = -2x (\)

           Using these three equations as linear inequalities, we can check any pair of fractional
           coordinates and find which field it is in. (As a reminder, fields are numbered clockwise
           starting from 0 at the top.)

           For our example (0.333, -0.417), we see that:
               y + x/2 >= 0, so our field is above (-), and is either 5, 0, or 1.
               y +  2x >= 0, so our field is to the right of (\), and is either 0 or 1.
               y -   x <  0, so our field is to the right of (/), and is 1.
        */
        let x_diff = q - hex.x() as f32;
        let y_diff = r - hex.y() as f32;
        let mut i = 0;

        if y_diff + x_diff / 2.0 >= 0.0 {
            i |= 0b100;
        }
        if y_diff + x_diff * 2.0 >= 0.0 {
            i |= 0b010;
        }
        if y_diff - x_diff >= 0.0 {
            i |= 0b001;
        }

        // Using a lookup table because nested ifs are too confusing
        const INVALID: u32 = 6;
        let field_lookup = [3, 4, 2, INVALID, INVALID, 5, 1, 0];

        Some(hex.to_field(field_lookup[i]))
    } else {
        None
    }
}

// Algorithm from http://www.redblobgames.com/grids/hexagons/#rounding
fn round_hex_coord(x: f32, y: f32) -> Option<HexCoord> {
    let z = -x - y;

    let mut rx = x.round();
    let mut ry = y.round();
    let rz = z.round();

    let x_diff = (rx - x).abs();
    let y_diff = (ry - y).abs();
    let z_diff = (rz - z).abs();

    if x_diff > y_diff && x_diff > z_diff {
        rx = -ry - rz;
    } else if y_diff > z_diff {
        ry = -rx - rz;
    }

    let rx = rx as i32;
    let ry = ry as i32;

    HexCoord::try_new(rx, ry)
}
<|endoftext|>


<commit_before>use std::io::net::ip::{SocketAddr, Ipv4Addr};
use std::io::net::get_host_addresses;
use std::io::{self, Buffer, Reader, Writer, IoResult, BufferedStream, standard_error};
use std::mem;
use std::collections::BTreeMap;
use url::Url;

#[cfg(test)]
use test::Bencher;
#[cfg(test)]
use serialize::json::ToJson;

use nonce::Nonce;
use message::{WSMessage, WSHeader, WS_MASK, WS_LEN, WS_LEN16, WS_LEN64, WS_OPCONT, WS_OPCODE, WS_FIN};
use stream::NetworkStream;

pub struct WebSocket<S = NetworkStream> {
    stream: Option<BufferedStream<S>>,
    connected: bool,
    pub remote_addr: Option<SocketAddr>,
    pub url: Url,
    use_ssl: bool,
}

impl WebSocket {
    pub fn new(url: Url) -> IoResult<WebSocket> {
        let addr = match try!(url.domain()
            .map(|h| get_host_addresses(h)
                 .map(|v| v.into_iter().find(|&a| {
                     match a {
                         Ipv4Addr(..) => true,
                         _ => false
                     }
                 })))
            .unwrap_or(Err(standard_error(io::InvalidInput)))) {
                Some(a) => a,
                None => return Err(standard_error(io::FileNotFound))
            };

        let use_ssl = url.scheme[] == "wss";

        let port = match url.port() {
            Some(p) => p,
            None if use_ssl => 443,
            _ => 80
        };

        Ok(WebSocket {
            stream: None,
            connected: false,
            remote_addr: Some(SocketAddr{ ip: addr, port: port }),
            url: url,
            use_ssl: use_ssl,
        })
    }

    #[allow(unused_variable)]
    fn try_connect(&mut self) -> IoResult<()> {
        self.stream = Some(BufferedStream::new(try!(self.remote_addr.map(|ref a| NetworkStream::connect(format!("{}:{}", a.ip, a.port)[], self.use_ssl))
                               .unwrap_or_else(|| Err(standard_error(io::InvalidInput))))));
        Ok(())
    }

    fn write_request(&mut self, nonce: &str) -> IoResult<()> {
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::NotConnected)) };

        try!(s.write(format!("GET {} HTTP/1.1\r\n", self.url.serialize_path().unwrap_or("/".to_string())).as_bytes()));
        try!(s.write(format!("Host: {}\r\n", self.url.host().unwrap()).as_bytes()));
        try!(s.write(format!("Origin: {}\r\n", self.url.serialize_no_fragment()).as_bytes()));
        try!(s.write(format!("Sec-WebSocket-Key: {}\r\n", nonce).as_bytes()));

        try!(s.write(b"Upgrade: websocket\r\n"));
        try!(s.write(b"Connection: Upgrade\r\n"));
        try!(s.write(b"Sec-WebSocket-Protocol: char, superchat\r\n"));
        try!(s.write(b"Sec-WebSocket-Version: 13\r\n"));
        try!(s.write(b"\r\n"));

        s.flush()
    }

    fn read_response(&mut self, nonce: &str) -> IoResult<()> {
        let spaces: &[_] = &[' ', '\t', '\r', '\n'];
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::NotConnected)) };
        let status = try!(s.read_line())[].splitn(2, ' ').nth(1).and_then(|s| s.parse::<uint>());

        match status {
            Some(101) => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        let headers = s.lines().map(|r| r.unwrap_or("\r\n".to_string())) .take_while(|l| l[] != "\r\n")
            .map(|s| s[].splitn(1, ':').map(|s| s.trim_chars(spaces).to_string()).collect::<Vec<String>>())
            .map(|p| (p[0].to_string(), p[1].to_string()))
            .collect::<BTreeMap<String, String>>();

        try!(s.flush());

        let response = headers.find(&"Sec-WebSocket-Accept".to_string());
        match response {
            Some(r) if nonce == r[] => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        Ok(())
    }

    pub fn connect(&mut self) -> IoResult<()> {
        let mut nonce = Nonce::new();

        try!(self.try_connect());
        try!(self.write_request(nonce[]));

        nonce = nonce.encode();
        try!(self.read_response(nonce[]));

        self.connected = true;

        Ok(())
    }

    fn read_header(&mut self) -> IoResult<WSHeader> {
        // XXX: this is a bug, WSHeader should accept u16
        Ok(WSHeader::from_bits_truncate(try!(self.read_be_u16())))
    }

    fn read_length(&mut self, header: &WSHeader) -> IoResult<uint> {
        let wslen = *header & WS_LEN;
        if wslen == WS_LEN16 { self.read_be_u16().map(|v| v as uint) }
        else if wslen == WS_LEN64 { self.read_be_u64().map(|v| v as uint) }
        else { Ok(wslen.bits() as uint) }
    }

    pub fn read_message(&mut self) -> IoResult<WSMessage> {
        let header = try!(self.read_header());
        let len = try!(self.read_length(&header));

        let data = if header.contains(WS_MASK) {
            WebSocket::unmask_data(try!(self.read_exact(len)), try!(self.read_be_u32()))
        } else {
            try!(self.read_exact(len))
        };

        Ok(WSMessage { header: header, data: data })
    }

    fn unmask_data(data: Vec<u8>, mask: u32) -> Vec<u8> {
        data.iter().enumerate().map(|(i, b)| *b ^ (mask >> ((i % 4) << 3) & 0xff) as u8).collect::<Vec<u8>>()
    }

    // TODO: send_message(&mut self, &WSMessage) -> IoResult<()>

    pub fn iter(&mut self) -> WSMessages {
        WSMessages { sock: self }
    }
}

impl Reader for WebSocket {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
        match self.stream {
            Some(ref mut s) => s.read(buf),
            None => Err(standard_error(io::NotConnected))
        }
    }
}

impl Writer for WebSocket {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.write(buf),
            None => Err(standard_error(io::NotConnected))
        }
    }

    fn flush(&mut self) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.flush(),
            None => Err(standard_error(io::NotConnected))
        }
    }
}

impl Buffer for WebSocket {
    fn fill_buf<'a>(&'a mut self) -> IoResult<&'a [u8]> {
        match self.stream {
            Some(ref mut s) => s.fill_buf(),
            None => Err(standard_error(io::NotConnected))
        }
    }

    fn consume(&mut self, amt: uint) {
        match self.stream {
            Some(ref mut s) => s.consume(amt),
            None => ()
        }
    }
}

pub struct WSMessages<'a> {
    sock: &'a mut WebSocket
}

pub struct WSDefragMessages<'a> {
    underlying: &'a mut WSMessages<'a>,
    buffer: WSMessage
}

impl<'a> WSMessages<'a> {
    pub fn defrag(&'a mut self) -> WSDefragMessages<'a> {
        WSDefragMessages{ underlying: self, buffer: WSMessage{ header: WSHeader::empty(), data: Vec::new() } }
    }
}

impl<'a> Iterator for WSMessages<'a> {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        self.sock.read_message().ok()
    }
}

impl<'a> WSDefragMessages<'a> {
    fn popbuf(&mut self) -> Option<WSMessage> {
        if self.buffer.data.is_empty() {
            None
        } else {
            let buf = WSMessage{ header: WSHeader::empty(), data: Vec::new() };
            mem::swap(&mut self.buffer, &mut buf);
            Some(buf)
        }
    }

    fn swapbuf(&mut self, msg: &mut WSMessage) -> WSMessage {
        mem::swap(&mut self.buffer, msg);
        return msg;
    }
}

impl<'a> Iterator for WSDefragMessages<'a> {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        loop {
            match self.underlying.next() {
                None => return self.popbuf(),
                Some(msg) => if msg.header.contains(WS_FIN) {
                    if msg.header & WS_OPCODE == WS_OPCONT {
                        self.buffer.push(msg);
                        return self.popbuf();
                    } else {
                        return Some(msg);
                    }

                } else {
                    if msg.header & WS_OPCODE == WS_OPCONT {
                        self.buffer.push(msg);
                    } else {
                        return self.swapbuf(&mut msg);
                    }
                }
            }
        }
    }
}

#[bench]
#[allow(dead_code)]
fn test_connect(b: &mut Bencher) {
    let url = Url::parse("wss://stream.pushbullet.com/websocket/").unwrap();
    let mut ws = WebSocket::new(url).unwrap();

    match ws.connect() {
        Err(e) => panic!("error: {}", e),
        _ => ()
    }
    let msg = ws.read_message().unwrap();
    println!("received: {} {} {}", msg, msg.to_string(), msg.to_json());
    for msg in ws.iter() {
        println!("{}", msg.to_string());
    }
}
<commit_msg>fix return type<commit_after>use std::io::net::ip::{SocketAddr, Ipv4Addr};
use std::io::net::get_host_addresses;
use std::io::{self, Buffer, Reader, Writer, IoResult, BufferedStream, standard_error};
use std::mem;
use std::collections::BTreeMap;
use url::Url;

#[cfg(test)]
use test::Bencher;
#[cfg(test)]
use serialize::json::ToJson;

use nonce::Nonce;
use message::{WSMessage, WSHeader, WS_MASK, WS_LEN, WS_LEN16, WS_LEN64, WS_OPCONT, WS_OPCODE, WS_FIN};
use stream::NetworkStream;

pub struct WebSocket<S = NetworkStream> {
    stream: Option<BufferedStream<S>>,
    connected: bool,
    pub remote_addr: Option<SocketAddr>,
    pub url: Url,
    use_ssl: bool,
}

impl WebSocket {
    pub fn new(url: Url) -> IoResult<WebSocket> {
        let addr = match try!(url.domain()
            .map(|h| get_host_addresses(h)
                 .map(|v| v.into_iter().find(|&a| {
                     match a {
                         Ipv4Addr(..) => true,
                         _ => false
                     }
                 })))
            .unwrap_or(Err(standard_error(io::InvalidInput)))) {
                Some(a) => a,
                None => return Err(standard_error(io::FileNotFound))
            };

        let use_ssl = url.scheme[] == "wss";

        let port = match url.port() {
            Some(p) => p,
            None if use_ssl => 443,
            _ => 80
        };

        Ok(WebSocket {
            stream: None,
            connected: false,
            remote_addr: Some(SocketAddr{ ip: addr, port: port }),
            url: url,
            use_ssl: use_ssl,
        })
    }

    #[allow(unused_variable)]
    fn try_connect(&mut self) -> IoResult<()> {
        self.stream = Some(BufferedStream::new(try!(self.remote_addr.map(|ref a| NetworkStream::connect(format!("{}:{}", a.ip, a.port)[], self.use_ssl))
                               .unwrap_or_else(|| Err(standard_error(io::InvalidInput))))));
        Ok(())
    }

    fn write_request(&mut self, nonce: &str) -> IoResult<()> {
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::NotConnected)) };

        try!(s.write(format!("GET {} HTTP/1.1\r\n", self.url.serialize_path().unwrap_or("/".to_string())).as_bytes()));
        try!(s.write(format!("Host: {}\r\n", self.url.host().unwrap()).as_bytes()));
        try!(s.write(format!("Origin: {}\r\n", self.url.serialize_no_fragment()).as_bytes()));
        try!(s.write(format!("Sec-WebSocket-Key: {}\r\n", nonce).as_bytes()));

        try!(s.write(b"Upgrade: websocket\r\n"));
        try!(s.write(b"Connection: Upgrade\r\n"));
        try!(s.write(b"Sec-WebSocket-Protocol: char, superchat\r\n"));
        try!(s.write(b"Sec-WebSocket-Version: 13\r\n"));
        try!(s.write(b"\r\n"));

        s.flush()
    }

    fn read_response(&mut self, nonce: &str) -> IoResult<()> {
        let spaces: &[_] = &[' ', '\t', '\r', '\n'];
        let s = match self.stream { Some(ref mut s) => s, None => return Err(standard_error(io::NotConnected)) };
        let status = try!(s.read_line())[].splitn(2, ' ').nth(1).and_then(|s| s.parse::<uint>());

        match status {
            Some(101) => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        let headers = s.lines().map(|r| r.unwrap_or("\r\n".to_string())) .take_while(|l| l[] != "\r\n")
            .map(|s| s[].splitn(1, ':').map(|s| s.trim_chars(spaces).to_string()).collect::<Vec<String>>())
            .map(|p| (p[0].to_string(), p[1].to_string()))
            .collect::<BTreeMap<String, String>>();

        try!(s.flush());

        let response = headers.find(&"Sec-WebSocket-Accept".to_string());
        match response {
            Some(r) if nonce == r[] => (),
            _ => return Err(standard_error(io::InvalidInput))
        }

        Ok(())
    }

    pub fn connect(&mut self) -> IoResult<()> {
        let mut nonce = Nonce::new();

        try!(self.try_connect());
        try!(self.write_request(nonce[]));

        nonce = nonce.encode();
        try!(self.read_response(nonce[]));

        self.connected = true;

        Ok(())
    }

    fn read_header(&mut self) -> IoResult<WSHeader> {
        // XXX: this is a bug, WSHeader should accept u16
        Ok(WSHeader::from_bits_truncate(try!(self.read_be_u16())))
    }

    fn read_length(&mut self, header: &WSHeader) -> IoResult<uint> {
        let wslen = *header & WS_LEN;
        if wslen == WS_LEN16 { self.read_be_u16().map(|v| v as uint) }
        else if wslen == WS_LEN64 { self.read_be_u64().map(|v| v as uint) }
        else { Ok(wslen.bits() as uint) }
    }

    pub fn read_message(&mut self) -> IoResult<WSMessage> {
        let header = try!(self.read_header());
        let len = try!(self.read_length(&header));

        let data = if header.contains(WS_MASK) {
            WebSocket::unmask_data(try!(self.read_exact(len)), try!(self.read_be_u32()))
        } else {
            try!(self.read_exact(len))
        };

        Ok(WSMessage { header: header, data: data })
    }

    fn unmask_data(data: Vec<u8>, mask: u32) -> Vec<u8> {
        data.iter().enumerate().map(|(i, b)| *b ^ (mask >> ((i % 4) << 3) & 0xff) as u8).collect::<Vec<u8>>()
    }

    // TODO: send_message(&mut self, &WSMessage) -> IoResult<()>

    pub fn iter(&mut self) -> WSMessages {
        WSMessages { sock: self }
    }
}

impl Reader for WebSocket {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<uint> {
        match self.stream {
            Some(ref mut s) => s.read(buf),
            None => Err(standard_error(io::NotConnected))
        }
    }
}

impl Writer for WebSocket {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.write(buf),
            None => Err(standard_error(io::NotConnected))
        }
    }

    fn flush(&mut self) -> IoResult<()> {
        match self.stream {
            Some(ref mut s) => s.flush(),
            None => Err(standard_error(io::NotConnected))
        }
    }
}

impl Buffer for WebSocket {
    fn fill_buf<'a>(&'a mut self) -> IoResult<&'a [u8]> {
        match self.stream {
            Some(ref mut s) => s.fill_buf(),
            None => Err(standard_error(io::NotConnected))
        }
    }

    fn consume(&mut self, amt: uint) {
        match self.stream {
            Some(ref mut s) => s.consume(amt),
            None => ()
        }
    }
}

pub struct WSMessages<'a> {
    sock: &'a mut WebSocket
}

pub struct WSDefragMessages<'a> {
    underlying: &'a mut WSMessages<'a>,
    buffer: WSMessage
}

impl<'a> WSMessages<'a> {
    pub fn defrag(&'a mut self) -> WSDefragMessages<'a> {
        WSDefragMessages{ underlying: self, buffer: WSMessage{ header: WSHeader::empty(), data: Vec::new() } }
    }
}

impl<'a> Iterator for WSMessages<'a> {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        self.sock.read_message().ok()
    }
}

impl<'a> WSDefragMessages<'a> {
    fn popbuf(&mut self) -> Option<WSMessage> {
        if self.buffer.data.is_empty() {
            None
        } else {
            let buf = WSMessage{ header: WSHeader::empty(), data: Vec::new() };
            mem::swap(&mut self.buffer, &mut buf);
            Some(buf)
        }
    }

    fn swapbuf(&mut self, msg: &mut WSMessage) -> WSMessage {
        mem::swap(&mut self.buffer, msg);
        return msg;
    }
}

impl<'a> Iterator for WSDefragMessages<'a> {
    type Item = WSMessage;
    fn next(&mut self) -> Option<WSMessage> {
        loop {
            match self.underlying.next() {
                None => return self.popbuf(),
                Some(msg) => if msg.header.contains(WS_FIN) {
                    if msg.header & WS_OPCODE == WS_OPCONT {
                        self.buffer.push(msg);
                        return self.popbuf();
                    } else {
                        return Some(msg);
                    }

                } else {
                    if msg.header & WS_OPCODE == WS_OPCONT {
                        self.buffer.push(msg);
                    } else {
                        return Some(self.swapbuf(&mut msg));
                    }
                }
            }
        }
    }
}

#[bench]
#[allow(dead_code)]
fn test_connect(b: &mut Bencher) {
    let url = Url::parse("wss://stream.pushbullet.com/websocket/").unwrap();
    let mut ws = WebSocket::new(url).unwrap();

    match ws.connect() {
        Err(e) => panic!("error: {}", e),
        _ => ()
    }
    let msg = ws.read_message().unwrap();
    println!("received: {} {} {}", msg, msg.to_string(), msg.to_json());
    for msg in ws.iter() {
        println!("{}", msg.to_string());
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2019 DDN. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

use crate::ActionResponse;
use bootstrap_components::{bs_button, bs_modal};
use futures::Future;
use iml_duration_picker::duration_picker;
use iml_environment::csrf_token;
use iml_utils::dispatch_custom_event;
use seed::{attrs, class, div, dom_types::At, fetch, form, h4, i, input, label, prelude::*};

#[derive(Debug, serde::Serialize)]
pub struct StratagemScan {
    pub filesystem: u32,
    pub report_duration: Option<u64>,
    pub purge_duration: Option<u64>,
}

#[derive(Debug, Default)]
pub struct Model {
    pub data: Option<StratagemScan>,
    pub open: bool,
    pub disabled: bool,
    pub is_locked: bool,
    pub report_config: iml_duration_picker::Model,
    pub purge_config: iml_duration_picker::Model,
}

impl Model {
    /// Validates the input fields for the duration picker.
    /// It would be much better if we relied on HTML5 validation,
    /// but we need a solution to https://github.com/David-OConnor/seed/issues/82 first.
    fn validate(&mut self) {
        let check = self
            .report_config
            .value
            .and_then(|r| self.purge_config.value.map(|p| r >= p))
            .unwrap_or(false);

        if check {
            self.report_config.validation_message =
                Some("Report duration must be less than Purge duration.".into());
        }
    }
}

#[derive(Clone, Debug)]
pub enum Msg {
    ScanStratagem(u32),
    OpenModal,
    CloseModal,
    ReportConfig(iml_duration_picker::Msg),
    PurgeConfig(iml_duration_picker::Msg),
    StratagemScanned(fetch::FetchObject<ActionResponse>),
}

pub fn update(msg: Msg, model: &mut Model, orders: &mut impl Orders<Msg>) {
    match msg {
        Msg::ScanStratagem(filesystem) => {
            model.disabled = true;

            let x = StratagemScan {
                filesystem,
                purge_duration: model.purge_config.value_as_ms(),
                report_duration: model.report_config.value_as_ms(),
            };

            orders.perform_cmd(scan(&x));
        }
        Msg::StratagemScanned(fetch_object) => {
            model.open = false;

            match fetch_object.response() {
                Ok(response) => {
                    dispatch_custom_event("show_command_modal", &response.data);
                }
                Err(fail_reason) => {
                    model.disabled = false;
                    log::error!("Fetch error: {:#?}", fail_reason);
                }
            }
        }
        Msg::ReportConfig(msg) => {
            iml_duration_picker::update(msg, &mut model.report_config);
            model.validate();
        }
        Msg::PurgeConfig(msg) => {
            iml_duration_picker::update(msg, &mut model.purge_config);
            model.validate();
        }
        Msg::OpenModal => {
            model.open = true;
        }
        Msg::CloseModal => {
            model.open = false;
        }
    }
}

fn scan(x: &StratagemScan) -> impl Future<Item = Msg, Error = Msg> {
    seed::fetch::Request::new("/api/run_stratagem/")
        .method(seed::fetch::Method::Post)
        .header(
            "X-CSRFToken",
            &csrf_token().expect("Couldn't get csrf token."),
        )
        .send_json(x)
        .fetch_json(Msg::StratagemScanned)
}

pub fn view(fs_id: u32, model: &Model) -> Vec<Node<Msg>> {
    let mut scan_now_button = bs_button::btn(
        class![bs_button::BTN_PRIMARY],
        vec![
            Node::new_text("Scan Filesystem Now"),
            i![class!["far", "fa-chart-bar"]],
        ],
    )
    .add_style("margin-left", px(15));

    if !model.disabled && !model.is_locked {
        scan_now_button = scan_now_button.add_listener(simple_ev(Ev::Click, Msg::OpenModal));
    } else {
        scan_now_button = scan_now_button.add_attr(At::Disabled.as_str(), "disabled");
    }

    let mut xs = vec![scan_now_button];

    if model.open {
        xs.append(&mut scan_modal(fs_id, &model));
    }

    xs
}

fn scan_modal(fs_id: u32, model: &Model) -> Vec<Node<Msg>> {
    let mut scan_button = bs_button::btn(
        class![bs_button::BTN_SUCCESS],
        vec![Node::new_text(if model.disabled {
            "Scanning..."
        } else if model.is_locked {
            "Locked 🔒"
        } else {
            "Scan Now"
        })],
    )
    .add_listener(mouse_ev(Ev::Click, move |_| Msg::ScanStratagem(fs_id)));

    if model.disabled || model.is_locked {
        scan_button = scan_button.add_attr("disabled", true);
    }

    vec![
        bs_modal::backdrop(),
        bs_modal::modal(vec![
            bs_modal::header(vec![h4![
                "Scan Filesystem Now ",
                i![class!["far", "fa-chart-bar"]]
            ]]),
            bs_modal::body(vec![form![
                div![
                    class!["form-group"],
                    label!["Generate report for files older than"],
                    div![
                        class!["input-group"],
                        duration_picker(
                            &model.report_config,
                            input![attrs! {At::AutoFocus => true, At::Placeholder => "Optional"}]
                        )
                        .map_message(Msg::ReportConfig)
                    ],
                ],
                div![
                    class!["form-group"],
                    label!["Purge Files older than"],
                    div![
                        class!["input-group"],
                        duration_picker(
                            &model.purge_config,
                            input![attrs! {At::Placeholder => "Optional"}]
                        )
                        .map_message(Msg::PurgeConfig)
                    ],
                ],
            ]]),
            bs_modal::footer(vec![
                scan_button,
                bs_button::btn(
                    class![bs_button::BTN_DEFAULT],
                    vec![
                        Node::new_text("Close"),
                        i![class!["far", "fa-times-circle"]],
                    ],
                )
                .add_listener(simple_ev(Ev::Click, Msg::CloseModal)),
            ]),
        ]),
    ]
}
<commit_msg>Update Lock (#1173)<commit_after>// Copyright (c) 2019 DDN. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

use crate::ActionResponse;
use bootstrap_components::{bs_button, bs_modal};
use futures::Future;
use iml_duration_picker::duration_picker;
use iml_environment::csrf_token;
use iml_utils::dispatch_custom_event;
use seed::{attrs, class, div, dom_types::At, fetch, form, h4, i, input, label, prelude::*, span};

#[derive(Debug, serde::Serialize)]
pub struct StratagemScan {
    pub filesystem: u32,
    pub report_duration: Option<u64>,
    pub purge_duration: Option<u64>,
}

#[derive(Debug, Default)]
pub struct Model {
    pub data: Option<StratagemScan>,
    pub open: bool,
    pub disabled: bool,
    pub is_locked: bool,
    pub report_config: iml_duration_picker::Model,
    pub purge_config: iml_duration_picker::Model,
}

impl Model {
    /// Validates the input fields for the duration picker.
    /// It would be much better if we relied on HTML5 validation,
    /// but we need a solution to https://github.com/David-OConnor/seed/issues/82 first.
    fn validate(&mut self) {
        let check = self
            .report_config
            .value
            .and_then(|r| self.purge_config.value.map(|p| r >= p))
            .unwrap_or(false);

        if check {
            self.report_config.validation_message =
                Some("Report duration must be less than Purge duration.".into());
        }
    }
}

#[derive(Clone, Debug)]
pub enum Msg {
    ScanStratagem(u32),
    OpenModal,
    CloseModal,
    ReportConfig(iml_duration_picker::Msg),
    PurgeConfig(iml_duration_picker::Msg),
    StratagemScanned(fetch::FetchObject<ActionResponse>),
}

pub fn update(msg: Msg, model: &mut Model, orders: &mut impl Orders<Msg>) {
    match msg {
        Msg::ScanStratagem(filesystem) => {
            model.disabled = true;

            let x = StratagemScan {
                filesystem,
                purge_duration: model.purge_config.value_as_ms(),
                report_duration: model.report_config.value_as_ms(),
            };

            orders.perform_cmd(scan(&x));
        }
        Msg::StratagemScanned(fetch_object) => {
            model.open = false;

            match fetch_object.response() {
                Ok(response) => {
                    dispatch_custom_event("show_command_modal", &response.data);
                }
                Err(fail_reason) => {
                    model.disabled = false;
                    log::error!("Fetch error: {:#?}", fail_reason);
                }
            }
        }
        Msg::ReportConfig(msg) => {
            iml_duration_picker::update(msg, &mut model.report_config);
            model.validate();
        }
        Msg::PurgeConfig(msg) => {
            iml_duration_picker::update(msg, &mut model.purge_config);
            model.validate();
        }
        Msg::OpenModal => {
            model.open = true;
        }
        Msg::CloseModal => {
            model.open = false;
        }
    }
}

fn scan(x: &StratagemScan) -> impl Future<Item = Msg, Error = Msg> {
    seed::fetch::Request::new("/api/run_stratagem/")
        .method(seed::fetch::Method::Post)
        .header(
            "X-CSRFToken",
            &csrf_token().expect("Couldn't get csrf token."),
        )
        .send_json(x)
        .fetch_json(Msg::StratagemScanned)
}

pub fn view(fs_id: u32, model: &Model) -> Vec<Node<Msg>> {
    let mut scan_now_button = bs_button::btn(
        class![bs_button::BTN_PRIMARY],
        vec![
            Node::new_text("Scan Filesystem Now"),
            i![class!["far", "fa-chart-bar"]],
        ],
    )
    .add_style("margin-left", px(15));

    if !model.disabled && !model.is_locked {
        scan_now_button = scan_now_button.add_listener(simple_ev(Ev::Click, Msg::OpenModal));
    } else {
        scan_now_button = scan_now_button.add_attr(At::Disabled.as_str(), "disabled");
    }

    let mut xs = vec![scan_now_button];

    if model.open {
        xs.append(&mut scan_modal(fs_id, &model));
    }

    xs
}

fn scan_modal(fs_id: u32, model: &Model) -> Vec<Node<Msg>> {
    let mut scan_button = bs_button::btn(
        class![bs_button::BTN_SUCCESS],
        vec![if model.disabled {
            Node::new_text("Scanning...")
        } else if model.is_locked {
            span!["Locked", i![class!["fas fa-lock"]]]
        } else {
            Node::new_text("Scan Now")
        }],
    )
    .add_listener(mouse_ev(Ev::Click, move |_| Msg::ScanStratagem(fs_id)));

    if model.disabled || model.is_locked {
        scan_button = scan_button.add_attr("disabled", true);
    }

    vec![
        bs_modal::backdrop(),
        bs_modal::modal(vec![
            bs_modal::header(vec![h4![
                "Scan Filesystem Now ",
                i![class!["far", "fa-chart-bar"]]
            ]]),
            bs_modal::body(vec![form![
                div![
                    class!["form-group"],
                    label!["Generate report for files older than"],
                    div![
                        class!["input-group"],
                        duration_picker(
                            &model.report_config,
                            input![attrs! {At::AutoFocus => true, At::Placeholder => "Optional"}]
                        )
                        .map_message(Msg::ReportConfig)
                    ],
                ],
                div![
                    class!["form-group"],
                    label!["Purge Files older than"],
                    div![
                        class!["input-group"],
                        duration_picker(
                            &model.purge_config,
                            input![attrs! {At::Placeholder => "Optional"}]
                        )
                        .map_message(Msg::PurgeConfig)
                    ],
                ],
            ]]),
            bs_modal::footer(vec![
                scan_button,
                bs_button::btn(
                    class![bs_button::BTN_DEFAULT],
                    vec![
                        Node::new_text("Close"),
                        i![class!["far", "fa-times-circle"]],
                    ],
                )
                .add_listener(simple_ev(Ev::Click, Msg::CloseModal)),
            ]),
        ]),
    ]
}
<|endoftext|>


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

//! This file contains code for interfacing with souffle (a dialect and
//! implementation of Datalog).
use std::fs;
use std::process::Command;

use crate::{
    ast::*,
    parsing::astconstructionvisitor,
    souffle::{universe_translation::*, lowering_ast_datalog::*, souffle_emitter},
};

/// Given a filename containing policy code this function parses it,
/// constructs an AST, translates the AST to the datalogIR, and emits a string
/// contianing souffle code.
pub fn emit_souffle(filename: &str) {
    let source = fs::read_to_string(filename).expect("failed to read input in emit_souffle");
    let prog = astconstructionvisitor::parse_program(&source[..]);
    let universe_handled_prog = UniverseHandlingPass::handle_universes(&prog);
    let dlir_prog = LoweringToDatalogPass::lower(&universe_handled_prog);
    let souffle_code = souffle_emitter::emit_program(&dlir_prog, &None);
    println!("souffle code from {}, \n{}", &filename, souffle_code);
}

/// Given a filename containing policy code this function parses it,
/// translates it to datalogIR, and emits souffle code to a new file with the
/// same name as the input but with the .dl extension in the out_dir
pub fn input_to_souffle_file(filename: &str, in_dir: &str, out_dir: &str) {
    let source = fs::read_to_string(&format!("{}/{}", in_dir, filename))
        .expect("failed to read input in input_to_souffle_file");
    let prog = astconstructionvisitor::parse_program(&source[..]);
    let universe_handled_prog = UniverseHandlingPass::handle_universes(&prog);
    let dlir_prog = LoweringToDatalogPass::lower(&universe_handled_prog);
    let souffle_code = souffle_emitter::emit_program(&dlir_prog, &None);
    fs::write(&format!("{}/{}.dl", out_dir, filename), souffle_code)
        .expect("failed to write output to file");
}

/// Given a parsed AstProgram, this function emits souffle code to a file.
/// This is needed in addition to the above function since there are other
/// passes that work on the highest level IR besides the one used to emit the
/// main code.
pub fn ast_to_souffle_file(prog: &AstProgram, filename: &str,
                           out_dir: &str, decl_skip: &Option<Vec<String>>) {
    let universe_handled_prog = UniverseHandlingPass::handle_universes(&prog);
    let dlir_prog = LoweringToDatalogPass::lower(&universe_handled_prog);
    let souffle_code = souffle_emitter::emit_program(&dlir_prog, decl_skip);
    fs::write(&format!("{}/{}.dl", out_dir, filename), souffle_code)
        .expect("failed to write output to file");
}

/// The function this function calls the souffle command on a .dl file. CSVs
/// generated by souffle are placed in outdir.
pub fn run_souffle(filename: &str, outdir: &str) {
    let souffle_command = match std::env::var("SOUFFLE_BIN") {
        Ok(val) => val,
        Err(e) => "souffle".to_string()
    };
    Command::new(souffle_command)
        .arg(&filename)
        .arg(&format!("-D{}/", outdir))
        .spawn()
        .expect("had an error in run_souffle()")
        .wait_with_output()
        .expect("had an error with souffle outputs");
}

/// The function true if the contents of a file are empty. This
/// is used to check if queries are true. Queries are translated into souffle
/// assertions that contain one entry if they are true and are empty
/// if they are false. These queries are emitted as CSV files by souffle.
pub fn is_file_empty(filename: &str) -> bool {
    let contents = fs::read_to_string(filename).expect("is_file_empty failed to read");
    contents == ""
}
<commit_msg>Add path information to Raksha errors (#466)<commit_after>/*
 * Copyright 2021 The Raksha Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! This file contains code for interfacing with souffle (a dialect and
//! implementation of Datalog).
use std::fs;
use std::process::Command;

use crate::{
    ast::*,
    parsing::astconstructionvisitor,
    souffle::{universe_translation::*, lowering_ast_datalog::*, souffle_emitter},
};

/// Given a filename containing policy code this function parses it,
/// constructs an AST, translates the AST to the datalogIR, and emits a string
/// contianing souffle code.
pub fn emit_souffle(filename: &str) {
    let source = fs::read_to_string(filename).expect("failed to read input in emit_souffle");
    let prog = astconstructionvisitor::parse_program(&source[..]);
    let universe_handled_prog = UniverseHandlingPass::handle_universes(&prog);
    let dlir_prog = LoweringToDatalogPass::lower(&universe_handled_prog);
    let souffle_code = souffle_emitter::emit_program(&dlir_prog, &None);
    println!("souffle code from {}, \n{}", &filename, souffle_code);
}

/// Given a filename containing policy code this function parses it,
/// translates it to datalogIR, and emits souffle code to a new file with the
/// same name as the input but with the .dl extension in the out_dir
pub fn input_to_souffle_file(filename: &str, in_dir: &str, out_dir: &str) {
    let in_path = format!("{}/{}", in_dir, filename);
    let source = fs::read_to_string(&in_path)
        .unwrap_or_else(|err|panic!("failed to read input in input_to_souffle_file: {}\n{}", &in_path, err));
    let prog = astconstructionvisitor::parse_program(&source[..]);
    let universe_handled_prog = UniverseHandlingPass::handle_universes(&prog);
    let dlir_prog = LoweringToDatalogPass::lower(&universe_handled_prog);
    let souffle_code = souffle_emitter::emit_program(&dlir_prog, &None);
    let out_path = format!("{}/{}.dl", out_dir, filename);
    fs::write(&out_path, souffle_code)
        .unwrap_or_else(|err|panic!("failed to read output in input_to_souffle_file: {}\n{}", &out_path, err));
}

/// Given a parsed AstProgram, this function emits souffle code to a file.
/// This is needed in addition to the above function since there are other
/// passes that work on the highest level IR besides the one used to emit the
/// main code.
pub fn ast_to_souffle_file(prog: &AstProgram, filename: &str,
                           out_dir: &str, decl_skip: &Option<Vec<String>>) {
    let universe_handled_prog = UniverseHandlingPass::handle_universes(&prog);
    let dlir_prog = LoweringToDatalogPass::lower(&universe_handled_prog);
    let souffle_code = souffle_emitter::emit_program(&dlir_prog, decl_skip);
    fs::write(&format!("{}/{}.dl", out_dir, filename), souffle_code)
        .expect("failed to write output to file");
}

/// The function this function calls the souffle command on a .dl file. CSVs
/// generated by souffle are placed in outdir.
pub fn run_souffle(filename: &str, outdir: &str) {
    let souffle_command = match std::env::var("SOUFFLE_BIN") {
        Ok(val) => val,
        Err(e) => "souffle".to_string()
    };
    Command::new(souffle_command)
        .arg(&filename)
        .arg(&format!("-D{}/", outdir))
        .spawn()
        .expect("had an error in run_souffle()")
        .wait_with_output()
        .expect("had an error with souffle outputs");
}

/// The function true if the contents of a file are empty. This
/// is used to check if queries are true. Queries are translated into souffle
/// assertions that contain one entry if they are true and are empty
/// if they are false. These queries are emitted as CSV files by souffle.
pub fn is_file_empty(filename: &str) -> bool {
    let contents = fs::read_to_string(filename).expect("is_file_empty failed to read");
    contents == ""
}
<|endoftext|>


<commit_before>//! The base traits that are used for serializing and deserializing rust types.
//!
//! `BerSerialize` implements serialization, and `BerDeserialize` implements deserialization.
//! Both traits depend upon the `Asn1Info` trait being implemented, to provide
//! metadata about the type (like ASN.1 tag and ASN.1 type).
//!
//! For both traits, only the `(de)serialize_value` functions should need overridden
//! to get custom behaviour. If you need to change how tags are handled, you may need
//! to override `_serialize_enc` or `_deserialize_with_tag`.
use std::io;

use ::Asn1Info;
use tag;
use err;
use ber::enc;

/// Provides the methods required to serialize this Rust type into an ASN.1 stream.
///
/// When implementing this for a simple primitive type, implementing `serialize_value`
/// should be all that's required. For more complex, structured types you may need to
/// implement `_serialize_enc` (this is called first by `serialize_enc`).
pub trait BerSerialize: Asn1Info {
  /// Serialize a value into ASN.1 data as DER.
  fn serialize<W: io::Write>(&self, writer: &mut W) -> Result<(), err::EncodeError> {
    self.serialize_enc(enc::DER, writer)
  }

  /// Serialize a value into ASN.1 data using a specific set of encoding rules.
  fn serialize_enc<E: enc::BerEncRules, W: io::Write>
      (&self, e: E, writer: &mut W) -> Result<(), err::EncodeError> {
    debug!("Encoding the type {}", Self::asn1_type());
    if let Some(r) = self._serialize_enc(e, writer) {
      return r;
    }

    let tag = match Self::asn1_tag() {
      Some(tag) => tag,
      None => {
        debug!("Type does not have a tag");
        return self.serialize_value(e, writer);
      },
    };

    try!(tag.write_tag(writer));

    // If this is indefinite length and constructed, write the data directly.
    if E::len_rules() == enc::LenEnc::Indefinite &&
       tag.constructed {
      debug!("Using indefinite length");
      try!(tag::Len::Indef.write_len(writer));
      try!(self.serialize_value(e, writer));
      try!(tag::Len::write_indef_end(writer));
    // Otherwise write to a Vec first, so we can write the length.
    } else {
      debug!("Using definite length encoding");
      let mut bytes: Vec<u8> = Vec::new();
      try!(self.serialize_value(e, &mut bytes));
      try!(tag::Len::write_len(Some(bytes.len() as tag::LenNum).into(), writer));
      try!(writer.write_all(&bytes));
    }

    Ok(())
  }

  /// An empty method that is called first by `serialize_enc` to allow custom
  /// handling, without losing normal serialization behaviour.
  ///
  /// Return `Some(())` to indicate stop normal behaviour, or `None` to continue.
  fn _serialize_enc<E: enc::BerEncRules, W: io::Write>
      (&self, e: E, writer: &mut W) -> Option<Result<(), err::EncodeError>> {
    let _ = (e, writer);
    None
  }

  /// Serialise a value into ASN.1 data, without a tag (implicit tagging).
  fn serialize_value<E: enc::BerEncRules, W: io::Write>
    (&self, e: E, writer: &mut W) -> Result<(), err::EncodeError>;
}

/// Provides the methods required to deserialize this Rust type from an ASN.1 stream.
///
/// When implementing this for a simple primitive type, implementing `deserialize_value`
/// should be all that's required. For more complex, structued types you may need to
/// implement `_deserialize_with_tag` (this is called first by `deserialize_with_tag).
pub trait BerDeserialize: Asn1Info + Sized {
  /// Deserialize ASN.1 data into a Rust value, accepting any valid BER.
  fn deserialize<I: Iterator<Item=io::Result<u8>>>(reader: &mut I) -> Result<Self, err::DecodeError> {
    Self::deserialize_enc(enc::BER, reader)
  }

  /// Deserialize ASN.1 data into a Rust value, using a specific set of encoding rules.
  fn deserialize_enc<E: enc::BerEncRules, I: Iterator<Item=io::Result<u8>>>
      (e: E, reader: &mut I) -> Result<Self, err::DecodeError> {
    let (tag, len) = try!(tag::read_taglen(reader));
    Self::deserialize_with_tag(e, reader, tag, len)
  }

  /// Deserialize ASN.1 data into a Rust value, using a specific set of encoding rules, and
  /// also providing the decoded tag and length.
  ///
  /// This function assumes the next bytes to decode are
  /// the BER length of this element.
  fn deserialize_with_tag<E: enc::BerEncRules, I: Iterator<Item=io::Result<u8>>>
      (e: E, reader: &mut I, tag: tag::Tag, len: tag::Len) -> Result<Self, err::DecodeError> {
    debug!("Decoding the type {}", Self::asn1_type());
    if let Some(r) = Self::_deserialize_with_tag(e, reader, tag, len) {
      return r;
    }

    if Some(tag) != Self::asn1_tag() {
      if let Some(our_tag) = Self::asn1_tag() {
        warn!("Expected tag {}, but found tag {}", tag, our_tag);
      } else {
        warn!("Expected tag {}, but found no tag", tag);
      }
      return Err(err::DecodeError::TagTypeMismatch);
    }

    // Handle any indefinite length error conditions.
    if len == tag::Len::Indef {
      // Return an error if the encoding rules only allow definite length
      // encoding.
      if E::len_rules() == enc::LenEnc::Definite {
        warn!("Encountered indefinite length encoding, but encoding rules don't allow this");
        return Err(err::DecodeError::IndefiniteLen);
      // If this element is primitve, the length isn't allowed to be indefinite length.
      } else if !tag.constructed {
        warn!("Encountered indefinite length encoding, but this is a primitive element");
        return Err(err::DecodeError::PrimIndef)
      }
    }
    // Read the main data.
    let item: Self = try!(Self::deserialize_value(e, reader, len));

    // If this is encoded with an indefinte length, try to read the end octets.
    if len == tag::Len::Indef {
      try!(tag::Len::read_indef_end(reader));
    }

    Ok(item)
  }

  /// An empty method that is called first by `deserialize_with_tag` to allow
  /// custom handling, without losing normal deserialization behaviour.
  ///
  /// Return `Some(..)` to return that value, or `None` to use normal behaviour.
  fn _deserialize_with_tag<E: enc::BerEncRules, I: Iterator<Item=io::Result<u8>>>
      (e: E, reader: &mut I, tag: tag::Tag, len: tag::Len) -> Option<Result<Self, err::DecodeError>> {
    let _ = (e, reader, tag, len);
    None
  }

  /// Deserialize an ASN.1 value from a BER stream, after having the tag and length
  /// decoded.
  ///
  /// The data length must be explicitly passed to this function. For primitive types,
  /// an error will be returned if this length is Indefinite.
  fn deserialize_value<E: enc::BerEncRules, I: Iterator<Item=io::Result<u8>>>
    (e: E, reader: &mut I, len: tag::Len) -> Result<Self, err::DecodeError> {
    let (_, _, _) = (e, reader, len);
    panic!("deserialize_value must be implemented for types that require it");
  }
}
<commit_msg>Update log message in asn1-cereal<commit_after>//! The base traits that are used for serializing and deserializing rust types.
//!
//! `BerSerialize` implements serialization, and `BerDeserialize` implements deserialization.
//! Both traits depend upon the `Asn1Info` trait being implemented, to provide
//! metadata about the type (like ASN.1 tag and ASN.1 type).
//!
//! For both traits, only the `(de)serialize_value` functions should need overridden
//! to get custom behaviour. If you need to change how tags are handled, you may need
//! to override `_serialize_enc` or `_deserialize_with_tag`.
use std::io;

use ::Asn1Info;
use tag;
use err;
use ber::enc;

/// Provides the methods required to serialize this Rust type into an ASN.1 stream.
///
/// When implementing this for a simple primitive type, implementing `serialize_value`
/// should be all that's required. For more complex, structured types you may need to
/// implement `_serialize_enc` (this is called first by `serialize_enc`).
pub trait BerSerialize: Asn1Info {
  /// Serialize a value into ASN.1 data as DER.
  fn serialize<W: io::Write>(&self, writer: &mut W) -> Result<(), err::EncodeError> {
    self.serialize_enc(enc::DER, writer)
  }

  /// Serialize a value into ASN.1 data using a specific set of encoding rules.
  fn serialize_enc<E: enc::BerEncRules, W: io::Write>
      (&self, e: E, writer: &mut W) -> Result<(), err::EncodeError> {
    debug!("Encoding the type {}", Self::asn1_type());
    if let Some(r) = self._serialize_enc(e, writer) {
      return r;
    }

    let tag = match Self::asn1_tag() {
      Some(tag) => tag,
      None => {
        debug!("Skipping encoding of empty tag");
        return self.serialize_value(e, writer);
      },
    };

    try!(tag.write_tag(writer));

    // If this is indefinite length and constructed, write the data directly.
    if E::len_rules() == enc::LenEnc::Indefinite &&
       tag.constructed {
      debug!("Using indefinite length");
      try!(tag::Len::Indef.write_len(writer));
      try!(self.serialize_value(e, writer));
      try!(tag::Len::write_indef_end(writer));
    // Otherwise write to a Vec first, so we can write the length.
    } else {
      debug!("Using definite length encoding");
      let mut bytes: Vec<u8> = Vec::new();
      try!(self.serialize_value(e, &mut bytes));
      try!(tag::Len::write_len(Some(bytes.len() as tag::LenNum).into(), writer));
      try!(writer.write_all(&bytes));
    }

    Ok(())
  }

  /// An empty method that is called first by `serialize_enc` to allow custom
  /// handling, without losing normal serialization behaviour.
  ///
  /// Return `Some(())` to indicate stop normal behaviour, or `None` to continue.
  fn _serialize_enc<E: enc::BerEncRules, W: io::Write>
      (&self, e: E, writer: &mut W) -> Option<Result<(), err::EncodeError>> {
    let _ = (e, writer);
    None
  }

  /// Serialise a value into ASN.1 data, without a tag (implicit tagging).
  fn serialize_value<E: enc::BerEncRules, W: io::Write>
    (&self, e: E, writer: &mut W) -> Result<(), err::EncodeError>;
}

/// Provides the methods required to deserialize this Rust type from an ASN.1 stream.
///
/// When implementing this for a simple primitive type, implementing `deserialize_value`
/// should be all that's required. For more complex, structued types you may need to
/// implement `_deserialize_with_tag` (this is called first by `deserialize_with_tag).
pub trait BerDeserialize: Asn1Info + Sized {
  /// Deserialize ASN.1 data into a Rust value, accepting any valid BER.
  fn deserialize<I: Iterator<Item=io::Result<u8>>>(reader: &mut I) -> Result<Self, err::DecodeError> {
    Self::deserialize_enc(enc::BER, reader)
  }

  /// Deserialize ASN.1 data into a Rust value, using a specific set of encoding rules.
  fn deserialize_enc<E: enc::BerEncRules, I: Iterator<Item=io::Result<u8>>>
      (e: E, reader: &mut I) -> Result<Self, err::DecodeError> {
    let (tag, len) = try!(tag::read_taglen(reader));
    Self::deserialize_with_tag(e, reader, tag, len)
  }

  /// Deserialize ASN.1 data into a Rust value, using a specific set of encoding rules, and
  /// also providing the decoded tag and length.
  ///
  /// This function assumes the next bytes to decode are
  /// the BER length of this element.
  fn deserialize_with_tag<E: enc::BerEncRules, I: Iterator<Item=io::Result<u8>>>
      (e: E, reader: &mut I, tag: tag::Tag, len: tag::Len) -> Result<Self, err::DecodeError> {
    debug!("Decoding the type {}", Self::asn1_type());
    if let Some(r) = Self::_deserialize_with_tag(e, reader, tag, len) {
      return r;
    }

    if Some(tag) != Self::asn1_tag() {
      if let Some(our_tag) = Self::asn1_tag() {
        warn!("Expected tag {}, but found tag {}", tag, our_tag);
      } else {
        warn!("Expected tag {}, but found no tag", tag);
      }
      return Err(err::DecodeError::TagTypeMismatch);
    }

    // Handle any indefinite length error conditions.
    if len == tag::Len::Indef {
      // Return an error if the encoding rules only allow definite length
      // encoding.
      if E::len_rules() == enc::LenEnc::Definite {
        warn!("Encountered indefinite length encoding, but encoding rules don't allow this");
        return Err(err::DecodeError::IndefiniteLen);
      // If this element is primitve, the length isn't allowed to be indefinite length.
      } else if !tag.constructed {
        warn!("Encountered indefinite length encoding, but this is a primitive element");
        return Err(err::DecodeError::PrimIndef)
      }
    }
    // Read the main data.
    let item: Self = try!(Self::deserialize_value(e, reader, len));

    // If this is encoded with an indefinte length, try to read the end octets.
    if len == tag::Len::Indef {
      try!(tag::Len::read_indef_end(reader));
    }

    Ok(item)
  }

  /// An empty method that is called first by `deserialize_with_tag` to allow
  /// custom handling, without losing normal deserialization behaviour.
  ///
  /// Return `Some(..)` to return that value, or `None` to use normal behaviour.
  fn _deserialize_with_tag<E: enc::BerEncRules, I: Iterator<Item=io::Result<u8>>>
      (e: E, reader: &mut I, tag: tag::Tag, len: tag::Len) -> Option<Result<Self, err::DecodeError>> {
    let _ = (e, reader, tag, len);
    None
  }

  /// Deserialize an ASN.1 value from a BER stream, after having the tag and length
  /// decoded.
  ///
  /// The data length must be explicitly passed to this function. For primitive types,
  /// an error will be returned if this length is Indefinite.
  fn deserialize_value<E: enc::BerEncRules, I: Iterator<Item=io::Result<u8>>>
    (e: E, reader: &mut I, len: tag::Len) -> Result<Self, err::DecodeError> {
    let (_, _, _) = (e, reader, len);
    panic!("deserialize_value must be implemented for types that require it");
  }
}
<|endoftext|>