<commit_before>extern crate termion;
extern crate voodoo;

use std::io::{Write};

use termion::event::{Key, Event, MouseEvent};
use termion::input::{TermRead};

use voodoo::color::ColorValue;
use voodoo::window::{Point, TermCell, Window};

const LEVEL_DESCR: [&'static str; 22] = [
    "                                                          ",
    "                                                          ",
    "                                                          ",
    "                                                          ",
    "          .........................                       ",
    "          .........................                       ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "   o................             .....                    ",
    "   o................             .....                    ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "          .........................                       ",
    "          .........................                       ",
    "                                                          ",
    "                                                          ",
    "                                                          ",
];

struct Program {
    position: Point,
    name: String,
    abilities: Vec<String>,
}

impl Program {
    fn new(position: Point, name: &str) -> Program {
        Program {
            position: position,
            name: name.to_owned(),
            abilities: vec![],
        }
    }

    fn render(&self) -> (Point, TermCell) {
        let mut tc: TermCell = '◘'.into();
        tc.bg = Some(ColorValue::Green);
        (self.position, tc)
    }
}

struct Level {
    layout: Vec<String>,
    player_programs: Vec<Program>,
}

impl Level {
    fn new(description: &[&str; 22]) -> Level {
        let mut layout = Vec::new();
        for s in description.iter() {
            layout.push(s.to_string());
        }
        Level {
            layout: layout,
            player_programs: Vec::new(),
        }
    }

    fn display_for(&self, y: usize, x: usize) -> Option<TermCell> {
        Self::convert(self.layout[y].chars().nth(x).unwrap())
    }

    // TODO: need char -> Tile -> DisplayChar

    fn convert(c: char) -> Option<TermCell> {
        match c {
            '.' => Some('·'.into()),
            'o' => Some('O'.into()),
            _ => None,
        }
    }

    fn display(&self, map: &mut Window) {
        for (y, line) in self.layout.iter().enumerate() {
            let y = y + 1;
            for (x, tile) in line.chars().enumerate() {
                let x = x + 1;
                match Self::convert(tile) {
                    Some(c) => map.put_at(Point::new(x as u16, y as u16), c),
                    None => {},
                }
            }
        }

        for program in self.player_programs.iter() {
            let (point, c) = program.render();
            map.put_at(point, c);
        }
    }
}

enum UiState {
    Unselected,
    Selected,
}

enum UiEvent {
    Click(Point),
}

enum GameState {
    Setup,
    PlayerTurn,
    AITurn,
    AIExecute,
}

struct UiModelView {

}

impl UiState {
    fn next(self, event: UiEvent, level: &mut Level, info: &mut Window, map: &mut Window) -> UiState {
        use UiEvent::*;
        use UiState::*;

        match (self, event) {
            (Unselected, Click(p)) => {
                for program in level.player_programs.iter() {
                    if intersects(&program, p) {
                        let (_, mut tc) = program.render();
                        tc.bg = Some(ColorValue::Blue);
                        map.put_at(p, tc);
                        info.print_at(Point::new(2, 2), &program.name);
                        return Selected;
                    }
                }
                Unselected
            }
            (Selected, Click(_)) => {
                for program in level.player_programs.iter() {
                    let (p, mut tc) = program.render();
                    map.put_at(p, tc);
                }
                info.border();
                Unselected
            }
        }
    }
}

fn intersects(program: &Program, point: Point) -> bool {
    program.position == point
}

fn main() {
    use voodoo::terminal::{Mode, Terminal};
    let mut level = Level::new(&LEVEL_DESCR);
    level.player_programs.push(Program::new(Point::new(5, 11), "Hack"));
    level.player_programs.push(Program::new(Point::new(5, 12), "Hack"));

    let mut terminal = Terminal::new();
    terminal.cursor(Mode::Disabled);
    terminal.clear_color(ColorValue::Black);
    let Terminal { ref mut stdin, ref mut stdout } = terminal;

    stdout.flush().unwrap();

    let mut info = voodoo::window::Window::new(Point::new(0, 0), 20, 24);
    let mut map = voodoo::window::Window::new(Point::new(20, 0), 60, 24);
    info.border();
    map.border();
    level.display(&mut map);
    info.refresh(stdout);
    map.refresh(stdout);

    let mut ui_state = UiState::Unselected;

    for c in stdin.events() {
        let evt = c.unwrap();
        match evt {
            Event::Key(Key::Char('q')) => break,
            Event::Mouse(me) => {
                match me {
                    MouseEvent::Press(_, x, y) => {
                        if let Some(p) = map.position.from_global_frame(Point::new(x, y)) {
                            ui_state = ui_state.next(
                                UiEvent::Click(p),
                                &mut level,
                                &mut info,
                                &mut map
                            );
                        }
                    },
                    _ => (),
                }
            }
            _ => {}
        }
        info.refresh(stdout);
        map.refresh(stdout);
    }
 }
<commit_msg>refactor: move rendering code to info/map view structs<commit_after>extern crate termion;
extern crate voodoo;

use std::io::{Write};

use termion::event::{Key, Event, MouseEvent};
use termion::input::{TermRead};

use voodoo::color::ColorValue;
use voodoo::window::{Point, TermCell, Window};

const LEVEL_DESCR: [&'static str; 22] = [
    "                                                          ",
    "                                                          ",
    "                                                          ",
    "                                                          ",
    "          .........................                       ",
    "          .........................                       ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "   o................             .....                    ",
    "   o................             .....                    ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "          ..                     ..                       ",
    "          .........................                       ",
    "          .........................                       ",
    "                                                          ",
    "                                                          ",
    "                                                          ",
];

struct Program {
    position: Point,
    name: String,
    abilities: Vec<String>,
}

impl Program {
    fn new(position: Point, name: &str) -> Program {
        Program {
            position: position,
            name: name.to_owned(),
            abilities: vec![],
        }
    }

    fn render(&self) -> (Point, TermCell) {
        let mut tc: TermCell = '◘'.into();
        tc.bg = Some(ColorValue::Green);
        (self.position, tc)
    }
}

struct Level {
    layout: Vec<String>,
    player_programs: Vec<Program>,
}

impl Level {
    fn new(description: &[&str; 22]) -> Level {
        let mut layout = Vec::new();
        for s in description.iter() {
            layout.push(s.to_string());
        }
        Level {
            layout: layout,
            player_programs: Vec::new(),
        }
    }

    // TODO: need char -> Tile -> DisplayChar

    fn convert(c: char) -> Option<TermCell> {
        match c {
            '.' => Some('·'.into()),
            'o' => Some('O'.into()),
            _ => None,
        }
    }
}

enum UiState {
    Unselected,
    Selected,
}

enum UiEvent {
    Click(Point),
}

enum GameState {
    Setup,
    PlayerTurn,
    AITurn,
    AIExecute,
}

struct InfoView {
    window: Window,
}

impl InfoView {
    fn new(window: Window) -> InfoView {
        InfoView {
            window: window,
        }
    }

    fn refresh(&mut self, stdout: &mut std::io::Stdout) {
        self.window.refresh(stdout);
    }

    fn clear(&mut self) {
        for col in 2..self.window.width - 2 {
            self.window.put_at(Point::new(col, 2), ' ');
        }
    }

    fn display_program(&mut self, program: &Program) {
        self.window.print_at(Point::new(2, 2), &program.name);
    }
}

struct MapView {
    window: Window,
    highlight: Option<Point>,
}

impl MapView {
    fn new(window: Window) -> MapView {
        MapView {
            window: window,
            highlight: None,
        }
    }

    fn from_global_frame(&self, p: Point) -> Option<Point> {
        self.window.position.from_global_frame(p)
    }

    fn display(&mut self, level: &Level) {
        for (y, line) in level.layout.iter().enumerate() {
            let y = y + 1;
            for (x, tile) in line.chars().enumerate() {
                let x = x + 1;
                match Level::convert(tile) {
                    Some(c) => self.window.put_at(Point::new(x as u16, y as u16), c),
                    None => {},
                }
            }
        }

        for program in level.player_programs.iter() {
            let (point, mut c) = program.render();
            if let Some(p) = self.highlight {
                if p == point {
                    c.bg = Some(ColorValue::Blue);
                }
            }
            self.window.put_at(point, c);
        }
    }

    fn refresh(&mut self, stdout: &mut std::io::Stdout) {
        self.window.refresh(stdout);
    }

    fn highlight(&mut self, p: Point) {
        self.highlight = Some(p);
    }

    fn clear_highlight(&mut self) {
        self.highlight = None;
    }
}

struct UiModelView {
    info: InfoView,
    map: MapView,
}

impl UiState {
    fn next(self, event: UiEvent, level: &mut Level, mv: &mut UiModelView) -> UiState {
        use UiEvent::*;
        use UiState::*;

        let UiModelView { ref mut info, ref mut map } = *mv;

        match (self, event) {
            (Unselected, Click(p)) => {
                for program in level.player_programs.iter() {
                    if intersects(&program, p) {
                        map.highlight(p);
                        info.display_program(program);
                        return Selected;
                    }
                }
                Unselected
            }
            (Selected, Click(_)) => {
                for program in level.player_programs.iter() {
                    let (p, mut tc) = program.render();
                    map.clear_highlight();
                }
                info.clear();
                Unselected
            }
        }
    }
}

fn intersects(program: &Program, point: Point) -> bool {
    program.position == point
}

fn main() {
    use voodoo::terminal::{Mode, Terminal};
    let mut level = Level::new(&LEVEL_DESCR);
    level.player_programs.push(Program::new(Point::new(5, 11), "Hack"));
    level.player_programs.push(Program::new(Point::new(5, 12), "Hack"));

    let mut terminal = Terminal::new();
    terminal.cursor(Mode::Disabled);
    terminal.clear_color(ColorValue::Black);
    let Terminal { ref mut stdin, ref mut stdout } = terminal;

    stdout.flush().unwrap();

    let mut info = voodoo::window::Window::new(Point::new(0, 0), 20, 24);
    let mut map = voodoo::window::Window::new(Point::new(20, 0), 60, 24);
    info.border();
    map.border();

    let mut info_view = InfoView::new(info);
    let mut map_view = MapView::new(map);
    info_view.refresh(stdout);
    map_view.display(&level);
    map_view.refresh(stdout);

    let mut ui_modelview = UiModelView {
        info: info_view,
        map: map_view,
    };
    let mut ui_state = UiState::Unselected;

    for c in stdin.events() {
        let evt = c.unwrap();
        match evt {
            Event::Key(Key::Char('q')) => break,
            Event::Mouse(me) => {
                match me {
                    MouseEvent::Press(_, x, y) => {
                        if let Some(p) = ui_modelview.map.from_global_frame(Point::new(x, y)) {
                            ui_state = ui_state.next(
                                UiEvent::Click(p),
                                &mut level,
                                &mut ui_modelview,
                            );
                        }
                    },
                    _ => (),
                }
            }
            _ => {}
        }
        ui_modelview.info.refresh(stdout);
        ui_modelview.map.display(&level);
        ui_modelview.map.refresh(stdout);
    }
 }
