<commit_before>extern crate getopts;
extern crate log;
extern crate time;

use std::net::SocketAddr;
use std::net::ToSocketAddrs;
use std::env;

use getopts::Options;

enum Protocol {
    ProtocolBuffer,
    Avro,
}

struct Settings {
    addr: SocketAddr,
    log_level: log::LogLevel,
    protocol: Protocol,
    db: String,
    config: Option<String>,
    profile: bool,
}

struct Bucket {
    name: String,
    perSecond: u64,
    purpose: Option<String>,
    size: u64,
    until: Option<time::Tm>,
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let mut opts = Options::new();
    opts.optopt("d", "db", "set output file name", "DATABASE");
    opts.optopt("p", "port", "Port to bind [9231]", "PORT");
    opts.optopt("h", "hostname", "Hostname to bind [0.0.0.0]", "HOSTNAME");
    opts.optopt("c", "config-file", "Configuration file", "CONFIG");
    opts.optflag("", "avro", "Use the avro protocol");
    opts.optflag("", "profile", "");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };

    let port: u16 = matches.opt_str("p").and_then(|x| x.parse().ok()).unwrap_or(9231);
    let settings = Settings {
        addr: (&*matches.opt_str("h").unwrap_or("0.0.0.0".to_owned()), port).to_socket_addrs().unwrap().next().unwrap(),
        log_level: log::LogLevel::Info,
        protocol: if matches.opt_present("avro") { Protocol::Avro } else { Protocol::ProtocolBuffer },
        db: matches.opt_str("d").unwrap(),
        config: matches.opt_str("c"),
        profile: matches.opt_present("profile"),
    };
}
<commit_msg>Allow custom log level<commit_after>extern crate getopts;
extern crate log;
extern crate time;

use std::net::SocketAddr;
use std::net::ToSocketAddrs;
use std::env;

use getopts::Options;
use log::LogLevel;

enum Protocol {
    ProtocolBuffer,
    Avro,
}

struct Settings {
    addr: SocketAddr,
    log_level: LogLevel,
    protocol: Protocol,
    db: String,
    config: Option<String>,
    profile: bool,
}

struct Bucket {
    name: String,
    perSecond: u64,
    purpose: Option<String>,
    size: u64,
    until: Option<time::Tm>,
}

fn main() {
    let args: Vec<String> = env::args().collect();

    let mut opts = Options::new();
    opts.optopt("d", "db", "set output file name", "DATABASE");
    opts.optopt("l", "log-level", "Log level [INFO]", "LOG-LEVEL");
    opts.optopt("p", "port", "Port to bind [9231]", "PORT");
    opts.optopt("h", "hostname", "Hostname to bind [0.0.0.0]", "HOSTNAME");
    opts.optopt("c", "config-file", "Configuration file", "CONFIG");
    opts.optflag("", "avro", "Use the avro protocol");
    opts.optflag("", "profile", "");

    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };

    let port: u16 = matches.opt_str("p").and_then(|x| x.parse().ok()).unwrap_or(9231);
    let settings = Settings {
        addr: (&*matches.opt_str("h").unwrap_or("0.0.0.0".to_owned()), port).to_socket_addrs().unwrap().next().unwrap(),
        protocol: if matches.opt_present("avro") { Protocol::Avro } else { Protocol::ProtocolBuffer },
        db: matches.opt_str("d").unwrap(),
        config: matches.opt_str("c"),
        profile: matches.opt_present("profile"),
        log_level: match &*matches.opt_str("l").unwrap_or("info".to_owned()) {
            "fatal" => LogLevel::Error,
            "error" => LogLevel::Error,
            "warn" => LogLevel::Warn,
            "debug" => LogLevel::Debug,
            "trace" => LogLevel::Trace,
            _ => LogLevel::Info,
        }
    };
}
<|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;

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

use websocket::{Message, Sender, Receiver, DataFrame};
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>,
}

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
            ::std::io::copy(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        })?;

        // Actual data transfer happens here
        ::std::io::copy(&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(),
        }
    )
}

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 peer = self.accept_client()?;
                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)
}

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())
    }
}

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 => 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.1")
        .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
  tcp:<socket address>               Connect to TCP
  l-tcp:<socket address>             Listen TCP connections
  -                                  stdin/stdout
  (more to be implemented)

Examples:
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to secure web socket once.
    Like netcat, but for websocket.
  websocat ws://localhost:1234/ tcp:localhost:1235
    Connect both to websocket and to TCP and exchange data.
  websocat l-tcp::::9559 ws://echo.websocket.org/
    Listen TCPv6 port 9959 on address :: and forward
    all connections to a public loopback websocket

Specify listening part first, unless you want websocat to serve once.
"#)
        .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>Move code a bit<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;

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

use websocket::{Message, Sender, Receiver, DataFrame};
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>,
}

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
            ::std::io::copy(&mut reader1, &mut writer2)?;
            writer2.write(b"")?; // signal close
            Ok(())
        })?;

        // Actual data transfer happens here
        ::std::io::copy(&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())
    }
}

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 peer = self.accept_client()?;
                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 => 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.1")
        .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
  tcp:<socket address>               Connect to TCP
  l-tcp:<socket address>             Listen TCP connections
  -                                  stdin/stdout
  (more to be implemented)

Examples:
  websocat - wss://myserver/mysocket
    Connect stdin/stdout to secure web socket once.
    Like netcat, but for websocket.
  websocat ws://localhost:1234/ tcp:localhost:1235
    Connect both to websocket and to TCP and exchange data.
  websocat l-tcp::::9559 ws://echo.websocket.org/
    Listen TCPv6 port 9959 on address :: and forward
    all connections to a public loopback websocket

Specify listening part first, unless you want websocat to serve once.
"#)
        .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|>


<commit_before>extern crate bencode;
extern crate rand;
extern crate bittorrent;

use std::{env, thread};
use std::thread::JoinHandle;
use std::sync::mpsc::{channel, Sender};
use std::sync::Arc;
use std::ops::Deref;
use bencode::{deserialize_file, Bencode};
use bittorrent::metadata::{MetadataDict, Metadata};
use bittorrent::bt_messages::Message;
use bittorrent::tracker::{get_http_tracker_peers, PEER_ID_PREFIX};
use bittorrent::peer::{connect_to_peer, gen_rand_peer_id};
use bittorrent::default_handler::{Handler, DefaultHandler};

/// Sets up a sink pool. it functions as an Actor
/// atm, rust doesn't support HKTs
fn init (mut handler: DefaultHandler) -> (Sender<(Message, Arc<String>)>, JoinHandle<()>) {
    let (tx, rx) = channel();
    let sink = thread::spawn(move|| {
        loop {
            let (message, cell): (Message, Arc<String>) = rx.recv().unwrap();
            handler.handle(message, cell.deref());
        }
    });
    (tx, sink)
}

/// Sets up a transmission based on a single torrent
fn init_torrent (tx: &Sender<(Message, Arc<String>)>, metadata: &Metadata, listen_port: u32, bytes_dled: u32) {
    let peer_id = gen_rand_peer_id(PEER_ID_PREFIX);
    let peers = match get_http_tracker_peers(&peer_id, metadata, listen_port, bytes_dled) {
        Some(peers) => peers,
        _ => panic!("cannot get peers from tracker")
    };

    println!("got {} peers", peers.len());

    for peer in peers {
        let child_meta = metadata.clone();
        let peer_id = peer_id.clone();
        let tx = tx.clone();
        thread::spawn(move || {
            match connect_to_peer(peer, &child_meta, &peer_id) {
                Ok((ref peer_id, ref mut reader)) => {
                    let peer_id_str = peer_id.iter().map(|x| *x as char).collect::<String>();
                    let c = Arc::new(peer_id_str);
                    loop {
                        match reader.wait_for_message() {
                            Ok(message) => {
                                //cannot figure out how to get the borrow checker to stop yelling
                                //at me about sending a string borrow... for now send an owned
                                //string. not even sure if its possible not to
                                let _ = tx.send((message, c.clone()));
                            },
                            Err(_) => {
                                println!("error waiting for message");
                            }
                        }
                    }
                },
                Err(e) => {
                    println!("{:?}", e);
                }
            };
        });
    }
}

fn main () {
    let path = env::args().nth(1)
                          .unwrap_or_else(||panic!("no path to torrent provided"));

    let content = deserialize_file(path).unwrap_or_else(||panic!("unable to parse bencoded metadata"));

    assert_eq!(content.len(), 1);

    let metadata = match content.first() {
        Some(&Bencode::Dict(ref dict)) => dict.to_metadata(),
        _ => panic!("no valid information in torrent file")
    }.unwrap();

    let (tx, sink) = init(DefaultHandler::new());

    //for now initialize torrents inline with main
    init_torrent(&tx, &metadata, 6887, 0);

    //block until the sink shuts down
    let _ = sink.join();
}
<commit_msg>improved concurrency pattern<commit_after>extern crate bencode;
extern crate bittorrent;

