<commit_before>//! Unix terminal formatter and drain for slog-rs
//!
//! ```
//! #[macro_use]
//! extern crate slog;
//! extern crate slog_term;
//!
//! use slog::*;
//!
//! fn main() {
//!     let root = Logger::root(slog_term::streamer().build(), o!("build-id" => "8dfljdf"));
//! }
//! ```
#![warn(missing_docs)]

#[macro_use]
extern crate slog;
extern crate isatty;
extern crate ansi_term;
extern crate chrono;

use std::{io, fmt, sync};

use ansi_term::Colour;
use isatty::{stderr_isatty, stdout_isatty};

use slog::Record;
use slog::{Level, OwnedKeyValueList};
use slog::format::Format as SlogFormat;
use slog::format::{Decorator, RecordDecorator};

/// Formatting mode
pub enum FormatMode {
    /// Compact logging format
    Compact,
    /// Full logging format
    Full,
}

/// Full formatting with optional color support
pub struct Format<D : Decorator> {
    mode : FormatMode,
    decorator : D,
    history : sync::Mutex<Vec<usize>>,
}

impl<D : Decorator> Format<D> {
    /// New Format format that prints using color
    pub fn new(mode : FormatMode, d : D) -> Self {
        Format {
            decorator: d,
            mode: mode,
            history : sync::Mutex::new(vec!()),
        }
    }

    fn print_msg_header(&self, io: &mut io::Write, rd : &D::RecordDecorator, info : &Record) -> io::Result<()> {
        let ts = chrono::Local::now();
        try!(rd.fmt_timestamp(
            io,
            format_args!("{} ", ts.format("%b %d %H:%M:%S%.3f"))
        ));
        try!(rd.fmt_level(
            io,
            format_args!("{} ", info.level().as_short_str())
        ));

        try!(rd.fmt_msg(
            io,
            format_args!("{}", info.msg()
            )
        ));
        Ok(())
    }
    fn format_full(&self,
                   io: &mut io::Write,
                   info: &Record,
                   logger_values: &OwnedKeyValueList)
                   -> io::Result<()> {

        let r_decorator = self.decorator.decorate(info);

        try!(self.print_msg_header(io, &r_decorator, info));
        let mut serializer = Serializer::new(io, r_decorator);

        for &(ref k, ref v) in logger_values.iter() {
            try!(serializer.print_comma());
            try!(v.serialize(info, k, &mut serializer));
        }

        for &(k, v) in info.values() {
            try!(serializer.print_comma());
            try!(v.serialize(info, k, &mut serializer));
        }
        let (mut io, _decorator_r) = serializer.finish();

        let _ = try!(write!(io, "\n"));

        Ok(())
    }

    fn format_compact(&self,
                      io: &mut io::Write,
                      info: &Record,
                      logger_values: &OwnedKeyValueList)
                      -> io::Result<()> {

        let r_decorator = self.decorator.decorate(info);
        let mut ser = Serializer::new(io, r_decorator);

        let indent = try!(self.format_recurse(&mut ser, info, logger_values));

        try!(self.print_indent(&mut ser.io, indent));

        try!(self.print_msg_header(&mut ser.io, &ser.decorator, info));

        for &(k, v) in info.values() {
            try!(ser.print_comma());
            try!(v.serialize(info, k, &mut ser));
        }
        try!(write!(&mut ser.io, "\n"));

        Ok(())
    }

    fn print_indent<W: io::Write>(
        &self,
        io: &mut W,
        indent : usize) -> io::Result<()> {
        for _ in 0..indent {
            try!(write!(io, "  "));
        }
        Ok(())
    }

    // record in the history, and check if should print
    // given set of values
    fn should_print(&self, address : usize, indent : usize) -> bool {
        let mut history = self.history.lock().unwrap();
        if history.len() <= indent {
            debug_assert_eq!(history.len(), indent);
            history.push(address);
            true
        } else {
            let should = history[indent] != address;
            history[indent] = address;
            should
        }
    }

    fn format_recurse<W: io::Write>(&self,
                                    ser : &mut Serializer<W, D::RecordDecorator>,
                                    info: &slog::Record,
                                    logger_values: &slog::OwnedKeyValueList)
                                    -> io::Result<usize> {
        let mut indent = if logger_values.parent().is_none() {
            0
        } else {
            try!(self.format_recurse(ser, info, logger_values.parent().as_ref().unwrap()))
        };

        if !logger_values.values().is_empty() {
            if self.should_print(logger_values.values() as *const _ as usize, indent) {
                try!(self.print_indent(&mut ser.io, indent));
                let mut clean = true;
                for &(ref k, ref v) in logger_values.values() {
                    if !clean {
                        try!(ser.print_comma());
                    }
                    try!(v.serialize(info, k, ser));
                    clean = false;
                }
                try!(write!(&mut ser.io, "\n"));
            }
            indent += 1
        }

        Ok(indent)
    }

}

fn severity_to_color(lvl: Level) -> u8 {
    match lvl {
        Level::Critical => 1,
        Level::Error => 9,
        Level::Warning => 3,
        Level::Info => 2,
        Level::Debug => 6,
        Level::Trace => 4,
    }
}

/// Decorator that can colors things
pub struct ColorDecorator {
    use_color : bool,
}

impl ColorDecorator {
    /// New decorator that does color records
    pub fn new_colored() -> Self {
        ColorDecorator { use_color: true}
    }
    /// New decorator that does not color records
    pub fn new_plain() -> Self {
        ColorDecorator { use_color: true}
    }
}

/// Decorator that colors things
pub struct ColorRecordDecorator {
    level_color : Option<Colour>,
    key_style: Option<ansi_term::Style>,
}

impl Decorator for ColorDecorator {

    type RecordDecorator = ColorRecordDecorator;

    fn decorate(&self, record : &Record) -> ColorRecordDecorator {
        if self.use_color {
            ColorRecordDecorator {
                level_color: Some(Colour::Fixed(severity_to_color(record.level()))),
                key_style: Some(ansi_term::Style::new().bold()),
            }
        } else {
            ColorRecordDecorator {
                level_color: None,
                key_style: None,
            }
        }
    }
}

impl RecordDecorator for ColorRecordDecorator {

    fn fmt_level(&self, io : &mut io::Write, args : fmt::Arguments) -> io::Result<()> {
        if let Some(level_color) = self.level_color {
            write!(io, "{}", level_color.paint(format!("{}", args)))
        } else {
            io.write_fmt(args)
        }
    }

    fn fmt_msg(&self, io : &mut io::Write, args : fmt::Arguments) -> io::Result<()> {
        if let Some(key_style) = self.key_style {
            write!(io, "{}", key_style.paint(format!("{}", args)))
        } else {
            io.write_fmt(args)
        }
    }

    fn fmt_key(&self, io : &mut io::Write, args : fmt::Arguments) -> io::Result<()> {
        self.fmt_msg(io, args)
    }
}

struct Serializer<W, D : RecordDecorator> {
    io: W,
    decorator : D,
}

impl<W: io::Write, D: RecordDecorator> Serializer<W, D> {
    fn new(io: W, d : D) -> Self {
        Serializer {
            io: io,
            decorator: d,
        }
    }

    fn print_comma(&mut self) -> io::Result<()> {
        try!(self.decorator.fmt_separator(
            &mut self.io,
            format_args!(", ")
        ));
        Ok(())
    }

    fn finish(self) -> (W, D) {
        (self.io, self.decorator)
    }
}

macro_rules! s(
    ($s:expr, $k:expr, $v:expr) => {
        try!($s.decorator.fmt_key(&mut $s.io, format_args!("{}", $k)));
        try!($s.decorator.fmt_separator(&mut $s.io, format_args!(": ")));
        try!($s.decorator.fmt_value(&mut $s.io, format_args!("{}", $v)));
    };
);

impl<W: io::Write, D : RecordDecorator> slog::ser::Serializer for Serializer<W, D> {
    fn emit_none(&mut self, key: &str) -> io::Result<()> {
        s!(self, key, "None");
        Ok(())
    }
    fn emit_unit(&mut self, key: &str) -> io::Result<()> {
        s!(self, key, "()");
        Ok(())
    }

    fn emit_bool(&mut self, key: &str, val: bool) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }

    fn emit_char(&mut self, key: &str, val: char) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }

    fn emit_usize(&mut self, key: &str, val: usize) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_isize(&mut self, key: &str, val: isize) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }

    fn emit_u8(&mut self, key: &str, val: u8) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i8(&mut self, key: &str, val: i8) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_u16(&mut self, key: &str, val: u16) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i16(&mut self, key: &str, val: i16) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_u32(&mut self, key: &str, val: u32) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i32(&mut self, key: &str, val: i32) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_f32(&mut self, key: &str, val: f32) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_u64(&mut self, key: &str, val: u64) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i64(&mut self, key: &str, val: i64) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_f64(&mut self, key: &str, val: f64) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_str(&mut self, key: &str, val: &str) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
}

impl<D : Decorator+Send+Sync> SlogFormat for Format<D> {
    fn format(&self,
              io: &mut io::Write,
              info: &Record,
              logger_values: &OwnedKeyValueList)
              -> io::Result<()> {
        match self.mode {
            FormatMode::Compact => self.format_compact(io, info, logger_values),
            FormatMode::Full => self.format_full(io, info, logger_values),
        }
    }
}

/// Streamer builder
pub struct StreamerBuilder {
    color: Option<bool>, // None = auto
    stdout: bool,
    async: bool,
    mode: FormatMode,
}

impl StreamerBuilder {

    /// New `StreamerBuilder`
    pub fn new() -> Self {
        StreamerBuilder {
            color: None,
            stdout: true,
            async: false,
            mode: FormatMode::Full,
        }
    }

    /// Force colored output
    pub fn color(mut self) -> Self {
        self.color = Some(true);
        self
    }

    /// Force plain output
    pub fn plain(mut self) -> Self {
        self.color = Some(false);
        self
    }

    /// Auto detect color (default)
    pub fn auto_color(mut self) -> Self {
        self.color = None;
        self
    }

    /// Output to stderr
    pub fn stderr(mut self) -> Self {
        self.stdout = false;
        self
    }

    /// Output to stdout (default)
    pub fn stdout(mut self) -> Self {
        self.stdout = true;
        self
    }

    /// Output using full mode
    pub fn full(mut self) -> Self {
        self.mode = FormatMode::Full;
        self
    }

    /// Output using compact mode (default)
    pub fn compact(mut self) -> Self {
        self.mode = FormatMode::Compact;
        self
    }

    /// Use asynchronous streamer
    pub fn async(mut self) -> Self {
        self.async = true;
        self
    }

    /// Use synchronous streamer (default)
    pub fn sync(mut self) -> Self {
        self.async = false;
        self
    }

    /// Build the streamer
    pub fn build(self) -> Box<slog::Drain> {
        let color = self.color.unwrap_or(
            if self.stdout {
                stdout_isatty()
            } else {
                stderr_isatty()
            });

        let format = Format::new(
            self.mode,
            ColorDecorator{use_color: color}
        );
        let io = if self.stdout {
            Box::new(io::stdout()) as Box<io::Write+Send>
        } else {
            Box::new(io::stderr()) as Box<io::Write+Send>
        };

        if self.async {
            Box::new(
                slog::AsyncStreamer::new(
                    io ,
                    format,
                ))
        } else {
            Box::new(
                slog::Streamer::new(
                    io ,
                    format,
                ))
        }
    }
}

/// Build `Streamer`/`AsyncStreamer`
pub fn streamer() -> StreamerBuilder {
    StreamerBuilder::new()
}