<|endoftext|>


<commit_before>extern crate clap;
extern crate codemap;
extern crate mask;

use clap::App;
use clap::Arg;
use codemap::CodeMap;
use std::fs::File;
use std::path::Path;
use std::io::prelude::*;

fn print_tokens(map: &CodeMap, tokens: &Vec<codemap::Spanned<mask::lexer::Token>>) {
  let mut indent = 0;

  for token in tokens {
    match token.node {
      mask::lexer::Token::End | mask::lexer::Token::EOF => {
        println!("{:?}", token.node);
        for _ in 0..indent {
          print!(" ");
        }
      }

      mask::lexer::Token::Enter => {
        indent += 2;
        println!("{:?}", token.node);
        for _ in 0..indent {
          print!(" ");
        }
      }

      mask::lexer::Token::Exit => {
        indent -= 2;
        print!("{:?} ", token.node);
      }

      _ => {
        let span = map.look_up_span(token.span);
        print!("{:?} ", token.node);
        //print!("{}:{}:{}: {:?} ", span.file.name(), span.begin.line, span.begin.column, token.node);
        //print!("{:?}<{:?}> ", token.node, );
      }
    }
  }
}

fn main() {
  let argv = App::new("mask")
    .version("0.0.1")
    .author("John Weachock <jweachock@gmail.com>")
    .about("A programming language.")
    .arg(
      Arg::with_name("code")
        .short("c")
        .long("code")
        .value_name("CODE")
        .help("Mask code to execute. Overrides any specified module")
        .takes_value(true),
    )
    .arg(
      Arg::with_name("path")
        .takes_value(true)
        .index(1)
        .help("Mask module to execute"),
    )
    .get_matches();

  let mut map = CodeMap::new();

  let file = match argv.value_of("code") {
    Some(x) => map.add_file(String::from("_stdin"), x.to_string()),
    None => {
      let path_name = argv.value_of("path").unwrap_or("test.ms");
      let path = Path::new(&path_name);
      let mut file = match File::open(path) {
        Err(why) => panic!("Couldn't open {}: {}", path.display(), why),
        Ok(file) => file,
      };

      let mut contents = String::new();
      match file.read_to_string(&mut contents) {
        Err(why) => panic!("Couldn't read {}: {}", path.display(), why),
        Ok(_) => map.add_file(path_name.to_string(), contents.to_string()),
      }
    }
  };

  let tokens = mask::lexer::lex(&file);
  print_tokens(&map, &tokens);
  let ast = mask::parser::parse(tokens);
  println!("AST: {:?}", ast);
}
<commit_msg>main: Add dummy REPL<commit_after>extern crate clap;
extern crate codemap;
extern crate mask;

use clap::App;
use clap::Arg;
use codemap::CodeMap;
use std::fs::File;
use std::io::Write;
use std::io::prelude::*;
use std::io;
use std::path::Path;

fn print_tokens(map: &CodeMap, tokens: &Vec<codemap::Spanned<mask::lexer::Token>>) {
  let mut indent = 0;

  for token in tokens {
    match token.node {
      mask::lexer::Token::End | mask::lexer::Token::EOF => {
        println!("{:?}", token.node);
        for _ in 0..indent {
          print!(" ");
        }
      }

      mask::lexer::Token::Enter => {
        indent += 2;
        println!("{:?}", token.node);
        for _ in 0..indent {
          print!(" ");
        }
      }

      mask::lexer::Token::Exit => {
        indent -= 2;
        print!("{:?} ", token.node);
      }

      _ => {
        let span = map.look_up_span(token.span);
        print!("{:?} ", token.node);
        //print!("{}:{}:{}: {:?} ", span.file.name(), span.begin.line, span.begin.column, token.node);
        //print!("{:?}<{:?}> ", token.node, );
      }
    }
  }
}

fn main() {
  let argv = App::new("mask")
    .version("0.0.1")
    .author("John Weachock <jweachock@gmail.com>")
    .about("A programming language.")
    .arg(
      Arg::with_name("code")
        .short("c")
        .long("code")
        .value_name("CODE")
        .help("Mask code to execute. Overrides any specified module")
        .takes_value(true),
    )
    .arg(
      Arg::with_name("path")
        .takes_value(true)
        .index(1)
        .help("Mask module to execute"),
    )
    .get_matches();

  let mut map = CodeMap::new();

  if let Some(source) = argv.value_of("code") {
    let file = map.add_file(String::from("_stdin"), source.to_string());

    // FIXME this code is duplicated a lot, but that's because there's no
    // "module" component in the compiler yet
    let tokens = mask::lexer::lex(&file);
    let ast = mask::parser::parse(tokens);
    println!("AST: {:?}", ast);
  } else if let Some(filename) = argv.value_of("path") {
    let path = Path::new(&filename);
    let mut file = match File::open(path) {
      Ok(file) => file,
      Err(why) => panic!("Couldn't open {}: {}", path.display(), why),
    };

    let mut contents = String::new();
    let cm_file = match file.read_to_string(&mut contents) {
      Ok(_) => map.add_file(filename.to_string(), contents.to_string()),
      Err(why) => panic!("Couldn't read {}: {}", path.display(), why),
    };

    // see FIXME above
    let tokens = mask::lexer::lex(&cm_file);
    let ast = mask::parser::parse(tokens);
    println!("AST: {:?}", ast);
  } else {
    println!("Starting REPL");
    loop {
      let mut buffer = String::new();
      print!("> ");
      io::stdout().flush();

      match io::stdin().read_line(&mut buffer) {
        Ok(nbytes) => {
          if nbytes == 0 || buffer == "quit\n" {
            println!("");
            break;
          }

          let file = map.add_file(String::from("_stdin"), buffer);

          // see FIXME above
          let tokens = mask::lexer::lex(&file);
          let ast = mask::parser::parse(tokens);
          println!("AST: {:?}", ast);
        }

        Err(why) => {
          panic!("Unable to read line: {}", why);
        }
      }
    }
  }
}
<|endoftext|>


<commit_before>use byteorder::{ByteOrder, LittleEndian};

pub struct Stream {
    chacha20: ChaCha20,
    counter: u32,
    block: [u8; 64],
    block_index: u8,
}

/// A ChaCha20 iterator that can be used for encryption and decryption.
///
/// Initialized with a 32-byte key and 12-byte nonce. If reusing a key for encryption, be sure to
/// generate a unique nonce so that a given (key, nonce) pair is never used twice.
///
/// # Overflow Behavior
///
/// Iterating past 256 GB of the bytestream will cause the block counter to overflow.
impl Stream {
    pub fn new(key: &[u8], nonce: &[u8]) -> Self {
        let mut stream = Self {
            chacha20: ChaCha20::new(key),
            counter: 0,
            block: [0; 64],
            block_index: 0,
        };
        stream.chacha20.set_nonce(nonce);
        stream.chacha20.write_block(0, &mut stream.block);
        stream
    }

    // not sure how legal it is to encrypt without starting with a fresh block
    pub fn encrypt(&mut self, message: &[u8], ciphertext: &mut [u8]) {
        self.xor(message, ciphertext)
    }

    pub fn decrypt(&mut self, ciphertext: &[u8], message: &mut [u8]) {
        self.xor(ciphertext, message)
    }

    fn xor(&mut self, input: &[u8], output: &mut [u8]) {
        for (x, (y, z)) in output.iter_mut().zip(self.take(input.len()).zip(input)) {
            *x = y ^ z;
        }
    }
}

// doing this is a bit strange but it seemed interesting
impl Iterator for Stream {
    type Item = u8;

    fn next(&mut self) -> Option<Self::Item> {
        if self.block_index == 64 {
            self.counter += 1; // could overflow
            self.chacha20.write_block(self.counter, &mut self.block);
            self.block_index = 0;
        }
        let byte = self.block[self.block_index as usize];
        self.block_index += 1;
        Some(byte)
    }
}

pub(crate) struct ChaCha20 {
    state: [u32; 16],
}

impl ChaCha20 {
    pub(crate) fn new(key: &[u8]) -> Self {
        assert_eq!(32, key.len());
        let mut chacha20 = Self { state: [0; 16] };
        chacha20.setup_state(key);
        chacha20
    }

    pub(crate) fn set_nonce(&mut self, nonce: &[u8]) {
        assert_eq!(12, nonce.len());
        LittleEndian::read_u32_into(nonce, &mut self.state[13..]);
    }

    pub(crate) fn write_block(&self, counter: u32, output: &mut [u8]) {
        assert_eq!(64, output.len());
        let mut state = [0; 16];
        self.transform_state(&mut state, counter);
        Self::serialize_block(state, output)
    }

    fn setup_state(&mut self, key: &[u8]) {
        self.state[0] = 0x6170_7865;
        self.state[1] = 0x3320_646e;
        self.state[2] = 0x7962_2d32;
        self.state[3] = 0x6b20_6574;
        LittleEndian::read_u32_into(key, &mut self.state[4..12]);
    }

    fn transform_state(&self, state: &mut [u32; 16], counter: u32) {
        *state = self.state;
        state[12] = counter;
        for _ in 0..10 {
            Self::inner_block(state);
        }
        for (x, &y) in state.iter_mut().zip(&self.state) {
            *x = x.wrapping_add(y);
        }
        state[12] = state[12].wrapping_add(counter);
    }

