<commit_before>#[macro_use]
extern crate log;
#[macro_use]
extern crate nom;

extern crate bip_bencode;
extern crate bip_metainfo;
extern crate bip_utracker;
extern crate bip_util;
extern crate chrono;
extern crate hyper;
extern crate pnet;
extern crate pnet_macros_support;
extern crate pretty_env_logger;
extern crate rustyline;
extern crate url;

mod btclient;
mod packet;

use bip_bencode::Bencode;
use bip_metainfo::MetainfoFile;
use bip_utracker::contact::CompactPeersV4;
use chrono::{TimeZone, UTC};
use hyper::Client;
use hyper::header::Connection;
use packet::peer_pkt::{MutablePeerHandshakePacket, MutablePeerMessagePacket};
use pnet::packet::Packet;
use rustyline::Editor;
use rustyline::error::ReadlineError;
use url::Url;

use std::fs::File;
use std::io::prelude::*;
use std::net::TcpStream;
use std::str;
use std::string::String;
use std::time::{SystemTime, UNIX_EPOCH};
use std::thread;

const HISTORY_FILE: &'static str = ".rustyline.history";
const PEER_HANDSHAKE_STRUCT_SZ: usize = 68;
const PEER_REQ_PKT_SZ: usize = 14;

/// Print general information about the torrent.
fn print_metainfo_overview(bytes: &[u8]) {
    let metainfo = MetainfoFile::from_bytes(bytes).unwrap();
    let info = metainfo.info();
    let info_hash_hex = metainfo.info_hash()
        .as_ref()
        .iter()
        .map(|b| format!("{:02x}", b))
        .fold(String::new(), |mut acc, nex| {
            acc.push_str(&nex);
            acc
        });
    let utc_creation_date = metainfo.creation_date().map(|c| UTC.timestamp(c, 0));

    println!("------Metainfo File Overview-----");

    println!("InfoHash: {}", info_hash_hex);
    println!("Main Tracker: {:?}", metainfo.main_tracker());
    println!("Comment: {:?}", metainfo.comment());
    println!("Creator: {:?}", metainfo.created_by());
    println!("Creation Date: {:?}", utc_creation_date);

    println!("Directory: {:?}", info.directory());
    println!("Piece Length: {:?}", info.piece_length());
    println!("Number Of Pieces: {}", info.pieces().count());
    println!("Number Of Files: {}", info.files().count());
    println!("Total File Size: {}",
             info.files().fold(0, |acc, nex| acc + nex.length()));
}

fn connect_to_tracker(metainfo: MetainfoFile,
                      peer_id: &str,
                      port: u16)
                      -> Result<(Vec<String>, String), String> {
    debug!("connecting to tracker: {:?}", metainfo.main_tracker());

    let info_hash = metainfo.info_hash();
    // TODO figure can this conversion to url-encoded form be done safely?
    let info_hash_str = unsafe { str::from_utf8_unchecked(info_hash.as_ref()) };

    // TODO this needs to be calculated based on what we have
    let total_len = metainfo.info().files().fold(0, |acc, nex| acc + nex.length());

    let mut url = Url::parse(metainfo.main_tracker().unwrap()).unwrap();
    url.query_pairs_mut()
        .append_pair("info_hash", info_hash_str)
        .append_pair("peer_id", peer_id)
        .append_pair("port", &(port.to_string()))
        // TODO parametrize this
        .append_pair("uploaded", "0")
        // TODO parametrize this
        .append_pair("downloaded", "0")
        // TODO see note on total_len above
        .append_pair("left", &(total_len.to_string()))
        .append_pair("compact", "1")
        .append_pair("event", "started")
        .append_pair("supportcrypto", "0");
    trace!("URL {:?}", url);

    let client = Client::new();
    let mut res = client.get(url).header(Connection::close()).send().unwrap();
    let mut buffer = Vec::new();
    res.read_to_end(&mut buffer).unwrap();
    debug!("{:?}", res);
    let bencode = Bencode::decode(&buffer).unwrap();
    trace!("{:?}", bencode);
    let (_, peers) = CompactPeersV4::from_bytes(bencode.dict()
            .unwrap()
            .lookup("peers")
            .unwrap()
            .bytes()
            .unwrap())
        .unwrap();
    trace!("{:?}", peers);
    let mut ip_ports: Vec<String> = Vec::new();
    debug!("Peer list received:");
    for peer in peers.iter() {
        debug!("{:?}", peer);
        ip_ports.push(peer.to_string());
    }
    Ok((ip_ports, info_hash_str.to_string()))
}

fn peer_connections(peer_ip_ports: Vec<String>,
                    info_hash: &str,
                    peer_id: &str)
                    -> Result<(), String> {
    for peer_ip_port in peer_ip_ports {
        let peer_ip_port_cl = peer_ip_port.clone();
        let info_hash_cl = info_hash.to_string().clone();
        let peer_id_cl = peer_id.to_string().clone();
        thread::spawn(move || {
            match handshake_peer(&peer_ip_port_cl, &info_hash_cl, &peer_id_cl) {
                Ok(mut stream) => {
                    debug!("Communicating with peer {:?}", peer_ip_port_cl);
                    let mut buff = [0; PEER_HANDSHAKE_STRUCT_SZ];
                    match stream.read(&mut buff) {
			Ok(_) => {
				let mut buf = vec![0u8; PEER_REQ_PKT_SZ];
				let mut peer_msg_pkt  = MutablePeerMessagePacket::new(&mut buf).unwrap();
				peer_msg_pkt.set_len(13);
				peer_msg_pkt.set_id(6);
				//TODO fill length properly
				let index_begin_length = vec![0, 0, 10];
				peer_msg_pkt.set_payload(&index_begin_length);
				match stream.write(peer_msg_pkt.packet()) {
					Ok(_) => {
						let mut buf_read = vec![0; 2048];
						match stream.read(&mut buf_read) {
							Ok(bytes_read) => {debug!("Bytes read: {:?}", bytes_read);},
							Err(e) => error!("Read failed! {:?}", e),
						}
					},
					Err(e) => error!("Write to stream failed! {:?}", e),
				}
			},
			Err(e) => error!("Reading from the stream failed! {:?}", e),
			}
                }
                Err(_) => {
                    error!("Closing thread with peer {:?}", peer_ip_port_cl);
                }
            }
        });
    }
    Ok(())
}

fn handshake_peer(peer_ip_port: &str, info_hash: &str, peer_id: &str) -> Result<TcpStream, String> {
    let mut buf = vec![0u8; PEER_HANDSHAKE_STRUCT_SZ];
    let mut ph = MutablePeerHandshakePacket::new(&mut buf).unwrap();
    ph.set_pstrlen("BitTorrent protocol".len() as u8);
    ph.set_pstr(&String::from("BitTorrent protocol").into_bytes());
    ph.set_reserved(&[0; 8]);
    ph.set_info_hash(info_hash.as_bytes());
    ph.set_peer_id(peer_id.as_bytes());
    match TcpStream::connect(peer_ip_port) {
        Ok(mut stream) => {
            debug!("Connection to peer {:?} successful!", peer_ip_port);
            match stream.write(ph.packet()) {
                Ok(_) => {
                    debug!("Sending message successful!");
                    trace!("{:?}", ph.packet());
                    Ok((stream))
                }
                Err(_) => {
                    error!("Sending message failed!");
                    Err("Sending message failed!".to_owned())
                }
            }
        }
        Err(e) => {
            error!("Connection to peer {:?} failed! {:?}", peer_ip_port, e);
            Err("Connection to peer failed!".to_owned())
        }
    }
}

fn main() {
    pretty_env_logger::init();

    let mut rl = Editor::<()>::new();
    if rl.load_history(HISTORY_FILE).is_err() {
        info!("No previous history!");
    }

    let now = SystemTime::now();
    let duration = now.duration_since(UNIX_EPOCH).unwrap();
    let client_id = "-bittorrent-rs-".to_owned() + &format!("{}", duration.as_secs() % 100_000);

    loop {
        let readline = rl.readline("> ");
        match readline {
            Ok(cmd) => {
                rl.add_history_entry(&cmd);

                let cmd = cmd.trim().split(' ').collect::<Vec<&str>>();

                match cmd[0] {
                    "parse" => {
                        if cmd.len() != 2 {
                            error!("usage: parse <torrent file>");
                        } else {
                            let path = cmd[1];
                            debug!("MetainfoFile: {:?}", path);

                            match File::open(path) {
                                Ok(mut f) => {
                                    let mut bytes: Vec<u8> = Vec::new();
                                    f.read_to_end(&mut bytes).unwrap();
                                    print_metainfo_overview(&bytes);
                                }
                                Err(e) => error!("{:?}", e),
                            }
                        }
                    }
                    "connect" => {
                        if cmd.len() != 2 {
                            error!("usage: connect <torrent file>");
                        } else {
                            let path = cmd[1];
                            match File::open(path) {
                                Ok(mut f) => {
                                    let mut bytes: Vec<u8> = Vec::new();
                                    f.read_to_end(&mut bytes).unwrap();

                                    // TODO: generate peer ID
                                    let result =
                                        connect_to_tracker(MetainfoFile::from_bytes(&bytes)
                                                               .unwrap(),
                                                           &client_id,
                                                           6882)
                                            .unwrap();
                                    peer_connections(result.0, &result.1, &client_id).unwrap();
                                }
                                Err(e) => error!("{:?}", e),
                            }
                        }
                    }
                    _ => {}
                }
            }
            Err(ReadlineError::Interrupted) |
            Err(ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                error!("Error: {:?}", err);
                break;
            }
        }
    }
    rl.save_history(HISTORY_FILE).unwrap();
}
<commit_msg>rustfmt<commit_after>#[macro_use]
extern crate log;
#[macro_use]
extern crate nom;

extern crate bip_bencode;
extern crate bip_metainfo;
extern crate bip_utracker;
extern crate bip_util;
extern crate chrono;
extern crate hyper;
extern crate pnet;
extern crate pnet_macros_support;
extern crate pretty_env_logger;
extern crate rustyline;
extern crate url;

mod btclient;
mod packet;

use bip_bencode::Bencode;
use bip_metainfo::MetainfoFile;
use bip_utracker::contact::CompactPeersV4;
use chrono::{TimeZone, UTC};
use hyper::Client;
use hyper::header::Connection;
use packet::peer_pkt::{MutablePeerHandshakePacket, MutablePeerMessagePacket};
use pnet::packet::Packet;
use rustyline::Editor;
use rustyline::error::ReadlineError;
use url::Url;