use std::{env, thread};
use std::thread::JoinHandle;
use std::sync::mpsc::{channel, Sender};
use std::sync::{Arc, Mutex};
use std::net::TcpStream;
use std::ops::{Deref, DerefMut};
use bencode::{deserialize_file, Bencode};
use bittorrent::metadata::{MetadataDict, Metadata};
use bittorrent::bt_messages::Message;
use bittorrent::buffered_reader::BufferedReader;
use bittorrent::tracker::{get_http_tracker_peers, PEER_ID_PREFIX};
use bittorrent::peer::{connect_to_peer, gen_rand_peer_id};
use bittorrent::default_handler::{Handler, DefaultHandler};

struct PeerState {
    id: String,
    chan: Sender<u32>
}

/// Sets up a sink pool. it functions as an Actor
/// atm, rust doesn't support HKTs
fn init <'a> (mut handler: DefaultHandler) -> (Sender<(Message, Arc<Mutex<PeerState>>)>, JoinHandle<()>) {
    let (tx, rx) = channel();
    let sink = thread::spawn(move|| {
        loop {
            let (message, cell): (Message, Arc<Mutex<PeerState>>) = rx.recv().unwrap();
            let f = cell.deref().lock().unwrap();
            let g = f.deref();
            g.chan.send(3);

            println!("{:?}", g.id);
            handler.handle(message);
        }
    });
    (tx, sink)
}

/// Sets up a transmission based on a single torrent
fn init_torrent (tx: &Sender<(Message, Arc<Mutex<PeerState>>)>, metadata: &Metadata, listen_port: u32, bytes_dled: u32) {
    let peer_id = gen_rand_peer_id(PEER_ID_PREFIX);
    let peers = match get_http_tracker_peers(&peer_id, metadata, listen_port, bytes_dled) {
        Some(peers) => peers,
        _ => panic!("cannot get peers from tracker")
    };

    println!("got {} peers", peers.len());

    for peer in peers {
        let child_meta = metadata.clone();
        let peer_id = peer_id.clone();
        let tx = tx.clone();
        thread::spawn(move || {
            match connect_to_peer(peer, &child_meta, &peer_id) {
                Ok((peer_id, mut reader)) => {
                    let peer_id_str = peer_id.iter().map(|x| *x as char).collect::<String>();
                    //requests need a response - use channels to accomplish this
                    let (btx, brx) = channel();
                    let arc = Arc::new(Mutex::new(PeerState {id: peer_id_str, chan: btx}));
                    loop {
                        //we can't just block in a loop - we'll never have a chance to send out
                        //outgoing messages over TCP
                        match reader.wait_for_message() {
                            Ok(message) => {
                                //this is an ask - expects a message
                                let _ = tx.send((message, arc.clone()));
                                match brx.recv().unwrap() {
                                    _ => println!("freed")
                                }
                            },
                            Err(_) => {
                                println!("error waiting for message");
                            }
                        };
                    }
                },
                Err(e) => {
                    println!("{:?}", e);
                }
            };
        });
    }
}

fn main () {
    let path = env::args().nth(1)
                          .unwrap_or_else(||panic!("no path to torrent provided"));

    let content = deserialize_file(path).unwrap_or_else(||panic!("unable to parse bencoded metadata"));

    assert_eq!(content.len(), 1);

    let metadata = match content.first() {
        Some(&Bencode::Dict(ref dict)) => dict.to_metadata(),
        _ => panic!("no valid information in torrent file")
    }.unwrap();

    let (tx, sink) = init(DefaultHandler::new());

    //for now initialize torrents inline with main
    init_torrent(&tx, &metadata, 6887, 0);

    //block until the sink shuts down
    let _ = sink.join();
}
<|endoftext|>


<commit_before>#![warn(trivial_casts, trivial_numeric_casts)]
#![warn(unused_qualifications)]
#![warn(unused_results)]

extern crate ansi_term;
extern crate datetime;
extern crate getopts;
extern crate libc;
extern crate locale;
extern crate natord;
extern crate num_cpus;
extern crate number_prefix;
extern crate scoped_threadpool;
extern crate term_grid;
extern crate unicode_width;
extern crate users;
extern crate zoneinfo_compiled;

#[cfg(feature="git")] extern crate git2;
#[macro_use] extern crate lazy_static;

use std::env;
use std::path::{Component, Path};
use std::process;

use dir::Dir;
use file::File;
use options::{Options, View};

mod dir;
mod feature;
mod file;
mod filetype;
mod options;
mod output;
mod term;

struct Exa {
    options: Options,
}

impl Exa {
    fn run(&mut self, mut args_file_names: Vec<String>) {
        let mut files = Vec::new();
        let mut dirs = Vec::new();

        if args_file_names.is_empty() {
            args_file_names.push(".".to_owned());
        }

        for file_name in args_file_names.iter() {
            match File::from_path(Path::new(&file_name), None) {
                Err(e) => {
                    println!("{}: {}", file_name, e);
                },
                Ok(f) => {
                    if f.is_directory() && !self.options.dir_action.treat_dirs_as_files() {
                        match f.to_dir(self.options.should_scan_for_git()) {
                            Ok(d) => dirs.push(d),
                            Err(e) => println!("{}: {}", file_name, e),
                        }
                    }
                    else {
                        files.push(f);
                    }
                },
            }
        }

        let no_files = files.is_empty();
        if !no_files {
            self.print_files(None, files);
        }

        let is_only_dir = dirs.len() == 1;
        self.print_dirs(dirs, no_files, is_only_dir);
    }

    fn print_dirs(&self, dir_files: Vec<Dir>, mut first: bool, is_only_dir: bool) {
        for dir in dir_files {

            // Put a gap between directories, or between the list of files and the
            // first directory.
            if first {
                first = false;
            }
            else {
                print!("\n");
            }

            if !is_only_dir {
                println!("{}:", dir.path.display());
            }

            let mut children = Vec::new();
            for file in dir.files() {
                match file {
                    Ok(file)       => children.push(file),
                    Err((path, e)) => println!("[{}: {}]", path.display(), e),
                }
            };

            self.options.filter.filter_files(&mut children);
            self.options.filter.sort_files(&mut children);

            if let Some(recurse_opts) = self.options.dir_action.recurse_options() {
                let depth = dir.path.components().filter(|&c| c != Component::CurDir).count() + 1;
                if !recurse_opts.tree && !recurse_opts.is_too_deep(depth) {

                    let mut child_dirs = Vec::new();
                    for child_dir in children.iter().filter(|f| f.is_directory()) {
                        match child_dir.to_dir(false) {
                            Ok(d)  => child_dirs.push(d),
                            Err(e) => println!("{}: {}", child_dir.path.display(), e),
                        }
                    }

                    self.print_files(Some(&dir), children);

                    if !child_dirs.is_empty() {
                        self.print_dirs(child_dirs, false, false);
                    }

                    continue;
                }
            }

            self.print_files(Some(&dir), children);

        }
    }

    fn print_files(&self, dir: Option<&Dir>, files: Vec<File>) {
        match self.options.view {
            View::Grid(g)         => g.view(&files),
            View::Details(d)      => d.view(dir, files),
            View::GridDetails(gd) => gd.view(dir, files),
            View::Lines(l)        => l.view(files),
        }
    }
}

fn main() {
    let args: Vec<String> = env::args().skip(1).collect();

    match Options::getopts(&args) {
        Ok((options, paths)) => {
            let mut exa = Exa { options: options };
            exa.run(paths);
        },
        Err(e) => {
            println!("{}", e);
            process::exit(e.error_code());
        },
    };
}
<commit_msg>Fix bug where the directory name was not printed<commit_after>#![warn(trivial_casts, trivial_numeric_casts)]
#![warn(unused_qualifications)]
#![warn(unused_results)]

extern crate ansi_term;
extern crate datetime;
extern crate getopts;
extern crate libc;
extern crate locale;
extern crate natord;
extern crate num_cpus;
extern crate number_prefix;
extern crate scoped_threadpool;
extern crate term_grid;
extern crate unicode_width;
extern crate users;
extern crate zoneinfo_compiled;

#[cfg(feature="git")] extern crate git2;
#[macro_use] extern crate lazy_static;

use std::env;
use std::path::{Component, Path};
use std::process;

use dir::Dir;
use file::File;
use options::{Options, View};