    fn inner_block(state: &mut [u32; 16]) {
        Self::quarter_round(state, 0, 4, 8, 12);
        Self::quarter_round(state, 1, 5, 9, 13);
        Self::quarter_round(state, 2, 6, 10, 14);
        Self::quarter_round(state, 3, 7, 11, 15);
        Self::quarter_round(state, 0, 5, 10, 15);
        Self::quarter_round(state, 1, 6, 11, 12);
        Self::quarter_round(state, 2, 7, 8, 13);
        Self::quarter_round(state, 3, 4, 9, 14);
    }

    fn quarter_round(state: &mut [u32; 16], i: usize, j: usize, k: usize, l: usize) {
        let mut a = state[i];
        let mut b = state[j];
        let mut c = state[k];
        let mut d = state[l];
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(16);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(12);
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(8);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(7);
        state[i] = a;
        state[j] = b;
        state[k] = c;
        state[l] = d;
    }

    fn serialize_block(block: [u32; 16], output: &mut [u8]) {
        LittleEndian::write_u32_into(&block, output);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use test_helpers::*;

    const KEY: &str = "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f";
    const NONCE: &[u8] = &[0, 0, 0, 0x09, 0, 0, 0, 0x4a, 0, 0, 0, 0];
    const SETUP_STATE: [u32; 16] = [
        0x61707865,
        0x3320646e,
        0x79622d32,
        0x6b206574,
        0x03020100,
        0x07060504,
        0x0b0a0908,
        0x0f0e0d0c,
        0x13121110,
        0x17161514,
        0x1b1a1918,
        0x1f1e1d1c,
        0x00000000,
        0x09000000,
        0x4a000000,
        0x00000000,
    ];
    const FINAL_STATE: [u32; 16] = [
        0xe4e7f110,
        0x15593bd1,
        0x1fdd0f50,
        0xc47120a3,
        0xc7f4d1c7,
        0x0368c033,
        0x9aaa2204,
        0x4e6cd4c3,
        0x466482d2,
        0x09aa9f07,
        0x05d7c214,
        0xa2028bd9,
        0xd19c12b5,
        0xb94e16de,
        0xe883d0cb,
        0x4e3c50a2,
    ];

    #[test]
    fn test_encrypt() {
        let key = &h2b(KEY);
        let nonce = &[0, 0, 0, 0, 0, 0, 0, 0x4a, 0, 0, 0, 0];
        let message = "Ladies and Gentlemen of the class of '99: If I could offer you only one \
            tip for the future, sunscreen would be it.";
        let ciphertext = h2b(
            "6e2e359a2568f98041ba0728dd0d6981e97e7aec1d4360c20a27afccfd9fae0b\
             f91b65c5524733ab8f593dabcd62b3571639d624e65152ab8f530c359f0861d8\
             07ca0dbf500d6a6156a38e088a22b65e52bc514d16ccf806818ce91ab7793736\
             5af90bbf74a35be6b40b8eedf2785e42874d",
        );
        let encrypted_message = &mut vec![0; message.len()];
        let decrypted_ciphertext = &mut vec![0; ciphertext.len()];

        let mut stream = Stream::new(key, nonce);
        stream.nth(64 - 1);
        stream.encrypt(message.as_bytes(), encrypted_message);
        assert_eq!(&ciphertext, encrypted_message);

        stream = Stream::new(key, nonce);
        stream.nth(64 - 1);
        stream.decrypt(&ciphertext, decrypted_ciphertext);
        assert_eq!(message.as_bytes(), decrypted_ciphertext.as_slice());
    }

    #[test]
    fn test_new() {
        let mut chacha20 = ChaCha20::new(&h2b(KEY));
        chacha20.set_nonce(NONCE);
        assert_eq!(SETUP_STATE, chacha20.state);
    }

    fn check_serialized_block(block: &[u8]) {
        let block_one = h2b(
            "10f1e7e4d13b5915500fdd1fa32071c4c7d1f4c733c068030422aa9ac3d46c4e\
             d2826446079faa0914c2d705d98b02a2b5129cd1de164eb9cbd083e8a2503c4e",
        );
        assert_eq!(block_one, block.to_vec());
    }

    #[test]
    fn test_get_block() {
        let mut chacha20 = ChaCha20::new(&h2b(KEY));
        chacha20.set_nonce(NONCE);
        let block = &mut [0; 64];
        chacha20.write_block(1, block);
        check_serialized_block(block);
    }

    #[test]
    fn test_transform_state() {
        let chacha20 = ChaCha20 { state: SETUP_STATE };
        let mut state = [0; 16];
        chacha20.transform_state(&mut state, 1);
        assert_eq!(FINAL_STATE, state);
    }

    #[test]
    fn test_inner_block() {
        let mut state = SETUP_STATE;
        state[12] = 1;
        for _ in 0..10 {
            ChaCha20::inner_block(&mut state);
        }
        assert_eq!(
            [
                0x837778ab,
                0xe238d763,
                0xa67ae21e,
                0x5950bb2f,
                0xc4f2d0c7,
                0xfc62bb2f,
                0x8fa018fc,
                0x3f5ec7b7,
                0x335271c2,
                0xf29489f3,
                0xeabda8fc,
                0x82e46ebd,
                0xd19c12b4,
                0xb04e16de,
                0x9e83d0cb,
                0x4e3c50a2,
            ],
            state
        );
    }

    #[test]
    fn test_quarter_round() {
        let mut state = [0; 16];
        state[0] = 0x11111111;
        state[1] = 0x01020304;
        state[2] = 0x9b8d6f43;
        state[3] = 0x01234567;
        ChaCha20::quarter_round(&mut state, 0, 1, 2, 3);
        assert_eq!([0xea2a92f4, 0xcb1cf8ce, 0x4581472e, 0x5881c4bb], state[..4]);

        state = [
            0x879531e0,
            0xc5ecf37d,
            0x516461b1,
            0xc9a62f8a,
            0x44c20ef3,
            0x3390af7f,
            0xd9fc690b,
            0x2a5f714c,
            0x53372767,
            0xb00a5631,
            0x974c541a,
            0x359e9963,
            0x5c971061,
            0x3d631689,
            0x2098d9d6,
            0x91dbd320,
        ];
        ChaCha20::quarter_round(&mut state, 2, 7, 8, 13);
        assert_eq!(
            [
                0x879531e0,
                0xc5ecf37d,
                0xbdb886dc,
                0xc9a62f8a,
                0x44c20ef3,
                0x3390af7f,
                0xd9fc690b,
                0xcfacafd2,
                0xe46bea80,
                0xb00a5631,
                0x974c541a,
                0x359e9963,
                0x5c971061,
                0xccc07c79,
                0x2098d9d6,
                0x91dbd320,
            ],
            state
        );
    }

    #[test]
    fn test_serialize_block() {
        let block = &mut [0; 64];
        ChaCha20::serialize_block(FINAL_STATE, block);
        check_serialized_block(block);
    }
}
<commit_msg>chacha20: documentation<commit_after>//! Module for the ChaCha20 stream cipher.
use byteorder::{ByteOrder, LittleEndian};

/// A ChaCha20 iterator that can be used as a stream cipher or pseudorandom generator.
///
/// Initialized with a 32-byte key and 12-byte nonce. If reusing a key for encryption, be sure to
/// generate a unique nonce so that a given (key, nonce) pair is never used twice.
///
/// # Overflow Behavior
///
/// Iterating past 256 GB of the bytestream will cause the block counter to overflow.
pub struct Stream {
    chacha20: ChaCha20,
    counter: u32,
    block: [u8; 64],
    block_index: u8,
}

impl Stream {
    /// Initializes a ChaCha20 stream given a key and nonce.
    ///
    /// # Panics
    ///
    /// Panics if `key.len()` is not equal to 32 or if `nonce.len()` is not equal to 12.
    pub fn new(key: &[u8], nonce: &[u8]) -> Self {
        let mut stream = Self {
            chacha20: ChaCha20::new(key),
            counter: 0,
            block: [0; 64],
            block_index: 0,
        };
        stream.chacha20.set_nonce(nonce);
        stream.chacha20.write_block(0, &mut stream.block);
        stream
    }

    /// Encrypts a message into a ciphertext.
    ///
    /// # Panics
    ///
    /// Panics if `message.len()` is not equal to `ciphertext.len()`.
    pub fn encrypt(&mut self, message: &[u8], ciphertext: &mut [u8]) {
        self.xor(message, ciphertext)
    }

    /// Decrypts a ciphertext into a message.
    ///
    /// # Panics
    ///
    /// Panics if `message.len()` is not equal to `ciphertext.len()`.
    pub fn decrypt(&mut self, ciphertext: &[u8], message: &mut [u8]) {
        self.xor(ciphertext, message)
    }

    fn xor(&mut self, input: &[u8], output: &mut [u8]) {
        assert_eq!(input.len(), output.len());
        for (x, (y, z)) in output.iter_mut().zip(self.take(input.len()).zip(input)) {
            *x = y ^ z;
        }
    }
}

impl Iterator for Stream {
    type Item = u8;

    fn next(&mut self) -> Option<Self::Item> {
        if self.block_index == 64 {
            self.counter += 1; // could overflow
            self.chacha20.write_block(self.counter, &mut self.block);
            self.block_index = 0;
        }
        let byte = self.block[self.block_index as usize];
        self.block_index += 1;
        Some(byte)
    }
}

pub(crate) struct ChaCha20 {
    state: [u32; 16],
}

impl ChaCha20 {
    pub(crate) fn new(key: &[u8]) -> Self {
        assert_eq!(32, key.len());
        let mut chacha20 = Self { state: [0; 16] };
        chacha20.setup_state(key);
        chacha20
    }

    pub(crate) fn set_nonce(&mut self, nonce: &[u8]) {
        assert_eq!(12, nonce.len());
        LittleEndian::read_u32_into(nonce, &mut self.state[13..]);
    }

    pub(crate) fn write_block(&self, counter: u32, output: &mut [u8]) {
        assert_eq!(64, output.len());
        let mut state = [0; 16];
        self.transform_state(&mut state, counter);
        Self::serialize_block(state, output)
    }

    fn setup_state(&mut self, key: &[u8]) {
        self.state[0] = 0x6170_7865;
        self.state[1] = 0x3320_646e;
        self.state[2] = 0x7962_2d32;
        self.state[3] = 0x6b20_6574;
        LittleEndian::read_u32_into(key, &mut self.state[4..12]);
    }