use std::fs::File;
use std::io::prelude::*;
use std::net::TcpStream;
use std::str;
use std::string::String;
use std::time::{SystemTime, UNIX_EPOCH};
use std::thread;

const HISTORY_FILE: &'static str = ".rustyline.history";
const PEER_HANDSHAKE_STRUCT_SZ: usize = 68;
const PEER_REQ_PKT_SZ: usize = 14;

/// Print general information about the torrent.
fn print_metainfo_overview(bytes: &[u8]) {
    let metainfo = MetainfoFile::from_bytes(bytes).unwrap();
    let info = metainfo.info();
    let info_hash_hex = metainfo.info_hash()
        .as_ref()
        .iter()
        .map(|b| format!("{:02x}", b))
        .fold(String::new(), |mut acc, nex| {
            acc.push_str(&nex);
            acc
        });
    let utc_creation_date = metainfo.creation_date().map(|c| UTC.timestamp(c, 0));

    println!("------Metainfo File Overview-----");

    println!("InfoHash: {}", info_hash_hex);
    println!("Main Tracker: {:?}", metainfo.main_tracker());
    println!("Comment: {:?}", metainfo.comment());
    println!("Creator: {:?}", metainfo.created_by());
    println!("Creation Date: {:?}", utc_creation_date);

    println!("Directory: {:?}", info.directory());
    println!("Piece Length: {:?}", info.piece_length());
    println!("Number Of Pieces: {}", info.pieces().count());
    println!("Number Of Files: {}", info.files().count());
    println!("Total File Size: {}",
             info.files().fold(0, |acc, nex| acc + nex.length()));
}

fn connect_to_tracker(metainfo: MetainfoFile,
                      peer_id: &str,
                      port: u16)
                      -> Result<(Vec<String>, String), String> {
    debug!("connecting to tracker: {:?}", metainfo.main_tracker());

    let info_hash = metainfo.info_hash();
    // TODO figure can this conversion to url-encoded form be done safely?
    let info_hash_str = unsafe { str::from_utf8_unchecked(info_hash.as_ref()) };

    // TODO this needs to be calculated based on what we have
    let total_len = metainfo.info().files().fold(0, |acc, nex| acc + nex.length());

    let mut url = Url::parse(metainfo.main_tracker().unwrap()).unwrap();
    url.query_pairs_mut()
        .append_pair("info_hash", info_hash_str)
        .append_pair("peer_id", peer_id)
        .append_pair("port", &(port.to_string()))
        // TODO parametrize this
        .append_pair("uploaded", "0")
        // TODO parametrize this
        .append_pair("downloaded", "0")
        // TODO see note on total_len above
        .append_pair("left", &(total_len.to_string()))
        .append_pair("compact", "1")
        .append_pair("event", "started")
        .append_pair("supportcrypto", "0");
    trace!("URL {:?}", url);

    let client = Client::new();
    let mut res = client.get(url).header(Connection::close()).send().unwrap();
    let mut buffer = Vec::new();
    res.read_to_end(&mut buffer).unwrap();
    debug!("{:?}", res);
    let bencode = Bencode::decode(&buffer).unwrap();
    trace!("{:?}", bencode);
    let (_, peers) = CompactPeersV4::from_bytes(bencode.dict()
            .unwrap()
            .lookup("peers")
            .unwrap()
            .bytes()
            .unwrap())
        .unwrap();
    trace!("{:?}", peers);
    let mut ip_ports: Vec<String> = Vec::new();
    debug!("Peer list received:");
    for peer in peers.iter() {
        debug!("{:?}", peer);
        ip_ports.push(peer.to_string());
    }
    Ok((ip_ports, info_hash_str.to_string()))
}

fn peer_connections(peer_ip_ports: Vec<String>,
                    info_hash: &str,
                    peer_id: &str)
                    -> Result<(), String> {
    for peer_ip_port in peer_ip_ports {
        let peer_ip_port_cl = peer_ip_port.clone();
        let info_hash_cl = info_hash.to_string().clone();
        let peer_id_cl = peer_id.to_string().clone();
        thread::spawn(move || {
            match handshake_peer(&peer_ip_port_cl, &info_hash_cl, &peer_id_cl) {
                Ok(mut stream) => {
                    debug!("Communicating with peer {:?}", peer_ip_port_cl);
                    let mut buff = [0; PEER_HANDSHAKE_STRUCT_SZ];
                    match stream.read(&mut buff) {
                        Ok(_) => {
                            let mut buf = vec![0u8; PEER_REQ_PKT_SZ];
                            let mut peer_msg_pkt = MutablePeerMessagePacket::new(&mut buf).unwrap();
                            peer_msg_pkt.set_len(13);
                            peer_msg_pkt.set_id(6);
                            // TODO fill length properly
                            let index_begin_length = vec![0, 0, 10];
                            peer_msg_pkt.set_payload(&index_begin_length);
                            match stream.write(peer_msg_pkt.packet()) {
                                Ok(_) => {
                                    let mut buf_read = vec![0; 2048];
                                    match stream.read(&mut buf_read) {
                                        Ok(bytes_read) => {
                                            debug!("Bytes read: {:?}", bytes_read);
                                        }
                                        Err(e) => error!("Read failed! {:?}", e),
                                    }
                                }
                                Err(e) => error!("Write to stream failed! {:?}", e),
                            }
                        }
                        Err(e) => error!("Reading from the stream failed! {:?}", e),
                    }
                }
                Err(_) => {
                    error!("Closing thread with peer {:?}", peer_ip_port_cl);
                }
            }
        });
    }
    Ok(())
}

fn handshake_peer(peer_ip_port: &str, info_hash: &str, peer_id: &str) -> Result<TcpStream, String> {
    let mut buf = vec![0u8; PEER_HANDSHAKE_STRUCT_SZ];
    let mut ph = MutablePeerHandshakePacket::new(&mut buf).unwrap();
    ph.set_pstrlen("BitTorrent protocol".len() as u8);
    ph.set_pstr(&String::from("BitTorrent protocol").into_bytes());
    ph.set_reserved(&[0; 8]);
    ph.set_info_hash(info_hash.as_bytes());
    ph.set_peer_id(peer_id.as_bytes());
    match TcpStream::connect(peer_ip_port) {
        Ok(mut stream) => {
            debug!("Connection to peer {:?} successful!", peer_ip_port);
            match stream.write(ph.packet()) {
                Ok(_) => {
                    debug!("Sending message successful!");
                    trace!("{:?}", ph.packet());
                    Ok((stream))
                }
                Err(_) => {
                    error!("Sending message failed!");
                    Err("Sending message failed!".to_owned())
                }
            }
        }
        Err(e) => {
            error!("Connection to peer {:?} failed! {:?}", peer_ip_port, e);
            Err("Connection to peer failed!".to_owned())
        }
    }
}

fn main() {
    pretty_env_logger::init();

    let mut rl = Editor::<()>::new();
    if rl.load_history(HISTORY_FILE).is_err() {
        info!("No previous history!");
    }

    let now = SystemTime::now();
    let duration = now.duration_since(UNIX_EPOCH).unwrap();
    let client_id = "-bittorrent-rs-".to_owned() + &format!("{}", duration.as_secs() % 100_000);

    loop {
        let readline = rl.readline("> ");
        match readline {
            Ok(cmd) => {
                rl.add_history_entry(&cmd);

                let cmd = cmd.trim().split(' ').collect::<Vec<&str>>();

                match cmd[0] {
                    "parse" => {
                        if cmd.len() != 2 {
                            error!("usage: parse <torrent file>");
                        } else {
                            let path = cmd[1];
                            debug!("MetainfoFile: {:?}", path);

                            match File::open(path) {
                                Ok(mut f) => {
                                    let mut bytes: Vec<u8> = Vec::new();
                                    f.read_to_end(&mut bytes).unwrap();
                                    print_metainfo_overview(&bytes);
                                }
                                Err(e) => error!("{:?}", e),
                            }
                        }
                    }
                    "connect" => {
                        if cmd.len() != 2 {
                            error!("usage: connect <torrent file>");
                        } else {
                            let path = cmd[1];
                            match File::open(path) {
                                Ok(mut f) => {
                                    let mut bytes: Vec<u8> = Vec::new();
                                    f.read_to_end(&mut bytes).unwrap();

                                    // TODO: generate peer ID
                                    let result =
                                        connect_to_tracker(MetainfoFile::from_bytes(&bytes)
                                                               .unwrap(),
                                                           &client_id,
                                                           6882)
                                            .unwrap();
                                    peer_connections(result.0, &result.1, &client_id).unwrap();
                                }
                                Err(e) => error!("{:?}", e),
                            }
                        }
                    }
                    _ => {}
                }
            }
            Err(ReadlineError::Interrupted) |
            Err(ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                error!("Error: {:?}", err);
                break;
            }
        }
    }
    rl.save_history(HISTORY_FILE).unwrap();
}
<|endoftext|>


<commit_before>//  Copyright 2015 Nathanael Merrill
//
//  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.
extern crate rustty;
extern crate gst;
extern crate argparse;
extern crate regex;

mod playlist;
mod default_ui;

use argparse::{ArgumentParser, Store, List};
use default_ui::*;
use regex::Regex;
use gst::ElementT;