mod dir;
mod feature;
mod file;
mod filetype;
mod options;
mod output;
mod term;

struct Exa {
    options: Options,
}

impl Exa {
    fn run(&mut self, mut args_file_names: Vec<String>) {
        let mut files = Vec::new();
        let mut dirs = Vec::new();

        if args_file_names.is_empty() {
            args_file_names.push(".".to_owned());
        }

        for file_name in args_file_names.iter() {
            match File::from_path(Path::new(&file_name), None) {
                Err(e) => {
                    println!("{}: {}", file_name, e);
                },
                Ok(f) => {
                    if f.is_directory() && !self.options.dir_action.treat_dirs_as_files() {
                        match f.to_dir(self.options.should_scan_for_git()) {
                            Ok(d) => dirs.push(d),
                            Err(e) => println!("{}: {}", file_name, e),
                        }
                    }
                    else {
                        files.push(f);
                    }
                },
            }
        }

        let no_files = files.is_empty();
        let is_only_dir = dirs.len() == 1 && no_files;

        if !no_files {
            self.print_files(None, files);
        }

        self.print_dirs(dirs, no_files, is_only_dir);
    }

    fn print_dirs(&self, dir_files: Vec<Dir>, mut first: bool, is_only_dir: bool) {
        for dir in dir_files {

            // Put a gap between directories, or between the list of files and the
            // first directory.
            if first {
                first = false;
            }
            else {
                print!("\n");
            }

            if !is_only_dir {
                println!("{}:", dir.path.display());
            }

            let mut children = Vec::new();
            for file in dir.files() {
                match file {
                    Ok(file)       => children.push(file),
                    Err((path, e)) => println!("[{}: {}]", path.display(), e),
                }
            };

            self.options.filter.filter_files(&mut children);
            self.options.filter.sort_files(&mut children);

            if let Some(recurse_opts) = self.options.dir_action.recurse_options() {
                let depth = dir.path.components().filter(|&c| c != Component::CurDir).count() + 1;
                if !recurse_opts.tree && !recurse_opts.is_too_deep(depth) {

                    let mut child_dirs = Vec::new();
                    for child_dir in children.iter().filter(|f| f.is_directory()) {
                        match child_dir.to_dir(false) {
                            Ok(d)  => child_dirs.push(d),
                            Err(e) => println!("{}: {}", child_dir.path.display(), e),
                        }
                    }

                    self.print_files(Some(&dir), children);

                    if !child_dirs.is_empty() {
                        self.print_dirs(child_dirs, false, false);
                    }

                    continue;
                }
            }

            self.print_files(Some(&dir), children);

        }
    }

    fn print_files(&self, dir: Option<&Dir>, files: Vec<File>) {
        match self.options.view {
            View::Grid(g)         => g.view(&files),
            View::Details(d)      => d.view(dir, files),
            View::GridDetails(gd) => gd.view(dir, files),
            View::Lines(l)        => l.view(files),
        }
    }
}

fn main() {
    let args: Vec<String> = env::args().skip(1).collect();

    match Options::getopts(&args) {
        Ok((options, paths)) => {
            let mut exa = Exa { options: options };
            exa.run(paths);
        },
        Err(e) => {
            println!("{}", e);
            process::exit(e.error_code());
        },
    };
}
<|endoftext|>


<commit_before>#![feature(box_patterns,box_syntax)]
#![feature(scheme)]
#![feature(compile)]
#![feature(convert)]

//! Seax
//! ----
//!
//! Command-line application for [Seax](hawkweisman.me/seax), a VM-based
//! platform for executing programs in functional languages.

extern crate rustc_serialize;
extern crate docopt;
extern crate regex;

extern crate seax_svm as svm;
extern crate seax_scheme as scheme;

#[macro_use]
extern crate log;

use docopt::Docopt;
use regex::Regex;

use std::io;
use std::io::{Write, Read, BufRead,BufReader};
use std::error::Error;
use std::fs::File;
use std::path::PathBuf;
use std::convert::AsRef;

#[allow(dead_code)]
static USAGE: &'static str = "
Usage:
    seax repl [-vd]
    seax [-vd] <file>
    seax compile [-vd] file

Options:
    -v, --verbose   Enable verbose mode
    -d, --debug     Enable debug mode
";

#[derive(RustcDecodable)]
#[allow(dead_code)]
struct Args {
    cmd_repl: bool,
    cmd_compile: bool,
    arg_file: String,
    flag_verbose: bool,
    flag_debug: bool,
}

mod loggers;

#[allow(dead_code)]
fn main() {
    let args: Args = Docopt::new(USAGE)
                .and_then(|d| d.decode())
                .unwrap_or_else(|e| e.exit());

    let ext_re = Regex::new(r".+?(?P<ext>\.[^.]*$|$)").unwrap();

    if args.flag_verbose {
        let _ = log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Debug);
            Box::new(loggers::DebugLogger)
        });
    } else {
        let _ = log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Info);
            Box::new(loggers::DefaultLogger)
        });
    };

    if args.cmd_repl {
        let stdin = BufReader::new(io::stdin());
        let mut stdout = io::stdout();

        print!("scheme> ");
        let _ = stdout.flush();

        for line in stdin.lines() {
            match line.map_err(|error| String::from(error.description()) )
                .and_then(  |ref code| scheme::compile(code) )
                .and_then(  |program | svm::eval_program(program, args.flag_debug) ) {
                    Ok(result)  => println!("===> {:?}",result),
                    Err(why)    => error!("{}", why)
                };
            print!("scheme> ");
            let _ = stdout.flush();
        }
    } else if args.cmd_compile {
        unimplemented!()
    } else {
        match ext_re
            .captures(args.arg_file.as_ref())
            .and_then(|c| c.name("ext")) {
            Some(".scm")   => { // interpret scheme
                debug!("Interpreting Scheme file {}", args.arg_file);
                let path = PathBuf::from(args.arg_file.as_str());
                match File::open(&path)
                    .map_err(|error    | String::from(error.description()) )
                    .and_then(|mut file| {
                        let mut s = String::new();
                        file.read_to_string(&mut s).map(|_| s)
                            .map_err(|error| String::from(error.description()) ) })
                    .and_then(  |ref code| scheme::compile(code) )
                    .and_then(  |program | svm::eval_program(program, args.flag_debug) ) {
                        Ok(result)  => println!("===> {:?}",result),
                        Err(why)    => error!("{}", why)
                };
            },
            _              => {
                debug!("Executing binary {}", args.arg_file);
                unimplemented!()
            } // bin, figure out
        }
    }
}
<commit_msg>Add crate attributes<commit_after>#![crate_name = "seax"]
#![crate_type = "bin"]
#![feature(box_patterns,box_syntax)]
#![feature(scheme)]
#![feature(compile)]
#![feature(convert)]

//! Seax
//! ----
//!
//! Command-line application for [Seax](hawkweisman.me/seax), a VM-based
//! platform for executing programs in functional languages.

extern crate rustc_serialize;
extern crate docopt;
extern crate regex;

extern crate seax_svm as svm;
extern crate seax_scheme as scheme;

#[macro_use]
extern crate log;

use docopt::Docopt;
use regex::Regex;

use std::io;
use std::io::{Write, Read, BufRead,BufReader};
use std::error::Error;
use std::fs::File;
use std::path::PathBuf;
use std::convert::AsRef;

#[allow(dead_code)]
static USAGE: &'static str = "
Usage:
    seax repl [-vd]
    seax [-vd] <file>
    seax compile [-vd] file

Options:
    -v, --verbose   Enable verbose mode
    -d, --debug     Enable debug mode
";

#[derive(RustcDecodable)]
#[allow(dead_code)]
struct Args {
    cmd_repl: bool,
    cmd_compile: bool,
    arg_file: String,
    flag_verbose: bool,
    flag_debug: bool,
}

mod loggers;