    fn transform_state(&self, state: &mut [u32; 16], counter: u32) {
        *state = self.state;
        state[12] = counter;
        for _ in 0..10 {
            Self::inner_block(state);
        }
        for (x, &y) in state.iter_mut().zip(&self.state) {
            *x = x.wrapping_add(y);
        }
        state[12] = state[12].wrapping_add(counter);
    }

    fn inner_block(state: &mut [u32; 16]) {
        Self::quarter_round(state, 0, 4, 8, 12);
        Self::quarter_round(state, 1, 5, 9, 13);
        Self::quarter_round(state, 2, 6, 10, 14);
        Self::quarter_round(state, 3, 7, 11, 15);
        Self::quarter_round(state, 0, 5, 10, 15);
        Self::quarter_round(state, 1, 6, 11, 12);
        Self::quarter_round(state, 2, 7, 8, 13);
        Self::quarter_round(state, 3, 4, 9, 14);
    }

    fn quarter_round(state: &mut [u32; 16], i: usize, j: usize, k: usize, l: usize) {
        let mut a = state[i];
        let mut b = state[j];
        let mut c = state[k];
        let mut d = state[l];
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(16);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(12);
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(8);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(7);
        state[i] = a;
        state[j] = b;
        state[k] = c;
        state[l] = d;
    }