fn main() {
    let mut regex = String::new();
    let mut songs : Vec<String> = vec![]; // TODO: add with capacity!
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Rust Fucking Simple Music Player");
        ap.refer(&mut regex)
            .add_option(&["-r", "--regex"], Store, "Use a regual expression");
        ap.refer(&mut songs)
            .add_argument("arguments", List, "Songs to play");
        ap.parse_args_or_exit();
    }

    if regex != "" {
        let re = Regex::new(&regex).expect("regex invalid");
        songs.retain(|i| re.is_match(i));
    }

    let mut playlist = playlist::Playlist::new(songs);

    let mut ui = UI::new(&playlist);

    gst::init();
    let mut playbin = gst::PlayBin::new("audio_player")
        .expect("Couldn't create playlist");
    let mut main_loop = gst::MainLoop::new();

    let mut bus;
    let bus_receiver;

    let song = match playlist.get_next_song() {
        Some(a) => gst::filename_to_uri(a).expect("filename error"),
        None => panic!("can't get song"),
    };
    playbin.set_uri(&song);
    bus = playbin.bus().expect("Couldn't get pipeline bus");
    bus_receiver = bus.receiver();
    main_loop.spawn();
    playbin.play();

    let mut song_buffered = false;

    'outer: loop {
        loop {
            match bus_receiver.try_recv() {
                Ok(message) => {
                    match message.parse(){
                        gst::Message::ErrorParsed{ref error, ..} => {
                            panic!("error msg from element `{}`: {}, quit",
                                     message.src_name(), error.message());
                        }
                        gst::Message::Eos(ref _msg) => {
                            break 'outer;
                        }
                        gst::Message::StreamStart(ref _msg) => {
                            song_buffered = false;
                            playlist.go_to_next();
                        }
                        _ => {
                            //println!("msg of type `{}` from element `{}`",
                            //         message.type_name(), message.src_name());
                        }
                    }
                }
                Err(err) => {
                    match err {
                        std::sync::mpsc::TryRecvError::Empty => break,
                        std::sync::mpsc::TryRecvError::Disconnected => break,
                    }
                }
            }
        }
        let stream_dir = playbin.duration_s();
        let stream_pos = playbin.position_s();

        if stream_dir.is_some() && stream_pos.is_some() && !song_buffered {
            if stream_dir.unwrap() - stream_pos.unwrap() < 3.0 {
                match playlist.get_next_song() {
                    Some(a) => {
                        let song = gst::filename_to_uri(a).expect("URI error");
                        playbin.set_uri(&song);
                        song_buffered = true;
                    }
                    None => {
                        //println!("All songs played");
                    }
                };
            }
        }

        let stream_dir = match stream_dir {
            Some(a) => a as i32,
            None => 0,
        };
        let stream_pos = match stream_pos {
            Some(a) => a as i32,
            None => 0,
        };

        match ui.manage_ui(&playlist, stream_pos, stream_dir) {
            UIResult::PlayPause => {
                println!("play/pause");
            }
            UIResult::Previous => {
                println!("previous");
            }
            UIResult::Next => {
                println!("next");
            }
            UIResult::Exit => {
                println!("exit");
                break;
            }
            UIResult::Error => {
                println!("error");
                break;
            }
            UIResult::NA => {}
        }
    }
    main_loop.quit();
}
<commit_msg>Implement play/pause, next, and previous commands.<commit_after>//  Copyright 2015 Nathanael Merrill
//
//  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.
extern crate rustty;
extern crate gst;
extern crate argparse;
extern crate regex;

mod playlist;
mod default_ui;

use argparse::{ArgumentParser, Store, List};
use default_ui::*;
use regex::Regex;
use gst::ElementT;

fn main() {
    let mut regex = String::new();
    let mut songs : Vec<String> = vec![]; // TODO: add with capacity!
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Rust Fucking Simple Music Player");
        ap.refer(&mut regex)
            .add_option(&["-r", "--regex"], Store, "Use a regual expression");
        ap.refer(&mut songs)
            .add_argument("arguments", List, "Songs to play");
        ap.parse_args_or_exit();
    }

    if regex != "" {
        let re = Regex::new(&regex).expect("regex invalid");
        songs.retain(|i| re.is_match(i));
    }

    let mut playlist = playlist::Playlist::new(songs);

    let mut ui = UI::new(&playlist);

    gst::init();
    let mut playbin = gst::PlayBin::new("audio_player")
        .expect("Couldn't create playlist");
    let mut main_loop = gst::MainLoop::new();

    let mut bus;
    let bus_receiver;

    let song = match playlist.get_next_song() {
        Some(a) => gst::filename_to_uri(a).expect("filename error"),
        None => panic!("can't get song"),
    };
    playbin.set_uri(&song);
    bus = playbin.bus().expect("Couldn't get pipeline bus");
    bus_receiver = bus.receiver();
    main_loop.spawn();
    playbin.play();

    let mut song_buffered = false;

    'outer: loop {
        loop {
            match bus_receiver.try_recv() {
                Ok(message) => {
                    match message.parse(){
                        gst::Message::ErrorParsed{ref error, ..} => {
                            panic!("error msg from element `{}`: {}, quit",
                                     message.src_name(), error.message());
                        }
                        gst::Message::Eos(ref _msg) => {
                            break 'outer;
                        }
                        gst::Message::StreamStart(ref _msg) => {
                            song_buffered = false;
                            playlist.go_to_next();
                        }
                        _ => {
                            //println!("msg of type `{}` from element `{}`",
                            //         message.type_name(), message.src_name());
                        }
                    }
                }
                Err(err) => {
                    match err {
                        std::sync::mpsc::TryRecvError::Empty => break,
                        std::sync::mpsc::TryRecvError::Disconnected => break,
                    }
                }
            }
        }
        let stream_dir = playbin.duration_s();
        let stream_pos = playbin.position_s();

        if stream_dir.is_some() && stream_pos.is_some() && !song_buffered {
            if stream_dir.unwrap() - stream_pos.unwrap() < 3.0 {
                match playlist.get_next_song() {
                    Some(a) => {
                        let song = gst::filename_to_uri(a).expect("URI error");
                        playbin.set_uri(&song);
                        song_buffered = true;
                    }
                    None => {
                        //println!("All songs played");
                    }
                };
            }
        }

        let stream_dir = match stream_dir {
            Some(a) => a as i32,
            None => 0,
        };
        let stream_pos = match stream_pos {
            Some(a) => a as i32,
            None => 0,
        };

        match ui.manage_ui(&playlist, stream_pos, stream_dir) {
            UIResult::PlayPause => {
                if playbin.is_paused() {
                    playbin.play();
                }
                else {
                    playbin.pause();
                }
            }
            UIResult::Previous => {
                playbin.set_state(gst::ffi::GstState::GST_STATE_NULL);
                playlist.go_to_prev();
                playlist.go_to_prev();
                match playlist.get_next_song() {
                    Some(a) => {
                        let song = gst::filename_to_uri(a).expect("URI error");
                        playbin.set_uri(&song);
                        song_buffered = true;
                    }
                    None => {
                        println!("All songs played");
                        break;
                    }
                };
                playbin.play();

            }
            UIResult::Next => {
                playbin.set_state(gst::ffi::GstState::GST_STATE_NULL);
                match playlist.get_next_song() {
                    Some(a) => {
                        let song = gst::filename_to_uri(a).expect("URI error");
                        playbin.set_uri(&song);
                        song_buffered = true;
                    }
                    None => {
                        println!("All songs played");
                        break;
                    }
                };
                playbin.play();
            }
            UIResult::Exit => {
                break;
            }
            UIResult::Error => {
                break;
            }
            UIResult::NA => {}
        }
    }
    main_loop.quit();
}
<|endoftext|>


<commit_before> #[macro_use]
 #[macro_use]
extern crate clap;
extern crate yaml_rust;
extern crate time;

mod models;
mod parser;
mod mayberef;

use clap::{Arg, App};
use std::fs::{File, read_dir};
use std::io::{BufReader, BufRead, Read, Write, stdin};
use std::io;
use std::path::{Path};
use std::collections::HashMap;
use models::*;

fn get_line() -> String {
    let stdin = stdin();
    let mut buffer = String::new();
    stdin.read_line(&mut buffer).unwrap();
    buffer
}

fn find_most_common_words(corpus: &Path) -> Vec<(String, u32)> {
    let words = files(corpus).flat_map(|file| read_words(file));

    let mut word_counts = HashMap::new();
    for word in words {
        *word_counts.entry(word).or_insert(0) += 1;
    }

    let mut counts: Vec<_> = word_counts.into_iter().collect();
    counts.sort_by(|a, b| b.1.cmp(&a.1));

    counts
}

fn save_words(path: &Path, words: &Vec<(String, u32)>) -> io::Result<()> {
    let mut out = try!(File::create(path));
    for &(ref word, count) in words {
        try!(out.write_all(format!("{}, {}\n", &word, count).as_bytes()));
    }
    Ok(())
}

fn load_most_common_words(filename: &str, num: usize) -> Vec<String> {
    let file = File::open(&Path::new(filename)).unwrap();
    let reader = BufReader::new(file);
    reader.lines().filter_map(|line|
        line.ok().and_then(|l| l.split(',').next().map(|s| s.to_string()))
    )
    .take(num)
    .collect()
}

fn sentences<T: Read + 'static>(reader: BufReader<T>) -> Box<Iterator<Item = String>> {
    Box::new(reader.split('.' as u8)
                   .filter_map(|v| {
                        let lowercase = v.unwrap();
                        String::from_utf8(lowercase).ok().map(|s| s.to_lowercase())
                    }))
}

fn read_words<R: BufRead + 'static>(reader: R) -> Box<Iterator<Item = String>> {
    Box::new(reader.lines()
                   .filter_map(|line| line.ok())
                   .flat_map(|line| {
                       line.split(|c| match c {
                                'a'...'z' | 'A'...'Z' => false,
                                _ => true,
                            })
                           .filter(|word| !word.is_empty())
                           .map(|word| word.to_lowercase())
                           .collect::<Vec<_>>()
                           .into_iter()
                   }))
}

fn files(path: &Path) -> Box<Iterator<Item = BufReader<File>>> {
    Box::new(read_dir(path)
            .unwrap()
            .map(|path| path.unwrap().path())
            .filter(|path| path.to_str().unwrap().ends_with(".txt"))
            .map(|path| {
                let file = File::open(&path).unwrap();
                BufReader::new(file)
            }))
}

fn create_model(corpus: &Path, words: Vec<String>) -> LanguageModelBuilder {
    let start_time = time::get_time();
    let mut builder = LanguageModelBuilder::new(10, words);

    let mut num_words = 0;

    for sentence in files(corpus).flat_map(sentences) {
        let mut acc = builder.new_sentence();
        for word in sentence.split(|c| match c {
            'a'...'z' => false,
            _ => true,
        }).filter(|w| !w.is_empty())
        {
            num_words += 1;
            acc.add_word(word);
            if num_words % 10_000_000 == 0 {
                println!("Loaded {} million words in {} seconds",
                         num_words / 1_000_000,
                         time::get_time().sec - start_time.sec);
            }
        }
    }

    let end_time = time::get_time();
    println!("Model built in {}s", end_time.sec - start_time.sec);
    builder
}