#[allow(dead_code)]
fn main() {
    let args: Args = Docopt::new(USAGE)
                .and_then(|d| d.decode())
                .unwrap_or_else(|e| e.exit());

    let ext_re = Regex::new(r".+?(?P<ext>\.[^.]*$|$)").unwrap();

    if args.flag_verbose {
        let _ = log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Debug);
            Box::new(loggers::DebugLogger)
        });
    } else {
        let _ = log::set_logger(|max_log_level| {
            max_log_level.set(log::LogLevelFilter::Info);
            Box::new(loggers::DefaultLogger)
        });
    };

    if args.cmd_repl {
        let stdin = BufReader::new(io::stdin());
        let mut stdout = io::stdout();

        print!("scheme> ");
        let _ = stdout.flush();

        for line in stdin.lines() {
            match line.map_err(|error| String::from(error.description()) )
                .and_then(  |ref code| scheme::compile(code) )
                .and_then(  |program | svm::eval_program(program, args.flag_debug) ) {
                    Ok(result)  => println!("===> {:?}",result),
                    Err(why)    => error!("{}", why)
                };
            print!("scheme> ");
            let _ = stdout.flush();
        }
    } else if args.cmd_compile {
        unimplemented!()
    } else {
        match ext_re
            .captures(args.arg_file.as_ref())
            .and_then(|c| c.name("ext")) {
            Some(".scm")   => { // interpret scheme
                debug!("Interpreting Scheme file {}", args.arg_file);
                let path = PathBuf::from(args.arg_file.as_str());
                match File::open(&path)
                    .map_err(|error    | String::from(error.description()) )
                    .and_then(|mut file| {
                        let mut s = String::new();
                        file.read_to_string(&mut s).map(|_| s)
                            .map_err(|error| String::from(error.description()) ) })
                    .and_then(  |ref code| scheme::compile(code) )
                    .and_then(  |program | svm::eval_program(program, args.flag_debug) ) {
                        Ok(result)  => println!("===> {:?}",result),
                        Err(why)    => error!("{}", why)
                };
            },
            _              => {
                debug!("Executing binary {}", args.arg_file);
                unimplemented!()
            } // bin, figure out
        }
    }
}
<|endoftext|>


<commit_before>extern crate mmap;
extern crate time;
extern crate timely;
extern crate columnar;
extern crate docopt;

use docopt::Docopt;

use std::thread;

use timely::progress::timestamp::RootTimestamp;
use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::example_shared::*;
use timely::example_shared::operators::*;
use timely::communication::*;
use timely::communication::pact::Exchange;

use timely::networking::initialize_networking;
use timely::networking::initialize_networking_from_file;

use timely::drain::DrainExt;

use columnar::Columnar;

mod typedrw;
mod graphmap;
use graphmap::GraphMMap;

mod sorting;
use sorting::{SegmentList, radix_sort_32};

static USAGE: &'static str = "
Usage: pagerank <source> [options] [<arguments>...]

Options:
    -w <arg>, --workers <arg>    number of workers per process [default: 1]
    -p <arg>, --processid <arg>  identity of this process      [default: 0]
    -n <arg>, --processes <arg>  number of processes involved  [default: 1]
    -h <arg>, --hosts <arg>      file containing list of host:port for workers
";

#[derive(Copy, Clone, Debug)]
struct NodeRank { node: u32, rank: f32, }
impl Columnar for NodeRank { type Stack = Vec<NodeRank>; }

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    let source = args.get_str("<source>").to_owned();

    let workers: u64 = if let Ok(threads) = args.get_str("-w").parse() { threads }
                       else { panic!("invalid setting for --workers: {}", args.get_str("-t")) };
    let process_id: u64 = if let Ok(proc_id) = args.get_str("-p").parse() { proc_id }
                          else { panic!("invalid setting for --processid: {}", args.get_str("-p")) };
    let processes: u64 = if let Ok(processes) = args.get_str("-n").parse() { processes }
                         else { panic!("invalid setting for --processes: {}", args.get_str("-n")) };

    println!("Starting pagerank dataflow with");
    println!("\tworkers:\t{}", workers);
    println!("\tprocesses:\t{}", processes);
    println!("\tprocessid:\t{}", process_id);

    // vector holding communicators to use; one per local worker.
    if processes > 1 {
        println!("Initializing BinaryCommunicator");

        let hosts = args.get_str("-h");
        let communicators = if hosts != "" {
            initialize_networking_from_file(hosts, process_id, workers).ok().expect("error initializing networking")
        }
        else {
            let addresses = (0..processes).map(|index| format!("localhost:{}", 2101 + index).to_string()).collect();
            initialize_networking(addresses, process_id, workers).ok().expect("error initializing networking")
        };

        pagerank_spawn(communicators, source);
    }
    else if workers > 1 { pagerank_spawn(ProcessCommunicator::new_vector(workers), source); }
    else { pagerank_spawn(vec![ThreadCommunicator], source); };
}

fn pagerank_spawn<C>(communicators: Vec<C>, filename: String)
where C: Communicator+Send {
    let mut guards = Vec::new();
    let workers = communicators.len();
    for communicator in communicators.into_iter() {
        let filename = filename.clone();
        guards.push(thread::Builder::new().name(format!("timely worker {}", communicator.index()))
                                          .spawn(move || pagerank_thread(communicator, filename, workers))
                                          .unwrap());
    }

    for guard in guards { guard.join().unwrap(); }
}

// returns [src/peers] degrees, (dst, deg) pairs, and a list of [src/peers] endpoints
fn transpose(mut edges: Vec<Vec<(u32, u32)>>, peers: usize, nodes: usize) -> (Vec<u32>, Vec<(u32, u32)>, Vec<u32>)  {

    let mut deg = vec![0; (nodes as usize / peers) + 1];
    for list in &edges {
        for &(s, _) in list {
            deg[s as usize / peers] += 1;
        }
    }

    // let start = time::precise_time_s();
    radix_sort_32(&mut edges, &mut Vec::new(), &|&(_,d)| d);
    // println!("sorted in {}s", time::precise_time_s() - start);

    let mut rev = Vec::<(u32,u32)>::with_capacity(deg.len());
    let mut trn = Vec::with_capacity(edges.len() * 1024);
    for list in edges {
        for (s,d) in list {
            let len = rev.len();
            if (len == 0) || (rev[len-1].0 < d) {
                rev.push((d, 0u32));
            }

            let len = rev.len();
            rev[len-1].1 += 1;
            trn.push(s / peers as u32);
        }
    }

    return (deg, rev, trn);
}

// pagerank dataflow graph has a set of edges as input, and a binary vertex that for each epoch of
// received edges initiates an iterative subcomputation to compute the pagerank.

fn pagerank_thread<C>(communicator: C, filename: String, _workers: usize)
where C: Communicator {
    let index = communicator.index() as usize;
    let peers = communicator.peers() as usize;

    let mut root = GraphRoot::new(communicator);

    let start = time::precise_time_s();
    let mut going = start;

    let nodes = GraphMMap::new(&filename).nodes();

    let mut segments = SegmentList::new(1024); // list of edge segments

    let mut src = vec![];   // holds ranks
    let mut deg = vec![];   // holds source degrees
    let mut rev = vec![];   // holds (dst, deg) pairs
    let mut trn = vec![];   // holds transposed sources

    let mut accounted = 0.0;

    let mut input = root.subcomputation(|builder| {

        let (input, edges) = builder.new_input::<(u32, u32)>();
        let (cycle, ranks) = builder.loop_variable::<NodeRank>(RootTimestamp::new(20), Local(1));

        edges.binary_notify(&ranks,
                            Exchange::new(|x: &(u32,u32)| x.0 as u64),
                            Exchange::new(|x: &NodeRank| x.node as u64),
                            format!("pagerank"),
                            vec![RootTimestamp::new(0)],
                            move |input1, input2, output, notificator| {

            // receive incoming edges
            while let Some((_index, data)) = input1.pull() {
                segments.push(data.drain_temp());
            }

            while let Some((iter, _)) = notificator.next() {
                // if the very first iteration, prepare some stuff
                if iter.inner == 0 {
                    // println!("received edges in {}s", time::precise_time_s() - start);
                    // (deg, rev, trn) = transpose(srcs, dsts, peers);
                    let (a, b, c) = transpose(segments.finalize(), peers, nodes);
                    deg = a;
                    rev = b;
                    trn = c;
                    src = vec![0.0f32; deg.len()];
                    if index == 0 { println!("transformed; {} x {} shape", deg.len(), rev.len()); }
                }

                if iter.inner == 10 { going = time::precise_time_s(); }
                if iter.inner == 20 && index == 0 { println!("avg: {}", (time::precise_time_s() - going) / 10.0); println!("rcv: {}", accounted / 20.0); }

                let iter_start = time::precise_time_s();

                for s in 0..src.len() { src[s] = 0.15 + 0.85 * src[s] / deg[s] as f32; }

                let mut trn_slice = &trn[..];
                let mut rev_slice = &rev[..];
                while rev_slice.len() > 0 {
                    // TODO: session should just flush...
                    let mut session = output.session(&iter);
                    let next = std::cmp::min(100_000, rev_slice.len());
                    for &(dst, deg) in &rev_slice[..next] {
                        let mut accum = 0.0;
                        for &s in &trn_slice[..deg as usize] {
                            accum += src[s as usize];
                        }
                        trn_slice = &trn_slice[deg as usize..];
                        session.give(NodeRank { node: dst, rank: accum });
                    }
                    rev_slice = &rev_slice[next..];
                }

                for s in &mut src { *s = 0.0; }

                if index == 0 { println!("iter: {} took {}s", iter.inner, time::precise_time_s() - iter_start); }
            }

            let start = time::precise_time_s();
            while let Some((iter, data)) = input2.pull() {
                notificator.notify_at(&iter);
                for x in data.drain_temp() {
                    src[x.node as usize / peers] += x.rank;
                }
            }
            accounted += time::precise_time_s() - start;
        })
        .connect_loop(cycle);

        input
    });

    {
        let graph = GraphMMap::new(&filename);

        let mut edges = Vec::new();
        for node in 0..graph.nodes() {
            if node % peers == index {
                for dst in graph.edges(node) {
                    edges.push((node as u32, *dst as u32));
                }
                if edges.len() > 1024 {
                    input.send_at(0, edges.drain_temp());
                    root.step();
                }
            }
        }

        input.send_at(0, edges.drain_temp());
    }
    input.close();
    while root.step() { }

    if index == 0 { println!("elapsed: {}", time::precise_time_s() - start); }
}

    // .unary_notify(
    //     Exchange::new(move |x: &(u32, f32)| (workers * (index / workers)) as u64 + (x.0 as u64 % workers as u64)),
    //     format!("Aggregation"),
    //     vec![],
    //     move |input, output, iterator| {
    //         while let Some((iter, data)) = input.pull() {
    //             iterator.notify_at(&iter);
    //             for (node, rank) in data.drain_temp() {
    //                 acc[node as usize / workers] += rank;
    //             }
    //         }
    //
    //         while let Some((item, _)) = iterator.next() {
    //
    //             output.give_at(&item, acc.drain_temp().enumerate().filter(|x| x.1 != 0.0)
    //                                      .map(|(u,f)| (((u * workers + local_index) as u32), f)));
    //
    //             for _ in 0..(1 + (nodes/workers)) { acc.push(0.0); }
    //             assert!(acc.len() == (1 + (nodes/workers)));
    //         }
    //     }
    // )