    fn serialize_block(block: [u32; 16], output: &mut [u8]) {
        LittleEndian::write_u32_into(&block, output);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use test_helpers::*;

    const KEY: &str = "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f";
    const NONCE: &[u8] = &[0, 0, 0, 0x09, 0, 0, 0, 0x4a, 0, 0, 0, 0];
    const SETUP_STATE: [u32; 16] = [
        0x61707865,
        0x3320646e,
        0x79622d32,
        0x6b206574,
        0x03020100,
        0x07060504,
        0x0b0a0908,
        0x0f0e0d0c,
        0x13121110,
        0x17161514,
        0x1b1a1918,
        0x1f1e1d1c,
        0x00000000,
        0x09000000,
        0x4a000000,
        0x00000000,
    ];
    const FINAL_STATE: [u32; 16] = [
        0xe4e7f110,
        0x15593bd1,
        0x1fdd0f50,
        0xc47120a3,
        0xc7f4d1c7,
        0x0368c033,
        0x9aaa2204,
        0x4e6cd4c3,
        0x466482d2,
        0x09aa9f07,
        0x05d7c214,
        0xa2028bd9,
        0xd19c12b5,
        0xb94e16de,
        0xe883d0cb,
        0x4e3c50a2,
    ];

    #[test]
    fn test_encrypt() {
        let key = &h2b(KEY);
        let nonce = &[0, 0, 0, 0, 0, 0, 0, 0x4a, 0, 0, 0, 0];
        let message = "Ladies and Gentlemen of the class of '99: If I could offer you only one \
            tip for the future, sunscreen would be it.";
        let ciphertext = h2b(
            "6e2e359a2568f98041ba0728dd0d6981e97e7aec1d4360c20a27afccfd9fae0b\
             f91b65c5524733ab8f593dabcd62b3571639d624e65152ab8f530c359f0861d8\
             07ca0dbf500d6a6156a38e088a22b65e52bc514d16ccf806818ce91ab7793736\
             5af90bbf74a35be6b40b8eedf2785e42874d",
        );
        let encrypted_message = &mut vec![0; message.len()];
        let decrypted_ciphertext = &mut vec![0; ciphertext.len()];

        let mut stream = Stream::new(key, nonce);
        stream.nth(64 - 1);
        stream.encrypt(message.as_bytes(), encrypted_message);
        assert_eq!(&ciphertext, encrypted_message);

        stream = Stream::new(key, nonce);
        stream.nth(64 - 1);
        stream.decrypt(&ciphertext, decrypted_ciphertext);
        assert_eq!(message.as_bytes(), decrypted_ciphertext.as_slice());
    }

    #[test]
    fn test_new() {
        let mut chacha20 = ChaCha20::new(&h2b(KEY));
        chacha20.set_nonce(NONCE);
        assert_eq!(SETUP_STATE, chacha20.state);
    }

    fn check_serialized_block(block: &[u8]) {
        let block_one = h2b(
            "10f1e7e4d13b5915500fdd1fa32071c4c7d1f4c733c068030422aa9ac3d46c4e\
             d2826446079faa0914c2d705d98b02a2b5129cd1de164eb9cbd083e8a2503c4e",
        );
        assert_eq!(block_one, block.to_vec());
    }

    #[test]
    fn test_get_block() {
        let mut chacha20 = ChaCha20::new(&h2b(KEY));
        chacha20.set_nonce(NONCE);
        let block = &mut [0; 64];
        chacha20.write_block(1, block);
        check_serialized_block(block);
    }

    #[test]
    fn test_transform_state() {
        let chacha20 = ChaCha20 { state: SETUP_STATE };
        let mut state = [0; 16];
        chacha20.transform_state(&mut state, 1);
        assert_eq!(FINAL_STATE, state);
    }

    #[test]
    fn test_inner_block() {
        let mut state = SETUP_STATE;
        state[12] = 1;
        for _ in 0..10 {
            ChaCha20::inner_block(&mut state);
        }
        assert_eq!(
            [
                0x837778ab,
                0xe238d763,
                0xa67ae21e,
                0x5950bb2f,
                0xc4f2d0c7,
                0xfc62bb2f,
                0x8fa018fc,
                0x3f5ec7b7,
                0x335271c2,
                0xf29489f3,
                0xeabda8fc,
                0x82e46ebd,
                0xd19c12b4,
                0xb04e16de,
                0x9e83d0cb,
                0x4e3c50a2,
            ],
            state
        );
    }

    #[test]
    fn test_quarter_round() {
        let mut state = [0; 16];
        state[0] = 0x11111111;
        state[1] = 0x01020304;
        state[2] = 0x9b8d6f43;
        state[3] = 0x01234567;
        ChaCha20::quarter_round(&mut state, 0, 1, 2, 3);
        assert_eq!([0xea2a92f4, 0xcb1cf8ce, 0x4581472e, 0x5881c4bb], state[..4]);

        state = [
            0x879531e0,
            0xc5ecf37d,
            0x516461b1,
            0xc9a62f8a,
            0x44c20ef3,
            0x3390af7f,
            0xd9fc690b,
            0x2a5f714c,
            0x53372767,
            0xb00a5631,
            0x974c541a,
            0x359e9963,
            0x5c971061,
            0x3d631689,
            0x2098d9d6,
            0x91dbd320,
        ];
        ChaCha20::quarter_round(&mut state, 2, 7, 8, 13);
        assert_eq!(
            [
                0x879531e0,
                0xc5ecf37d,
                0xbdb886dc,
                0xc9a62f8a,
                0x44c20ef3,
                0x3390af7f,
                0xd9fc690b,
                0xcfacafd2,
                0xe46bea80,
                0xb00a5631,
                0x974c541a,
                0x359e9963,
                0x5c971061,
                0xccc07c79,
                0x2098d9d6,
                0x91dbd320,
            ],
            state
        );
    }

    #[test]
    fn test_serialize_block() {
        let block = &mut [0; 64];
        ChaCha20::serialize_block(FINAL_STATE, block);
        check_serialized_block(block);
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
extern crate bytes;
extern crate futures;
extern crate tokio_core;
extern crate net2;

use std::net::SocketAddr;
use std::time::Duration;

use net2::UdpBuilder;
use net2::unix::UnixUdpBuilderExt;

use std::thread;
use std::sync::atomic::{ATOMIC_USIZE_INIT, AtomicUsize, Ordering};

use tokio_core::reactor::{Core, Interval};
use tokio_core::net::{UdpSocket, UdpCodec};

use bytes::{Bytes, IntoBuf, Buf};
use futures::{Stream, Sink, AsyncSink};

use clap::{App, Arg};

pub static CHUNK_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;
pub static DROP_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;
pub static ERR_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;

struct DgramCodec {
    addr: SocketAddr,
}

impl DgramCodec {
    pub fn new(addr: SocketAddr) -> Self {
        DgramCodec { addr: addr }
    }
}

impl UdpCodec for DgramCodec {
    type In = Bytes;
    type Out = Bytes;

    fn decode(&mut self,
              _src: &SocketAddr,
              buf: &[u8])
              -> ::std::result::Result<Self::In, ::std::io::Error> {
        Ok(buf.into())
    }

    fn encode(&mut self, msg: Self::Out, buf: &mut Vec<u8>) -> SocketAddr {
        use std::io::Read;
        msg.into_buf().reader().read_to_end(buf).unwrap();
        self.addr
    }
}

fn main() {
    let matches = App::new("MEGA UDP DUPLICATOR")
        .version("0.0.0.0.0.0.1alpha")
        .author("Sergey Noskov <snoskov@avito.ru>")
        .about("UDP multiplexor: copies packets from listened address to specified backends")
        .arg(Arg::with_name("listen")
             .short("l")
             .long("listen")
             .value_name("IP:PORT")
             .takes_value(true)
             .help("ip and port to listen to")
            )
        .arg(Arg::with_name("backend")
             .short("b")
             .long("backend")
             .value_name("IP:PORT")
             .multiple(true)
             .takes_value(true)
             .help("IP and port of a backend to forward data to. Can be specified multiple times.")
            )
        .arg(Arg::with_name("bufsize")
             .help("Internal queue buffer size")
             .short("s")
             .long("size")
             .takes_value(true)
             .default_value("1048576")
            )
        .arg(Arg::with_name("nthreads")
             .short("n")
             .long("nthreads")
             .takes_value(true)
             .default_value("0")
            )
        .arg(Arg::with_name("interval")
             .short("t")
             .long("interval")
             .takes_value(true)
             .default_value("1000")
             .help("Stats printing interval (in milliseconds)")
            )
        .get_matches();

    let listen = value_t!(matches.value_of("listen"), SocketAddr).unwrap_or_else(|e| e.exit());
    let backends = values_t!(matches.values_of("backend"), SocketAddr).unwrap_or_else(|e| e.exit());
    let nthreads = value_t!(matches.value_of("nthreads"), usize).unwrap_or_else(|e| e.exit());
    let bufsize = value_t!(matches.value_of("bufsize"), usize).unwrap_or_else(|e| e.exit());
    let interval = value_t!(matches.value_of("interval"), u64).unwrap_or_else(|e| e.exit());

    // Init chunk counter
    // In the main thread start a reporting timer
    let mut core = Core::new().unwrap();
    let handle = core.handle();

    let timer = Interval::new(Duration::from_millis(interval), &handle).unwrap();
    let timer = timer.for_each(|()|{
        let interval = interval as f64 / 1000f64;
        let chunks = CHUNK_COUNTER.load(Ordering::Relaxed) as f64 / interval;
        let drops = DROP_COUNTER.load(Ordering::Relaxed) as f64 / interval;
        let errors = ERR_COUNTER.load(Ordering::Relaxed) as f64 / interval;

    println!("chunks/drops/errors: {:.2}/{:.2}/{:.2}", chunks, drops, errors);
    CHUNK_COUNTER.store(0, Ordering::Relaxed);
    DROP_COUNTER.store(0, Ordering::Relaxed);
    ERR_COUNTER.store(0, Ordering::Relaxed);
    Ok(())
    });

    let socket = UdpBuilder::new_v4().unwrap();
    socket.reuse_address(true).unwrap();
    socket.reuse_port(true).unwrap();
    let socket = socket.bind(&listen).unwrap();

    // Create sockets from backend addresses
    let senders = backends
        .iter()
        .map(|_| {
                 let socket = UdpBuilder::new_v4().unwrap();
                 let bind: SocketAddr = "0.0.0.0:0".parse().unwrap();
                 socket.bind(&bind).unwrap()
             })
        .collect::<Vec<_>>();

    for _ in 0..nthreads {
        // create clone of each sender socket for thread
        let senders = senders
            .iter()
            .map(|socket| socket.try_clone().unwrap())
            .collect::<Vec<_>>();

        let socket = socket.try_clone().unwrap();
        let mut backends = backends.clone();
        thread::spawn(move || {
            // each thread runs it's own core
            let mut core = Core::new().unwrap();
            let handle = core.handle();

            // make senders into framed UDP
            let mut senders = senders
                .into_iter()
                .map(|socket| {
                    let sender = UdpSocket::from_socket(socket, &handle).unwrap();
                    // socket order becomes reversed but we don't care about this
                    // because the order is same everywhere
                    // and anyways, we push same data to every socket
                    let sender = sender
                        .framed(DgramCodec::new(backends.pop().unwrap()))
                        .buffer(bufsize);
                    sender
                })
                .collect::<Vec<_>>();

            // create server now
            let socket = UdpSocket::from_socket(socket, &handle).unwrap();
            let server = socket.framed(DgramCodec::new(listen)).buffer(bufsize);

            let server = server.for_each(move |buf| {
                CHUNK_COUNTER.fetch_add(1, Ordering::Relaxed);
                senders
                    .iter_mut()
                    .map(|sender| {
                        let res = sender.start_send(buf.clone());
                        match res {
                            Err(_) => {
                                ERR_COUNTER.fetch_add(1, Ordering::Relaxed);
                            }
                            Ok(AsyncSink::NotReady(_)) => {
                                DROP_COUNTER.fetch_add(1, Ordering::Relaxed);
                            }
                            Ok(AsyncSink::Ready) => {
                                let res = sender.poll_complete();
                                if res.is_err() {
                                    ERR_COUNTER.fetch_add(1, Ordering::Relaxed);
                                }
                            }
                        }
                    })
                    .last();
                Ok(())
            });

            core.run(server).unwrap();
        });
    }
    core.run(timer).unwrap();
}
<commit_msg>Fixed joky version description<commit_after>#[macro_use]
extern crate clap;
extern crate bytes;
extern crate futures;
extern crate tokio_core;
extern crate net2;

use std::net::SocketAddr;
use std::time::Duration;

use net2::UdpBuilder;
use net2::unix::UnixUdpBuilderExt;

use std::thread;
use std::sync::atomic::{ATOMIC_USIZE_INIT, AtomicUsize, Ordering};

use tokio_core::reactor::{Core, Interval};
use tokio_core::net::{UdpSocket, UdpCodec};

use bytes::{Bytes, IntoBuf, Buf};
use futures::{Stream, Sink, AsyncSink};

use clap::{App, Arg};

pub static CHUNK_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;
pub static DROP_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;
pub static ERR_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;

struct DgramCodec {
    addr: SocketAddr,
}

impl DgramCodec {
    pub fn new(addr: SocketAddr) -> Self {
        DgramCodec { addr: addr }
    }
}

impl UdpCodec for DgramCodec {
    type In = Bytes;
    type Out = Bytes;

    fn decode(&mut self,
              _src: &SocketAddr,
              buf: &[u8])
        -> ::std::result::Result<Self::In, ::std::io::Error> {
            Ok(buf.into())
        }

    fn encode(&mut self, msg: Self::Out, buf: &mut Vec<u8>) -> SocketAddr {
        use std::io::Read;
        msg.into_buf().reader().read_to_end(buf).unwrap();
        self.addr
    }
}

fn main() {
    let matches = App::new("Simple UDP duplicator")
        .version("0.1.0")
        .author("Sergey Noskov <snoskov@avito.ru>")
        .about("UDP multiplexor: copies packets from listened address to specified backends dropping packets on any error")
        .arg(Arg::with_name("listen")
             .short("l")
             .long("listen")
             .value_name("IP:PORT")
             .takes_value(true)
             .help("ip and port to listen to")
            )
        .arg(Arg::with_name("backend")
             .short("b")
             .long("backend")
             .value_name("IP:PORT")
             .multiple(true)
             .takes_value(true)
             .help("IP and port of a backend to forward data to. Can be specified multiple times.")
            )
        .arg(Arg::with_name("bufsize")
             .help("Internal queue buffer size")
             .short("s")
             .long("size")
             .takes_value(true)
             .default_value("1048576")
            )
        .arg(Arg::with_name("nthreads")
             .short("n")
             .long("nthreads")
             .takes_value(true)
             .default_value("0")
            )
        .arg(Arg::with_name("interval")
             .short("t")
             .long("interval")
             .takes_value(true)
             .default_value("1000")
             .help("Stats printing interval (in milliseconds)")
            )
        .get_matches();

    let listen = value_t!(matches.value_of("listen"), SocketAddr).unwrap_or_else(|e| e.exit());
    let backends = values_t!(matches.values_of("backend"), SocketAddr).unwrap_or_else(|e| e.exit());
    let nthreads = value_t!(matches.value_of("nthreads"), usize).unwrap_or_else(|e| e.exit());
    let bufsize = value_t!(matches.value_of("bufsize"), usize).unwrap_or_else(|e| e.exit());
    let interval = value_t!(matches.value_of("interval"), u64).unwrap_or_else(|e| e.exit());

    // Init chunk counter
    // In the main thread start a reporting timer
    let mut core = Core::new().unwrap();
    let handle = core.handle();

    let timer = Interval::new(Duration::from_millis(interval), &handle).unwrap();
    let timer = timer.for_each(|()|{
        let interval = interval as f64 / 1000f64;
        let chunks = CHUNK_COUNTER.load(Ordering::Relaxed) as f64 / interval;
        let drops = DROP_COUNTER.load(Ordering::Relaxed) as f64 / interval;
        let errors = ERR_COUNTER.load(Ordering::Relaxed) as f64 / interval;

        println!("chunks/drops/errors: {:.2}/{:.2}/{:.2}", chunks, drops, errors);
        CHUNK_COUNTER.store(0, Ordering::Relaxed);
        DROP_COUNTER.store(0, Ordering::Relaxed);
        ERR_COUNTER.store(0, Ordering::Relaxed);
        Ok(())
    });

    let socket = UdpBuilder::new_v4().unwrap();
    socket.reuse_address(true).unwrap();
    socket.reuse_port(true).unwrap();
    let socket = socket.bind(&listen).unwrap();

    // Create sockets from backend addresses
    let senders = backends
        .iter()
        .map(|_| {
            let socket = UdpBuilder::new_v4().unwrap();
            let bind: SocketAddr = "0.0.0.0:0".parse().unwrap();
            socket.bind(&bind).unwrap()
        })
    .collect::<Vec<_>>();

    for _ in 0..nthreads {
        // create clone of each sender socket for thread
        let senders = senders
            .iter()
            .map(|socket| socket.try_clone().unwrap())
            .collect::<Vec<_>>();

        let socket = socket.try_clone().unwrap();
        let mut backends = backends.clone();
        thread::spawn(move || {
            // each thread runs it's own core
            let mut core = Core::new().unwrap();
            let handle = core.handle();

            // make senders into framed UDP
            let mut senders = senders
                .into_iter()
                .map(|socket| {
                    let sender = UdpSocket::from_socket(socket, &handle).unwrap();
                    // socket order becomes reversed but we don't care about this
                    // because the order is same everywhere
                    // and anyways, we push same data to every socket
                    let sender = sender
                        .framed(DgramCodec::new(backends.pop().unwrap()))
                        .buffer(bufsize);
                    sender
                })
            .collect::<Vec<_>>();

            // create server now
            let socket = UdpSocket::from_socket(socket, &handle).unwrap();
            let server = socket.framed(DgramCodec::new(listen)).buffer(bufsize);

            let server = server.for_each(move |buf| {
                CHUNK_COUNTER.fetch_add(1, Ordering::Relaxed);
                senders
                    .iter_mut()
                    .map(|sender| {
                        let res = sender.start_send(buf.clone());
                        match res {
                            Err(_) => {
                                ERR_COUNTER.fetch_add(1, Ordering::Relaxed);
                            }
                            Ok(AsyncSink::NotReady(_)) => {
                                DROP_COUNTER.fetch_add(1, Ordering::Relaxed);
                            }
                            Ok(AsyncSink::Ready) => {
                                let res = sender.poll_complete();
                                if res.is_err() {
                                    ERR_COUNTER.fetch_add(1, Ordering::Relaxed);
                                }
                            }
                        }
                    })
                .last();
                Ok(())
            });

            core.run(server).unwrap();
        });
    }
    core.run(timer).unwrap();
}
<|endoftext|>


<commit_before>
#[derive(Copy, Clone)]
enum InstructionSet {
	PSH,
	ADD,
	POP,
	SET,
	HLT,
	VAL{value: i32},
}

struct VM {
	ip: usize,
	program: Vec<InstructionSet>,
	running: bool,
	sp: isize,
	stack: [InstructionSet; 256],
}

impl VM {
	pub fn new(prog: Vec<InstructionSet>) -> VM {
		VM {
			ip : 0,
			program: prog,
			running: true,
			sp: -1,
			stack:[InstructionSet::VAL{value:0}; 256],
		}
	}

	fn fetch(&mut self) -> InstructionSet {
		let instr = self.program[self.ip];
		instr
	}

	fn eval(&mut self, instr: InstructionSet) {
		use InstructionSet::*;
		match instr {
			HLT => self.running = false,
			PSH => {
				self.sp += 1;
				self.ip += 1; // move to operand location
				self.stack[self.sp as usize] = self.program[self.ip];
				//[DEBUG]
				println!("PSH {0}", match self.program[self.ip] {
					InstructionSet::VAL{value} => value,
					_ => 0,});
			},
			POP => {
				let val_popped = match self.stack[self.sp as usize] {
					InstructionSet::VAL{value} => value,
					_ => 0,
				};
				self.sp -= 1;
				println!("POP {0}", val_popped);
			},
			ADD => {
				// first we pop the stack and store it as a
				let a = match self.stack[self.sp as usize] {
					InstructionSet::VAL{value} => value,
					_ => 0,
				};
				self.sp -= 1;

				// we then pop the top of the stack and store it as b
				let b = match self.stack[self.sp as usize] {
					InstructionSet::VAL{value} => value,
					_ => 0,
				};
				self.sp -= 1;

				//[DEBUG]
				println!("ADD {0}; {1}", a, b);

				// we then add the result and push it to the stack
				let result = b + a;
				self.sp += 1; // increment stack pointer **before**
				self.stack[self.sp as usize] = InstructionSet::VAL{value: result}; // set the value to the top of the stack
			},
			_ => println!("[ERROR]: Unkown instruction"),
		}
	}
}

fn main() {
	let program = vec![
		InstructionSet::PSH, InstructionSet::VAL{value: 5},
		InstructionSet::PSH, InstructionSet::VAL{value: 6},
		InstructionSet::ADD,
		InstructionSet::POP,
		InstructionSet::HLT,
	];

	let mut vm = VM::new(program);
	loop {
		let instr = vm.fetch();
		vm.eval(instr);
		vm.ip += 1;
		if !vm.running {
			break;
		}
	}
}
<commit_msg>More cleaning<commit_after>
#[derive(Copy, Clone)]
enum InstructionSet {
	PSH,
	ADD,
	POP,
	SET,
	HLT,
	VAL{value: i32},
}

struct VM {
	ip: usize,
	program: Vec<InstructionSet>,
	running: bool,
	sp: isize,
	stack: [InstructionSet; 256],
}

impl VM {
	pub fn new(prog: Vec<InstructionSet>) -> VM {
		VM {
			ip : 0,
			program: prog,
			running: true,
			sp: -1,
			stack:[InstructionSet::VAL{value:0}; 256],
		}
	}

	fn fetch(&mut self) -> InstructionSet {
		let instr = self.program[self.ip];
		instr
	}

	fn eval(&mut self, instr: InstructionSet) {
		use InstructionSet::*;
		match instr {
			HLT => self.running = false,
			PSH => {
				self.sp += 1;
				self.ip += 1; // move to operand location
				self.stack[self.sp as usize] = self.program[self.ip];
				//[DEBUG]
				println!("PSH {0}", match self.program[self.ip] {
					VAL{value} => value,
					_ => 0,});
			},
			POP => {
				let val_popped = match self.stack[self.sp as usize] {
					VAL{value} => value,
					_ => 0,
				};
				self.sp -= 1;
				println!("POP {0}", val_popped);
			},
			ADD => {
				// first we pop the stack and store it as a
				let a = match self.stack[self.sp as usize] {
					VAL{value} => value,
					_ => 0,
				};
				self.sp -= 1;

				// we then pop the top of the stack and store it as b
				let b = match self.stack[self.sp as usize] {
					VAL{value} => value,
					_ => 0,
				};
				self.sp -= 1;

				//[DEBUG]
				println!("ADD {0}; {1}", a, b);

				// we then add the result and push it to the stack
				let result = b + a;
				self.sp += 1; // increment stack pointer **before**
				self.stack[self.sp as usize] = VAL{value: result}; // set the value to the top of the stack
			},
			_ => println!("[ERROR]: Unkown instruction"),
		}
	}
}

fn main() {
	use InstructionSet::*;
	let program = vec![
		PSH, VAL{value: 5},
		PSH, VAL{value: 6},
		ADD,
		POP,
		HLT,
	];

	let mut vm = VM::new(program);
	loop {
		let instr = vm.fetch();
		vm.eval(instr);
		vm.ip += 1;
		if !vm.running {
			break;
		}
	}
}
<|endoftext|>


<commit_before>#![recursion_limit = "1024"] // error_chain

extern crate websocket;
extern crate env_logger;
#[macro_use]
extern crate log;
#[macro_use]
extern crate error_chain;
extern crate url;
extern crate clap;

const BUFSIZ : usize = 8192;

use std::thread;
use std::io::{stdin,stdout};

use websocket::{Message, Sender, Receiver, DataFrame, Server as WsServer};
use websocket::message::Type;
use websocket::client::request::Url;
use websocket::Client;

use std::borrow::Borrow;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Write, Read};

error_chain! {
    foreign_links {
        ::std::io::Error, Io;
        log::SetLoggerError, Log;
        ::url::ParseError, Url;
        ::websocket::result::WebSocketError, Ws;
        ::std::env::VarError, Ev;
    }
    errors {
        InvalidSpecifier(t : String) {
            description("invalid specifier")
            display("Invalid client or server specifier `{}`", t)
        }
    }
}

// Initialize logger with default "info" log level:
fn init_logger() -> Result<()> {
    let mut builder = env_logger::LogBuilder::new();
    builder.filter(None, log::LogLevelFilter::Info);
    if ::std::env::var("RUST_LOG").is_ok() {
       builder.parse(&::std::env::var("RUST_LOG")?);
    }
    builder.init()?;
    Ok(())
}

struct SenderWrapper<T: Sender> (T);

impl<T: Sender> ::std::io::Write for SenderWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
        let message = Message::binary(buf);
        let ret;
        let len = buf.len();
        if len > 0 {
            debug!("Sending message of {} bytes", len);
            ret = self.0.send_message(&message);
        } else {
            // Interpret zero length buffer is request
            // to close communication

            debug!("Sending the closing message");
            ret = self.0.send_message(&Message::close());
        }
        ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;
        Ok(len)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        Ok(())
    }
}

struct ReceiverWrapper<T: Receiver<DataFrame>> (T);

impl<T:Receiver<DataFrame>> ::std::io::Read for ReceiverWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> ::std::io::Result<usize> {
        let ret = self.0.recv_message();
        let msg : Message = ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;