fn main() {
    let matches = App::new("gauntlet")
        .version("0.0.1")
        .author("James Moughan <jamougha@gmail.com>")
        .about("Implementation of GloVe algorithm")
        .arg(Arg::with_name("CORPUS")
            .short("c")
            .long("corpus")
            .help("Sets a directory to search for the corpus")
            .takes_value(true))
        .arg(Arg::with_name("LOAD")
            .short("l")
            .long("load")
            .help("Loads a pre-saved language model")
            .takes_value(true))
        .arg(Arg::with_name("SAVE")
            .short("s")
            .long("save")
            .help("Generates a model from the corpus specified and saves it")
            .takes_value(true))
        .arg(Arg::with_name("WORDS")
            .short("w")
            .long("save_words")
            .help("Saves or the vocabulary list to the specified file")
            .takes_value(true))
        .arg(Arg::with_name("NUM_WORDS")
            .short("n")
            .long("num_words")
            .help("The maximum number of words to use in the vocabulary list, defaults to 30000")
            .takes_value(true))
        .get_matches();

    let (load, save, corpus) = (matches.value_of("LOAD"), matches.value_of("SAVE"), matches.value_of("CORPUS"));
    let builder = match (load, save, corpus) {
        (Some(l), None, None) => LanguageModelBuilder::load(Path::new(&l)).expect("Couldn't load model"),
        (Some(_), _, _) => { println!("You must specify either a model to load or a corpus directory location"); return; }
        (_, save, Some(corpus)) => {
            let corpus = Path::new(corpus);
            let words = find_most_common_words(corpus);
            if let Some(wordfile) = matches.value_of("WORDS") {
                if let Err(e) = save_words(Path::new(wordfile), &words) {
                    println!("Couldn't save vocabulary list: {}", e);
                }
            }
            let builder = create_model(&corpus, words.into_iter().map(|x| x.0).collect());
            if let Some(save) = save {
                if let Err(e) = builder.save(Path::new(save)) {
                    println!("Couldn't save model: {}", e);
                }
            }
            builder
        }
        _ => panic!("what you want?")
    };

    let start_time = time::get_time();
    let model = builder.build();
    println!("Model built in {}s", time::get_time().sec - start_time.sec);

    loop {
        println!("");
        let input = get_line();
        if input.starts_with(":q") {
            break;
        }

        match parser::parse(input.trim_matches(|c: char| c.is_whitespace()), &model) {
            Ok(word_vec) => {
                let nearest = model.nearest_words(&word_vec);
                println!(" = {:?}", word_vec);
                println!("-------------");
                for word in nearest.iter().take(20) {
                    println!("{:?}, {}", word, word_vec.distance(word));
                }
            }
            Err(e) => println!("{:?}", e)
        }

    }

}
<commit_msg>Clean up error handling and allow loading the vocab list<commit_after> #[macro_use]
 #[macro_use]
extern crate clap;
extern crate yaml_rust;
extern crate time;

mod models;
mod parser;
mod mayberef;

use clap::{Arg, App};
use std::fs::{File, read_dir};
use std::io::{BufReader, BufRead, Read, Write, stdin};
use std::io;
use std::path::{Path};
use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use models::*;

fn get_line() -> String {
    let stdin = stdin();
    let mut buffer = String::new();
    stdin.read_line(&mut buffer).unwrap();
    buffer
}

fn find_most_common_words(corpus: &Path) -> Vec<(String, u32)> {
    let words = files(corpus).flat_map(|file| read_words(file));

    let mut word_counts = HashMap::new();
    for word in words {
        *word_counts.entry(word).or_insert(0) += 1;
    }

    let mut counts: Vec<_> = word_counts.into_iter().collect();
    counts.sort_by(|a, b| b.1.cmp(&a.1));

    counts
}

fn save_words(path: &Path, words: &Vec<(String, u32)>) -> io::Result<()> {
    let mut out = try!(File::create(path));
    for &(ref word, count) in words {
        try!(out.write_all(format!("{}, {}\n", &word, count).as_bytes()));
    }
    Ok(())
}

#[derive(Debug)]
struct FormatError;
impl fmt::Display for FormatError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Invalid format for vocabulary file")
    }
}

impl Error for FormatError {
    fn description(&self) -> &str {
        "The format of the vocabulary file was Invalid"
    }
    fn cause(&self) -> Option<&Error> {
        None
    }
}

fn load_most_common_words(filename: &str, num: usize) -> Result<Vec<(String, u32)>, Box<Error>> {
    let file = try!(File::open(&Path::new(filename)));
    let reader = BufReader::new(file);
    let mut words = vec!();
    for line in reader.lines().take(num) {
        if let Ok(line) = line {
            let mut columns = line.split(',');
            let word = columns.next().map(|w| w.to_string());
            let num = columns.next().map(|n| n.parse());
            match (word, num) {
                (Some(word), Some(num)) => words.push((word, try!(num))),
                _ => return Err(Box::new(FormatError)),
            }
        } else {
            return Err(Box::new(FormatError));
        }
    }
    Ok(words)
}

fn sentences<T: Read + 'static>(reader: BufReader<T>) -> Box<Iterator<Item = String>> {
    Box::new(reader.split('.' as u8)
                   .filter_map(|v| {
                        let lowercase = v.unwrap();
                        String::from_utf8(lowercase).ok().map(|s| s.to_lowercase())
                    }))
}

fn read_words<R: BufRead + 'static>(reader: R) -> Box<Iterator<Item = String>> {
    Box::new(reader.lines()
                   .filter_map(|line| line.ok())
                   .flat_map(|line| {
                       line.split(|c| match c {
                                'a'...'z' | 'A'...'Z' => false,
                                _ => true,
                            })
                           .filter(|word| !word.is_empty())
                           .map(|word| word.to_lowercase())
                           .collect::<Vec<_>>()
                           .into_iter()
                   }))
}

fn files(path: &Path) -> Box<Iterator<Item = BufReader<File>>> {
    Box::new(read_dir(path)
            .unwrap()
            .map(|path| path.unwrap().path())
            .filter(|path| path.to_str().unwrap().ends_with(".txt"))
            .map(|path| {
                let file = File::open(&path).unwrap();
                BufReader::new(file)
            }))
}

fn create_model(corpus: &Path, words: Vec<String>) -> LanguageModelBuilder {
    let start_time = time::get_time();
    let mut builder = LanguageModelBuilder::new(10, words);

    let mut num_words = 0;

    for sentence in files(corpus).flat_map(sentences) {
        let mut acc = builder.new_sentence();
        for word in sentence.split(|c| match c {
            'a'...'z' => false,
            _ => true,
        }).filter(|w| !w.is_empty())
        {
            num_words += 1;
            acc.add_word(word);
            if num_words % 10_000_000 == 0 {
                println!("Loaded {} million words in {} seconds",
                         num_words / 1_000_000,
                         time::get_time().sec - start_time.sec);
            }
        }
    }

    let end_time = time::get_time();
    println!("Model built in {}s", end_time.sec - start_time.sec);
    builder
}

fn main() {
    let matches = App::new("gauntlet")
        .version("0.0.1")
        .author("James Moughan <jamougha@gmail.com>")
        .about("Implementation of GloVe algorithm")
        .arg(Arg::with_name("CORPUS")
            .short("c")
            .long("corpus")
            .help("Sets a directory to search for the corpus")
            .takes_value(true))
        .arg(Arg::with_name("LOAD")
            .short("l")
            .long("load")
            .help("Loads a pre-saved language model")
            .takes_value(true))
        .arg(Arg::with_name("SAVE")
            .short("s")
            .long("save")
            .help("Generates a model from the corpus specified and saves it")
            .takes_value(true))
        .arg(Arg::with_name("WORDS")
            .short("w")
            .long("save_words")
            .help("Saves the vocabulary list to the specified file")
            .takes_value(true))
        .arg(Arg::with_name("LOAD_WORDS")
            .short("W")
            .long("load_words")
            .help("Loads the vocabulary list from the specified file")
            .takes_value(true))
        .arg(Arg::with_name("NUM_WORDS")
            .short("n")
            .long("num_words")
            .help("The maximum number of words to use in the vocabulary list, defaults to 30000")
            .takes_value(true))
        .get_matches();

    let (load, save, corpus) = (matches.value_of("LOAD"), matches.value_of("SAVE"), matches.value_of("CORPUS"));
    let builder = match (load, save, corpus) {
        (Some(l), None, None) => LanguageModelBuilder::load(Path::new(&l)).expect("Couldn't load model"),
        (Some(_), _, _) => { println!("You must specify either a model to load or a corpus directory location"); return; }
        (_, save, Some(corpus)) => {
            let corpus = Path::new(corpus);
            let words = matches.value_of("LOAD_WORDS").map(|file|
                load_most_common_words(file, 30000)
                   .unwrap_or_else(|e| panic!("Couldn't read vocabulary file: {}", e))
            ).unwrap_or(find_most_common_words(corpus));

            if let Some(wordfile) = matches.value_of("WORDS") {
                if let Err(e) = save_words(Path::new(wordfile), &words) {
                    println!("Couldn't save vocabulary list: {}", e);
                }
            }
            let builder = create_model(&corpus, words.into_iter().map(|x| x.0).collect());
            if let Some(save) = save {
                if let Err(e) = builder.save(Path::new(save)) {
                    println!("Couldn't save model: {}", e);
                }
            }
            builder
        }
        _ => panic!("what you want?")
    };

    let start_time = time::get_time();
    let model = builder.build();
    println!("Model built in {}s", time::get_time().sec - start_time.sec);

    loop {
        println!("");
        let input = get_line();
        if input.starts_with(":q") {
            break;
        }

        match parser::parse(input.trim_matches(|c: char| c.is_whitespace()), &model) {
            Ok(word_vec) => {
                let nearest = model.nearest_words(&word_vec);
                println!(" = {:?}", word_vec);
                println!("-------------");
                for word in nearest.iter().take(20) {
                    println!("{:?}, {}", word, word_vec.distance(word));
                }
            }
            Err(e) => println!("{:?}", e)
        }

    }

}
<|endoftext|>


<commit_before>extern crate websocat;

extern crate futures;
extern crate tokio_core;
extern crate tokio_stdin_stdout;

extern crate env_logger;

#[macro_use]
extern crate structopt;

use structopt::StructOpt;

use tokio_core::reactor::{Core};
use websocat::{spec, WebsocatConfiguration, Options};

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