/*
/// Short for `Format::colored()`
fn format_colored() -> Format<ColorDecorator> {
    Format::new(
        FormatMode::Compact,
        ColorDecorator{use_color: true}
    )
}

/// Short for `Format::plain()`
fn format_plain() -> Format<ColorDecorator> {
    Format::new(
        FormatMode::Compact,
        ColorDecorator{use_color: false}
    )
}

fn auto_format(is_atty : bool) -> Format<ColorDecorator> {
    if is_atty {
        format_colored()
    } else {
        format_plain()
    }
}

/// Drain to `stdout`
///
/// Automatically using color if printing to tty
pub fn stdout() -> Streamer<io::Stdout, Format<ColorDecorator>> {
    Streamer::new(
        io::stdout(),
        auto_format(stdout_isatty()),
    )
}

/// Drain to `stderr`
///
/// Automatically using color if output goes to tty
pub fn stderr() -> Streamer<io::Stderr, Format<ColorDecorator>> {
    Streamer::new(
        io::stderr(),
        auto_format(stderr_isatty())
    )
}

/// Asynchronous drain to `stdout`
///
/// Automatically using color if printing to tty
pub fn async_stdout() -> AsyncStreamer<Format<ColorDecorator>> {
    AsyncStreamer::new(
        io::stdout(),
        auto_format(stdout_isatty()),
    )
}

/// Asynchronos drain to `stderr`
///
/// Automatically using color if output goes to tty
pub fn async_stderr() -> AsyncStreamer<Format<ColorDecorator>> {
    AsyncStreamer::new(
        io::stderr(),
        auto_format(stderr_isatty()),
    )
}
*/
<commit_msg>Remove dead code.<commit_after>//! Unix terminal formatter and drain for slog-rs
//!
//! ```
//! #[macro_use]
//! extern crate slog;
//! extern crate slog_term;
//!
//! use slog::*;
//!
//! fn main() {
//!     let root = Logger::root(slog_term::streamer().build(), o!("build-id" => "8dfljdf"));
//! }
//! ```
#![warn(missing_docs)]

#[macro_use]
extern crate slog;
extern crate isatty;
extern crate ansi_term;
extern crate chrono;

use std::{io, fmt, sync};

use ansi_term::Colour;
use isatty::{stderr_isatty, stdout_isatty};

use slog::Record;
use slog::{Level, OwnedKeyValueList};
use slog::format::Format as SlogFormat;
use slog::format::{Decorator, RecordDecorator};

/// Formatting mode
pub enum FormatMode {
    /// Compact logging format
    Compact,
    /// Full logging format
    Full,
}

/// Full formatting with optional color support
pub struct Format<D : Decorator> {
    mode : FormatMode,
    decorator : D,
    history : sync::Mutex<Vec<usize>>,
}

impl<D : Decorator> Format<D> {
    /// New Format format that prints using color
    pub fn new(mode : FormatMode, d : D) -> Self {
        Format {
            decorator: d,
            mode: mode,
            history : sync::Mutex::new(vec!()),
        }
    }

    fn print_msg_header(&self, io: &mut io::Write, rd : &D::RecordDecorator, info : &Record) -> io::Result<()> {
        let ts = chrono::Local::now();
        try!(rd.fmt_timestamp(
            io,
            format_args!("{} ", ts.format("%b %d %H:%M:%S%.3f"))
        ));
        try!(rd.fmt_level(
            io,
            format_args!("{} ", info.level().as_short_str())
        ));

        try!(rd.fmt_msg(
            io,
            format_args!("{}", info.msg()
            )
        ));
        Ok(())
    }
    fn format_full(&self,
                   io: &mut io::Write,
                   info: &Record,
                   logger_values: &OwnedKeyValueList)
                   -> io::Result<()> {

        let r_decorator = self.decorator.decorate(info);

        try!(self.print_msg_header(io, &r_decorator, info));
        let mut serializer = Serializer::new(io, r_decorator);

        for &(ref k, ref v) in logger_values.iter() {
            try!(serializer.print_comma());
            try!(v.serialize(info, k, &mut serializer));
        }

        for &(k, v) in info.values() {
            try!(serializer.print_comma());
            try!(v.serialize(info, k, &mut serializer));
        }
        let (mut io, _decorator_r) = serializer.finish();

        let _ = try!(write!(io, "\n"));

        Ok(())
    }

    fn format_compact(&self,
                      io: &mut io::Write,
                      info: &Record,
                      logger_values: &OwnedKeyValueList)
                      -> io::Result<()> {

        let r_decorator = self.decorator.decorate(info);
        let mut ser = Serializer::new(io, r_decorator);

        let indent = try!(self.format_recurse(&mut ser, info, logger_values));

        try!(self.print_indent(&mut ser.io, indent));

        try!(self.print_msg_header(&mut ser.io, &ser.decorator, info));

        for &(k, v) in info.values() {
            try!(ser.print_comma());
            try!(v.serialize(info, k, &mut ser));
        }
        try!(write!(&mut ser.io, "\n"));

        Ok(())
    }

    fn print_indent<W: io::Write>(
        &self,
        io: &mut W,
        indent : usize) -> io::Result<()> {
        for _ in 0..indent {
            try!(write!(io, "  "));
        }
        Ok(())
    }

    // record in the history, and check if should print
    // given set of values
    fn should_print(&self, address : usize, indent : usize) -> bool {
        let mut history = self.history.lock().unwrap();
        if history.len() <= indent {
            debug_assert_eq!(history.len(), indent);
            history.push(address);
            true
        } else {
            let should = history[indent] != address;
            history[indent] = address;
            should
        }
    }

    fn format_recurse<W: io::Write>(&self,
                                    ser : &mut Serializer<W, D::RecordDecorator>,
                                    info: &slog::Record,
                                    logger_values: &slog::OwnedKeyValueList)
                                    -> io::Result<usize> {
        let mut indent = if logger_values.parent().is_none() {
            0
        } else {
            try!(self.format_recurse(ser, info, logger_values.parent().as_ref().unwrap()))
        };

        if !logger_values.values().is_empty() {
            if self.should_print(logger_values.values() as *const _ as usize, indent) {
                try!(self.print_indent(&mut ser.io, indent));
                let mut clean = true;
                for &(ref k, ref v) in logger_values.values() {
                    if !clean {
                        try!(ser.print_comma());
                    }
                    try!(v.serialize(info, k, ser));
                    clean = false;
                }
                try!(write!(&mut ser.io, "\n"));
            }
            indent += 1
        }

        Ok(indent)
    }

}

fn severity_to_color(lvl: Level) -> u8 {
    match lvl {
        Level::Critical => 1,
        Level::Error => 9,
        Level::Warning => 3,
        Level::Info => 2,
        Level::Debug => 6,
        Level::Trace => 4,
    }
}

/// Decorator that can colors things
pub struct ColorDecorator {
    use_color : bool,
}

impl ColorDecorator {
    /// New decorator that does color records
    pub fn new_colored() -> Self {
        ColorDecorator { use_color: true}
    }
    /// New decorator that does not color records
    pub fn new_plain() -> Self {
        ColorDecorator { use_color: true}
    }
}

/// Decorator that colors things
pub struct ColorRecordDecorator {
    level_color : Option<Colour>,
    key_style: Option<ansi_term::Style>,
}

impl Decorator for ColorDecorator {

    type RecordDecorator = ColorRecordDecorator;

    fn decorate(&self, record : &Record) -> ColorRecordDecorator {
        if self.use_color {
            ColorRecordDecorator {
                level_color: Some(Colour::Fixed(severity_to_color(record.level()))),
                key_style: Some(ansi_term::Style::new().bold()),
            }
        } else {
            ColorRecordDecorator {
                level_color: None,
                key_style: None,
            }
        }
    }
}

impl RecordDecorator for ColorRecordDecorator {

    fn fmt_level(&self, io : &mut io::Write, args : fmt::Arguments) -> io::Result<()> {
        if let Some(level_color) = self.level_color {
            write!(io, "{}", level_color.paint(format!("{}", args)))
        } else {
            io.write_fmt(args)
        }
    }

    fn fmt_msg(&self, io : &mut io::Write, args : fmt::Arguments) -> io::Result<()> {
        if let Some(key_style) = self.key_style {
            write!(io, "{}", key_style.paint(format!("{}", args)))
        } else {
            io.write_fmt(args)
        }
    }

    fn fmt_key(&self, io : &mut io::Write, args : fmt::Arguments) -> io::Result<()> {
        self.fmt_msg(io, args)
    }
}

struct Serializer<W, D : RecordDecorator> {
    io: W,
    decorator : D,
}

impl<W: io::Write, D: RecordDecorator> Serializer<W, D> {
    fn new(io: W, d : D) -> Self {
        Serializer {
            io: io,
            decorator: d,
        }
    }

    fn print_comma(&mut self) -> io::Result<()> {
        try!(self.decorator.fmt_separator(
            &mut self.io,
            format_args!(", ")
        ));
        Ok(())
    }

    fn finish(self) -> (W, D) {
        (self.io, self.decorator)
    }
}

macro_rules! s(
    ($s:expr, $k:expr, $v:expr) => {
        try!($s.decorator.fmt_key(&mut $s.io, format_args!("{}", $k)));
        try!($s.decorator.fmt_separator(&mut $s.io, format_args!(": ")));
        try!($s.decorator.fmt_value(&mut $s.io, format_args!("{}", $v)));
    };
);

impl<W: io::Write, D : RecordDecorator> slog::ser::Serializer for Serializer<W, D> {
    fn emit_none(&mut self, key: &str) -> io::Result<()> {
        s!(self, key, "None");
        Ok(())
    }
    fn emit_unit(&mut self, key: &str) -> io::Result<()> {
        s!(self, key, "()");
        Ok(())
    }

    fn emit_bool(&mut self, key: &str, val: bool) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }

    fn emit_char(&mut self, key: &str, val: char) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }

    fn emit_usize(&mut self, key: &str, val: usize) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_isize(&mut self, key: &str, val: isize) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }

    fn emit_u8(&mut self, key: &str, val: u8) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i8(&mut self, key: &str, val: i8) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_u16(&mut self, key: &str, val: u16) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i16(&mut self, key: &str, val: i16) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_u32(&mut self, key: &str, val: u32) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i32(&mut self, key: &str, val: i32) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_f32(&mut self, key: &str, val: f32) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_u64(&mut self, key: &str, val: u64) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_i64(&mut self, key: &str, val: i64) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_f64(&mut self, key: &str, val: f64) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
    fn emit_str(&mut self, key: &str, val: &str) -> io::Result<()> {
        s!(self, key, val);
        Ok(())
    }
}

impl<D : Decorator+Send+Sync> SlogFormat for Format<D> {
    fn format(&self,
              io: &mut io::Write,
              info: &Record,
              logger_values: &OwnedKeyValueList)
              -> io::Result<()> {
        match self.mode {
            FormatMode::Compact => self.format_compact(io, info, logger_values),
            FormatMode::Full => self.format_full(io, info, logger_values),
        }
    }
}

/// Streamer builder
pub struct StreamerBuilder {
    color: Option<bool>, // None = auto
    stdout: bool,
    async: bool,
    mode: FormatMode,
}

impl StreamerBuilder {

    /// New `StreamerBuilder`
    pub fn new() -> Self {
        StreamerBuilder {
            color: None,
            stdout: true,
            async: false,
            mode: FormatMode::Full,
        }
    }

    /// Force colored output
    pub fn color(mut self) -> Self {
        self.color = Some(true);
        self
    }

    /// Force plain output
    pub fn plain(mut self) -> Self {
        self.color = Some(false);
        self
    }

    /// Auto detect color (default)
    pub fn auto_color(mut self) -> Self {
        self.color = None;
        self
    }

    /// Output to stderr
    pub fn stderr(mut self) -> Self {
        self.stdout = false;
        self
    }

    /// Output to stdout (default)
    pub fn stdout(mut self) -> Self {
        self.stdout = true;
        self
    }

    /// Output using full mode
    pub fn full(mut self) -> Self {
        self.mode = FormatMode::Full;
        self
    }

    /// Output using compact mode (default)
    pub fn compact(mut self) -> Self {
        self.mode = FormatMode::Compact;
        self
    }

    /// Use asynchronous streamer
    pub fn async(mut self) -> Self {
        self.async = true;
        self
    }

    /// Use synchronous streamer (default)
    pub fn sync(mut self) -> Self {
        self.async = false;
        self
    }