        match msg.opcode {
            Type::Close => {
                Ok(0)
            }
            Type::Ping => {
                // Sender used to be in a separate thread with a channel
                // now there's no channel, so trickier to combine ping replies
                // and usual data exchange
                unimplemented!();
            }
            _ => {
                let msgpayload : &[u8] = msg.payload.borrow();
                let len = msgpayload.len();
                debug!("Received message of {} bytes", len);

                assert!(buf.len() >= len);

                buf[0..len].copy_from_slice(msgpayload);

                Ok(len)
            }
        }
    }
}

struct Peer<R, W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    reader: R,
    writer: W,
}

type IPeer = Peer<Box<Read+Send>, Box<Write+Send>>;

struct DataExchangeSession<R1, R2, W1, W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    peer1: Peer<R1, W1>,
    peer2: Peer<R2, W2>,
}

// Derived from https://doc.rust-lang.org/src/std/up/src/libstd/io/util.rs.html#46-61
pub fn copy_with_flushes<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> ::std::io::Result<u64>
    where R: Read, W: Write
{
    let mut buf = [0; BUFSIZ];
    let mut written = 0;
    loop {
        let len = match reader.read(&mut buf) {
            Ok(0) => return Ok(written),
            Ok(len) => len,
            Err(ref e) if e.kind() == IoErrorKind::Interrupted => continue,
            Err(ref e) if e.kind() == IoErrorKind::WouldBlock => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(&buf[..len])?;
        writer.flush()?;
        written += len as u64;
    }
}

impl<R1,R2,W1,W2> DataExchangeSession<R1,R2,W1,W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    fn data_exchange(self) -> Result<()> {

        let mut reader1 = self.peer1.reader;
        let mut writer1 = self.peer1.writer;
        let mut reader2 = self.peer2.reader;
        let mut writer2 = self.peer2.writer;

        let receive_loop = thread::Builder::new().spawn(move || -> Result<()> {
            // Actual data transfer happens here
            copy_with_flushes(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        })?;

        // Actual data transfer happens here
        copy_with_flushes(&mut reader2, &mut writer1)?;
        writer1.write(b"")?; // Signal close

        debug!("Waiting for receiver side to exit");

        receive_loop.join().map_err(|x|format!("{:?}",x))?
    }
}

fn get_websocket_peer(urlstr: &str) -> Result<
        Peer<
            ReceiverWrapper<websocket::client::Receiver<websocket::WebSocketStream>>,
            SenderWrapper<websocket::client::Sender<websocket::WebSocketStream>>>
        > {
    let url = Url::parse(urlstr)?;

    info!("Connecting to {}", url);

    let request = Client::connect(url)?;

    let response = request.send()?; // Send the request and retrieve a response

    info!("Validating response...");

    response.validate()?; // Validate the response

    info!("Successfully connected");

    let (sender, receiver) = response.begin().split();

    let peer = Peer {
        reader : ReceiverWrapper(receiver),
        writer : SenderWrapper(sender),
    };
    Ok(peer)
}

fn get_tcp_peer(addr: &str) -> Result<
        Peer<
            ::std::net::TcpStream,
            ::std::net::TcpStream,
        >> {
    let sock = ::std::net::TcpStream::connect(addr)?;

    let peer = Peer {
        reader : sock.try_clone()?,
        writer : sock.try_clone()?,
    };
    info!("Connected to TCP {}", addr);
    Ok(peer)
}

fn get_stdio_peer() -> Result<Peer<std::io::Stdin, std::io::Stdout>> {
    Ok(
        Peer {
            reader : stdin(),
            writer : stdout(),
        }
    )
}

struct TcpServer(::std::net::TcpListener);

impl TcpServer {
    fn new(addr: &str) -> Result<Self> {
        Ok(TcpServer(::std::net::TcpListener::bind(addr)?))
    }
}

impl Server for TcpServer {
    fn accept_client(&mut self) -> Result<IPeer> {
        let (sock, addr) = self.0.accept()?;
        info!("TCP client connection from {}", addr);
        let peer = Peer {
            reader : sock.try_clone()?,
            writer : sock.try_clone()?,
        };
        Ok(peer.upcast())
    }
}

struct WebsockServer<'a>(WsServer<'a>);