#[derive(StructOpt, Debug)]
#[structopt(after_help = "
Basic examples:
  Connect stdin/stdout to a websocket:
    websocat - ws://echo.websocket.org/

  Listen websocket and redirect it to a TCP port:
    websocat ws-l:127.0.0.1:8080 tcp:127.0.0.1:5678

  See more examples with the --long-help option

Short list of specifiers (see --long-help):
  ws:// wss:// - inetd: ws-listen: inetd-ws: tcp: tcp-l: ws-c:
  autoreconnect: reuse: mirror: threadedstdio: clogged:
  literal: literalreply: assert: udp-connect: open-async:
  readfile: writefile: open-fd: unix-connect: unix-listen:
  unix-dgram: abstract-connect: abstract-listen:
  exec: sh-c:
")]
struct Opt {
    /// First, listening/connecting specifier. See --long-help for info about specifiers.
    s1: String,
    /// Second, connecting specifier
    s2: String,

    #[structopt(short = "u", long = "unidirectional", help="Inhibit copying data from right specifier to left")]
    unidirectional: bool,
    #[structopt(short = "U", long = "unidirectional-reverse", help="Inhibit copying data from left specifier to right")]
    unidirectional_reverse: bool,

    #[structopt(long = "exit-on-eof", help="Close a data transfer direction if the other one reached EOF")]
    exit_on_eof: bool,

    #[structopt(short = "t", long = "text", help="Send text WebSocket messages instead of binary")]
    websocket_text_mode: bool,

    #[structopt(long="oneshot", help="Serve only once")]
    oneshot: bool,

    #[structopt(long="long-help", help="Show full help aboput specifiers and examples")]
    longhelp: bool,

    #[structopt(long="dump-spec", help="Instead of running, dump the specifiers representation to stdout")]
    dumpspec: bool,

    #[structopt(long="protocol", help="Specify Sec-WebSocket-Protocol: header")]
    websocket_protocol: Option<String>,

    #[structopt(long="udp-oneshot", help="udp-listen: replies only one packet per client")]
    udp_oneshot_mode: bool,

    #[structopt(long="unlink", help="Unlink listening UNIX socket before binding to it")]
    unlink_unix_socket:bool,

    #[structopt(long="exec-args", raw(allow_hyphen_values=r#"true"#), help="Arguments for the `exec:` specifier. Must be the last option, everything after it gets into the exec args list.")]
    exec_args: Vec<String>,

    #[structopt(long="ws-c-uri", help="URI to use for ws-c: specifier", default_value="ws://0.0.0.0/")]
    ws_c_uri: String,
}