    /// Build the streamer
    pub fn build(self) -> Box<slog::Drain> {
        let color = self.color.unwrap_or(
            if self.stdout {
                stdout_isatty()
            } else {
                stderr_isatty()
            });

        let format = Format::new(
            self.mode,
            ColorDecorator{use_color: color}
        );
        let io = if self.stdout {
            Box::new(io::stdout()) as Box<io::Write+Send>
        } else {
            Box::new(io::stderr()) as Box<io::Write+Send>
        };

        if self.async {
            Box::new(
                slog::AsyncStreamer::new(
                    io ,
                    format,
                ))
        } else {
            Box::new(
                slog::Streamer::new(
                    io ,
                    format,
                ))
        }
    }
}

/// Build `Streamer`/`AsyncStreamer`
pub fn streamer() -> StreamerBuilder {
    StreamerBuilder::new()
}
<|endoftext|>


<commit_before>use failure::Error;
use ggez::{graphics, Context, GameResult};
use ggez_goodies::scene::{Scene, SceneSwitch};
use lazy_static::lazy_static;
use warmy::{LogicalKey, Store};

use crate::game::{Game, SceneState};
use crate::input::{Axis, Button, InputEvent};
use crate::objects::TextureAtlas;
use crate::scenes::FSceneSwitch;
use crate::text::Text;

use super::menu::Menu;
use super::transition::FadeStyle;
use super::Fade;

struct MenuImage {
    sheet: &'static str,
    image: usize,
    x: u32,
    y: u32,
}

const ARROW_POSITIONS: [u32; 4] = [85, 110, 152, 172];