<commit_msg>added recv time accounting<commit_after>extern crate mmap;
extern crate time;
extern crate timely;
extern crate columnar;
extern crate docopt;

use docopt::Docopt;

use std::thread;

use timely::progress::timestamp::RootTimestamp;
use timely::progress::scope::Scope;
use timely::progress::nested::Summary::Local;
use timely::example_shared::*;
use timely::example_shared::operators::*;
use timely::communication::*;
use timely::communication::pact::Exchange;

use timely::networking::initialize_networking;
use timely::networking::initialize_networking_from_file;

use timely::drain::DrainExt;

use columnar::Columnar;

mod typedrw;
mod graphmap;
use graphmap::GraphMMap;

mod sorting;
use sorting::{SegmentList, radix_sort_32};

static USAGE: &'static str = "
Usage: pagerank <source> [options] [<arguments>...]

Options:
    -w <arg>, --workers <arg>    number of workers per process [default: 1]
    -p <arg>, --processid <arg>  identity of this process      [default: 0]
    -n <arg>, --processes <arg>  number of processes involved  [default: 1]
    -h <arg>, --hosts <arg>      file containing list of host:port for workers
";

#[derive(Copy, Clone, Debug)]
struct NodeRank { node: u32, rank: f32, }
impl Columnar for NodeRank { type Stack = Vec<NodeRank>; }

fn main () {
    let args = Docopt::new(USAGE).and_then(|dopt| dopt.parse()).unwrap_or_else(|e| e.exit());

    let source = args.get_str("<source>").to_owned();

    let workers: u64 = if let Ok(threads) = args.get_str("-w").parse() { threads }
                       else { panic!("invalid setting for --workers: {}", args.get_str("-t")) };
    let process_id: u64 = if let Ok(proc_id) = args.get_str("-p").parse() { proc_id }
                          else { panic!("invalid setting for --processid: {}", args.get_str("-p")) };
    let processes: u64 = if let Ok(processes) = args.get_str("-n").parse() { processes }
                         else { panic!("invalid setting for --processes: {}", args.get_str("-n")) };

    println!("Starting pagerank dataflow with");
    println!("\tworkers:\t{}", workers);
    println!("\tprocesses:\t{}", processes);
    println!("\tprocessid:\t{}", process_id);

    // vector holding communicators to use; one per local worker.
    if processes > 1 {
        println!("Initializing BinaryCommunicator");

        let hosts = args.get_str("-h");
        let communicators = if hosts != "" {
            initialize_networking_from_file(hosts, process_id, workers).ok().expect("error initializing networking")
        }
        else {
            let addresses = (0..processes).map(|index| format!("localhost:{}", 2101 + index).to_string()).collect();
            initialize_networking(addresses, process_id, workers).ok().expect("error initializing networking")
        };

        pagerank_spawn(communicators, source);
    }
    else if workers > 1 { pagerank_spawn(ProcessCommunicator::new_vector(workers), source); }
    else { pagerank_spawn(vec![ThreadCommunicator], source); };
}

fn pagerank_spawn<C>(communicators: Vec<C>, filename: String)
where C: Communicator+Send {
    let mut guards = Vec::new();
    let workers = communicators.len();
    for communicator in communicators.into_iter() {
        let filename = filename.clone();
        guards.push(thread::Builder::new().name(format!("timely worker {}", communicator.index()))
                                          .spawn(move || pagerank_thread(communicator, filename, workers))
                                          .unwrap());
    }

    for guard in guards { guard.join().unwrap(); }
}

// returns [src/peers] degrees, (dst, deg) pairs, and a list of [src/peers] endpoints
fn transpose(mut edges: Vec<Vec<(u32, u32)>>, peers: usize, nodes: usize) -> (Vec<u32>, Vec<(u32, u32)>, Vec<u32>)  {

    let mut deg = vec![0; (nodes as usize / peers) + 1];
    for list in &edges {
        for &(s, _) in list {
            deg[s as usize / peers] += 1;
        }
    }

    // let start = time::precise_time_s();
    radix_sort_32(&mut edges, &mut Vec::new(), &|&(_,d)| d);
    // println!("sorted in {}s", time::precise_time_s() - start);

    let mut rev = Vec::<(u32,u32)>::with_capacity(deg.len());
    let mut trn = Vec::with_capacity(edges.len() * 1024);
    for list in edges {
        for (s,d) in list {
            let len = rev.len();
            if (len == 0) || (rev[len-1].0 < d) {
                rev.push((d, 0u32));
            }

            let len = rev.len();
            rev[len-1].1 += 1;
            trn.push(s / peers as u32);
        }
    }

    return (deg, rev, trn);
}

// pagerank dataflow graph has a set of edges as input, and a binary vertex that for each epoch of
// received edges initiates an iterative subcomputation to compute the pagerank.