fn longhelp() {
    println!("(see also usual --help message)

Full list of specifiers:
  `-` -- Stdin/stdout
    Read input from console, print to console.
    Can be specified only one time.
    Aliases: `stdio:`, `inetd:`

    `inetd:` also disables logging to stderr.

    Example: like `cat(1)`.
      websocat - -

    Example: for inetd mode
      websocat inetd: literal:$'Hello, world.\n'

  `ws://<url>`, `wss://<url>` -- WebSocket client
    Example: forward port 4554 to a websocket
      websocat tcp-l:127.0.0.1:4554 wss://127.0.0.1/some_websocket

  `ws-listen:<spec>` - Listen for websocket connections
    A combining specifier, but given IPv4 address as argument auto-inserts `tcp-l:`
    Aliases: `listen-ws:` `ws-l:` `l-ws:`

    Example:
        websocat ws-l:127.0.0.1:8808 -

    Example: the same, but more verbose:
        websocat ws-l:tcp-l:127.0.0.1:8808 reuse:-

  `inetd-ws:` - Alias of `ws-l:inetd:`

    Example of inetd.conf line:
      1234 stream tcp nowait myuser  /opt/websocat websocat inetd-ws: tcp:127.0.0.1:22

  `tcp:<hostport>` - connect to specified TCP host and port
    Aliases: `tcp-connect:`,`connect-tcp:`,`c-tcp:`,`tcp-c:`

    Example: like netcat
      websocat - tcp:127.0.0.1:22

    Example: IPv6
      websocat ws-l:0.0.0.0:8084 tcp:[::1]:22

  `tcp-l:<hostport>` - listen TCP port on specified address
    Aliases: `l-tcp:`  `tcp-listen:` `listen-tcp:`

    Example: echo server
      websocat tcp-l:0.0.0.0:1441 mirror:

  `exec:<program_path> --exec-args <arguments...> --`
    Execute a program (subprocess) directly, without a subshell.

    Example: date server
      websocat -U ws-l:127.0.0.1:5667 exec:date

    Example: pinger
      websocat -U ws-l:127.0.0.1:5667 exec:ping --exec-args 127.0.0.1 -c 1 --

  `sh-c:<command line>` - start subprocess though 'sh -c' or `cmd /C`

    Example: unauthenticated shell
      websocat --exit-on-eof ws-l:127.0.0.1:5667 sh-c:'bash -i 2>&1'

  `udp:<hostport>` - send and receive packets to specified UDP socket
    Aliases: `udp-connect:` `connect-udp:` `c-udp:` `udp-c:`

  `udp-listen:<hostport>` - bind to socket on host and port
    Aliases: `udp-l:`, `l-udp:`, `listen-udp:`

    Note that it is not a multiconnect specifier: entire lifecycle
    of the UDP socket is the same connection.

    Packets get sent to the most recent seen peer.
    If no peers are seen yet, it waits for the first packet.

    File a feature request on Github if you want proper DNS-like request-reply UDP mode here.

  `ws-connect:<spec>` - low-level WebSocket connector
    A combining specifier. Underlying specifier is should be after the colon.
    URL and Host: header being sent are independent from underlying specifier
    Aliases: `ws-c:` `c-ws:` `connect-ws:`

    Example: connect to echo server in more explicit way
      websocat --ws-c-uri=ws://echo.websocket.org/ - ws-c:tcp:174.129.224.73:80

  `autoreconnect:<spec>` - Auto-reconnector
    Re-establish underlying specifier on any error or EOF

    Example: keep connecting to the port or spin 100% CPU trying if it is closed.
      websocat - autoreconnect:tcp:127.0.0.1:5445

    TODO: implement timeouts

  `reuse:<spec>` - Reuse one connection for serving multiple clients
    Better suited for unidirectional connections

    Example (unreliable): don't disconnect SSH when websocket reconnects
      websocat ws-l:[::]:8088 reuse:tcp:127.0.0.1:22

  `threadedstdio:` - Stdin/stdout, spawning a thread
    Like `-`, but forces threaded mode instead of async mode
    Use when standard input is not `epoll(7)`-able.
    Replaces `-` when `no_unix_stdio` Cargo feature is activated

  `mirror:` - Simply copy output to input
    Similar to `exec:cat`.

  `open-async:<path>` - Open file for read and write and use it like a socket
    Not for regular files, see readfile: and writefile: instead.

    Example:
      websocat - open-async:/dev/null

  `open-fd:<number>` - Use specified file descriptor like a socket

  `unix-connect:<path>` - Connect to UNIX socket
    Aliases: `unix:`, `connect-unix:`, `unix-c:`, `c-unix:`

  `unix-listen:<path>` - Listen for connections on a UNIX socket
    Aliases: `unix-l:`, `listen-unix:`, `l-unix:`

    Example: with nginx
      umask 0000
      websocat --unlink ws-l:unix-l:/tmp/wstest tcp:[::]:22

    Nginx config:

    location /ws {{
        proxy_read_timeout 7d;
        proxy_send_timeout 7d;
        #proxy_pass http://localhost:3012;
        proxy_pass http://unix:/tmp/wstest;
        proxy_http_version 1.1;
        proxy_set_header Upgrade $http_upgrade;
        proxy_set_header Connection \"upgrade\";
    }}

      TODO: chmod?

  `unix-dgram:<path>:<path>` - Send packets to one path, receive from the other

  `abstract-connect:<string>` - Connect to Linux abstract-namespaced socket
    Aliases: `abstract-c:`, `connect-abstract:`, `c-abstract:`, `abstract:`

  `abstract-listen:<path>` - Listen for connections on Linux abstract-namespaced socket
    Aliases: `abstract-l:`, `listen-abstract:`, `l-abstract:`

  `readfile:<path>` - synchronously read files
    Blocking on operations with the file pauses the whole process

    Example:
      websocat ws-l:127.0.0.1:8000 readfile:hello.json

  `write:<path>` - synchronously write files
    Blocking on operations with the file pauses the whole process
    Files are opened in overwrite mode.

    Example:
      websocat ws-l:127.0.0.1:8000 reuse:writefile:log.txt

  `clogged:` - Do nothing
    Don't read or write any bytes. Keep connections hanging.

  `literal:<string>` - Output a string, discard input.
    Ignore all input, use specified string as output.

  `literalreply:<string>` - Reply with this string for each input packet

  `assert:<string>` - Check the input.
    Read entire input and panic the program if the input is not equal
    to the specified string.

  TODO:
  `exec:`, `seqpacket:`

More examples:
  Wacky mode:
    websocat ws-l:ws-l:ws-c:- tcp:127.0.0.1:5678

    Connect to a websocket using stdin/stdout as a transport,
    then accept a websocket connection over this previous websocket as a transport,
    then connect to a websocket using previous step as a transport,
    then forward resulting connection to the TCP port.

    (Excercise to the reader: manage to actually connect to it).
");
}

fn run() -> Result<()> {
    if std::env::args().nth(1).unwrap_or_default() == "--long-help" {
        longhelp();
        return Ok(());
    }

    let cmd = Opt::from_args();

    if cmd.longhelp {
        longhelp();
        return Ok(());
    }

    if false
    //    || cmd.oneshot
    {
        Err("This mode is not implemented")?
    }

    let opts = {
        macro_rules! opts {
            ($($o:ident)*) => {
                Options {
                    $($o : cmd.$o,)*
                }
            }
        }
        opts!(
            websocket_text_mode
            websocket_protocol
            udp_oneshot_mode
            unidirectional
            unidirectional_reverse
            exit_on_eof
            oneshot
            unlink_unix_socket
            exec_args
            ws_c_uri
        )
    };

    let s1 = spec(&cmd.s1)?;
    let s2 = spec(&cmd.s2)?;

    let mut websocat = WebsocatConfiguration { opts, s1, s2 };

    while let Some(concern) = websocat.get_concern() {
        use websocat::ConfigurationConcern::*;
        if concern == StdinToStdout {
            if cmd.dumpspec {
                println!("cat mode");
                return Ok(())
            }

            // Degenerate mode: just copy stdin to stdout and call it a day
            ::std::io::copy(&mut ::std::io::stdin(), &mut ::std::io::stdout())?;
            return Ok(())
        }

        if concern == DegenerateMode {
            if cmd.dumpspec {
                println!("noop");
            }
            return Ok(())
        }

        if concern == StdioConflict {
            Err("Too many usages of stdin/stdout")?;
        }

        if concern == NeedsStdioReuser {
            eprintln!("Warning: replies on stdio get directed at random connected client");
            websocat = websocat.auto_install_reuser();
            continue;
        }

        if concern == NeedsStdioReuser2 {
            websocat = websocat.auto_install_reuser();
            continue;
        }

        if concern == MultipleReusers {
            eprintln!("Specifier dump: {:?} {:?}", websocat.s1, websocat.s2);
            Err("Multiple reusers is not allowed")?;
        }
        break;
    }

    if cmd.dumpspec {
        println!("{:?}", websocat.s1);
        println!("{:?}", websocat.s2);
        println!("{:?}", websocat.opts);
        return Ok(())
    }

    let mut core = Core::new()?;

    let prog = websocat.serve(core.handle(), std::rc::Rc::new(|e| {
        eprintln!("websocat: {}", e);
    }));
    core.run(prog).map_err(|()|"error running".to_string())?;
    Ok(())
}

fn main() {
    env_logger::init();
    let r = run();

    if let Err(e) = r {
        eprintln!("websocat: {}", e);
        ::std::process::exit(1);
    }
}
<commit_msg>Modify TODO list<commit_after>extern crate websocat;

extern crate futures;
extern crate tokio_core;
extern crate tokio_stdin_stdout;

extern crate env_logger;

#[macro_use]
extern crate structopt;

use structopt::StructOpt;

use tokio_core::reactor::{Core};
use websocat::{spec, WebsocatConfiguration, Options};

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

#[derive(StructOpt, Debug)]
#[structopt(after_help = "
Basic examples:
  Connect stdin/stdout to a websocket:
    websocat - ws://echo.websocket.org/

  Listen websocket and redirect it to a TCP port:
    websocat ws-l:127.0.0.1:8080 tcp:127.0.0.1:5678

  See more examples with the --long-help option

Short list of specifiers (see --long-help):
  ws:// wss:// - inetd: ws-listen: inetd-ws: tcp: tcp-l: ws-c:
  autoreconnect: reuse: mirror: threadedstdio: clogged:
  literal: literalreply: assert: udp-connect: open-async:
  readfile: writefile: open-fd: unix-connect: unix-listen:
  unix-dgram: abstract-connect: abstract-listen:
  exec: sh-c:
")]
struct Opt {
    /// First, listening/connecting specifier. See --long-help for info about specifiers.
    s1: String,
    /// Second, connecting specifier
    s2: String,

    #[structopt(short = "u", long = "unidirectional", help="Inhibit copying data from right specifier to left")]
    unidirectional: bool,
    #[structopt(short = "U", long = "unidirectional-reverse", help="Inhibit copying data from left specifier to right")]
    unidirectional_reverse: bool,

    #[structopt(long = "exit-on-eof", help="Close a data transfer direction if the other one reached EOF")]
    exit_on_eof: bool,

    #[structopt(short = "t", long = "text", help="Send text WebSocket messages instead of binary")]
    websocket_text_mode: bool,

    #[structopt(long="oneshot", help="Serve only once")]
    oneshot: bool,

    #[structopt(long="long-help", help="Show full help aboput specifiers and examples")]
    longhelp: bool,

    #[structopt(long="dump-spec", help="Instead of running, dump the specifiers representation to stdout")]
    dumpspec: bool,

    #[structopt(long="protocol", help="Specify Sec-WebSocket-Protocol: header")]
    websocket_protocol: Option<String>,

    #[structopt(long="udp-oneshot", help="udp-listen: replies only one packet per client")]
    udp_oneshot_mode: bool,

    #[structopt(long="unlink", help="Unlink listening UNIX socket before binding to it")]
    unlink_unix_socket:bool,

    #[structopt(long="exec-args", raw(allow_hyphen_values=r#"true"#), help="Arguments for the `exec:` specifier. Must be the last option, everything after it gets into the exec args list.")]
    exec_args: Vec<String>,

    #[structopt(long="ws-c-uri", help="URI to use for ws-c: specifier", default_value="ws://0.0.0.0/")]
    ws_c_uri: String,
}

fn longhelp() {
    println!("(see also usual --help message)

Full list of specifiers:
  `-` -- Stdin/stdout
    Read input from console, print to console.
    Can be specified only one time.
    Aliases: `stdio:`, `inetd:`

    `inetd:` also disables logging to stderr.

    Example: like `cat(1)`.
      websocat - -

    Example: for inetd mode
      websocat inetd: literal:$'Hello, world.\n'

  `ws://<url>`, `wss://<url>` -- WebSocket client
    Example: forward port 4554 to a websocket
      websocat tcp-l:127.0.0.1:4554 wss://127.0.0.1/some_websocket

  `ws-listen:<spec>` - Listen for websocket connections
    A combining specifier, but given IPv4 address as argument auto-inserts `tcp-l:`
    Aliases: `listen-ws:` `ws-l:` `l-ws:`

    Example:
        websocat ws-l:127.0.0.1:8808 -

    Example: the same, but more verbose:
        websocat ws-l:tcp-l:127.0.0.1:8808 reuse:-

  `inetd-ws:` - Alias of `ws-l:inetd:`

    Example of inetd.conf line:
      1234 stream tcp nowait myuser  /opt/websocat websocat inetd-ws: tcp:127.0.0.1:22

  `tcp:<hostport>` - connect to specified TCP host and port
    Aliases: `tcp-connect:`,`connect-tcp:`,`c-tcp:`,`tcp-c:`

    Example: like netcat
      websocat - tcp:127.0.0.1:22

    Example: IPv6
      websocat ws-l:0.0.0.0:8084 tcp:[::1]:22

  `tcp-l:<hostport>` - listen TCP port on specified address
    Aliases: `l-tcp:`  `tcp-listen:` `listen-tcp:`

    Example: echo server
      websocat tcp-l:0.0.0.0:1441 mirror:

  `exec:<program_path> --exec-args <arguments...> --`
    Execute a program (subprocess) directly, without a subshell.

    Example: date server
      websocat -U ws-l:127.0.0.1:5667 exec:date

    Example: pinger
      websocat -U ws-l:127.0.0.1:5667 exec:ping --exec-args 127.0.0.1 -c 1 --

  `sh-c:<command line>` - start subprocess though 'sh -c' or `cmd /C`

    Example: unauthenticated shell
      websocat --exit-on-eof ws-l:127.0.0.1:5667 sh-c:'bash -i 2>&1'

  `udp:<hostport>` - send and receive packets to specified UDP socket
    Aliases: `udp-connect:` `connect-udp:` `c-udp:` `udp-c:`

  `udp-listen:<hostport>` - bind to socket on host and port
    Aliases: `udp-l:`, `l-udp:`, `listen-udp:`

    Note that it is not a multiconnect specifier: entire lifecycle
    of the UDP socket is the same connection.

    Packets get sent to the most recent seen peer.
    If no peers are seen yet, it waits for the first packet.

    File a feature request on Github if you want proper DNS-like request-reply UDP mode here.

  `ws-connect:<spec>` - low-level WebSocket connector
    A combining specifier. Underlying specifier is should be after the colon.
    URL and Host: header being sent are independent from underlying specifier
    Aliases: `ws-c:` `c-ws:` `connect-ws:`

    Example: connect to echo server in more explicit way
      websocat --ws-c-uri=ws://echo.websocket.org/ - ws-c:tcp:174.129.224.73:80

  `autoreconnect:<spec>` - Auto-reconnector
    Re-establish underlying specifier on any error or EOF

    Example: keep connecting to the port or spin 100% CPU trying if it is closed.
      websocat - autoreconnect:tcp:127.0.0.1:5445

    TODO: implement timeouts

  `reuse:<spec>` - Reuse one connection for serving multiple clients
    Better suited for unidirectional connections

    Example (unreliable): don't disconnect SSH when websocket reconnects
      websocat ws-l:[::]:8088 reuse:tcp:127.0.0.1:22

  `threadedstdio:` - Stdin/stdout, spawning a thread
    Like `-`, but forces threaded mode instead of async mode
    Use when standard input is not `epoll(7)`-able.
    Replaces `-` when `no_unix_stdio` Cargo feature is activated

  `mirror:` - Simply copy output to input
    Similar to `exec:cat`.

  `open-async:<path>` - Open file for read and write and use it like a socket
    Not for regular files, see readfile: and writefile: instead.

    Example:
      websocat - open-async:/dev/null

  `open-fd:<number>` - Use specified file descriptor like a socket

  `unix-connect:<path>` - Connect to UNIX socket
    Aliases: `unix:`, `connect-unix:`, `unix-c:`, `c-unix:`

  `unix-listen:<path>` - Listen for connections on a UNIX socket
    Aliases: `unix-l:`, `listen-unix:`, `l-unix:`

    Example: with nginx
      umask 0000
      websocat --unlink ws-l:unix-l:/tmp/wstest tcp:[::]:22

    Nginx config:

    location /ws {{
        proxy_read_timeout 7d;
        proxy_send_timeout 7d;
        #proxy_pass http://localhost:3012;
        proxy_pass http://unix:/tmp/wstest;
        proxy_http_version 1.1;
        proxy_set_header Upgrade $http_upgrade;
        proxy_set_header Connection \"upgrade\";
    }}

      TODO: chmod?

  `unix-dgram:<path>:<path>` - Send packets to one path, receive from the other

  `abstract-connect:<string>` - Connect to Linux abstract-namespaced socket
    Aliases: `abstract-c:`, `connect-abstract:`, `c-abstract:`, `abstract:`

  `abstract-listen:<path>` - Listen for connections on Linux abstract-namespaced socket
    Aliases: `abstract-l:`, `listen-abstract:`, `l-abstract:`

  `readfile:<path>` - synchronously read files
    Blocking on operations with the file pauses the whole process

    Example:
      websocat ws-l:127.0.0.1:8000 readfile:hello.json

  `write:<path>` - synchronously write files
    Blocking on operations with the file pauses the whole process
    Files are opened in overwrite mode.

    Example:
      websocat ws-l:127.0.0.1:8000 reuse:writefile:log.txt

  `clogged:` - Do nothing
    Don't read or write any bytes. Keep connections hanging.

  `literal:<string>` - Output a string, discard input.
    Ignore all input, use specified string as output.

  `literalreply:<string>` - Reply with this string for each input packet

  `assert:<string>` - Check the input.
    Read entire input and panic the program if the input is not equal
    to the specified string.

  TODO:
  --unix-seqpacket
  sctp:
  ssl:

More examples:
  Wacky mode:
    websocat ws-l:ws-l:ws-c:- tcp:127.0.0.1:5678

    Connect to a websocket using stdin/stdout as a transport,
    then accept a websocket connection over this previous websocket as a transport,
    then connect to a websocket using previous step as a transport,
    then forward resulting connection to the TCP port.

    (Excercise to the reader: manage to actually connect to it).
");
}

fn run() -> Result<()> {
    if std::env::args().nth(1).unwrap_or_default() == "--long-help" {
        longhelp();
        return Ok(());
    }

    let cmd = Opt::from_args();

    if cmd.longhelp {
        longhelp();
        return Ok(());
    }

    if false
    //    || cmd.oneshot
    {
        Err("This mode is not implemented")?
    }

    let opts = {
        macro_rules! opts {
            ($($o:ident)*) => {
                Options {
                    $($o : cmd.$o,)*
                }
            }
        }
        opts!(
            websocket_text_mode
            websocket_protocol
            udp_oneshot_mode
            unidirectional
            unidirectional_reverse
            exit_on_eof
            oneshot
            unlink_unix_socket
            exec_args
            ws_c_uri
        )
    };

    let s1 = spec(&cmd.s1)?;
    let s2 = spec(&cmd.s2)?;

    let mut websocat = WebsocatConfiguration { opts, s1, s2 };

    while let Some(concern) = websocat.get_concern() {
        use websocat::ConfigurationConcern::*;
        if concern == StdinToStdout {
            if cmd.dumpspec {
                println!("cat mode");
                return Ok(())
            }

            // Degenerate mode: just copy stdin to stdout and call it a day
            ::std::io::copy(&mut ::std::io::stdin(), &mut ::std::io::stdout())?;
            return Ok(())
        }

        if concern == DegenerateMode {
            if cmd.dumpspec {
                println!("noop");
            }
            return Ok(())
        }

        if concern == StdioConflict {
            Err("Too many usages of stdin/stdout")?;
        }

        if concern == NeedsStdioReuser {
            eprintln!("Warning: replies on stdio get directed at random connected client");
            websocat = websocat.auto_install_reuser();
            continue;
        }

        if concern == NeedsStdioReuser2 {
            websocat = websocat.auto_install_reuser();
            continue;
        }

        if concern == MultipleReusers {
            eprintln!("Specifier dump: {:?} {:?}", websocat.s1, websocat.s2);
            Err("Multiple reusers is not allowed")?;
        }
        break;
    }

    if cmd.dumpspec {
        println!("{:?}", websocat.s1);
        println!("{:?}", websocat.s2);
        println!("{:?}", websocat.opts);
        return Ok(())
    }

    let mut core = Core::new()?;

    let prog = websocat.serve(core.handle(), std::rc::Rc::new(|e| {
        eprintln!("websocat: {}", e);
    }));
    core.run(prog).map_err(|()|"error running".to_string())?;
    Ok(())
}

fn main() {
    env_logger::init();
    let r = run();

    if let Err(e) = r {
        eprintln!("websocat: {}", e);
        ::std::process::exit(1);
    }
}
<|endoftext|>


<commit_before>extern crate mio;
extern crate bytes;
extern crate log4rs;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;

use mio::tcp::*;
use mio::EventLoop;
use mio::{TryRead,TryWrite};
use mio::util::Slab;
use std::collections::VecDeque;
use clap::{Arg, App};

const SERVER: mio::Token = mio::Token(0);

#[derive(Debug)]
enum MiChatCommand {
    Broadcast(String),
}

struct MiChat {
    listener: TcpListener,
    connections: Slab<Connection>,
}

impl MiChat {
    fn new(listener: TcpListener) -> MiChat {
        MiChat {
            listener: listener,
            connections: Slab::new_starting_at(mio::Token(1), 1024),
        }
    }
    fn process_action(&mut self, msg: &MiChatCommand) {
        trace!("{:p}; got {:?}", self, msg);
        match msg {
            &MiChatCommand::Broadcast(ref s) => {
                for c in self.connections.iter_mut() {
                    c.enqueue(&s);
                }
            }
        }
    }
}

#[derive(Debug)]
struct Connection {
    socket: TcpStream,
    sock_status: mio::EventSet,
    token: mio::Token,
    read_buf: Vec<u8>,
    read_eof: bool,
    failed: bool,
    write_buf: Vec<u8>,
}

impl Connection {
    fn new(socket: TcpStream, token: mio::Token) -> Connection {
        Connection {
            socket: socket,
            sock_status: mio::EventSet::none(),
            token: token,
            read_buf: Vec::with_capacity(1024),
            write_buf: Vec::new(),
            read_eof: false,
            failed: false,
        }
    }
    // Event updates arrive here
    fn ready(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        self.sock_status.insert(events);
        info!("Connection::ready: {:?}; this time: {:?}; now: {:?}",
                self.socket.peer_addr(), events, self.sock_status);
    }

    // actions are processed here on down.

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>) {
        if self.sock_status.is_readable() {
            self.read();
            self.sock_status.remove(mio::EventSet::readable());
        }

        self.process_buffer(to_parent);

        if self.sock_status.is_writable() {
            self.write();
            self.sock_status.remove(mio::EventSet::writable());
        }

        if !self.is_closed() {
            self.reregister(event_loop)
        }
    }

    fn process_buffer(&mut self, to_parent: &mut VecDeque<MiChatCommand>) {
        let mut prev = 0;
        info!("{:?}: Read buffer: {:?}", self.socket.peer_addr(), self.read_buf);
        for n in self.read_buf.iter().enumerate()
                .filter_map(|(i, e)| if *e == '\n' as u8 { Some(i) } else { None } ) {
            info!("{:?}: Pos: {:?}; chunk: {:?}", self.socket.peer_addr(), n, &self.read_buf[prev..n]);
            let s = String::from_utf8_lossy(&self.read_buf[prev..n]).to_string();
            let cmd = MiChatCommand::Broadcast(s);
            info!("Send! {:?}", cmd);
            to_parent.push_back(cmd);
            prev = n+1;
        }
        let remainder = self.read_buf[prev..].to_vec();
        info!("{:?}: read Remainder: {}", self.socket.peer_addr(), remainder.len());
        self.read_buf = remainder;
    }

    fn read(&mut self) {
        let mut abuf = Vec::new();
        match self.socket.try_read_buf(&mut abuf) {
            Ok(Some(0)) => {
                info!("{:?}: EOF!", self.socket.peer_addr());
                self.read_eof = true
            },
            Ok(Some(n)) => {
                info!("{:?}: Read {}bytes", self.socket.peer_addr(), n);
                self.read_buf.extend(abuf);
            },
            Ok(None) => {
                info!("{:?}: Noop!", self.socket.peer_addr());
            },
            Err(e) => {
                error!("got an error trying to read; err={:?}", e);
                self.failed =true;
            }
        }
    }

    fn write(&mut self) {
        match self.socket.try_write(&mut self.write_buf) {
            Ok(Some(n)) => {
                info!("{:?}: Wrote {} of {} in buffer ({:?}...)", self.socket.peer_addr(), n,
                    self.write_buf.len(), &self.write_buf[0..2]);
                self.write_buf = self.write_buf[n..].to_vec();
                info!("{:?}: Now {:?}b", self.socket.peer_addr(), self.write_buf.len());
            },
            Ok(None) => {
                info!("Write unready");
            },
            Err(e) => {
                error!("got an error trying to write; err={:?}", e);
                self.failed = true;
            }
        }
    }

    fn reregister(&mut self, event_loop: &mut mio::EventLoop<MiChat>) {
        let mut flags = mio::EventSet::readable();
        if !self.write_buf.is_empty() {
            flags.insert(mio::EventSet::writable());
        }
        info!("Registering {:?} with {:?}", self, flags);

        event_loop.reregister(
                &self.socket,
                self.token,
                flags,
                mio::PollOpt::oneshot()).expect("EventLoop#reregister")
    }

    fn is_closed(&self) -> bool {
        self.failed || (self.read_eof && self.write_buf.is_empty())
    }

    fn enqueue(&mut self, s: &str) {
        self.write_buf.extend(s.as_bytes());
        self.write_buf.push('\n' as u8);
    }
}

impl mio::Handler for MiChat {
    type Timeout = ();
    type Message = ();
    fn ready(&mut self, event_loop: &mut mio::EventLoop<Self>, token: mio::Token, events: mio::EventSet) {
        info!("{:?}: {:?}", token, events);
        let mut next = VecDeque::new();

        match token {
            SERVER => {
                // Only receive readable events
                assert!(events.is_readable());

                info!("the listener socket is ready to accept a connection");
                match self.listener.accept() {
                    Ok(Some(socket)) => {
                        let token = self.connections
                            .insert_with(|token| Connection::new(socket, token))
                            .expect("token insert");

                        event_loop.register_opt(
                            &self.connections[token].socket,
                            token,
                            mio::EventSet::readable(),
                            mio::PollOpt::edge() | mio::PollOpt::oneshot())
                        .expect("event loop register");
                    }
                    Ok(None) => {
                        info!("the listener socket wasn't actually ready");
                    }
                    Err(e) => {
                        info!("listener.accept() errored: {}", e);
                        event_loop.shutdown();
                    }
                }
            }
            _ => {
                self.connections[token].ready(event_loop, events);
                info!("TODO next on {:?}: {:?}", token, next);
                if self.connections[token].is_closed() {
                    info!("Removing; {:?}", token);
                    self.connections.remove(token);
                }
            }
        }

        loop {
            for conn in self.connections.iter_mut() {
                conn.process_rules(event_loop, &mut next);
            }
            if next.is_empty() { break; }

            for action in next.iter() {
                self.process_action(action);
            }
            next.clear();
        }
    }
}

const LOG_FILE: &'static str = "log.toml";

fn main() {
    if let Err(e) = log4rs::init_file(LOG_FILE, Default::default()) {
        panic!("Could not init logger from file {}: {}", LOG_FILE, e);
    }
    let matches = App::new("michat")
        .arg(Arg::with_name("bind").short("l").takes_value(true).required(true))
        .get_matches();

    let address = value_t_or_exit!(matches.value_of("bind"), std::net::SocketAddr);
    let listener = TcpListener::bind(&address).expect("bind");

    let mut event_loop = mio::EventLoop::new().expect("Create event loop");
    event_loop.register(&listener, SERVER).expect("Register listener");

    info!("running michat listener at: {:?}", address);

    let mut service = MiChat::new(listener);
    event_loop.run(&mut service).expect("Run loop");
}
<commit_msg>Make the listener just another event handler.<commit_after>extern crate mio;
extern crate bytes;
extern crate log4rs;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;

use mio::tcp::*;
use mio::EventLoop;
use mio::{TryRead,TryWrite};
use mio::util::Slab;
use std::collections::VecDeque;
use clap::{Arg, App};

#[derive(Debug)]
enum MiChatCommand {
    Broadcast(String),
    NewConnection(TcpStream),
}

struct MiChat {
    connections: Slab<EventHandler>,
}

impl MiChat {
    fn new() -> MiChat {
        MiChat {
            connections: Slab::new(1024),
        }
    }

    fn listen(&mut self, event_loop: &mut mio::EventLoop<MiChat>, listener: TcpListener) {
        let l = EventHandler::Listener(Listener::new(listener));
        let server = self.connections.insert(l).expect("insert listener");
        if let &EventHandler::Listener(ref l) = &self.connections[server] {
            event_loop.register(&l.listener, server).expect("Register listener");
        }
    }

    fn process_action(&mut self, msg: MiChatCommand, event_loop: &mut mio::EventLoop<MiChat>) {
        trace!("{:p}; got {:?}", self, msg);
        match msg {
            MiChatCommand::Broadcast(s) => {
                for c in self.connections.iter_mut() {
                    if let &mut EventHandler::Conn(ref mut c) = c {
                        c.enqueue(&s);
                    }
                }
            },

            MiChatCommand::NewConnection(socket) => {
                let token = self.connections
                    .insert_with(|token| EventHandler::Conn(Connection::new(socket, token)))
                    .expect("token insert");

                if let EventHandler::Conn(ref c) = self.connections[token] {
                    event_loop.register_opt(
                            &c.socket,
                            token,
                            mio::EventSet::readable(),
                            mio::PollOpt::edge() | mio::PollOpt::oneshot())
                        .expect("event loop register");
                }
            }
        }
    }
}

#[derive(Debug)]
struct Connection {
    socket: TcpStream,
    sock_status: mio::EventSet,
    token: mio::Token,
    read_buf: Vec<u8>,
    read_eof: bool,
    failed: bool,
    write_buf: Vec<u8>,
}

impl Connection {
    fn new(socket: TcpStream, token: mio::Token) -> Connection {
        Connection {
            socket: socket,
            sock_status: mio::EventSet::none(),
            token: token,
            read_buf: Vec::with_capacity(1024),
            write_buf: Vec::new(),
            read_eof: false,
            failed: false,
        }
    }
    // Event updates arrive here
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        self.sock_status.insert(events);
        info!("Connection::handle_event: {:?}; this time: {:?}; now: {:?}",
                self.socket.peer_addr(), events, self.sock_status);
    }

    // actions are processed here on down.

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>) {
        if self.sock_status.is_readable() {
            self.read();
            self.sock_status.remove(mio::EventSet::readable());
        }

        self.process_buffer(to_parent);

        if self.sock_status.is_writable() {
            self.write();
            self.sock_status.remove(mio::EventSet::writable());
        }

        if !self.is_closed() {
            self.reregister(event_loop)
        }
    }

    fn process_buffer(&mut self, to_parent: &mut VecDeque<MiChatCommand>) {
        let mut prev = 0;
        info!("{:?}: Read buffer: {:?}", self.socket.peer_addr(), self.read_buf);
        for n in self.read_buf.iter().enumerate()
                .filter_map(|(i, e)| if *e == '\n' as u8 { Some(i) } else { None } ) {
            info!("{:?}: Pos: {:?}; chunk: {:?}", self.socket.peer_addr(), n, &self.read_buf[prev..n]);
            let s = String::from_utf8_lossy(&self.read_buf[prev..n]).to_string();
            let cmd = MiChatCommand::Broadcast(s);
            info!("Send! {:?}", cmd);
            to_parent.push_back(cmd);
            prev = n+1;
        }
        let remainder = self.read_buf[prev..].to_vec();
        info!("{:?}: read Remainder: {}", self.socket.peer_addr(), remainder.len());
        self.read_buf = remainder;
    }

    fn read(&mut self) {
        let mut abuf = Vec::new();
        match self.socket.try_read_buf(&mut abuf) {
            Ok(Some(0)) => {
                info!("{:?}: EOF!", self.socket.peer_addr());
                self.read_eof = true
            },
            Ok(Some(n)) => {
                info!("{:?}: Read {}bytes", self.socket.peer_addr(), n);
                self.read_buf.extend(abuf);
            },
            Ok(None) => {
                info!("{:?}: Noop!", self.socket.peer_addr());
            },
            Err(e) => {
                error!("got an error trying to read; err={:?}", e);
                self.failed =true;
            }
        }
    }

    fn write(&mut self) {
        match self.socket.try_write(&mut self.write_buf) {
            Ok(Some(n)) => {
                info!("{:?}: Wrote {} of {} in buffer", self.socket.peer_addr(), n,
                    self.write_buf.len());
                self.write_buf = self.write_buf[n..].to_vec();
                info!("{:?}: Now {:?}b", self.socket.peer_addr(), self.write_buf.len());
            },
            Ok(None) => {
                info!("Write unready");
            },
            Err(e) => {
                error!("got an error trying to write; err={:?}", e);
                self.failed = true;
            }
        }
    }

    fn reregister(&mut self, event_loop: &mut mio::EventLoop<MiChat>) {
        let mut flags = mio::EventSet::readable();
        if !self.write_buf.is_empty() {
            flags.insert(mio::EventSet::writable());
        }
        info!("Registering {:?} with {:?}", self, flags);

        event_loop.reregister(
                &self.socket,
                self.token,
                flags,
                mio::PollOpt::oneshot()).expect("EventLoop#reregister")
    }

    fn is_closed(&self) -> bool {
        self.failed || (self.read_eof && self.write_buf.is_empty())
    }

    fn enqueue(&mut self, s: &str) {
        self.write_buf.extend(s.as_bytes());
        self.write_buf.push('\n' as u8);
    }
}