impl<'a> WebsockServer<'a> {
    fn new(addr: &str) -> Result<Self> {
        Ok(WebsockServer(WsServer::bind(addr)?))
    }
}

impl<'a> Server for WebsockServer<'a> {
    fn accept_client(&mut self) -> Result<IPeer> {
        let connection = self.0.accept()?;
        info!("WebSocket client connection ...");
        let request = connection.read_request()?;
        request.validate()?;
        let response = request.accept(); // Form a response
        let mut client = response.send()?; // Send the response

        let ip = client.get_mut_sender()
            .get_mut()
            .peer_addr()
            .unwrap();

        info!("... from IP {}", ip);

        let (sender, receiver) = client.split();

        let peer = Peer {
            reader : ReceiverWrapper(receiver),
            writer : SenderWrapper(sender),
        };
        Ok(peer.upcast())
    }
}

impl<R,W> Peer<R,W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    fn upcast(self) -> IPeer  {
        Peer {
            reader: Box::new(self.reader) as Box<Read +Send>,
            writer: Box::new(self.writer) as Box<Write+Send>,
        }
    }
}

trait Server
{
    fn accept_client(&mut self) -> Result<IPeer>;

    fn start_serving(&mut self, spec2: &str, once: bool) -> Result<()> {
        let spec2s = spec2.to_string();
        let closure = move |peer, spec2 : String|{
            let spec2_ = get_peer_by_spec(spec2.as_str())?;
            let peer2 = match spec2_ {
                Spec::Server(mut x) => {
                    x.accept_client()?
                }
                Spec::Client(p1) => {
                    p1
                }
            };
            let des = DataExchangeSession {
                peer1 : peer,
                peer2 : peer2,
            };

            des.data_exchange()
        };
        if once {
            let peer = self.accept_client()?;
            closure(peer, spec2s)
        } else {
            let cl2 = ::std::sync::Arc::new(closure);
            loop {
                let ret = self.accept_client();
                let peer = match ret {
                    Ok(x) => x,
                    Err(er) => {
                        warn!("Can't accept client: {}", er);
                        continue;
                    }
                };
                let cl3 = cl2.clone();
                let spec2s2 = spec2s.clone();
                if let Err(x) = thread::Builder::new().spawn(move|| {
                    if let Err(x) = cl3(peer, spec2s2) {
                        warn!("Error while serving: {}", x);
                    }
                }) {
                    warn!("Error creating thread: {}", x);
                    thread::sleep(::std::time::Duration::from_millis(200));
                }
            }
        }
    }

    fn upcast(self) -> Box<Server+Send>
        where Self : Sized + Send + 'static
        { Box::new(self) as Box<Server+Send> }
}

fn main2(spec1: &str, spec2: &str, once: bool) -> Result<()> {
    let spec1_ = get_peer_by_spec(spec1)?;

    match spec1_ {
        Spec::Server(mut x) => {
            x.start_serving(spec2, once)
        }
        Spec::Client(p1) => {
            let spec2_ = get_peer_by_spec(spec2)?;

            let otherpeer = match spec2_ {
                Spec::Server(mut x) => {
                    let t = x.accept_client()?;
                    t
                }
                Spec::Client(p2) => {
                    p2
                }
            };

            let des = DataExchangeSession {
                peer1 : p1,
                peer2 : otherpeer,
            };

            des.data_exchange()
        }
    }
}

enum Spec {
    Server(Box<Server + Send>),
    Client(IPeer)
}

fn get_peer_by_spec(specifier: &str) -> Result<Spec> {
    use Spec::{Server,Client};
    match specifier {
        x if x == "-"               => Ok(Client(get_stdio_peer()?.upcast())),
        x if x.starts_with("ws:")   => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("wss:")  => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("tcp:")  => Ok(Client(get_tcp_peer(&x[4..])?.upcast())),
        x if x.starts_with("l-tcp:")  => Ok(Server(TcpServer::new(&x[6..])?.upcast())),
        x if x.starts_with("l-ws:")  => Ok(Server(WebsockServer::new(&x[5..])?.upcast())),
        x => Err(ErrorKind::InvalidSpecifier(x.to_string()).into()),
    }
}

fn try_main() -> Result<()> {
    //env_logger::init()?;
    init_logger()?;

    // setup command line arguments
    let matches = ::clap::App::new("websocat")
        .version("0.2")
        .author("Vitaly \"_Vi\" Shukela <vi0oss@gmail.com>")
        .about("Exchange binary data between binary websocket and something.\nSocat analogue with websockets.")
        .arg(::clap::Arg::with_name("spec1")
             .help("First specifier.")
             .required(true)
             .index(1))
        .arg(::clap::Arg::with_name("spec2")
             .help("Second specifier.")
             .required(true)
             .index(2))
        .after_help(r#"
Specifiers can be:
  ws[s]://<rest of websocket URL>   Connect to websocket
  l-ws:host:port                    Listen unencrypted websocket
  tcp:host:port                     Connect to TCP
  l-tcp:host:port                   Listen TCP connections
  -                                 stdin/stdout
  (more to be implemented)

Examples:
  websocat l-tcp:0.0.0.0:9559 ws://echo.websocket.org/
    Listen port 9959 on address :: and forward
    all connections to a public loopback websocket
  websocat l-ws:127.0.0.1:7878 tcp:127.0.0.1:1194
    Listen websocket and forward connections to local tcp
    Use nginx proxy for SSL if you want
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to a secure web socket.
    Like netcat, but for websocket.
    `ssh user@host -o ProxyHommand "websocat - ws://..."`
  websocat ws://localhost:1234/ tcp:localhost:1235
    Connect both to websocket and to TCP and exchange data.

Specify listening part first, unless you want websocat to serve once.

IPv6 supported, just use specs like `l-ws:::1:4567`

Web socket usage is not obligatory, you can use any specs on both sides.
"#)
        .get_matches();

    let spec1  = matches.value_of("spec1") .ok_or("no listener_spec" )?;
    let spec2 = matches.value_of("spec2").ok_or("no connector_spec")?;

    main2(spec1, spec2, false)?;

    debug!("Exited");
    Ok(())
}

fn main() {
    if let Err(x) = try_main() {
        let _ = writeln!(::std::io::stderr(), "{:?}", x);
    }
}

<commit_msg>Sec-Websocket-Protocol: binary<commit_after>#![recursion_limit = "1024"] // error_chain

extern crate websocket;
extern crate env_logger;
#[macro_use]
extern crate log;
#[macro_use]
extern crate error_chain;
extern crate url;
extern crate clap;

const BUFSIZ : usize = 8192;

use std::thread;
use std::io::{stdin,stdout};

use websocket::{Message, Sender, Receiver, DataFrame, Server as WsServer};
use websocket::message::Type;
use websocket::client::request::Url;
use websocket::Client;

use std::borrow::Borrow;
use std::io::{Error as IoError, ErrorKind as IoErrorKind, Write, Read};

error_chain! {
    foreign_links {
        ::std::io::Error, Io;
        log::SetLoggerError, Log;
        ::url::ParseError, Url;
        ::websocket::result::WebSocketError, Ws;
        ::std::env::VarError, Ev;
    }
    errors {
        InvalidSpecifier(t : String) {
            description("invalid specifier")
            display("Invalid client or server specifier `{}`", t)
        }
    }
}

// Initialize logger with default "info" log level:
fn init_logger() -> Result<()> {
    let mut builder = env_logger::LogBuilder::new();
    builder.filter(None, log::LogLevelFilter::Info);
    if ::std::env::var("RUST_LOG").is_ok() {
       builder.parse(&::std::env::var("RUST_LOG")?);
    }
    builder.init()?;
    Ok(())
}

struct SenderWrapper<T: Sender> (T);

impl<T: Sender> ::std::io::Write for SenderWrapper<T> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<usize> {
        let message = Message::binary(buf);
        let ret;
        let len = buf.len();
        if len > 0 {
            debug!("Sending message of {} bytes", len);
            ret = self.0.send_message(&message);
        } else {
            // Interpret zero length buffer is request
            // to close communication

            debug!("Sending the closing message");
            ret = self.0.send_message(&Message::close());
        }
        ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;
        Ok(len)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        Ok(())
    }
}

struct ReceiverWrapper<T: Receiver<DataFrame>> (T);

impl<T:Receiver<DataFrame>> ::std::io::Read for ReceiverWrapper<T> {
    fn read(&mut self, buf: &mut [u8]) -> ::std::io::Result<usize> {
        let ret = self.0.recv_message();
        let msg : Message = ret.map_err(|e|IoError::new(IoErrorKind::BrokenPipe, e))?;

        match msg.opcode {
            Type::Close => {
                Ok(0)
            }
            Type::Ping => {
                // Sender used to be in a separate thread with a channel
                // now there's no channel, so trickier to combine ping replies
                // and usual data exchange
                unimplemented!();
            }
            _ => {
                let msgpayload : &[u8] = msg.payload.borrow();
                let len = msgpayload.len();
                debug!("Received message of {} bytes", len);

                assert!(buf.len() >= len);

                buf[0..len].copy_from_slice(msgpayload);

                Ok(len)
            }
        }
    }
}

struct Peer<R, W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    reader: R,
    writer: W,
}

type IPeer = Peer<Box<Read+Send>, Box<Write+Send>>;