fn pagerank_thread<C>(communicator: C, filename: String, _workers: usize)
where C: Communicator {
    let index = communicator.index() as usize;
    let peers = communicator.peers() as usize;

    let mut root = GraphRoot::new(communicator);

    let start = time::precise_time_s();
    let mut going = start;

    let nodes = GraphMMap::new(&filename).nodes();

    let mut segments = SegmentList::new(1024); // list of edge segments

    let mut src = vec![];   // holds ranks
    let mut deg = vec![];   // holds source degrees
    let mut rev = vec![];   // holds (dst, deg) pairs
    let mut trn = vec![];   // holds transposed sources

    let mut accounted = 0.0;
    let mut accounted_inner = 0.0;

    let mut input = root.subcomputation(|builder| {

        let (input, edges) = builder.new_input::<(u32, u32)>();
        let (cycle, ranks) = builder.loop_variable::<NodeRank>(RootTimestamp::new(20), Local(1));

        edges.binary_notify(&ranks,
                            Exchange::new(|x: &(u32,u32)| x.0 as u64),
                            Exchange::new(|x: &NodeRank| x.node as u64),
                            format!("pagerank"),
                            vec![RootTimestamp::new(0)],
                            move |input1, input2, output, notificator| {

            // receive incoming edges
            while let Some((_index, data)) = input1.pull() {
                segments.push(data.drain_temp());
            }

            while let Some((iter, _)) = notificator.next() {
                // if the very first iteration, prepare some stuff
                if iter.inner == 0 {
                    // println!("received edges in {}s", time::precise_time_s() - start);
                    // (deg, rev, trn) = transpose(srcs, dsts, peers);
                    let (a, b, c) = transpose(segments.finalize(), peers, nodes);
                    deg = a;
                    rev = b;
                    trn = c;
                    src = vec![0.0f32; deg.len()];
                    if index == 0 { println!("transformed; {} x {} shape", deg.len(), rev.len()); }
                }

                if iter.inner == 10 { going = time::precise_time_s(); }
                if iter.inner == 20 && index == 0 { println!("avg: {}", (time::precise_time_s() - going) / 10.0);
                    println!("rcv: {}, {}", accounted_inner / 20.0, accounted / 20.0); }

                let iter_start = time::precise_time_s();

                for s in 0..src.len() { src[s] = 0.15 + 0.85 * src[s] / deg[s] as f32; }

                let mut trn_slice = &trn[..];
                let mut rev_slice = &rev[..];
                while rev_slice.len() > 0 {
                    // TODO: session should just flush...
                    let mut session = output.session(&iter);
                    let next = std::cmp::min(100_000, rev_slice.len());
                    for &(dst, deg) in &rev_slice[..next] {
                        let mut accum = 0.0;
                        for &s in &trn_slice[..deg as usize] {
                            accum += src[s as usize];
                        }
                        trn_slice = &trn_slice[deg as usize..];
                        session.give(NodeRank { node: dst, rank: accum });
                    }
                    rev_slice = &rev_slice[next..];
                }

                for s in &mut src { *s = 0.0; }

                if index == 0 { println!("iter: {} took {}s", iter.inner, time::precise_time_s() - iter_start); }
            }

            let start = time::precise_time_s();
            while let Some((iter, data)) = input2.pull() {
                let start_inner = time::precise_time_s();
                notificator.notify_at(&iter);
                for x in data.drain_temp() {
                    src[x.node as usize / peers] += x.rank;
                }
                accounted_inner += time::precise_time_s() -start_inner;
            }
            accounted += time::precise_time_s() - start;
        })
        .connect_loop(cycle);

        input
    });

    {
        let graph = GraphMMap::new(&filename);

        let mut edges = Vec::new();
        for node in 0..graph.nodes() {
            if node % peers == index {
                for dst in graph.edges(node) {
                    edges.push((node as u32, *dst as u32));
                }
                if edges.len() > 1024 {
                    input.send_at(0, edges.drain_temp());
                    root.step();
                }
            }
        }

        input.send_at(0, edges.drain_temp());
    }
    input.close();
    while root.step() { }

    if index == 0 { println!("elapsed: {}", time::precise_time_s() - start); }
}

    // .unary_notify(
    //     Exchange::new(move |x: &(u32, f32)| (workers * (index / workers)) as u64 + (x.0 as u64 % workers as u64)),
    //     format!("Aggregation"),
    //     vec![],
    //     move |input, output, iterator| {
    //         while let Some((iter, data)) = input.pull() {
    //             iterator.notify_at(&iter);
    //             for (node, rank) in data.drain_temp() {
    //                 acc[node as usize / workers] += rank;
    //             }
    //         }
    //
    //         while let Some((item, _)) = iterator.next() {
    //
    //             output.give_at(&item, acc.drain_temp().enumerate().filter(|x| x.1 != 0.0)
    //                                      .map(|(u,f)| (((u * workers + local_index) as u32), f)));
    //
    //             for _ in 0..(1 + (nodes/workers)) { acc.push(0.0); }
    //             assert!(acc.len() == (1 + (nodes/workers)));
    //         }
    //     }
    // )
<|endoftext|>


<commit_before>use std::io::prelude::*;
use std::io::stdin;
use std::error::Error;
use std::env;

/// Joins two blocks of text, such that row 1 of the first block will be
/// appended to row 1 of the second block and so on.
///
/// Input is read from stdin until EOF, output to stdout. The the first half of
/// the input lines will become the first block.
pub fn main() {
    run().unwrap();
}

fn run() -> Result<(), Box<Error>> {
    let separator: String = env::args().nth(1).unwrap_or("".to_string());
    let mut buffer = String::new();
    try!(stdin().read_to_string(&mut buffer));
    let lines: Vec<&str> = buffer.lines().collect();
    let lines_per_block = lines.len() / 2;
    let first = &lines[0..lines_per_block];
    let second = &lines[lines_per_block..];
    for (&left, &right) in first.iter().zip(second.iter()) {
        println!("{}{}{}", left, separator, right);
    }
    Ok(())
}
<commit_msg>Turn doc into module doc<commit_after>//! Joins two blocks of text, such that row 1 of the first block will be
//! appended to row 1 of the second block and so on.
//!
//! Input is read from stdin until EOF, output to stdout. The the first half of
//! the input lines will become the first block.

use std::io::prelude::*;
use std::io::stdin;
use std::error::Error;
use std::env;

fn main() {
    run().unwrap();
}

fn run() -> Result<(), Box<Error>> {
    let separator: String = env::args().nth(1).unwrap_or("".to_string());
    let mut buffer = String::new();
    try!(stdin().read_to_string(&mut buffer));
    let lines: Vec<&str> = buffer.lines().collect();
    let lines_per_block = lines.len() / 2;
    let first = &lines[0..lines_per_block];
    let second = &lines[lines_per_block..];
    for (&left, &right) in first.iter().zip(second.iter()) {
        println!("{}{}{}", left, separator, right);
    }
    Ok(())
}
<|endoftext|>


<commit_before>extern crate fdupe;
extern crate getopts;

use getopts::Options;
use std::env;
use std::path::Path;

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = Options::new();
    opts.optopt("n", "num", "set number of edit distace", "NUM");
    opts.optopt("d", "dir", "set directory where to look for duplicate files", "DIRECTORY");
    opts.optopt("t", "filter", "set the file name regex filter in the directory", "FILTER");
    opts.optflag("h", "help", "print this help menu");
    if args.len() == 1 {
        print_usage(&program, opts);
        return;
    }
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }
    let dir = match matches.opt_str("d") {
        Some(d) => {
            {
                let path = Path::new(&d);
                if !path.is_dir() {
                    panic!("{} is not a directory", d)
                }
            }
            d
        }
        None => { panic!("-d is a required option") }
    };
    let n_ed = match matches.opt_str("n") {
        Some(n) => {
            let n = match n.parse::<i32>() {
                Ok(n) => { n }
                Err(f) => { panic!(f.to_string()) }
            };
            if n < 0 {
                panic!("number of edit distance must be >= 0, got {}", n);
            }
            n
        }
        None => { panic!("-n is a required option") }
    };
    let filter = matches.opt_str("t");
    fdupe::fdupe::find_duplicates(&dir, n_ed, filter);
}
<commit_msg>Add err_and_exit function<commit_after>extern crate fdupe;
extern crate getopts;

use std::env;
use std::process;
use std::path::Path;
use getopts::Options;

fn print_usage(program: &str, opts: Options) {
    let brief = format!("Usage: {} [options]", program);
    print!("{}", opts.usage(&brief));
}

fn err_and_exit(message: &str) {
    println!("Error: {}", message);
    process::exit(1);
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();

    let mut opts = Options::new();
    opts.optopt("n", "num", "set number of edit distace", "NUM");
    opts.optopt("d", "dir", "set directory where to look for duplicate files", "DIRECTORY");
    opts.optopt("t", "filter", "set the file name regex filter in the directory", "FILTER");
    opts.optflag("h", "help", "print this help menu");
    if args.len() == 1 {
        print_usage(&program, opts);
        return;
    }
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => { m }
        Err(f) => { panic!(f.to_string()) }
    };
    if matches.opt_present("h") {
        print_usage(&program, opts);
        return;
    }
    let dir = match matches.opt_str("d") {
        Some(d) => {
            {
                let path = Path::new(&d);
                if !path.is_dir() {
                    panic!("{} is not a directory", d)
                }
            }
            d
        }
        None => { panic!("-d is a required option") }
    };
    let n_ed = match matches.opt_str("n") {
        Some(n) => {
            let n = match n.parse::<i32>() {
                Ok(n) => { n }
                Err(f) => { panic!(f.to_string()) }
            };
            if n < 0 {
                panic!("number of edit distance must be >= 0, got {}", n);
            }
            n
        }
        None => { panic!("-n is a required option") }
    };
    let filter = matches.opt_str("t");
    fdupe::fdupe::find_duplicates(&dir, n_ed, filter);
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
#[macro_use]
extern crate lazy_static;

extern crate libc;
extern crate scmp;
extern crate tempdir;

use std::fs;
use std::env;
use std::ffi::OsStr;
use std::ffi::CString;
use std::io::Error;
use std::os::raw::c_void;
use std::os::unix::process::CommandExt;
use std::os::unix::ffi::OsStrExt;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;
use scmp::*;

unsafe fn allow_syscall(ctx : *mut c_void, nr : u32) {
    assert!(seccomp_rule_add_array(ctx, SCMP_ACT_ALLOW, nr as i32, 0, std::ptr::null()) == 0);
}