#[derive(Debug)]
struct Listener {
    listener: TcpListener,
    sock_status: mio::EventSet,
}

impl Listener {
    fn new(listener: TcpListener) -> Listener {
        Listener {
            listener: listener,
            sock_status: mio::EventSet::none(),
        }
    }

    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        assert!(events.is_readable());
        self.sock_status.insert(events);
        info!("Listener::handle_event: {:?}; this time: {:?}; now: {:?}",
                self.listener.local_addr(), events, self.sock_status);
    }

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
            to_parent: &mut VecDeque<MiChatCommand>) {
        info!("the listener socket is ready to accept a connection");
        match self.listener.accept() {
            Ok(Some(socket)) => {
                let cmd = MiChatCommand::NewConnection(socket);
                to_parent.push_back(cmd);

/*
*/

            }
            Ok(None) => {
                info!("the listener socket wasn't actually ready");
            }
            Err(e) => {
                info!("listener.accept() errored: {}", e);
                event_loop.shutdown();
            }
        }
    }

    fn is_closed(&self) -> bool {
        false
    }
}

#[derive(Debug)]
enum EventHandler {
    Listener (Listener),
    Conn (Connection),
}

impl EventHandler {
    fn handle_event(&mut self, _event_loop: &mut mio::EventLoop<MiChat>, events: mio::EventSet) {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn.handle_event(_event_loop, events),
            &mut EventHandler::Listener(ref mut listener) => listener.handle_event(_event_loop, events)
        }
    }

    fn process_rules(&mut self, event_loop: &mut mio::EventLoop<MiChat>,
        to_parent: &mut VecDeque<MiChatCommand>) {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn.process_rules(event_loop, to_parent),
            &mut EventHandler::Listener(ref mut listener) => listener.process_rules(event_loop, to_parent)
        }
    }

    fn is_closed(&self) -> bool {
        match self {
            &EventHandler::Conn(ref conn) => conn.is_closed(),
            &EventHandler::Listener(ref listener) => listener.is_closed()
        }
    }

}