lazy_static! {
    static ref ON: Text = Text {
        string: "On".to_string(),
        font: "bold.f".to_string(),
        bordered: true,
        centered: false,
        x: 214,
        y: 108,
    };
    static ref OFF: Text = Text {
        string: "Off".to_string(),
        font: "bold.f".to_string(),
        bordered: true,
        centered: false,
        x: 214,
        y: 108,
    };
    static ref PLAYER_COUNT: Text = Text {
        string: "1".to_string(),
        font: "bold.f".to_string(),
        bordered: true,
        centered: false,
        x: 214,
        y: 83,
    };
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
enum MenuOption {
    Players = 0,
    Gore = 1,
    Practice = 2,
    SelectKnight = 3,
}

impl MenuOption {
    fn from_i32(n: i32) -> MenuOption {
        match n {
            0 => MenuOption::Players,
            1 => MenuOption::Gore,
            2 => MenuOption::Practice,
            3 => MenuOption::SelectKnight,
            _ => panic!("Something went really wrong selecting menu options"),
        }
    }
}

#[derive(Debug)]
pub struct MainMenuScene {
    menu: Menu,
    done: bool,
    fade_out_done: bool,
    selected_option: MenuOption,
}

impl MainMenuScene {
    pub fn new(ctx: &mut Context, store: &mut Store<Context>) -> Result<Self, Error> {
        let menu = Menu::new(ctx, store, "/menu.yaml")?;
        // .unwrap_or_else(|| panic!("error in menu.yaml, must have 'background")),
        Ok(Self {
            menu,
            done: false,
            fade_out_done: false,
            selected_option: MenuOption::Players,
        })
    }
}

impl MainMenuScene {
    fn draw_arrow(&mut self, game: &mut Game, ctx: &mut Context) -> GameResult<()> {
        // we don't want to save this image to the game cache as we only
        // want to apply the pallete to the the arrow. We've extended the palette
        // by 16 empty colours so to_rgba8 won't break.
        let atlas = game
            .store
            .get::<_, TextureAtlas>(&LogicalKey::new(&self.menu.screen.cursor.sheet), ctx)
            // TODO: raise error
            .expect("Couldn't find sel.cel yaml metadata");

        let atlas_dimension = atlas.borrow().image.width as u32;
        let arrow_image = graphics::Image::from_rgba8(
            ctx,
            atlas_dimension as u16,
            atlas_dimension as u16,
            &atlas.borrow().image.to_rgba8(&self.menu.palette),
        )?;
        let rect = atlas.borrow().rects[self.menu.screen.cursor.image];
        let texture_size = atlas.borrow().image.width as f32;
        let y = ARROW_POSITIONS[self.selected_option as usize];
        let draw_params = graphics::DrawParam {
            src: graphics::Rect {
                x: rect.x as f32 / texture_size,
                y: rect.y as f32 / texture_size,
                w: rect.w as f32 / texture_size,
                h: rect.h as f32 / texture_size,
            },
            dest: graphics::Point2::new(self.menu.screen.cursor.x as f32 * 3.0, y as f32 * 3.0),
            scale: graphics::Point2::new(3.0, 3.0),
            ..Default::default()
        };
        graphics::draw_ex(ctx, &arrow_image, draw_params)?;
        Ok(())
    }

    fn draw_gore_option(&mut self, game: &mut Game, ctx: &mut Context) -> GameResult<()> {
        let mut batch = if game.gore_on {
            ON.as_sprite_batch(ctx, game, &self.menu.palette, self.menu.palette_hash)
                .expect("error drawing ON")
        } else {
            OFF.as_sprite_batch(ctx, game, &self.menu.palette, self.menu.palette_hash)
                .expect("error drawing OFF")
        };

        graphics::draw_ex(
            ctx,
            &batch,
            graphics::DrawParam {
                dest: graphics::Point2::new(0.0, 0.0),
                scale: graphics::Point2::new(3.0, 3.0),
                ..Default::default()
            },
        )?;
        batch.clear();
        Ok(())
    }

    fn draw_player_count_option(&mut self, game: &mut Game, ctx: &mut Context) -> GameResult<()> {
        let mut text = PLAYER_COUNT.clone();
        text.string = game.num_players.to_string();
        let mut batch = text
            .as_sprite_batch(ctx, game, &self.menu.palette, self.menu.palette_hash)
            .expect("error drawing PLAYER_COUNT");
        graphics::draw_ex(
            ctx,
            &batch,
            graphics::DrawParam {
                dest: graphics::Point2::new(0.0, 0.0),
                scale: graphics::Point2::new(3.0, 3.0),
                ..Default::default()
            },
        )?;
        batch.clear();
        Ok(())
    }
}

impl Scene<Game, InputEvent> for MainMenuScene {
    fn update(&mut self, _game: &mut Game, _ctx: &mut Context) -> FSceneSwitch {
        if self.done {
            if !self.fade_out_done {
                SceneSwitch::push(Fade::new(274, 1, FadeStyle::Out))
            } else {
                SceneSwitch::Pop // shouldn't happen
            }
        } else {
            SceneSwitch::None
        }
    }

    fn draw(&mut self, game: &mut Game, ctx: &mut Context) -> GameResult<()> {
        graphics::set_background_color(ctx, graphics::Color::from((0, 0, 0, 255)));
        graphics::clear(ctx);
        self.menu.draw(game, ctx)?;
        self.draw_arrow(game, ctx)?;
        self.draw_gore_option(game, ctx)?;
        self.draw_player_count_option(game, ctx)?;
        Ok(())
    }

    fn name(&self) -> &str {
        "Main menu"
    }

    fn input(&mut self, gameworld: &mut Game, _event: InputEvent, started: bool) {
        let x = gameworld.input.get_axis_raw(Axis::Horz1) as i32;
        if gameworld.input.get_button_down(Button::Fire1) {
            match self.selected_option {
                MenuOption::Practice => {
                    self.done = true;
                    gameworld.next_scene = SceneState::Practice;
                }
                MenuOption::Gore => gameworld.gore_on = !gameworld.gore_on,
                MenuOption::Players => gameworld.num_players = (gameworld.num_players % 4) + 1,
                _ => (),
            }
        } else if x != 0 {
            match self.selected_option {
                MenuOption::Gore => gameworld.gore_on = !gameworld.gore_on,
                MenuOption::Players => {
                    if x > 0 {
                        gameworld.num_players = (gameworld.num_players % 4) + x;
                    } else {
                        gameworld.num_players = (gameworld.num_players + 3 + x) % 4 + 1;
                    }
                }
                _ => (),
            }
        }

        if !started {
            let y = gameworld.input.get_axis_raw(Axis::Vert1) as i32;
            if y > 0 {
                self.selected_option = MenuOption::from_i32((self.selected_option as i32 + y) % 4);
            } else if y < 0 {
                self.selected_option =
                    MenuOption::from_i32((self.selected_option as i32 + 4 + y) % 4);
            }
        }
    }
}
<commit_msg>ggez-v0.5: update main_menu drawing.<commit_after>use failure::Error;
use ggez::nalgebra::{Point2, Vector2};
use ggez::{graphics, Context, GameResult};
use ggez_goodies::scene::{Scene, SceneSwitch};
use lazy_static::lazy_static;
use warmy::{LogicalKey, Store};

use crate::game::{Game, SceneState};
use crate::input::{Axis, Button, InputEvent};
use crate::objects::TextureAtlas;
use crate::scenes::FSceneSwitch;
use crate::text::Text;

use super::menu::Menu;
use super::transition::FadeStyle;
use super::Fade;

struct MenuImage {
    sheet: &'static str,
    image: usize,
    x: u32,
    y: u32,
}

const ARROW_POSITIONS: [u32; 4] = [85, 110, 152, 172];

lazy_static! {
    static ref ON: Text = Text {
        string: "On".to_string(),
        font: "bold.f".to_string(),
        bordered: true,
        centered: false,
        x: 214,
        y: 108,
    };
    static ref OFF: Text = Text {
        string: "Off".to_string(),
        font: "bold.f".to_string(),
        bordered: true,
        centered: false,
        x: 214,
        y: 108,
    };
    static ref PLAYER_COUNT: Text = Text {
        string: "1".to_string(),
        font: "bold.f".to_string(),
        bordered: true,
        centered: false,
        x: 214,
        y: 83,
    };
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
enum MenuOption {
    Players = 0,
    Gore = 1,
    Practice = 2,
    SelectKnight = 3,
}

impl MenuOption {
    fn from_i32(n: i32) -> MenuOption {
        match n {
            0 => MenuOption::Players,
            1 => MenuOption::Gore,
            2 => MenuOption::Practice,
            3 => MenuOption::SelectKnight,
            _ => panic!("Something went really wrong selecting menu options"),
        }
    }
}

#[derive(Debug)]
pub struct MainMenuScene {
    menu: Menu,
    done: bool,
    fade_out_done: bool,
    selected_option: MenuOption,
}

impl MainMenuScene {
    pub fn new(ctx: &mut Context, store: &mut Store<Context>) -> Result<Self, Error> {
        let menu = Menu::new(ctx, store, "/menu.yaml")?;
        // .unwrap_or_else(|| panic!("error in menu.yaml, must have 'background")),
        Ok(Self {
            menu,
            done: false,
            fade_out_done: false,
            selected_option: MenuOption::Players,
        })
    }
}

impl MainMenuScene {
    fn draw_arrow(&mut self, game: &mut Game, ctx: &mut Context) -> GameResult<()> {
        // we don't want to save this image to the game cache as we only
        // want to apply the pallete to the the arrow. We've extended the palette
        // by 16 empty colours so to_rgba8 won't break.
        let atlas = game
            .store
            .get::<_, TextureAtlas>(&LogicalKey::new(&self.menu.screen.cursor.sheet), ctx)
            // TODO: raise error
            .expect("Couldn't find sel.cel yaml metadata");

        let atlas_dimension = atlas.borrow().image.width as u32;
        let arrow_image = graphics::Image::from_rgba8(
            ctx,
            atlas_dimension as u16,
            atlas_dimension as u16,
            &atlas.borrow().image.to_rgba8(&self.menu.palette),
        )?;
        let rect = atlas.borrow().rects[self.menu.screen.cursor.image];
        let texture_size = atlas.borrow().image.width as f32;
        let y = ARROW_POSITIONS[self.selected_option as usize];
        let draw_params = graphics::DrawParam::default()
            .src(graphics::Rect {
                x: rect.x as f32 / texture_size,
                y: rect.y as f32 / texture_size,
                w: rect.w as f32 / texture_size,
                h: rect.h as f32 / texture_size})
            .dest(Point2::new(self.menu.screen.cursor.x as f32 * 3.0, y as f32 * 3.0))
            .scale(Vector2::new(3.0, 3.0));
        graphics::draw(ctx, &arrow_image, draw_params)?;
        Ok(())
    }

    fn draw_gore_option(&mut self, game: &mut Game, ctx: &mut Context) -> GameResult<()> {
        let mut batch = if game.gore_on {
            ON.as_sprite_batch(ctx, game, &self.menu.palette, self.menu.palette_hash)
                .expect("error drawing ON")
        } else {
            OFF.as_sprite_batch(ctx, game, &self.menu.palette, self.menu.palette_hash)
                .expect("error drawing OFF")
        };

        graphics::draw(
            ctx,
            &batch,
            graphics::DrawParam::default().scale(Vector2::new(3.0, 3.0)),
        )?;
        batch.clear();
        Ok(())
    }

    fn draw_player_count_option(&mut self, game: &mut Game, ctx: &mut Context) -> GameResult<()> {
        let mut text = PLAYER_COUNT.clone();
        text.string = game.num_players.to_string();
        let mut batch = text
            .as_sprite_batch(ctx, game, &self.menu.palette, self.menu.palette_hash)
            .expect("error drawing PLAYER_COUNT");
        graphics::draw(
            ctx,
            &batch,
            graphics::DrawParam::default().scale(Vector2::new(3.0, 3.0)),
        )?;
        batch.clear();
        Ok(())
    }
}

impl Scene<Game, InputEvent> for MainMenuScene {
    fn update(&mut self, _game: &mut Game, _ctx: &mut Context) -> FSceneSwitch {
        if self.done {
            if !self.fade_out_done {
                SceneSwitch::push(Fade::new(274, 1, FadeStyle::Out))
            } else {
                SceneSwitch::Pop // shouldn't happen
            }
        } else {
            SceneSwitch::None
        }
    }

    fn draw(&mut self, game: &mut Game, ctx: &mut Context) -> GameResult<()> {
        graphics::clear(ctx, graphics::Color::from((0, 0, 0, 255)));
        self.menu.draw(game, ctx)?;
        self.draw_arrow(game, ctx)?;
        self.draw_gore_option(game, ctx)?;
        self.draw_player_count_option(game, ctx)?;
        Ok(())
    }

    fn name(&self) -> &str {
        "Main menu"
    }

    fn input(&mut self, gameworld: &mut Game, _event: InputEvent, started: bool) {
        let x = gameworld.input.get_axis_raw(Axis::Horz1) as i32;
        if gameworld.input.get_button_down(Button::Fire1) {
            match self.selected_option {
                MenuOption::Practice => {
                    self.done = true;
                    gameworld.next_scene = SceneState::Practice;
                }
                MenuOption::Gore => gameworld.gore_on = !gameworld.gore_on,
                MenuOption::Players => gameworld.num_players = (gameworld.num_players % 4) + 1,
                _ => (),
            }
        } else if x != 0 {
            match self.selected_option {
                MenuOption::Gore => gameworld.gore_on = !gameworld.gore_on,
                MenuOption::Players => {
                    if x > 0 {
                        gameworld.num_players = (gameworld.num_players % 4) + x;
                    } else {
                        gameworld.num_players = (gameworld.num_players + 3 + x) % 4 + 1;
                    }
                }
                _ => (),
            }
        }

        if !started {
            let y = gameworld.input.get_axis_raw(Axis::Vert1) as i32;
            if y > 0 {
                self.selected_option = MenuOption::from_i32((self.selected_option as i32 + y) % 4);
            } else if y < 0 {
                self.selected_option =
                    MenuOption::from_i32((self.selected_option as i32 + 4 + y) % 4);
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Taken from dummy behaviour to have a layer of communication which reacts with
//! the embedded state machine (and inner ui), also back to net services:
//! currently kademlia, mdns
//! https://docs.rs/libp2p/0.21.1/libp2p/swarm/struct.DummyBehaviour.html
use libp2p::swarm::{
    protocols_handler, NetworkBehaviour,
    NetworkBehaviourAction::{self, GenerateEvent},
    PollParameters, ProtocolsHandler,
};
use libp2p_core::{
    connection::{ConnectedPoint, ConnectionId},
    Multiaddr, PeerId,
};
use std::task::{Context, Poll};

use super::{
    sm::{self, AdbfStateChart, Events, Events::*, NewPeerData, States},
    ui_data::UiData,
};

/// Events going from StateMachine back to the net behavior
#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Clone, Debug)]
pub enum SMOutEvents {
    MyPathSearchRunning(bool),
}

//#[derive(Clone, Default)]
pub struct SMBehaviour {
    sm: sm::StateMachine<AdbfStateChart>,
    send_buffer: Vec<SMOutEvents>,
}
impl SMBehaviour {
    pub fn new(own_peer: PeerId, ui_data: UiData) -> Self {
        Self {
            sm: AdbfStateChart::init(AdbfStateChart::new(own_peer, ui_data)),
            send_buffer: vec![],
        }
    }

    pub fn mdns_new_peer(&mut self, peer_id: &PeerId, multi_addr: &Multiaddr) {
        let new_peer_event = GotANewPeer(NewPeerData {
            id: peer_id.clone(),
            addr: multi_addr.clone(),
        });
        self.process_and_react(new_peer_event);
    }

    pub fn mdns_remove(&mut self, peer_id: &PeerId) {
        let remove_peer_event = HaveToRemovePeer(peer_id.clone());
        self.process_and_react(remove_peer_event);
    }

    fn process_and_react(&mut self, event: Events) {
        let return_state = self.sm.process_event(event);
        match return_state {
            Ok(good_state) => match good_state {
                States::Start => (),                // nothing to do
                States::WaitingForPeerAction => (), // is just waiting
                States::SendKademliaOut => {
                    self.send_buffer
                        .push(SMOutEvents::MyPathSearchRunning(true));
                    self.sm.process_event(Done);
                }
            },
            Err(_process_without_valid_state_transition) => (), // this is normal in a state chart
        }
    }
}

impl NetworkBehaviour for SMBehaviour {
    type ProtocolsHandler = protocols_handler::DummyProtocolsHandler;
    type OutEvent = SMOutEvents;

    fn new_handler(&mut self) -> Self::ProtocolsHandler {
        protocols_handler::DummyProtocolsHandler::default()
    }
    fn addresses_of_peer(&mut self, _: &PeerId) -> Vec<Multiaddr> {
        Vec::new()
    }
    fn inject_connected(&mut self, _: &PeerId) {}
    fn inject_disconnected(&mut self, _: &PeerId) {}
    fn inject_connection_established(&mut self, _: &PeerId, _: &ConnectionId, _: &ConnectedPoint) {}
    fn inject_connection_closed(&mut self, _: &PeerId, _: &ConnectionId, _: &ConnectedPoint) {}

    fn inject_event(
        &mut self,
        _: PeerId,
        _: ConnectionId,
        _: <Self::ProtocolsHandler as ProtocolsHandler>::OutEvent,
    ) {
    }

    fn poll(
        &mut self,
        _: &mut Context,
        _: &mut impl PollParameters,
    ) -> Poll<
        NetworkBehaviourAction<
            <Self::ProtocolsHandler as ProtocolsHandler>::InEvent,
            Self::OutEvent,
        >,
    > {
        if let Some(item) = self.send_buffer.pop() {
            Poll::Ready(GenerateEvent(item))
        } else {
            Poll::Pending
        }
    }
}
<commit_msg>[net] using vecdeque for a queue instead of vector<commit_after>//! Taken from dummy behaviour to have a layer of communication which reacts with
//! the embedded state machine (and inner ui), also back to net services:
//! currently kademlia, mdns
//! https://docs.rs/libp2p/0.21.1/libp2p/swarm/struct.DummyBehaviour.html
use libp2p::swarm::{
    protocols_handler, NetworkBehaviour,
    NetworkBehaviourAction::{self, GenerateEvent},
    PollParameters, ProtocolsHandler,
};
use libp2p_core::{
    connection::{ConnectedPoint, ConnectionId},
    Multiaddr, PeerId,
};
use std::{
    collections::vec_deque::VecDeque,
    task::{Context, Poll},
};

use super::{
    sm::{self, AdbfStateChart, Events, Events::*, NewPeerData, States},
    ui_data::UiData,
};

/// Events going from StateMachine back to the net behavior
#[allow(non_camel_case_types)]
#[derive(Serialize, Deserialize, Clone, Debug)]
pub enum SMOutEvents {
    MyPathSearchRunning(bool),
}

//#[derive(Clone, Default)]
pub struct SMBehaviour {
    sm: sm::StateMachine<AdbfStateChart>,
    send_buffer: VecDeque<SMOutEvents>,
}
impl SMBehaviour {
    pub fn new(own_peer: PeerId, ui_data: UiData) -> Self {
        Self {
            sm: AdbfStateChart::init(AdbfStateChart::new(own_peer, ui_data)),
            send_buffer: VecDeque::new(),
        }
    }

    pub fn mdns_new_peer(&mut self, peer_id: &PeerId, multi_addr: &Multiaddr) {
        let new_peer_event = GotANewPeer(NewPeerData {
            id: peer_id.clone(),
            addr: multi_addr.clone(),
        });
        self.process_and_react(new_peer_event);
    }

    pub fn mdns_remove(&mut self, peer_id: &PeerId) {
        let remove_peer_event = HaveToRemovePeer(peer_id.clone());
        self.process_and_react(remove_peer_event);
    }

    fn process_and_react(&mut self, event: Events) {
        let return_state = self.sm.process_event(event);
        match return_state {
            Ok(good_state) => match good_state {
                States::Start => (),                // nothing to do
                States::WaitingForPeerAction => (), // is just waiting
                States::SendKademliaOut => {
                    self.send_buffer
                        .push_back(SMOutEvents::MyPathSearchRunning(true));
                    self.sm.process_event(Done);
                }
            },
            Err(_process_without_valid_state_transition) => (), // this is normal in a state chart
        }
    }
}

impl NetworkBehaviour for SMBehaviour {
    type ProtocolsHandler = protocols_handler::DummyProtocolsHandler;
    type OutEvent = SMOutEvents;

    fn new_handler(&mut self) -> Self::ProtocolsHandler {
        protocols_handler::DummyProtocolsHandler::default()
    }
    fn addresses_of_peer(&mut self, _: &PeerId) -> Vec<Multiaddr> {
        Vec::new()
    }
    fn inject_connected(&mut self, _: &PeerId) {}
    fn inject_disconnected(&mut self, _: &PeerId) {}
    fn inject_connection_established(&mut self, _: &PeerId, _: &ConnectionId, _: &ConnectedPoint) {}
    fn inject_connection_closed(&mut self, _: &PeerId, _: &ConnectionId, _: &ConnectedPoint) {}

    fn inject_event(
        &mut self,
        _: PeerId,
        _: ConnectionId,
        _: <Self::ProtocolsHandler as ProtocolsHandler>::OutEvent,
    ) {
    }

    fn poll(
        &mut self,
        _: &mut Context,
        _: &mut impl PollParameters,
    ) -> Poll<
        NetworkBehaviourAction<
            <Self::ProtocolsHandler as ProtocolsHandler>::InEvent,
            Self::OutEvent,
        >,
    > {
        if let Some(item) = self.send_buffer.pop_front() {
            Poll::Ready(GenerateEvent(item))
        } else {
            Poll::Pending
        }
    }
}
<|endoftext|>


<commit_before>use std::iter::repeat;
use std::collections::VecDeque;
use std::collections::BinaryHeap;
use std::cmp::Ordering;

fn max(a: u32, b: u32) -> u32 { if a > b { a } else { b } }

#[derive(Clone)]
struct Edge {
    vertex: u32,
    weight: f32,
}

impl Edge {
    pub fn new(vertex_: u32, weight_: f32) -> Edge {
        Edge{vertex: vertex_, weight: weight_,}
    }
}

impl Ord for Edge {
    fn cmp(&self, other: &Self) -> Ordering {
        self.partial_cmp(other).unwrap()
    }
}

impl Eq for Edge {}

impl PartialOrd for Edge {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        (self.weight, self.vertex).partial_cmp(&(other.weight, other.vertex))
    }
}

impl PartialEq for Edge {
    fn eq(&self, other: &Self) -> bool {
        (self.weight, &self.vertex) == (other.weight, &other.vertex)
    }
}

struct Graph {
    nodes: Vec<Vec<Edge>>,
}

impl Graph {
    pub fn new() -> Graph { Graph{nodes: Vec::new(),} }

    pub fn add_edge(&mut self, src: u32, dst: u32, weight: f32) {
        let len = self.nodes.len();

        if (max(src, dst)) > len as u32 {
            let new_len = (max(src, dst) + 1) as usize;
            self.nodes.extend(repeat(Vec::new()).take(new_len - len))
        }
        self.nodes[src as usize].push(Edge::new(dst, weight));
    }
    pub fn bfs(&self, src: u32) {
        let mut queue: VecDeque<u32> = VecDeque::new();
        let len = self.nodes.len();
        let mut visited = vec!(false ; len);

        queue.push_front(src);

        while let Some(current) = queue.pop_back() {
            if !visited[current as usize] {
                println!("current: {}" , current);
                visited[current as usize] = true;
            } else { continue ; }

            for n in &self.nodes[current as usize] {
                let neighbor: u32 = n.vertex;
                queue.push_front(neighbor);
            }
        }
    }
    pub fn dfs(&self, src: u32) {
        let mut stack: Vec<u32> = Vec::new();
        let len = self.nodes.len() as usize;
        let mut visited = vec!(false ; len);

        stack.push(src);

        while let Some(current) = stack.pop() {
            if !visited[current as usize] {
                println!("current: {}" , current);
                visited[current as usize] = true;
            } else { continue ; }

            for n in &self.nodes[current as usize] {
                let neighbor: u32 = n.vertex;
                stack.push(neighbor);
            }
        }
    }

    pub fn dijkstra(&self, src: u32, dst: u32) {
      let mut dist: Vec<f32> = Vec::new();
      let mut prev: Vec<u32> = Vec::new();
      let mut q = BinaryHeap::new();
      const max_weight: f32 = std::f32::MAX;

      // init dist, prev table
      for i in 0..self.nodes.len() {
        dist.push(max_weight);
        prev.push(0);
      }

      // We're at `start`, with a zero cost
      dist[src as usize] = 0.0;
      q.push(Edge::new(src, dist[src as usize]));

      //while (!q.is_empty()) {
      while let Some(u) = q.pop() {
        //let u: Edge = el.unwrap();

        // loop for all edges connected to
        for v in self.nodes[u.vertex as usize].iter() {

          let alt: f32 = dist[u.vertex as usize] + v.weight; // accumulate shortest dist from source

          if (alt < dist[v.vertex as usize]) {
            dist[v.vertex as usize] = alt; // keep the shortest dist from src to v
            prev[v.vertex as usize] = u.vertex;

            q.push(Edge::new(v.vertex, dist[v.vertex as usize])); // Add unvisited v into the Q to be processed
          }
        }
      }

      let mut shortestPath: Vec<u32> = Vec::new();
      let mut curr: u32 = dst;

      shortestPath.push(curr);

      while (curr != src) {
        curr = prev[curr as usize];
        shortestPath.push(curr);
      }

      shortestPath.reverse();

      for v in shortestPath.iter() {
        println!("{}", v);
      }
    }
}

fn main() {
    let mut g1 = Graph::new();
    g1.add_edge(0, 1, 1.0);
    g1.add_edge(0, 2, 4.0);
    g1.add_edge(1, 0, 1.0);
    g1.add_edge(1, 2, 2.0);
    g1.add_edge(1, 3, 6.0);
    g1.add_edge(2, 0, 4.0);
    g1.add_edge(2, 1, 2.0);
    g1.add_edge(2, 3, 3.0);
    g1.add_edge(3, 1, 6.0);
    g1.add_edge(3, 2, 3.0);

    g1.bfs(3);

    println!("");

    g1.dfs(3);

    println!("");

    g1.dijkstra(0, 3);
}<commit_msg>fixed compiler warnings<commit_after>use std::iter::repeat;
use std::collections::VecDeque;
use std::collections::BinaryHeap;
use std::cmp::Ordering;

fn max(a: u32, b: u32) -> u32 { if a > b { a } else { b } }

#[derive(Clone)]
struct Edge {
    vertex: u32,
    weight: f32,
}

impl Edge {
    pub fn new(vertex_: u32, weight_: f32) -> Edge {
        Edge{vertex: vertex_, weight: weight_,}
    }
}

impl Ord for Edge {
    fn cmp(&self, other: &Self) -> Ordering {
        self.partial_cmp(other).unwrap()
    }
}

impl Eq for Edge {}

impl PartialOrd for Edge {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        (self.weight, self.vertex).partial_cmp(&(other.weight, other.vertex))
    }
}

impl PartialEq for Edge {
    fn eq(&self, other: &Self) -> bool {
        (self.weight, &self.vertex) == (other.weight, &other.vertex)
    }
}

struct Graph {
    nodes: Vec<Vec<Edge>>,
}

impl Graph {
    pub fn new() -> Graph { Graph{nodes: Vec::new(),} }

    pub fn add_edge(&mut self, src: u32, dst: u32, weight: f32) {
        let len = self.nodes.len();

        if (max(src, dst)) > len as u32 {
            let new_len = (max(src, dst) + 1) as usize;
            self.nodes.extend(repeat(Vec::new()).take(new_len - len))
        }
        self.nodes[src as usize].push(Edge::new(dst, weight));
    }
    pub fn bfs(&self, src: u32) {
        let mut queue: VecDeque<u32> = VecDeque::new();
        let len = self.nodes.len();
        let mut visited = vec!(false ; len);

        queue.push_front(src);

        while let Some(current) = queue.pop_back() {
            if !visited[current as usize] {
                println!("current: {}" , current);
                visited[current as usize] = true;
            } else { continue ; }

            for n in &self.nodes[current as usize] {
                let neighbor: u32 = n.vertex;
                queue.push_front(neighbor);
            }
        }
    }
    pub fn dfs(&self, src: u32) {
        let mut stack: Vec<u32> = Vec::new();
        let len = self.nodes.len() as usize;
        let mut visited = vec!(false ; len);

        stack.push(src);

        while let Some(current) = stack.pop() {
            if !visited[current as usize] {
                println!("current: {}" , current);
                visited[current as usize] = true;
            } else { continue ; }

            for n in &self.nodes[current as usize] {
                let neighbor: u32 = n.vertex;
                stack.push(neighbor);
            }
        }
    }

    pub fn dijkstra(&self, src: u32, dst: u32) {
      let mut dist: Vec<f32> = Vec::new();
      let mut prev: Vec<u32> = Vec::new();
      let mut q = BinaryHeap::new();
      const MAX_WEIGHT: f32 = std::f32::MAX;

      // init dist, prev table
      for _ in 0..self.nodes.len() {
        dist.push(MAX_WEIGHT);
        prev.push(0);
      }

      // We're at `start`, with a zero cost
      dist[src as usize] = 0.0;
      q.push(Edge::new(src, dist[src as usize]));

      while let Some(u) = q.pop() {
        // loop for all edges connected to
        for v in self.nodes[u.vertex as usize].iter() {

          let alt: f32 = dist[u.vertex as usize] + v.weight; // accumulate shortest dist from source

          // The variable alt is the length of the path from the root node to the neighbor node v
          // if it were to go through u. If this path is shorter than the current shortest path
          // recorded for v, that current path is replaced with this alt path.
          if alt < dist[v.vertex as usize] {
            dist[v.vertex as usize] = alt; // keep the shortest dist from src to v
            prev[v.vertex as usize] = u.vertex;

            q.push(Edge::new(v.vertex, dist[v.vertex as usize])); // Add unvisited v into the Q to be processed
          }
        }
      }

      let mut shortest_path: Vec<u32> = Vec::new();
      let mut curr: u32 = dst;

      shortest_path.push(curr);

      while curr != src {
        curr = prev[curr as usize];
        shortest_path.push(curr);
      }

      shortest_path.reverse();

      for v in shortest_path.iter() {
        println!("{}", v);
      }
    }
}

fn main() {
    let mut g1 = Graph::new();
    g1.add_edge(0, 1, 1.0);
    g1.add_edge(0, 2, 4.0);
    g1.add_edge(1, 0, 1.0);
    g1.add_edge(1, 2, 2.0);
    g1.add_edge(1, 3, 6.0);
    g1.add_edge(2, 0, 4.0);
    g1.add_edge(2, 1, 2.0);
    g1.add_edge(2, 3, 3.0);
    g1.add_edge(3, 1, 6.0);
    g1.add_edge(3, 2, 3.0);

    g1.bfs(3);

    println!("");

    g1.dfs(3);

    println!("");

    g1.dijkstra(0, 3);
}<|endoftext|>


<commit_before>use std::collections::BTreeMap;

use toml::Value;

use error::RefErrorKind as REK;
use result::Result;

#[derive(Default)]
pub struct RefFlags {
    content_hashing:       bool,
    permission_tracking:   bool,
}

impl RefFlags {

    /// Read the RefFlags from a TOML document
    ///
    /// Assumes that the whole TOML tree is passed. So this looks up `ref.flags` to get the flags.
    /// It assumes that this is a Map with Key = <name of the setting> and Value = boolean.
    pub fn read(v: &Value) -> Result<RefFlags> {
        fn get_field(v: &Value, key: &str) -> Result<bool> {
            match v.lookup(key) {
                Some(&Value::Boolean(b)) => Ok(b),
                Some(_) => Err(REK::HeaderTypeError.into()),
                None    => Err(REK::HeaderFieldMissingError.into()),
            }
        }

        Ok(RefFlags {
            content_hashing:     try!(get_field(v, "ref.flags.content_hashing")),
            permission_tracking: try!(get_field(v, "ref.flags.permission_tracking")),
        })
    }

    /// Alias for `RefFlags::content_hashing()`
    pub fn is_often_moving(mut self, b: bool) -> RefFlags {
        self.with_content_hashing(b)
    }

    pub fn with_content_hashing(mut self, b: bool) -> RefFlags {
        self.content_hashing = b;
        self
    }

    pub fn with_permission_tracking(mut self, b: bool) -> RefFlags {
        self.permission_tracking = b;
        self
    }

    pub fn get_content_hashing(&self) -> bool {
        unimplemented!()
    }

    pub fn get_permission_tracking(&self) -> bool {
        unimplemented!()
    }

}

impl Into<Value> for RefFlags {

    /// Build a TOML::Value from this RefFlags object.
    ///
    /// Returns a Map which should be set in `ref.flags` in the header.
    fn into(self) -> Value {
        let mut btm = BTreeMap::new();
        btm.insert(String::from("content_hashing"),     Value::Boolean(self.content_hashing));
        btm.insert(String::from("permission_tracking"), Value::Boolean(self.permission_tracking));
        return Value::Table(btm)
    }

}

<commit_msg>Impl RefFlags::get_content_hashing()<commit_after>use std::collections::BTreeMap;

use toml::Value;

use error::RefErrorKind as REK;
use result::Result;

#[derive(Default)]
pub struct RefFlags {
    content_hashing:       bool,
    permission_tracking:   bool,
}

impl RefFlags {

    /// Read the RefFlags from a TOML document
    ///
    /// Assumes that the whole TOML tree is passed. So this looks up `ref.flags` to get the flags.
    /// It assumes that this is a Map with Key = <name of the setting> and Value = boolean.
    pub fn read(v: &Value) -> Result<RefFlags> {
        fn get_field(v: &Value, key: &str) -> Result<bool> {
            match v.lookup(key) {
                Some(&Value::Boolean(b)) => Ok(b),
                Some(_) => Err(REK::HeaderTypeError.into()),
                None    => Err(REK::HeaderFieldMissingError.into()),
            }
        }

        Ok(RefFlags {
            content_hashing:     try!(get_field(v, "ref.flags.content_hashing")),
            permission_tracking: try!(get_field(v, "ref.flags.permission_tracking")),
        })
    }

    /// Alias for `RefFlags::content_hashing()`
    pub fn is_often_moving(mut self, b: bool) -> RefFlags {
        self.with_content_hashing(b)
    }

    pub fn with_content_hashing(mut self, b: bool) -> RefFlags {
        self.content_hashing = b;
        self
    }

    pub fn with_permission_tracking(mut self, b: bool) -> RefFlags {
        self.permission_tracking = b;
        self
    }

    pub fn get_content_hashing(&self) -> bool {
        self.content_hashing
    }

    pub fn get_permission_tracking(&self) -> bool {
        unimplemented!()
    }

}

impl Into<Value> for RefFlags {

    /// Build a TOML::Value from this RefFlags object.
    ///
    /// Returns a Map which should be set in `ref.flags` in the header.
    fn into(self) -> Value {
        let mut btm = BTreeMap::new();
        btm.insert(String::from("content_hashing"),     Value::Boolean(self.content_hashing));
        btm.insert(String::from("permission_tracking"), Value::Boolean(self.permission_tracking));
        return Value::Table(btm)
    }

}

<|endoftext|>


<commit_before>// Copyright 2016 The Grin Developers
//
// 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.

//! Mining service, gathers transactions from the pool, assemble them in a
//! block and mine the block to produce a valid header with its proof-of-work.

use rand::{self, Rng};
use std::sync::{Arc, Mutex};
use std::ops::Deref;
use time;

use core::consensus;
use core::core;
use core::core::hash::{Hash, Hashed};
use core::pow;
use core::pow::cuckoo;
use chain;
use secp;

pub struct Miner {
	chain_head: Arc<Mutex<chain::Tip>>,
	chain_store: Arc<chain::ChainStore>,
}

impl Miner {
	/// Creates a new Miner. Needs references to the chain state and its
	/// storage.
	pub fn new(chain_head: Arc<Mutex<chain::Tip>>, chain_store: Arc<chain::ChainStore>) -> Miner {
		Miner {
			chain_head: chain_head,
			chain_store: chain_store,
		}
	}

	/// Starts the mining loop, building a new block on top of the existing
	/// chain anytime required and looking for PoW solution.
	pub fn run_loop(&self) {
		info!("Starting miner loop.");
		loop {
			// get the latest chain state and build a block on top of it
			let head: core::BlockHeader;
			let mut latest_hash: Hash;
			{
				head = self.chain_store.head_header().unwrap();
				latest_hash = self.chain_head.lock().unwrap().last_block_h;
			}
			let b = self.build_block(&head);
			let mut pow_header = pow::PowHeader::from_block(&b);

			// look for a pow for at most 2 sec on the same block (to give a chance to new
			// transactions) and as long as the head hasn't changed
			let deadline = time::get_time().sec + 2;
			let mut sol = None;
			debug!("Mining at Cuckoo{} for at most 2 secs.",
			       b.header.cuckoo_len);
			while head.hash() == latest_hash && time::get_time().sec < deadline {
				let pow_hash = pow_header.hash();
				let mut miner = cuckoo::Miner::new(pow_hash.to_slice(),
				                                   consensus::EASINESS,
				                                   b.header.cuckoo_len as u32);
				if let Ok(proof) = miner.mine() {
					if proof.to_target() <= b.header.target {
						sol = Some(proof);
						break;
					}
				}
				pow_header.nonce += 1;
				{
					latest_hash = self.chain_head.lock().unwrap().last_block_h;
				}
			}

			// if we found a solution, push our block out
			if let Some(proof) = sol {
				info!("Found valid proof of work, adding block {}.", b.hash());
				if let Err(e) = chain::process_block(&b, self.chain_store.clone(), chain::NONE) {
					error!("Error validating mined block: {:?}", e);
				}
			} else {
				debug!("No solution found, continuing...")
			}
		}
	}

	/// Builds a new block with the chain head as previous and eligible
	/// transactions from the pool.
	fn build_block(&self, head: &core::BlockHeader) -> core::Block {
		let mut now_sec = time::get_time().sec;
		let head_sec = head.timestamp.to_timespec().sec;
		if now_sec == head_sec {
			now_sec += 1;
		}
		let (target, cuckoo_len) =
			consensus::next_target(now_sec, head_sec, head.target, head.cuckoo_len);

		let mut rng = rand::OsRng::new().unwrap();
		let secp_inst = secp::Secp256k1::with_caps(secp::ContextFlag::Commit);
		// TODO get a new key from the user's wallet or something
		let skey = secp::key::SecretKey::new(&secp_inst, &mut rng);

		// TODO populate inputs and outputs from pool transactions
		core::Block {
			header: core::BlockHeader {
				height: head.height + 1,
				previous: head.hash(),
				timestamp: time::at(time::Timespec::new(now_sec, 0)),
				cuckoo_len: cuckoo_len,
				target: target,
				nonce: rng.gen(),
				..Default::default()
			},
			inputs: vec![],
			outputs: vec![],
			proofs: vec![],
		}
	}
}
<commit_msg>Several bug fixes to the miner loop. Now builds the block properly and updates the shared head state.<commit_after>// Copyright 2016 The Grin Developers
//
// 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.

//! Mining service, gathers transactions from the pool, assemble them in a
//! block and mine the block to produce a valid header with its proof-of-work.

use rand::{self, Rng};
use std::sync::{Arc, Mutex};
use std::ops::Deref;
use time;

use adapters::ChainToNetAdapter;
use core::consensus;
use core::core;
use core::core::hash::{Hash, Hashed};
use core::pow;
use core::pow::cuckoo;
use chain;
use secp;

pub struct Miner {
	chain_head: Arc<Mutex<chain::Tip>>,
	chain_store: Arc<chain::ChainStore>,
	/// chain adapter to net
	chain_adapter: Arc<ChainToNetAdapter>,
}

impl Miner {
	/// Creates a new Miner. Needs references to the chain state and its
	/// storage.
	pub fn new(chain_head: Arc<Mutex<chain::Tip>>,
	           chain_store: Arc<chain::ChainStore>,
	           chain_adapter: Arc<ChainToNetAdapter>)
	           -> Miner {
		Miner {
			chain_head: chain_head,
			chain_store: chain_store,
			chain_adapter: chain_adapter,
		}
	}

	/// Starts the mining loop, building a new block on top of the existing
	/// chain anytime required and looking for PoW solution.
	pub fn run_loop(&self) {
		info!("Starting miner loop.");
		loop {
			// get the latest chain state and build a block on top of it
			let head: core::BlockHeader;
			let mut latest_hash: Hash;
			{
				head = self.chain_store.head_header().unwrap();
				latest_hash = self.chain_head.lock().unwrap().last_block_h;
			}
			let mut b = self.build_block(&head);
			let mut pow_header = pow::PowHeader::from_block(&b);

			// look for a pow for at most 2 sec on the same block (to give a chance to new
			// transactions) and as long as the head hasn't changed
			let deadline = time::get_time().sec + 2;
			let mut sol = None;
			debug!("Mining at Cuckoo{} for at most 2 secs on block {}.",
			       b.header.cuckoo_len,
			       latest_hash);
			let mut iter_count = 0;
			while head.hash() == latest_hash && time::get_time().sec < deadline {
				let pow_hash = pow_header.hash();
				let mut miner = cuckoo::Miner::new(pow_hash.to_slice(),
				                                   consensus::EASINESS,
				                                   b.header.cuckoo_len as u32);
				if let Ok(proof) = miner.mine() {
					if proof.to_target() <= b.header.target {
						sol = Some(proof);
						break;
					}
				}
				pow_header.nonce += 1;
				{
					latest_hash = self.chain_head.lock().unwrap().last_block_h;
				}
				iter_count += 1;
			}

			// if we found a solution, push our block out
			if let Some(proof) = sol {
				info!("Found valid proof of work, adding block {}.", b.hash());
				b.header.pow = proof;
				b.header.nonce = pow_header.nonce;
				let res = chain::process_block(&b,
				                               self.chain_store.clone(),
				                               self.chain_adapter.clone(),
				                               chain::NONE);
				if let Err(e) = res {
					error!("Error validating mined block: {:?}", e);
				} else if let Ok(Some(tip)) = res {
					let chain_head = self.chain_head.clone();
					let mut head = chain_head.lock().unwrap();
					*head = tip;
				}
			} else {
				debug!("No solution found after {} iterations, continuing...",
				       iter_count)
			}
		}
	}

	/// Builds a new block with the chain head as previous and eligible
	/// transactions from the pool.
	fn build_block(&self, head: &core::BlockHeader) -> core::Block {
		let mut now_sec = time::get_time().sec;
		let head_sec = head.timestamp.to_timespec().sec;
		if now_sec == head_sec {
			now_sec += 1;
		}
		let (target, cuckoo_len) =
			consensus::next_target(now_sec, head_sec, head.target, head.cuckoo_len);

		let mut rng = rand::OsRng::new().unwrap();
		let secp_inst = secp::Secp256k1::with_caps(secp::ContextFlag::Commit);
		// TODO get a new key from the user's wallet or something
		let skey = secp::key::SecretKey::new(&secp_inst, &mut rng);

		// TODO populate inputs and outputs from pool transactions
		let mut b = core::Block::new(head, vec![], skey).unwrap();
		b.header.nonce = rng.gen();
		b.header.target = target;
		b.header.cuckoo_len = cuckoo_len;
		b.header.timestamp = time::at(time::Timespec::new(now_sec, 0));
		b
	}
}
<|endoftext|>


<commit_before>//! Client handshake machine.

use std::{
    io::{Read, Write},
    marker::PhantomData,
};

use http::{
    header::HeaderName, HeaderMap, Request as HttpRequest, Response as HttpResponse, StatusCode,
};
use httparse::Status;
use log::*;

use super::{
    derive_accept_key,
    headers::{FromHttparse, MAX_HEADERS},
    machine::{HandshakeMachine, StageResult, TryParse},
    HandshakeRole, MidHandshake, ProcessingResult,
};
use crate::{
    error::{Error, ProtocolError, Result, UrlError},
    protocol::{Role, WebSocket, WebSocketConfig},
};

/// Client request type.
pub type Request = HttpRequest<()>;

/// Client response type.
pub type Response = HttpResponse<()>;

/// Client handshake role.
#[derive(Debug)]
pub struct ClientHandshake<S> {
    verify_data: VerifyData,
    config: Option<WebSocketConfig>,
    _marker: PhantomData<S>,
}

impl<S: Read + Write> ClientHandshake<S> {
    /// Initiate a client handshake.
    pub fn start(
        stream: S,
        request: Request,
        config: Option<WebSocketConfig>,
    ) -> Result<MidHandshake<Self>> {
        if request.method() != http::Method::GET {
            return Err(Error::Protocol(ProtocolError::WrongHttpMethod));
        }

        if request.version() < http::Version::HTTP_11 {
            return Err(Error::Protocol(ProtocolError::WrongHttpVersion));
        }

        // Check the URI scheme: only ws or wss are supported
        let _ = crate::client::uri_mode(request.uri())?;

        // Convert and verify the `http::Request` and turn it into the request as per RFC.
        // Also extract the key from it (it must be present in a correct request).
        let (request, key) = generate_request(request)?;

        let machine = HandshakeMachine::start_write(stream, request);

        let client = {
            let accept_key = derive_accept_key(key.as_ref());
            ClientHandshake { verify_data: VerifyData { accept_key }, config, _marker: PhantomData }
        };

        trace!("Client handshake initiated.");
        Ok(MidHandshake { role: client, machine })
    }
}

impl<S: Read + Write> HandshakeRole for ClientHandshake<S> {
    type IncomingData = Response;
    type InternalStream = S;
    type FinalResult = (WebSocket<S>, Response);
    fn stage_finished(
        &mut self,
        finish: StageResult<Self::IncomingData, Self::InternalStream>,
    ) -> Result<ProcessingResult<Self::InternalStream, Self::FinalResult>> {
        Ok(match finish {
            StageResult::DoneWriting(stream) => {
                ProcessingResult::Continue(HandshakeMachine::start_read(stream))
            }
            StageResult::DoneReading { stream, result, tail } => {
                let result = self.verify_data.verify_response(result)?;
                debug!("Client handshake done.");
                let websocket =
                    WebSocket::from_partially_read(stream, tail, Role::Client, self.config);
                ProcessingResult::Done((websocket, result))
            }
        })
    }
}

/// Verifies and generates a client WebSocket request from the original request and extracts a WebSocket key from it.
fn generate_request(mut request: Request) -> Result<(Vec<u8>, String)> {
    let mut req = Vec::new();
    write!(
        req,
        "GET {path} {version:?}\r\n",
        path = request.uri().path_and_query().ok_or(Error::Url(UrlError::NoPathOrQuery))?.as_str(),
        version = request.version()
    )
    .unwrap();

    // Headers that must be present in a correct request.
    const KEY_HEADERNAME: &str = "Sec-WebSocket-Key";
    const WEBSOCKET_HEADERS: [&str; 5] =
        ["Host", "Connection", "Upgrade", "Sec-WebSocket-Version", KEY_HEADERNAME];

    // We must extract a WebSocket key from a properly formed request or fail if it's not present.
    let key = request
        .headers()
        .get(KEY_HEADERNAME)
        .ok_or_else(|| {
            Error::Protocol(ProtocolError::InvalidHeader(
                HeaderName::from_bytes(KEY_HEADERNAME.as_bytes()).unwrap(),
            ))
        })?
        .to_str()?
        .to_owned();

    // We must check that all necessary headers for a valid request are present. Note that we have to
    // deal with the fact that some apps seem to have a case-sensitive check for headers which is not
    // correct and should not considered the correct behavior, but it seems like some apps ignore it.
    // `http` by default writes all headers in lower-case which is fine (and does not violate the RFC)
    // but some servers seem to be poorely written and ignore RFC.
    //
    // See similar problem in `hyper`: https://github.com/hyperium/hyper/issues/1492
    let headers = request.headers_mut();
    for &header in &WEBSOCKET_HEADERS {
        let value = headers.remove(header).ok_or_else(|| {
            Error::Protocol(ProtocolError::InvalidHeader(
                HeaderName::from_bytes(header.as_bytes()).unwrap(),
            ))
        })?;
        write!(req, "{header}: {value}\r\n", header = header, value = value.to_str()?).unwrap();
    }

    // Now we must ensure that the headers that we've written once are not anymore present in the map.
    // If they do, then the request is invalid (some headers are duplicated there for some reason).
    let insensitive: Vec<String> =
        WEBSOCKET_HEADERS.iter().map(|h| h.to_ascii_lowercase()).collect();
    for (k, v) in headers {
        let mut name = k.as_str();

        // We have already written the necessary headers once (above) and removed them from the map.
        // If we encounter them again, then the request is considered invalid and error is returned.
        // Note that we can't use `.contains()`, since `&str` does not coerce to `&String` in Rust.
        if insensitive.iter().any(|x| x == name) {
            return Err(Error::Protocol(ProtocolError::InvalidHeader(k.clone())));
        }

        // Relates to the issue of some servers treating headers in a case-sensitive way, please see:
        // https://github.com/snapview/tungstenite-rs/pull/119 (original fix of the problem)
        if name == "sec-websocket-protocol" {
            name = "Sec-WebSocket-Protocol";
        }

        if name == "origin" {
            name = "Origin";
        }

        writeln!(req, "{}: {}\r", name, v.to_str()?).unwrap();
    }

    writeln!(req, "\r").unwrap();
    trace!("Request: {:?}", String::from_utf8_lossy(&req));
    Ok((req, key))
}

/// Information for handshake verification.
#[derive(Debug)]
struct VerifyData {
    /// Accepted server key.
    accept_key: String,
}

impl VerifyData {
    pub fn verify_response(&self, response: Response) -> Result<Response> {
        // 1. If the status code received from the server is not 101, the
        // client handles the response per HTTP [RFC2616] procedures. (RFC 6455)
        if response.status() != StatusCode::SWITCHING_PROTOCOLS {
            return Err(Error::Http(response.map(|_| None)));
        }

        let headers = response.headers();

        // 2. If the response lacks an |Upgrade| header field or the |Upgrade|
        // header field contains a value that is not an ASCII case-
        // insensitive match for the value "websocket", the client MUST
        // _Fail the WebSocket Connection_. (RFC 6455)
        if !headers
            .get("Upgrade")
            .and_then(|h| h.to_str().ok())
            .map(|h| h.eq_ignore_ascii_case("websocket"))
            .unwrap_or(false)
        {
            return Err(Error::Protocol(ProtocolError::MissingUpgradeWebSocketHeader));
        }
        // 3.  If the response lacks a |Connection| header field or the
        // |Connection| header field doesn't contain a token that is an
        // ASCII case-insensitive match for the value "Upgrade", the client
        // MUST _Fail the WebSocket Connection_. (RFC 6455)
        if !headers
            .get("Connection")
            .and_then(|h| h.to_str().ok())
            .map(|h| h.eq_ignore_ascii_case("Upgrade"))
            .unwrap_or(false)
        {
            return Err(Error::Protocol(ProtocolError::MissingConnectionUpgradeHeader));
        }
        // 4.  If the response lacks a |Sec-WebSocket-Accept| header field or
        // the |Sec-WebSocket-Accept| contains a value other than the
        // base64-encoded SHA-1 of ... the client MUST _Fail the WebSocket
        // Connection_. (RFC 6455)
        if !headers.get("Sec-WebSocket-Accept").map(|h| h == &self.accept_key).unwrap_or(false) {
            return Err(Error::Protocol(ProtocolError::SecWebSocketAcceptKeyMismatch));
        }
        // 5.  If the response includes a |Sec-WebSocket-Extensions| header
        // field and this header field indicates the use of an extension
        // that was not present in the client's handshake (the server has
        // indicated an extension not requested by the client), the client
        // MUST _Fail the WebSocket Connection_. (RFC 6455)
        // TODO

        // 6.  If the response includes a |Sec-WebSocket-Protocol| header field
        // and this header field indicates the use of a subprotocol that was
        // not present in the client's handshake (the server has indicated a
        // subprotocol not requested by the client), the client MUST _Fail
        // the WebSocket Connection_. (RFC 6455)
        // TODO

        Ok(response)
    }
}

impl TryParse for Response {
    fn try_parse(buf: &[u8]) -> Result<Option<(usize, Self)>> {
        let mut hbuffer = [httparse::EMPTY_HEADER; MAX_HEADERS];
        let mut req = httparse::Response::new(&mut hbuffer);
        Ok(match req.parse(buf)? {
            Status::Partial => None,
            Status::Complete(size) => Some((size, Response::from_httparse(req)?)),
        })
    }
}

impl<'h, 'b: 'h> FromHttparse<httparse::Response<'h, 'b>> for Response {
    fn from_httparse(raw: httparse::Response<'h, 'b>) -> Result<Self> {
        if raw.version.expect("Bug: no HTTP version") < /*1.*/1 {
            return Err(Error::Protocol(ProtocolError::WrongHttpMethod));
        }

        let headers = HeaderMap::from_httparse(raw.headers)?;

        let mut response = Response::new(());
        *response.status_mut() = StatusCode::from_u16(raw.code.expect("Bug: no HTTP status code"))?;
        *response.headers_mut() = headers;
        // TODO: httparse only supports HTTP 0.9/1.0/1.1 but not HTTP 2.0
        // so the only valid value we could get in the response would be 1.1.
        *response.version_mut() = http::Version::HTTP_11;

        Ok(response)
    }
}

/// Generate a random key for the `Sec-WebSocket-Key` header.
pub fn generate_key() -> String {
    // a base64-encoded (see Section 4 of [RFC4648]) value that,
    // when decoded, is 16 bytes in length (RFC 6455)
    let r: [u8; 16] = rand::random();
    base64::encode(&r)
}

#[cfg(test)]
mod tests {
    use super::{super::machine::TryParse, generate_key, generate_request, Response};
    use crate::client::IntoClientRequest;

    #[test]
    fn random_keys() {
        let k1 = generate_key();
        println!("Generated random key 1: {}", k1);
        let k2 = generate_key();
        println!("Generated random key 2: {}", k2);
        assert_ne!(k1, k2);
        assert_eq!(k1.len(), k2.len());
        assert_eq!(k1.len(), 24);
        assert_eq!(k2.len(), 24);
        assert!(k1.ends_with("=="));
        assert!(k2.ends_with("=="));
        assert!(k1[..22].find('=').is_none());
        assert!(k2[..22].find('=').is_none());
    }

    fn construct_expected(host: &str, key: &str) -> Vec<u8> {
        format!(
            "\
            GET /getCaseCount HTTP/1.1\r\n\
            Host: {host}\r\n\
            Connection: Upgrade\r\n\
            Upgrade: websocket\r\n\
            Sec-WebSocket-Version: 13\r\n\
            Sec-WebSocket-Key: {key}\r\n\
            \r\n",
            host = host,
            key = key
        )
        .into_bytes()
    }

    #[test]
    fn request_formatting() {
        let request = "ws://localhost/getCaseCount".into_client_request().unwrap();
        let (request, key) = generate_request(request).unwrap();
        let correct = construct_expected("localhost", &key);
        assert_eq!(&request[..], &correct[..]);
    }

    #[test]
    fn request_formatting_with_host() {
        let request = "wss://localhost:9001/getCaseCount".into_client_request().unwrap();
        let (request, key) = generate_request(request).unwrap();
        let correct = construct_expected("localhost:9001", &key);
        assert_eq!(&request[..], &correct[..]);
    }

    #[test]
    fn request_formatting_with_at() {
        let request = "wss://user:pass@localhost:9001/getCaseCount".into_client_request().unwrap();
        let (request, key) = generate_request(request).unwrap();
        let correct = construct_expected("localhost:9001", &key);
        assert_eq!(&request[..], &correct[..]);
    }

    #[test]
    fn response_parsing() {
        const DATA: &[u8] = b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n";
        let (_, resp) = Response::try_parse(DATA).unwrap().unwrap();
        assert_eq!(resp.status(), http::StatusCode::OK);
        assert_eq!(resp.headers().get("Content-Type").unwrap(), &b"text/html"[..],);
    }

    #[test]
    fn invalid_custom_request() {
        let request = http::Request::builder().method("GET").body(()).unwrap();
        assert!(generate_request(request).is_err());
    }
}
<commit_msg>feat: error responses return body upon non 101 status code<commit_after>//! Client handshake machine.

use std::{
    io::{Read, Write},
    marker::PhantomData,
};

use http::{
    header::HeaderName, HeaderMap, Request as HttpRequest, Response as HttpResponse, StatusCode,
};
use httparse::Status;
use log::*;

use super::{
    derive_accept_key,
    headers::{FromHttparse, MAX_HEADERS},
    machine::{HandshakeMachine, StageResult, TryParse},
    HandshakeRole, MidHandshake, ProcessingResult,
};
use crate::{
    error::{Error, ProtocolError, Result, UrlError},
    protocol::{Role, WebSocket, WebSocketConfig},
};

/// Client request type.
pub type Request = HttpRequest<()>;

/// Client response type.
pub type Response = HttpResponse<Option<String>>;

/// Client handshake role.
#[derive(Debug)]
pub struct ClientHandshake<S> {
    verify_data: VerifyData,
    config: Option<WebSocketConfig>,
    _marker: PhantomData<S>,
}

impl<S: Read + Write> ClientHandshake<S> {
    /// Initiate a client handshake.
    pub fn start(
        stream: S,
        request: Request,
        config: Option<WebSocketConfig>,
    ) -> Result<MidHandshake<Self>> {
        if request.method() != http::Method::GET {
            return Err(Error::Protocol(ProtocolError::WrongHttpMethod));
        }

        if request.version() < http::Version::HTTP_11 {
            return Err(Error::Protocol(ProtocolError::WrongHttpVersion));
        }

        // Check the URI scheme: only ws or wss are supported
        let _ = crate::client::uri_mode(request.uri())?;

        // Convert and verify the `http::Request` and turn it into the request as per RFC.
        // Also extract the key from it (it must be present in a correct request).
        let (request, key) = generate_request(request)?;

        let machine = HandshakeMachine::start_write(stream, request);

        let client = {
            let accept_key = derive_accept_key(key.as_ref());
            ClientHandshake { verify_data: VerifyData { accept_key }, config, _marker: PhantomData }
        };

        trace!("Client handshake initiated.");
        Ok(MidHandshake { role: client, machine })
    }
}

impl<S: Read + Write> HandshakeRole for ClientHandshake<S> {
    type IncomingData = Response;
    type InternalStream = S;
    type FinalResult = (WebSocket<S>, Response);
    fn stage_finished(
        &mut self,
        finish: StageResult<Self::IncomingData, Self::InternalStream>,
    ) -> Result<ProcessingResult<Self::InternalStream, Self::FinalResult>> {
        Ok(match finish {
            StageResult::DoneWriting(stream) => {
                ProcessingResult::Continue(HandshakeMachine::start_read(stream))
            }
            StageResult::DoneReading { stream, result, tail } => {
                let result = self.verify_data.verify_response(result, &tail)?;
                debug!("Client handshake done.");
                let websocket =
                    WebSocket::from_partially_read(stream, tail, Role::Client, self.config);
                ProcessingResult::Done((websocket, result))
            }
        })
    }
}

/// Verifies and generates a client WebSocket request from the original request and extracts a WebSocket key from it.
fn generate_request(mut request: Request) -> Result<(Vec<u8>, String)> {
    let mut req = Vec::new();
    write!(
        req,
        "GET {path} {version:?}\r\n",
        path = request.uri().path_and_query().ok_or(Error::Url(UrlError::NoPathOrQuery))?.as_str(),
        version = request.version()
    )
    .unwrap();

    // Headers that must be present in a correct request.
    const KEY_HEADERNAME: &str = "Sec-WebSocket-Key";
    const WEBSOCKET_HEADERS: [&str; 5] =
        ["Host", "Connection", "Upgrade", "Sec-WebSocket-Version", KEY_HEADERNAME];

    // We must extract a WebSocket key from a properly formed request or fail if it's not present.
    let key = request
        .headers()
        .get(KEY_HEADERNAME)
        .ok_or_else(|| {
            Error::Protocol(ProtocolError::InvalidHeader(
                HeaderName::from_bytes(KEY_HEADERNAME.as_bytes()).unwrap(),
            ))
        })?
        .to_str()?
        .to_owned();

    // We must check that all necessary headers for a valid request are present. Note that we have to
    // deal with the fact that some apps seem to have a case-sensitive check for headers which is not
    // correct and should not considered the correct behavior, but it seems like some apps ignore it.
    // `http` by default writes all headers in lower-case which is fine (and does not violate the RFC)
    // but some servers seem to be poorely written and ignore RFC.
    //
    // See similar problem in `hyper`: https://github.com/hyperium/hyper/issues/1492
    let headers = request.headers_mut();
    for &header in &WEBSOCKET_HEADERS {
        let value = headers.remove(header).ok_or_else(|| {
            Error::Protocol(ProtocolError::InvalidHeader(
                HeaderName::from_bytes(header.as_bytes()).unwrap(),
            ))
        })?;
        write!(req, "{header}: {value}\r\n", header = header, value = value.to_str()?).unwrap();
    }

    // Now we must ensure that the headers that we've written once are not anymore present in the map.
    // If they do, then the request is invalid (some headers are duplicated there for some reason).
    let insensitive: Vec<String> =
        WEBSOCKET_HEADERS.iter().map(|h| h.to_ascii_lowercase()).collect();
    for (k, v) in headers {
        let mut name = k.as_str();

        // We have already written the necessary headers once (above) and removed them from the map.
        // If we encounter them again, then the request is considered invalid and error is returned.
        // Note that we can't use `.contains()`, since `&str` does not coerce to `&String` in Rust.
        if insensitive.iter().any(|x| x == name) {
            return Err(Error::Protocol(ProtocolError::InvalidHeader(k.clone())));
        }

        // Relates to the issue of some servers treating headers in a case-sensitive way, please see:
        // https://github.com/snapview/tungstenite-rs/pull/119 (original fix of the problem)
        if name == "sec-websocket-protocol" {
            name = "Sec-WebSocket-Protocol";
        }

        if name == "origin" {
            name = "Origin";
        }

        writeln!(req, "{}: {}\r", name, v.to_str()?).unwrap();
    }

    writeln!(req, "\r").unwrap();
    trace!("Request: {:?}", String::from_utf8_lossy(&req));
    Ok((req, key))
}

/// Information for handshake verification.
#[derive(Debug)]
struct VerifyData {
    /// Accepted server key.
    accept_key: String,
}

impl VerifyData {
    pub fn verify_response(&self, mut response: Response, tail: &[u8]) -> Result<Response> {
        // 1. If the status code received from the server is not 101, the
        // client handles the response per HTTP [RFC2616] procedures. (RFC 6455)
        if response.status() != StatusCode::SWITCHING_PROTOCOLS {
            *response.body_mut() = Some(String::from_utf8_lossy(tail).to_string());
            return Err(Error::Http(response));
        }

        let headers = response.headers();

        // 2. If the response lacks an |Upgrade| header field or the |Upgrade|
        // header field contains a value that is not an ASCII case-
        // insensitive match for the value "websocket", the client MUST
        // _Fail the WebSocket Connection_. (RFC 6455)
        if !headers
            .get("Upgrade")
            .and_then(|h| h.to_str().ok())
            .map(|h| h.eq_ignore_ascii_case("websocket"))
            .unwrap_or(false)
        {
            return Err(Error::Protocol(ProtocolError::MissingUpgradeWebSocketHeader));
        }
        // 3.  If the response lacks a |Connection| header field or the
        // |Connection| header field doesn't contain a token that is an
        // ASCII case-insensitive match for the value "Upgrade", the client
        // MUST _Fail the WebSocket Connection_. (RFC 6455)
        if !headers
            .get("Connection")
            .and_then(|h| h.to_str().ok())
            .map(|h| h.eq_ignore_ascii_case("Upgrade"))
            .unwrap_or(false)
        {
            return Err(Error::Protocol(ProtocolError::MissingConnectionUpgradeHeader));
        }
        // 4.  If the response lacks a |Sec-WebSocket-Accept| header field or
        // the |Sec-WebSocket-Accept| contains a value other than the
        // base64-encoded SHA-1 of ... the client MUST _Fail the WebSocket
        // Connection_. (RFC 6455)
        if !headers.get("Sec-WebSocket-Accept").map(|h| h == &self.accept_key).unwrap_or(false) {
            return Err(Error::Protocol(ProtocolError::SecWebSocketAcceptKeyMismatch));
        }
        // 5.  If the response includes a |Sec-WebSocket-Extensions| header
        // field and this header field indicates the use of an extension
        // that was not present in the client's handshake (the server has
        // indicated an extension not requested by the client), the client
        // MUST _Fail the WebSocket Connection_. (RFC 6455)
        // TODO

        // 6.  If the response includes a |Sec-WebSocket-Protocol| header field
        // and this header field indicates the use of a subprotocol that was
        // not present in the client's handshake (the server has indicated a
        // subprotocol not requested by the client), the client MUST _Fail
        // the WebSocket Connection_. (RFC 6455)
        // TODO

        Ok(response)
    }
}

impl TryParse for Response {
    fn try_parse(buf: &[u8]) -> Result<Option<(usize, Self)>> {
        let mut hbuffer = [httparse::EMPTY_HEADER; MAX_HEADERS];
        let mut req = httparse::Response::new(&mut hbuffer);
        Ok(match req.parse(buf)? {
            Status::Partial => None,
            Status::Complete(size) => Some((size, Response::from_httparse(req)?)),
        })
    }
}

impl<'h, 'b: 'h> FromHttparse<httparse::Response<'h, 'b>> for Response {
    fn from_httparse(raw: httparse::Response<'h, 'b>) -> Result<Self> {
        if raw.version.expect("Bug: no HTTP version") < /*1.*/1 {
            return Err(Error::Protocol(ProtocolError::WrongHttpMethod));
        }

        let headers = HeaderMap::from_httparse(raw.headers)?;

        let mut response = Response::new(None);
        *response.status_mut() = StatusCode::from_u16(raw.code.expect("Bug: no HTTP status code"))?;
        *response.headers_mut() = headers;
        // TODO: httparse only supports HTTP 0.9/1.0/1.1 but not HTTP 2.0
        // so the only valid value we could get in the response would be 1.1.
        *response.version_mut() = http::Version::HTTP_11;

        Ok(response)
    }
}

/// Generate a random key for the `Sec-WebSocket-Key` header.
pub fn generate_key() -> String {
    // a base64-encoded (see Section 4 of [RFC4648]) value that,
    // when decoded, is 16 bytes in length (RFC 6455)
    let r: [u8; 16] = rand::random();
    base64::encode(&r)
}

#[cfg(test)]
mod tests {
    use super::{super::machine::TryParse, generate_key, generate_request, Response};
    use crate::client::IntoClientRequest;

    #[test]
    fn random_keys() {
        let k1 = generate_key();
        println!("Generated random key 1: {}", k1);
        let k2 = generate_key();
        println!("Generated random key 2: {}", k2);
        assert_ne!(k1, k2);
        assert_eq!(k1.len(), k2.len());
        assert_eq!(k1.len(), 24);
        assert_eq!(k2.len(), 24);
        assert!(k1.ends_with("=="));
        assert!(k2.ends_with("=="));
        assert!(k1[..22].find('=').is_none());
        assert!(k2[..22].find('=').is_none());
    }

    fn construct_expected(host: &str, key: &str) -> Vec<u8> {
        format!(
            "\
            GET /getCaseCount HTTP/1.1\r\n\
            Host: {host}\r\n\
            Connection: Upgrade\r\n\
            Upgrade: websocket\r\n\
            Sec-WebSocket-Version: 13\r\n\
            Sec-WebSocket-Key: {key}\r\n\
            \r\n",
            host = host,
            key = key
        )
        .into_bytes()
    }

    #[test]
    fn request_formatting() {
        let request = "ws://localhost/getCaseCount".into_client_request().unwrap();
        let (request, key) = generate_request(request).unwrap();
        let correct = construct_expected("localhost", &key);
        assert_eq!(&request[..], &correct[..]);
    }

    #[test]
    fn request_formatting_with_host() {
        let request = "wss://localhost:9001/getCaseCount".into_client_request().unwrap();
        let (request, key) = generate_request(request).unwrap();
        let correct = construct_expected("localhost:9001", &key);
        assert_eq!(&request[..], &correct[..]);
    }

    #[test]
    fn request_formatting_with_at() {
        let request = "wss://user:pass@localhost:9001/getCaseCount".into_client_request().unwrap();
        let (request, key) = generate_request(request).unwrap();
        let correct = construct_expected("localhost:9001", &key);
        assert_eq!(&request[..], &correct[..]);
    }

    #[test]
    fn response_parsing() {
        const DATA: &[u8] = b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n";
        let (_, resp) = Response::try_parse(DATA).unwrap().unwrap();
        assert_eq!(resp.status(), http::StatusCode::OK);
        assert_eq!(resp.headers().get("Content-Type").unwrap(), &b"text/html"[..],);
    }

    #[test]
    fn invalid_custom_request() {
        let request = http::Request::builder().method("GET").body(()).unwrap();
        assert!(generate_request(request).is_err());
    }
}
<|endoftext|>