fn setup_seccomp() -> Result<(), Error> {
    unsafe {
        let ctx0 = seccomp_init(SCMP_ACT_TRAP);
        assert!(!ctx0.is_null());
        allow_syscall(ctx0, __NR_execve);
        allow_syscall(ctx0, __NR_brk);
        allow_syscall(ctx0, __NR_access);
        allow_syscall(ctx0, __NR_mmap);
        allow_syscall(ctx0, __NR_open);
        allow_syscall(ctx0, __NR_read);
        allow_syscall(ctx0, __NR_write);
        allow_syscall(ctx0, __NR_stat);
        allow_syscall(ctx0, __NR_fstat);
        allow_syscall(ctx0, __NR_close);
        allow_syscall(ctx0, __NR_mprotect);
        allow_syscall(ctx0, __NR_arch_prctl);
        allow_syscall(ctx0, __NR_munmap);
        allow_syscall(ctx0, __NR_set_tid_address);
        allow_syscall(ctx0, __NR_set_robust_list);
        allow_syscall(ctx0, __NR_rt_sigaction);
        allow_syscall(ctx0, __NR_rt_sigprocmask);
        allow_syscall(ctx0, __NR_getrlimit);
        allow_syscall(ctx0, __NR_statfs);
        allow_syscall(ctx0, __NR_fcntl);
        allow_syscall(ctx0, __NR_geteuid);
        allow_syscall(ctx0, __NR_umask);
        allow_syscall(ctx0, __NR_pipe);
        allow_syscall(ctx0, __NR_clone);
        allow_syscall(ctx0, __NR_wait4);
        allow_syscall(ctx0, __NR_exit_group);
        allow_syscall(ctx0, __NR_dup);
        allow_syscall(ctx0, __NR_ioctl);
        allow_syscall(ctx0, __NR_getpid);
        allow_syscall(ctx0, __NR_gettid);
        allow_syscall(ctx0, __NR_tgkill);
        allow_syscall(ctx0, __NR_rt_sigreturn);
        allow_syscall(ctx0, __NR_socket);
        allow_syscall(ctx0, __NR_connect);
        allow_syscall(ctx0, __NR_lseek);
        allow_syscall(ctx0, __NR_openat);
        allow_syscall(ctx0, __NR_unlinkat);
        allow_syscall(ctx0, __NR_utimensat);
        allow_syscall(ctx0, __NR_lstat);
        allow_syscall(ctx0, __NR_fchmod);
        allow_syscall(ctx0, __NR_unlink);
        allow_syscall(ctx0, __NR_utime);
        allow_syscall(ctx0, __NR_futex);
        allow_syscall(ctx0, __NR_chmod);
        allow_syscall(ctx0, __NR_getuid);
        allow_syscall(ctx0, __NR_getgid);
        allow_syscall(ctx0, __NR_getppid);
        allow_syscall(ctx0, __NR_getcwd);
        allow_syscall(ctx0, __NR_getegid);
        allow_syscall(ctx0, __NR_getdents);
        allow_syscall(ctx0, __NR_readlink);
        allow_syscall(ctx0, __NR_mkdirat);
        allow_syscall(ctx0, __NR_mkdir);
        assert!(seccomp_load(ctx0) == 0);
    }
    return Ok(());
}

const RENAME_NOREPLACE : u64 = 1;

fn path_to_cstring(path : &Path) -> CString {
    return CString::new(path.as_os_str().as_bytes()).unwrap();
}

fn rename_noreplace(path0 : &Path, path1 : &Path) -> Result<(), Error> {
    unsafe {
        if libc::syscall(__NR_renameat2 as i64, libc::AT_FDCWD, path_to_cstring(path0).as_ptr(), libc::AT_FDCWD, path_to_cstring(path1).as_ptr(), RENAME_NOREPLACE, 0) == 0 {
            return Ok(());
        } else {
            return Err(Error::last_os_error());
        }
    }
}

fn rename_or_suffix(path0 : &Path, path1 : &Path) -> Result<PathBuf, Error> {
    let rv = rename_noreplace(path0, path1);
    if let Err(err) = rv {
        if err.raw_os_error().expect("Error fails to wrap errno") != libc::EEXIST {
            return Err(err);
        }
        for suffix in 1.. {
            let mut builder = path1.as_os_str().to_os_string();
            builder.push(format!(".{}", suffix));
            let rv = rename_noreplace(path0, Path::new(&builder));
            if let Err(err) = rv {
                if err.raw_os_error().expect("Error fails to wrap errno") != libc::EEXIST {
                    return Err(err);
                }
                continue;
            }
            return Ok(Path::new(&builder).to_path_buf());
        }
        return Ok(PathBuf::new()); // XXX Unreachable
    } else {
        return Ok(path1.to_path_buf());
    }
}

struct ArchiveType {
    extensions: Vec<&'static str>,
    extract_cmd: Vec<&'static str>,
}

lazy_static! {
    static ref ARCHIVE_TYPES : Vec<ArchiveType> = vec![
        ArchiveType {
            extensions: vec![
                ".tgz",
                ".tar.gz",
                ".tar.Z",
                ".tbz2",
                ".tar.bz2",
                ".tlz",
                ".tar.lz",
                ".tar.lzma",
                ".tar.lzo",
                ".tar.xz",
            ],
            extract_cmd: vec!["tar", "-xf"],
        },
        ArchiveType {
            extensions: vec![
                ".zip",
            ],
            extract_cmd: vec!["unzip", "--"],
        },
        ArchiveType {
            extensions: vec![
                ".part1.rar",
                ".rar",
            ],
            extract_cmd: vec!["unrar", "x", "--"],
        },
        ArchiveType {
            extensions: vec![
                ".7z",
            ],
            extract_cmd: vec!["7z", "x", "--"],
        },
    ];
}

fn main() {
    let matches = clap_app!(myapp =>
        (@arg ARCHIVE: +required ... "Archives to extract")
    ).get_matches();

    let mut found = false;

    for arch in matches.values_of_os("ARCHIVE").unwrap() {
        let arch = Path::new(arch);
        let arch_bytes = arch.as_os_str().as_bytes();
        println!("Extract {:?}", arch);
        for at in ARCHIVE_TYPES.iter() {
            for &ext in &at.extensions {
                let ext_bytes = ext.as_bytes();
                if !arch_bytes.ends_with(ext_bytes) {
                    continue;
                }
                found = true;

                let noext : &OsStr = OsStrExt::from_bytes(&arch_bytes[0..arch_bytes.len()-ext_bytes.len()]);
                let stripped =
                    if let Some(noext) = Path::new(noext).file_name() {
                        noext
                    } else {
                        OsStr::new("extracted")
                    };
                let cmd = at.extract_cmd[0];
                let tmpdir = tempdir::TempDir::new_in(
                    env::current_dir().expect("Unable to get the current directory"),
                    stripped.to_string_lossy().to_mut().as_str()).expect("Unable to create temporary directory");
                let status = Command::new(cmd)
                    .args(&at.extract_cmd[1..])
                    .arg(fs::canonicalize(&arch).expect("Unable to find archive"))
                    .current_dir(&tmpdir)
                    .before_exec(setup_seccomp)
                    .status()
                    .expect(format!("Failed to launch {}", cmd).as_str());
                assert!(status.success(), format!("{} wasn't successful", cmd));
                let extracted_path = tmpdir.into_path();
                let mut dir_list = fs::read_dir(&extracted_path).expect("Unable to read the extraction directory");
                if let Some(path) = dir_list.next() {
                    let path = path.expect("Unable to iterate the extraction directory");
                    if let Some(path1) = dir_list.next() {
                        path1.expect("Unable to iterate the extraction directory");
                        let newname = rename_or_suffix(&extracted_path, Path::new(stripped)).expect("Failed to rename");
                        println!("Extracted to {:?}", newname);
                    } else {
                        let newname = rename_or_suffix(&path.path(), Path::new(stripped)).expect("Failed to rename");
                        println!("Extracted to {:?}", newname);
                        fs::remove_dir(&extracted_path).expect("Failed to remove empty extraction directory");
                    }
                } else {
                    println!("No files were extracted");
                }
                break;
            }
        }
        if !found {
            println!("Can't handle {:?}", arch);
        }
    }
}

<commit_msg>Add madvise to the whitelist<commit_after>#[macro_use]
extern crate clap;
#[macro_use]
extern crate lazy_static;

extern crate libc;
extern crate scmp;
extern crate tempdir;

use std::fs;
use std::env;
use std::ffi::OsStr;
use std::ffi::CString;
use std::io::Error;
use std::os::raw::c_void;
use std::os::unix::process::CommandExt;
use std::os::unix::ffi::OsStrExt;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;
use scmp::*;