impl mio::Handler for MiChat {
    type Timeout = ();
    type Message = ();
    fn ready(&mut self, event_loop: &mut mio::EventLoop<Self>, token: mio::Token, events: mio::EventSet) {
        info!("{:?}: {:?}", token, events);
        match token {
            _ => {
                self.connections[token].handle_event(event_loop, events);
                if self.connections[token].is_closed() {
                    info!("Removing; {:?}", token);
                    self.connections.remove(token);
                }
            }
        }

        let mut parent_actions = VecDeque::new();
        loop {
            for conn in self.connections.iter_mut() {
                conn.process_rules(event_loop, &mut parent_actions);
            }
            // Anything left to process?
            if parent_actions.is_empty() { break; }

            for action in parent_actions.drain(..) {
                self.process_action(action, event_loop);
            }
        }
    }
}

const LOG_FILE: &'static str = "log.toml";

fn main() {
    if let Err(e) = log4rs::init_file(LOG_FILE, Default::default()) {
        panic!("Could not init logger from file {}: {}", LOG_FILE, e);
    }
    let matches = App::new("michat")
        .arg(Arg::with_name("bind").short("l").takes_value(true).required(true))
        .get_matches();

    let address = value_t_or_exit!(matches.value_of("bind"), std::net::SocketAddr);
    let listener = TcpListener::bind(&address).expect("bind");

    let mut event_loop = mio::EventLoop::new().expect("Create event loop");
    let mut service = MiChat::new();

    service.listen(&mut event_loop, listener);
    info!("running michat listener at: {:?}", address);
    event_loop.run(&mut service).expect("Run loop");
}
<|endoftext|>