struct DataExchangeSession<R1, R2, W1, W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    peer1: Peer<R1, W1>,
    peer2: Peer<R2, W2>,
}

// Derived from https://doc.rust-lang.org/src/std/up/src/libstd/io/util.rs.html#46-61
pub fn copy_with_flushes<R: ?Sized, W: ?Sized>(reader: &mut R, writer: &mut W) -> ::std::io::Result<u64>
    where R: Read, W: Write
{
    let mut buf = [0; BUFSIZ];
    let mut written = 0;
    loop {
        let len = match reader.read(&mut buf) {
            Ok(0) => return Ok(written),
            Ok(len) => len,
            Err(ref e) if e.kind() == IoErrorKind::Interrupted => continue,
            Err(ref e) if e.kind() == IoErrorKind::WouldBlock => continue,
            Err(e) => return Err(e),
        };
        writer.write_all(&buf[..len])?;
        writer.flush()?;
        written += len as u64;
    }
}

impl<R1,R2,W1,W2> DataExchangeSession<R1,R2,W1,W2>
    where R1 : Read  + Send + 'static,
          R2 : Read  + Send + 'static,
          W1 : Write + Send + 'static,
          W2 : Write + Send + 'static,
{
    fn data_exchange(self) -> Result<()> {

        let mut reader1 = self.peer1.reader;
        let mut writer1 = self.peer1.writer;
        let mut reader2 = self.peer2.reader;
        let mut writer2 = self.peer2.writer;

        let receive_loop = thread::Builder::new().spawn(move || -> Result<()> {
            // Actual data transfer happens here
            copy_with_flushes(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        })?;

        // Actual data transfer happens here
        copy_with_flushes(&mut reader2, &mut writer1)?;
        writer1.write(b"")?; // Signal close

        debug!("Waiting for receiver side to exit");

        receive_loop.join().map_err(|x|format!("{:?}",x))?
    }
}

fn get_websocket_peer(urlstr: &str) -> Result<
        Peer<
            ReceiverWrapper<websocket::client::Receiver<websocket::WebSocketStream>>,
            SenderWrapper<websocket::client::Sender<websocket::WebSocketStream>>>
        > {
    let url = Url::parse(urlstr)?;

    info!("Connecting to {}", url);

    let mut request = Client::connect(url)?;

    request.headers.set(
        ::websocket::header::WebSocketProtocol(
            vec!["binary".to_string()]
        )
    );

    let response = request.send()?; // Send the request and retrieve a response

    info!("Validating response...");

    response.validate()?; // Validate the response

    info!("Successfully connected");

    let (sender, receiver) = response.begin().split();

    let peer = Peer {
        reader : ReceiverWrapper(receiver),
        writer : SenderWrapper(sender),
    };
    Ok(peer)
}

fn get_tcp_peer(addr: &str) -> Result<
        Peer<
            ::std::net::TcpStream,
            ::std::net::TcpStream,
        >> {
    let sock = ::std::net::TcpStream::connect(addr)?;

    let peer = Peer {
        reader : sock.try_clone()?,
        writer : sock.try_clone()?,
    };
    info!("Connected to TCP {}", addr);
    Ok(peer)
}

fn get_stdio_peer() -> Result<Peer<std::io::Stdin, std::io::Stdout>> {
    Ok(
        Peer {
            reader : stdin(),
            writer : stdout(),
        }
    )
}

struct TcpServer(::std::net::TcpListener);

impl TcpServer {
    fn new(addr: &str) -> Result<Self> {
        Ok(TcpServer(::std::net::TcpListener::bind(addr)?))
    }
}

impl Server for TcpServer {
    fn accept_client(&mut self) -> Result<IPeer> {
        let (sock, addr) = self.0.accept()?;
        info!("TCP client connection from {}", addr);
        let peer = Peer {
            reader : sock.try_clone()?,
            writer : sock.try_clone()?,
        };
        Ok(peer.upcast())
    }
}

struct WebsockServer<'a>(WsServer<'a>);

impl<'a> WebsockServer<'a> {
    fn new(addr: &str) -> Result<Self> {
        Ok(WebsockServer(WsServer::bind(addr)?))
    }
}

impl<'a> Server for WebsockServer<'a> {
    fn accept_client(&mut self) -> Result<IPeer> {
        let connection = self.0.accept()?;
        info!("WebSocket client connection ...");
        let request = connection.read_request()?;
        request.validate()?;
        let response = request.accept(); // Form a response
        let mut client = response.send()?; // Send the response

        let ip = client.get_mut_sender()
            .get_mut()
            .peer_addr()
            .unwrap();

        info!("... from IP {}", ip);

        let (sender, receiver) = client.split();

        let peer = Peer {
            reader : ReceiverWrapper(receiver),
            writer : SenderWrapper(sender),
        };
        Ok(peer.upcast())
    }
}

impl<R,W> Peer<R,W>
    where R : Read + Send + 'static, W: Write + Send + 'static
{
    fn upcast(self) -> IPeer  {
        Peer {
            reader: Box::new(self.reader) as Box<Read +Send>,
            writer: Box::new(self.writer) as Box<Write+Send>,
        }
    }
}

trait Server
{
    fn accept_client(&mut self) -> Result<IPeer>;

    fn start_serving(&mut self, spec2: &str, once: bool) -> Result<()> {
        let spec2s = spec2.to_string();
        let closure = move |peer, spec2 : String|{
            let spec2_ = get_peer_by_spec(spec2.as_str())?;
            let peer2 = match spec2_ {
                Spec::Server(mut x) => {
                    x.accept_client()?
                }
                Spec::Client(p1) => {
                    p1
                }
            };
            let des = DataExchangeSession {
                peer1 : peer,
                peer2 : peer2,
            };

            des.data_exchange()
        };
        if once {
            let peer = self.accept_client()?;
            closure(peer, spec2s)
        } else {
            let cl2 = ::std::sync::Arc::new(closure);
            loop {
                let ret = self.accept_client();
                let peer = match ret {
                    Ok(x) => x,
                    Err(er) => {
                        warn!("Can't accept client: {}", er);
                        continue;
                    }
                };
                let cl3 = cl2.clone();
                let spec2s2 = spec2s.clone();
                if let Err(x) = thread::Builder::new().spawn(move|| {
                    if let Err(x) = cl3(peer, spec2s2) {
                        warn!("Error while serving: {}", x);
                    }
                }) {
                    warn!("Error creating thread: {}", x);
                    thread::sleep(::std::time::Duration::from_millis(200));
                }
            }
        }
    }

    fn upcast(self) -> Box<Server+Send>
        where Self : Sized + Send + 'static
        { Box::new(self) as Box<Server+Send> }
}

fn main2(spec1: &str, spec2: &str, once: bool) -> Result<()> {
    let spec1_ = get_peer_by_spec(spec1)?;

    match spec1_ {
        Spec::Server(mut x) => {
            x.start_serving(spec2, once)
        }
        Spec::Client(p1) => {
            let spec2_ = get_peer_by_spec(spec2)?;

            let otherpeer = match spec2_ {
                Spec::Server(mut x) => {
                    let t = x.accept_client()?;
                    t
                }
                Spec::Client(p2) => {
                    p2
                }
            };

            let des = DataExchangeSession {
                peer1 : p1,
                peer2 : otherpeer,
            };

            des.data_exchange()
        }
    }
}

enum Spec {
    Server(Box<Server + Send>),
    Client(IPeer)
}

fn get_peer_by_spec(specifier: &str) -> Result<Spec> {
    use Spec::{Server,Client};
    match specifier {
        x if x == "-"               => Ok(Client(get_stdio_peer()?.upcast())),
        x if x.starts_with("ws:")   => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("wss:")  => Ok(Client(get_websocket_peer(x)?.upcast())),
        x if x.starts_with("tcp:")  => Ok(Client(get_tcp_peer(&x[4..])?.upcast())),
        x if x.starts_with("l-tcp:")  => Ok(Server(TcpServer::new(&x[6..])?.upcast())),
        x if x.starts_with("l-ws:")  => Ok(Server(WebsockServer::new(&x[5..])?.upcast())),
        x => Err(ErrorKind::InvalidSpecifier(x.to_string()).into()),
    }
}

fn try_main() -> Result<()> {
    //env_logger::init()?;
    init_logger()?;

    // setup command line arguments
    let matches = ::clap::App::new("websocat")
        .version("0.2")
        .author("Vitaly \"_Vi\" Shukela <vi0oss@gmail.com>")
        .about("Exchange binary data between binary websocket and something.\nSocat analogue with websockets.")
        .arg(::clap::Arg::with_name("spec1")
             .help("First specifier.")
             .required(true)
             .index(1))
        .arg(::clap::Arg::with_name("spec2")
             .help("Second specifier.")
             .required(true)
             .index(2))
        .after_help(r#"
Specifiers can be:
  ws[s]://<rest of websocket URL>   Connect to websocket
  l-ws:host:port                    Listen unencrypted websocket
  tcp:host:port                     Connect to TCP
  l-tcp:host:port                   Listen TCP connections
  -                                 stdin/stdout
  (more to be implemented)

Examples:
  websocat l-tcp:0.0.0.0:9559 ws://echo.websocket.org/
    Listen port 9959 on address :: and forward
    all connections to a public loopback websocket
  websocat l-ws:127.0.0.1:7878 tcp:127.0.0.1:1194
    Listen websocket and forward connections to local tcp
    Use nginx proxy for SSL if you want
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to a secure web socket.
    Like netcat, but for websocket.
    `ssh user@host -o ProxyHommand "websocat - ws://..."`
  websocat ws://localhost:1234/ tcp:localhost:1235
    Connect both to websocket and to TCP and exchange data.

Specify listening part first, unless you want websocat to serve once.

IPv6 supported, just use specs like `l-ws:::1:4567`

Web socket usage is not obligatory, you can use any specs on both sides.
"#)
        .get_matches();

    let spec1  = matches.value_of("spec1") .ok_or("no listener_spec" )?;
    let spec2 = matches.value_of("spec2").ok_or("no connector_spec")?;

    main2(spec1, spec2, false)?;

    debug!("Exited");
    Ok(())
}

fn main() {
    if let Err(x) = try_main() {
        let _ = writeln!(::std::io::stderr(), "{:?}", x);
    }
}

<|endoftext|>