unsafe fn allow_syscall(ctx : *mut c_void, nr : u32) {
    assert!(seccomp_rule_add_array(ctx, SCMP_ACT_ALLOW, nr as i32, 0, std::ptr::null()) == 0);
}

fn setup_seccomp() -> Result<(), Error> {
    unsafe {
        let ctx0 = seccomp_init(SCMP_ACT_TRAP);
        assert!(!ctx0.is_null());
        allow_syscall(ctx0, __NR_execve);
        allow_syscall(ctx0, __NR_brk);
        allow_syscall(ctx0, __NR_access);
        allow_syscall(ctx0, __NR_mmap);
        allow_syscall(ctx0, __NR_open);
        allow_syscall(ctx0, __NR_read);
        allow_syscall(ctx0, __NR_write);
        allow_syscall(ctx0, __NR_stat);
        allow_syscall(ctx0, __NR_fstat);
        allow_syscall(ctx0, __NR_close);
        allow_syscall(ctx0, __NR_mprotect);
        allow_syscall(ctx0, __NR_arch_prctl);
        allow_syscall(ctx0, __NR_munmap);
        allow_syscall(ctx0, __NR_set_tid_address);
        allow_syscall(ctx0, __NR_set_robust_list);
        allow_syscall(ctx0, __NR_rt_sigaction);
        allow_syscall(ctx0, __NR_rt_sigprocmask);
        allow_syscall(ctx0, __NR_getrlimit);
        allow_syscall(ctx0, __NR_statfs);
        allow_syscall(ctx0, __NR_fcntl);
        allow_syscall(ctx0, __NR_geteuid);
        allow_syscall(ctx0, __NR_umask);
        allow_syscall(ctx0, __NR_pipe);
        allow_syscall(ctx0, __NR_clone);
        allow_syscall(ctx0, __NR_wait4);
        allow_syscall(ctx0, __NR_exit_group);
        allow_syscall(ctx0, __NR_dup);
        allow_syscall(ctx0, __NR_ioctl);
        allow_syscall(ctx0, __NR_getpid);
        allow_syscall(ctx0, __NR_gettid);
        allow_syscall(ctx0, __NR_tgkill);
        allow_syscall(ctx0, __NR_rt_sigreturn);
        allow_syscall(ctx0, __NR_socket);
        allow_syscall(ctx0, __NR_connect);
        allow_syscall(ctx0, __NR_lseek);
        allow_syscall(ctx0, __NR_openat);
        allow_syscall(ctx0, __NR_unlinkat);
        allow_syscall(ctx0, __NR_utimensat);
        allow_syscall(ctx0, __NR_lstat);
        allow_syscall(ctx0, __NR_fchmod);
        allow_syscall(ctx0, __NR_unlink);
        allow_syscall(ctx0, __NR_utime);
        allow_syscall(ctx0, __NR_futex);
        allow_syscall(ctx0, __NR_chmod);
        allow_syscall(ctx0, __NR_getuid);
        allow_syscall(ctx0, __NR_getgid);
        allow_syscall(ctx0, __NR_getppid);
        allow_syscall(ctx0, __NR_getcwd);
        allow_syscall(ctx0, __NR_getegid);
        allow_syscall(ctx0, __NR_getdents);
        allow_syscall(ctx0, __NR_readlink);
        allow_syscall(ctx0, __NR_mkdirat);
        allow_syscall(ctx0, __NR_mkdir);
        allow_syscall(ctx0, __NR_madvise);
        assert!(seccomp_load(ctx0) == 0);
    }
    return Ok(());
}

const RENAME_NOREPLACE : u64 = 1;

fn path_to_cstring(path : &Path) -> CString {
    return CString::new(path.as_os_str().as_bytes()).unwrap();
}

fn rename_noreplace(path0 : &Path, path1 : &Path) -> Result<(), Error> {
    unsafe {
        if libc::syscall(__NR_renameat2 as i64, libc::AT_FDCWD, path_to_cstring(path0).as_ptr(), libc::AT_FDCWD, path_to_cstring(path1).as_ptr(), RENAME_NOREPLACE, 0) == 0 {
            return Ok(());
        } else {
            return Err(Error::last_os_error());
        }
    }
}

fn rename_or_suffix(path0 : &Path, path1 : &Path) -> Result<PathBuf, Error> {
    let rv = rename_noreplace(path0, path1);
    if let Err(err) = rv {
        if err.raw_os_error().expect("Error fails to wrap errno") != libc::EEXIST {
            return Err(err);
        }
        for suffix in 1.. {
            let mut builder = path1.as_os_str().to_os_string();
            builder.push(format!(".{}", suffix));
            let rv = rename_noreplace(path0, Path::new(&builder));
            if let Err(err) = rv {
                if err.raw_os_error().expect("Error fails to wrap errno") != libc::EEXIST {
                    return Err(err);
                }
                continue;
            }
            return Ok(Path::new(&builder).to_path_buf());
        }
        return Ok(PathBuf::new()); // XXX Unreachable
    } else {
        return Ok(path1.to_path_buf());
    }
}

struct ArchiveType {
    extensions: Vec<&'static str>,
    extract_cmd: Vec<&'static str>,
}

lazy_static! {
    static ref ARCHIVE_TYPES : Vec<ArchiveType> = vec![
        ArchiveType {
            extensions: vec![
                ".tgz",
                ".tar.gz",
                ".tar.Z",
                ".tbz2",
                ".tar.bz2",
                ".tlz",
                ".tar.lz",
                ".tar.lzma",
                ".tar.lzo",
                ".tar.xz",
            ],
            extract_cmd: vec!["tar", "-xf"],
        },
        ArchiveType {
            extensions: vec![
                ".zip",
            ],
            extract_cmd: vec!["unzip", "--"],
        },
        ArchiveType {
            extensions: vec![
                ".part1.rar",
                ".rar",
            ],
            extract_cmd: vec!["unrar", "x", "--"],
        },
        ArchiveType {
            extensions: vec![
                ".7z",
            ],
            extract_cmd: vec!["7z", "x", "--"],
        },
    ];
}

fn main() {
    let matches = clap_app!(myapp =>
        (@arg ARCHIVE: +required ... "Archives to extract")
    ).get_matches();

    let mut found = false;

    for arch in matches.values_of_os("ARCHIVE").unwrap() {
        let arch = Path::new(arch);
        let arch_bytes = arch.as_os_str().as_bytes();
        println!("Extract {:?}", arch);
        for at in ARCHIVE_TYPES.iter() {
            for &ext in &at.extensions {
                let ext_bytes = ext.as_bytes();
                if !arch_bytes.ends_with(ext_bytes) {
                    continue;
                }
                found = true;

                let noext : &OsStr = OsStrExt::from_bytes(&arch_bytes[0..arch_bytes.len()-ext_bytes.len()]);
                let stripped =
                    if let Some(noext) = Path::new(noext).file_name() {
                        noext
                    } else {
                        OsStr::new("extracted")
                    };
                let cmd = at.extract_cmd[0];
                let tmpdir = tempdir::TempDir::new_in(
                    env::current_dir().expect("Unable to get the current directory"),
                    stripped.to_string_lossy().to_mut().as_str()).expect("Unable to create temporary directory");
                let status = Command::new(cmd)
                    .args(&at.extract_cmd[1..])
                    .arg(fs::canonicalize(&arch).expect("Unable to find archive"))
                    .current_dir(&tmpdir)
                    .before_exec(setup_seccomp)
                    .status()
                    .expect(format!("Failed to launch {}", cmd).as_str());
                assert!(status.success(), format!("{} wasn't successful", cmd));
                let extracted_path = tmpdir.into_path();
                let mut dir_list = fs::read_dir(&extracted_path).expect("Unable to read the extraction directory");
                if let Some(path) = dir_list.next() {
                    let path = path.expect("Unable to iterate the extraction directory");
                    if let Some(path1) = dir_list.next() {
                        path1.expect("Unable to iterate the extraction directory");
                        let newname = rename_or_suffix(&extracted_path, Path::new(stripped)).expect("Failed to rename");
                        println!("Extracted to {:?}", newname);
                    } else {
                        let newname = rename_or_suffix(&path.path(), Path::new(stripped)).expect("Failed to rename");
                        println!("Extracted to {:?}", newname);
                        fs::remove_dir(&extracted_path).expect("Failed to remove empty extraction directory");
                    }
                } else {
                    println!("No files were extracted");
                }
                break;
            }
        }
        if !found {
            println!("Can't handle {:?}", arch);
        }
    }
}

<|endoftext|>