<commit_before>use std::net::TcpStream;
use std::collections::HashMap;
use std::io;
use std::io::{Read,Write};
use std::ops::Deref;
use libc;

use unix_socket::UnixStream;
use rustc_serialize::hex::ToHex;
use dbus_serialize::types::{Value,BasicValue};
use dbus_serialize::decoder::DBusDecoder;

use message;
use message::{Message,HeaderFieldName,MessageBuf};
use demarshal::{demarshal,DemarshalError};
use marshal::Marshal;

trait StreamSocket : Read + Write { }

impl StreamSocket for TcpStream { }
impl StreamSocket for UnixStream { }

enum Socket {
    Tcp(TcpStream),
    Uds(UnixStream)
}

pub struct Connection {
    sock: Socket,
    next_serial: u32,
    queue: Vec<Message>,
}

#[derive(Debug)]
pub enum Error {
    Disconnected,
    IOError(io::Error),
    DemarshalError(DemarshalError),
    BadData,
    AuthFailed,
}

impl From<io::Error> for Error {
    fn from(x: io::Error) -> Self {
        Error::IOError(x)
    }
}

impl From<DemarshalError> for Error {
    fn from(x: DemarshalError) -> Self {
        Error::DemarshalError(x)
    }
}

fn read_exactly(sock: &mut StreamSocket, buf: &mut Vec<u8>, len: usize) -> Result<(),Error> {
    buf.truncate(0);
    buf.reserve(len);
    if try!(sock.take(len as u64).read_to_end(buf)) != len {
        return Err(Error::Disconnected);
    }
    Ok(())
}

fn read_line(sock: &mut StreamSocket) -> Result<String,Error> {
    let mut line = "".to_string();
    let mut last = '\0';

    loop {
        let mut buf = vec![0];
        match sock.read(&mut buf) {
            Ok(x) if x > 0 => (),
            _ => return Err(Error::Disconnected)
        };
        let chr = buf[0] as char;
        line.push(chr);
        if chr == '\n' && last == '\r' {
            break;
        }
        last = chr;
    }
    Ok(line)
}

impl Connection {
    fn get_sock(&mut self) -> &mut StreamSocket {
        match self.sock {
            Socket::Tcp(ref mut x) => x,
            Socket::Uds(ref mut x) => x
        }
    }

    fn auth_anonymous(&mut self) -> Result<(),Error> {
        let sock = self.get_sock();

        // Authenticate to the daemon
        let buf = vec![0];
        try!(sock.write_all(&buf));

        try!(sock.write_all(b"AUTH ANONYMOUS 6c69626462757320312e382e3132\r\n"));

        // Read response
        let resp = try!(read_line(sock));
        if !resp.starts_with("OK ") {
            return Err(Error::AuthFailed);
        }

        // Ready for action
        try!(sock.write_all(b"BEGIN\r\n"));
        Ok(())
    }

    fn auth_external(&mut self) -> Result<(),Error> {
        let sock = self.get_sock();

        // Authenticate to the daemon
        let buf = vec![0];
        try!(sock.write_all(&buf));

        let uid = unsafe {
            libc::funcs::posix88::unistd::getuid()
        };
        let uid_str = uid.to_string();
        let uid_hex = uid_str.into_bytes().to_hex();
        let cmd = "AUTH EXTERNAL ".to_string() + &uid_hex + "\r\n";
        try!(sock.write_all(&cmd.into_bytes()));

        // Read response
        let resp = try!(read_line(sock));
        if !resp.starts_with("OK ") {
            return Err(Error::AuthFailed);
        }

        // Ready for action
        try!(sock.write_all(b"BEGIN\r\n"));
        Ok(())
    }

    fn say_hello(&mut self) -> Result<(String),Error> {
        let mut msg = message::create_method_call("org.freedesktop.DBus",
                                                  "/org/freedesktop/DBus",
                                                  "org.freedesktop.DBus",
                                                  "Hello");
        match try!(self.call_sync(&mut msg)).get(0) {
            Some(&Value::BasicValue(BasicValue::String(ref x))) => Ok(x.to_string()),
            _ => Err(Error::BadData)
        }
    }

    pub fn connect_uds(addr: &str) -> Result<Connection,Error> {
        let sock = try!(UnixStream::connect(addr));
        let mut conn = Connection {
            sock: Socket::Uds(sock),
            queue: Vec::new(),
            next_serial: 1
        };

        try!(conn.auth_external());
        try!(conn.say_hello());
        Ok(conn)
    }

    pub fn connect_tcp(addr: &str) -> Result<Connection,Error> {
        let sock = try!(TcpStream::connect(addr));
        let mut conn = Connection {
            sock: Socket::Tcp(sock),
            queue: Vec::new(),
            next_serial: 1
        };

        try!(conn.auth_anonymous());
        try!(conn.say_hello());
        Ok(conn)
    }

    pub fn send(&mut self, mbuf: &mut MessageBuf) -> Result<u32, Error> {
        let mut msg = &mut mbuf.0;
        // A minimum header with no body is 16 bytes
        let mut len = msg.len() as u32;
        if len < 16 {
            return Err(Error::BadData);
        }

        // Get the current length from the message, which only include the length of the header.
        // That field should actually be the length of only the body, so update that now
        let old_len = message::get_length(msg);
        len -= old_len;
        let mut buf = Vec::new();
        len.dbus_encode(&mut buf);
        // Update the message with a correct serial number, as well
        let this_serial = self.next_serial;
        self.next_serial += 1;
        this_serial.dbus_encode(&mut buf);
        message::set_length(msg, &buf);

        let sock = self.get_sock();
        try!(sock.write_all(msg));
        Ok(this_serial)
    }

    pub fn call_sync(&mut self, mbuf: &mut MessageBuf) -> Result<Vec<Value>,Error> {
        // XXX: assert that this is a method call with reply
        let serial = try!(self.send(mbuf));
        // We need a local queue so that read_msg doesn't just give us
        // the same one over and over
        let mut queue = Vec::new();
        loop {
            let mut msg = try!(self.read_msg());
            match msg.headers.remove(&(HeaderFieldName::ReplySerial as u8)) {
                Some(Value::Variant(x)) => {
                    let obj : Value = *x.object;
                    let reply_serial : u32 = DBusDecoder::decode(obj).unwrap();
                    if reply_serial == serial {
                        // Move our queued messages into the Connection's queue
                        for _ in 0..queue.len() {
                            self.queue.push(queue.remove(0));
                        }
                        return Ok(msg.body);
                    }
                }
                _ => ()
            };
            queue.push(msg);
        }
    }

    pub fn read_msg(&mut self) -> Result<Message,Error> {
        match self.queue.get(0) {
            Some(_) => return Ok(self.queue.remove(0)),
            _ => ()
        };
        let mut buf = Vec::new();
        let sock = self.get_sock();

        // Read and demarshal the fixed portion of the header
        try!(read_exactly(sock, &mut buf, 12));
        let mut offset = 0;
        let mut sig = "(yyyyuu)".to_string();
        let header = match try!(demarshal(&mut buf, &mut offset, &mut sig)) {
            Value::Struct(x) => x,
            x => panic!("Demarshal didn't return what we asked for: {:?}", x)
        };

        let mut v = header.objects;
        let mut msg : Message = Default::default();
        let endian : u8 = DBusDecoder::decode(v.remove(0)).unwrap();
        if endian == 'B' as u8 {
            msg.big_endian = true;
        }
        msg.message_type = message::MessageType(DBusDecoder::decode(v.remove(0)).unwrap());
        msg.flags = DBusDecoder::decode::<u8>(v.remove(0)).unwrap();
        msg.version = DBusDecoder::decode::<u8>(v.remove(0)).unwrap();
        let body_len = DBusDecoder::decode::<u32>(v.remove(0)).unwrap();
        msg.serial = DBusDecoder::decode::<u32>(v.remove(0)).unwrap();

        // Read array length
        try!(read_exactly(sock, &mut buf, 4));
        // demarshal consumes the buf, so save a copy for when we demarshal the entire array
        let mut buf_copy = buf.clone();
        offset = 12;
        sig = "u".to_string();
        let data = demarshal(&mut buf, &mut offset, &mut sig).ok().unwrap();
        let arr_len = DBusDecoder::decode::<u32>(data).unwrap() as usize;

        // Make buf_copy big enough for the entire array, and fill it
        buf_copy.reserve(arr_len);
        if try!(sock.take(arr_len as u64).read_to_end(&mut buf_copy)) != arr_len {
            return Err(Error::Disconnected);
        };

        offset = 12;
        sig = "a(yv)".to_string();
        let header_fields = match try!(demarshal(&mut buf_copy, &mut offset, &mut sig)) {
            Value::Array(x) => x,
            x => panic!("Demarshal didn't return what we asked for: {:?}", x)
        };

        msg.headers = HashMap::new();
        for i in header_fields.objects {
            let mut st = match i {
                Value::Struct(x) => x,
                x => panic!("Demarshal didn't return what we asked for: {:?}", x)
            };
            let val = st.objects.remove(1);
            let code = DBusDecoder::decode::<u8>(st.objects.remove(0)).unwrap();
            msg.headers.insert(code, val);
        }

        // Read the padding, if any
        let trailing_pad = 8 - (offset % 8);
        if trailing_pad % 8 != 0 {
            try!(read_exactly(sock, &mut buf, trailing_pad));
        }

        // Finally, read the entire body
        if body_len > 0 {
            let v = match msg.headers.get(&(HeaderFieldName::Signature as u8)) {
                Some(&Value::Variant(ref x)) => x,
                _ => return Err(Error::DemarshalError(DemarshalError::BadSignature))
            };

            let sigval = match v.object.deref() {
                &Value::BasicValue(BasicValue::Signature(ref x)) => x,
                _ => return Err(Error::DemarshalError(DemarshalError::BadSignature))
            };

            let mut body = Vec::new();
            try!(read_exactly(sock, &mut body, body_len as usize));

            let mut sig = "(".to_string() + &sigval.0 + ")";
            offset = 0;
            let objs = match try!(demarshal(&mut body, &mut offset, &mut sig)) {
                Value::Struct(x) => x.objects,
                x => panic!("Didn't get a struct: {:?}", x)
            };
            for x in objs {
                msg.body.push(x);
            }
        }

        Ok(msg)
    }
}

#[test]
fn test_connect () {
    let mut conn = Connection::connect_uds("/var/run/dbus/system_bus_socket").unwrap();
    let mut msg = message::create_method_call("org.freedesktop.DBus", "/org/freedesktop/DBus",
                                          "org.freedesktop.DBus", "ListNames");
    let resp = conn.call_sync(&mut msg).unwrap();
    println!("ListNames: {:?}", resp);
}

#[cfg(tcp)]
#[test]
fn test_tcp () {
    let mut conn = Connection::connect_tcp("localhost:12345").unwrap();
    let mut msg = message::create_method_call("org.freedesktop.DBus", "/org/freedesktop/DBus",
                                          "org.freedesktop.DBus", "ListNames");
    conn.send(&mut msg).ok();
    let msg = conn.read_msg().unwrap();
    println!("{:?}", msg.body);
    //loop {
    //    conn.read_msg().unwrap();
    //}
}
<commit_msg>connection.rs: factor NUL byte out from auth mechanisms<commit_after>use std::net::TcpStream;
use std::collections::HashMap;
use std::io;
use std::io::{Read,Write};
use std::ops::Deref;
use libc;

use unix_socket::UnixStream;
use rustc_serialize::hex::ToHex;
use dbus_serialize::types::{Value,BasicValue};
use dbus_serialize::decoder::DBusDecoder;

use message;
use message::{Message,HeaderFieldName,MessageBuf};
use demarshal::{demarshal,DemarshalError};
use marshal::Marshal;

trait StreamSocket : Read + Write { }

impl StreamSocket for TcpStream { }
impl StreamSocket for UnixStream { }

enum Socket {
    Tcp(TcpStream),
    Uds(UnixStream)
}

pub struct Connection {
    sock: Socket,
    next_serial: u32,
    queue: Vec<Message>,
}

#[derive(Debug)]
pub enum Error {
    Disconnected,
    IOError(io::Error),
    DemarshalError(DemarshalError),
    BadData,
    AuthFailed,
}

impl From<io::Error> for Error {
    fn from(x: io::Error) -> Self {
        Error::IOError(x)
    }
}

impl From<DemarshalError> for Error {
    fn from(x: DemarshalError) -> Self {
        Error::DemarshalError(x)
    }
}

fn read_exactly(sock: &mut StreamSocket, buf: &mut Vec<u8>, len: usize) -> Result<(),Error> {
    buf.truncate(0);
    buf.reserve(len);
    if try!(sock.take(len as u64).read_to_end(buf)) != len {
        return Err(Error::Disconnected);
    }
    Ok(())
}

fn read_line(sock: &mut StreamSocket) -> Result<String,Error> {
    let mut line = "".to_string();
    let mut last = '\0';

    loop {
        let mut buf = vec![0];
        match sock.read(&mut buf) {
            Ok(x) if x > 0 => (),
            _ => return Err(Error::Disconnected)
        };
        let chr = buf[0] as char;
        line.push(chr);
        if chr == '\n' && last == '\r' {
            break;
        }
        last = chr;
    }
    Ok(line)
}

impl Connection {
    fn get_sock(&mut self) -> &mut StreamSocket {
        match self.sock {
            Socket::Tcp(ref mut x) => x,
            Socket::Uds(ref mut x) => x
        }
    }

    fn send_nul_byte(&mut self) -> Result<(),Error> {
        // Send NUL byte
        let sock = self.get_sock();
        let buf = vec![0];
        try!(sock.write_all(&buf));
        Ok(())
    }

    fn auth_anonymous(&mut self) -> Result<(),Error> {
        let sock = self.get_sock();

        try!(sock.write_all(b"AUTH ANONYMOUS 6c69626462757320312e382e3132\r\n"));

        // Read response
        let resp = try!(read_line(sock));
        if !resp.starts_with("OK ") {
            return Err(Error::AuthFailed);
        }

        // Ready for action
        try!(sock.write_all(b"BEGIN\r\n"));
        Ok(())
    }

    fn auth_external(&mut self) -> Result<(),Error> {
        let sock = self.get_sock();

        let uid = unsafe {
            libc::funcs::posix88::unistd::getuid()
        };
        let uid_str = uid.to_string();
        let uid_hex = uid_str.into_bytes().to_hex();
        let cmd = "AUTH EXTERNAL ".to_string() + &uid_hex + "\r\n";
        try!(sock.write_all(&cmd.into_bytes()));

        // Read response
        let resp = try!(read_line(sock));
        if !resp.starts_with("OK ") {
            return Err(Error::AuthFailed);
        }

        // Ready for action
        try!(sock.write_all(b"BEGIN\r\n"));
        Ok(())
    }

    fn say_hello(&mut self) -> Result<(String),Error> {
        let mut msg = message::create_method_call("org.freedesktop.DBus",
                                                  "/org/freedesktop/DBus",
                                                  "org.freedesktop.DBus",
                                                  "Hello");
        match try!(self.call_sync(&mut msg)).get(0) {
            Some(&Value::BasicValue(BasicValue::String(ref x))) => Ok(x.to_string()),
            _ => Err(Error::BadData)
        }
    }

    pub fn connect_uds(addr: &str) -> Result<Connection,Error> {
        let sock = try!(UnixStream::connect(addr));
        let mut conn = Connection {
            sock: Socket::Uds(sock),
            queue: Vec::new(),
            next_serial: 1
        };

        try!(conn.send_nul_byte());
        try!(conn.auth_external());
        try!(conn.say_hello());
        Ok(conn)
    }

    pub fn connect_tcp(addr: &str) -> Result<Connection,Error> {
        let sock = try!(TcpStream::connect(addr));
        let mut conn = Connection {
            sock: Socket::Tcp(sock),
            queue: Vec::new(),
            next_serial: 1
        };

        try!(conn.send_nul_byte());
        try!(conn.auth_anonymous());
        try!(conn.say_hello());
        Ok(conn)
    }

    pub fn send(&mut self, mbuf: &mut MessageBuf) -> Result<u32, Error> {
        let mut msg = &mut mbuf.0;
        // A minimum header with no body is 16 bytes
        let mut len = msg.len() as u32;
        if len < 16 {
            return Err(Error::BadData);
        }

        // Get the current length from the message, which only include the length of the header.
        // That field should actually be the length of only the body, so update that now
        let old_len = message::get_length(msg);
        len -= old_len;
        let mut buf = Vec::new();
        len.dbus_encode(&mut buf);
        // Update the message with a correct serial number, as well
        let this_serial = self.next_serial;
        self.next_serial += 1;
        this_serial.dbus_encode(&mut buf);
        message::set_length(msg, &buf);

        let sock = self.get_sock();
        try!(sock.write_all(msg));
        Ok(this_serial)
    }

    pub fn call_sync(&mut self, mbuf: &mut MessageBuf) -> Result<Vec<Value>,Error> {
        // XXX: assert that this is a method call with reply
        let serial = try!(self.send(mbuf));
        // We need a local queue so that read_msg doesn't just give us
        // the same one over and over
        let mut queue = Vec::new();
        loop {
            let mut msg = try!(self.read_msg());
            match msg.headers.remove(&(HeaderFieldName::ReplySerial as u8)) {
                Some(Value::Variant(x)) => {
                    let obj : Value = *x.object;
                    let reply_serial : u32 = DBusDecoder::decode(obj).unwrap();
                    if reply_serial == serial {
                        // Move our queued messages into the Connection's queue
                        for _ in 0..queue.len() {
                            self.queue.push(queue.remove(0));
                        }
                        return Ok(msg.body);
                    }
                }
                _ => ()
            };
            queue.push(msg);
        }
    }

    pub fn read_msg(&mut self) -> Result<Message,Error> {
        match self.queue.get(0) {
            Some(_) => return Ok(self.queue.remove(0)),
            _ => ()
        };
        let mut buf = Vec::new();
        let sock = self.get_sock();

        // Read and demarshal the fixed portion of the header
        try!(read_exactly(sock, &mut buf, 12));
        let mut offset = 0;
        let mut sig = "(yyyyuu)".to_string();
        let header = match try!(demarshal(&mut buf, &mut offset, &mut sig)) {
            Value::Struct(x) => x,
            x => panic!("Demarshal didn't return what we asked for: {:?}", x)
        };

        let mut v = header.objects;
        let mut msg : Message = Default::default();
        let endian : u8 = DBusDecoder::decode(v.remove(0)).unwrap();
        if endian == 'B' as u8 {
            msg.big_endian = true;
        }
        msg.message_type = message::MessageType(DBusDecoder::decode(v.remove(0)).unwrap());
        msg.flags = DBusDecoder::decode::<u8>(v.remove(0)).unwrap();
        msg.version = DBusDecoder::decode::<u8>(v.remove(0)).unwrap();
        let body_len = DBusDecoder::decode::<u32>(v.remove(0)).unwrap();
        msg.serial = DBusDecoder::decode::<u32>(v.remove(0)).unwrap();

        // Read array length
        try!(read_exactly(sock, &mut buf, 4));
        // demarshal consumes the buf, so save a copy for when we demarshal the entire array
        let mut buf_copy = buf.clone();
        offset = 12;
        sig = "u".to_string();
        let data = demarshal(&mut buf, &mut offset, &mut sig).ok().unwrap();
        let arr_len = DBusDecoder::decode::<u32>(data).unwrap() as usize;

        // Make buf_copy big enough for the entire array, and fill it
        buf_copy.reserve(arr_len);
        if try!(sock.take(arr_len as u64).read_to_end(&mut buf_copy)) != arr_len {
            return Err(Error::Disconnected);
        };

        offset = 12;
        sig = "a(yv)".to_string();
        let header_fields = match try!(demarshal(&mut buf_copy, &mut offset, &mut sig)) {
            Value::Array(x) => x,
            x => panic!("Demarshal didn't return what we asked for: {:?}", x)
        };

        msg.headers = HashMap::new();
        for i in header_fields.objects {
            let mut st = match i {
                Value::Struct(x) => x,
                x => panic!("Demarshal didn't return what we asked for: {:?}", x)
            };
            let val = st.objects.remove(1);
            let code = DBusDecoder::decode::<u8>(st.objects.remove(0)).unwrap();
            msg.headers.insert(code, val);
        }

        // Read the padding, if any
        let trailing_pad = 8 - (offset % 8);
        if trailing_pad % 8 != 0 {
            try!(read_exactly(sock, &mut buf, trailing_pad));
        }

        // Finally, read the entire body
        if body_len > 0 {
            let v = match msg.headers.get(&(HeaderFieldName::Signature as u8)) {
                Some(&Value::Variant(ref x)) => x,
                _ => return Err(Error::DemarshalError(DemarshalError::BadSignature))
            };

            let sigval = match v.object.deref() {
                &Value::BasicValue(BasicValue::Signature(ref x)) => x,
                _ => return Err(Error::DemarshalError(DemarshalError::BadSignature))
            };

            let mut body = Vec::new();
            try!(read_exactly(sock, &mut body, body_len as usize));

            let mut sig = "(".to_string() + &sigval.0 + ")";
            offset = 0;
            let objs = match try!(demarshal(&mut body, &mut offset, &mut sig)) {
                Value::Struct(x) => x.objects,
                x => panic!("Didn't get a struct: {:?}", x)
            };
            for x in objs {
                msg.body.push(x);
            }
        }

        Ok(msg)
    }
}

#[test]
fn test_connect () {
    let mut conn = Connection::connect_uds("/var/run/dbus/system_bus_socket").unwrap();
    let mut msg = message::create_method_call("org.freedesktop.DBus", "/org/freedesktop/DBus",
                                          "org.freedesktop.DBus", "ListNames");
    let resp = conn.call_sync(&mut msg).unwrap();
    println!("ListNames: {:?}", resp);
}

#[cfg(tcp)]
#[test]
fn test_tcp () {
    let mut conn = Connection::connect_tcp("localhost:12345").unwrap();
    let mut msg = message::create_method_call("org.freedesktop.DBus", "/org/freedesktop/DBus",
                                          "org.freedesktop.DBus", "ListNames");
    conn.send(&mut msg).ok();
    let msg = conn.read_msg().unwrap();
    println!("{:?}", msg.body);
    //loop {
    //    conn.read_msg().unwrap();
    //}
}
<|endoftext|>


<commit_before>use async_lapin::*;
use lapin::{
    executor::Executor, message::DeliveryResult, options::*, publisher_confirm::Confirmation,
    types::FieldTable, BasicProperties, Connection, ConnectionProperties, Result,
};
use log::info;
use std::{future::Future, pin::Pin};

#[derive(Debug)]
struct SmolExecutor;

impl Executor for SmolExecutor {
    fn spawn(&self, f: Pin<Box<dyn Future<Output = ()> + Send>>) -> Result<()> {
        smol::spawn(f).detach();
        Ok(())
    }
}

fn main() -> Result<()> {
    if std::env::var("RUST_LOG").is_err() {
        std::env::set_var("RUST_LOG", "info");
    }

    env_logger::init();

    let addr = std::env::var("AMQP_ADDR").unwrap_or_else(|_| "amqp://127.0.0.1:5672/%2f".into());

    smol::block_on(async {
        let conn = Connection::connect(
            &addr,
            ConnectionProperties::default()
                .with_async_io(SmolExecutor),
        )
        .await?;

        info!("CONNECTED");

        let channel_a = conn.create_channel().await?;
        let channel_b = conn.create_channel().await?;

        let queue = channel_a
            .queue_declare(
                "hello",
                QueueDeclareOptions::default(),
                FieldTable::default(),
            )
            .await?;

        info!("Declared queue {:?}", queue);

        let consumer = channel_b
            .basic_consume(
                "hello",
                "my_consumer",
                BasicConsumeOptions::default(),
                FieldTable::default(),
            )
            .await?;

        consumer
            .set_delegate(move |delivery: DeliveryResult| async move {
                let delivery = delivery.expect("error caught in in consumer");
                if let Some((channel, delivery)) = delivery {
                    channel
                        .basic_ack(delivery.delivery_tag, BasicAckOptions::default())
                        .await
                        .expect("failed to ack");
                }
            })
            .expect("set_delegate");

        let payload = b"Hello world!";

        loop {
            let confirm = channel_a
                .basic_publish(
                    "",
                    "hello",
                    BasicPublishOptions::default(),
                    payload.to_vec(),
                    BasicProperties::default(),
                )
                .await?
                .await?;
            assert_eq!(confirm, Confirmation::NotRequested);
        }
    })
}
<commit_msg>async-lapin: rustfmt<commit_after>use async_lapin::*;
use lapin::{
    executor::Executor, message::DeliveryResult, options::*, publisher_confirm::Confirmation,
    types::FieldTable, BasicProperties, Connection, ConnectionProperties, Result,
};
use log::info;
use std::{future::Future, pin::Pin};

#[derive(Debug)]
struct SmolExecutor;

impl Executor for SmolExecutor {
    fn spawn(&self, f: Pin<Box<dyn Future<Output = ()> + Send>>) -> Result<()> {
        smol::spawn(f).detach();
        Ok(())
    }
}

fn main() -> Result<()> {
    if std::env::var("RUST_LOG").is_err() {
        std::env::set_var("RUST_LOG", "info");
    }

    env_logger::init();

    let addr = std::env::var("AMQP_ADDR").unwrap_or_else(|_| "amqp://127.0.0.1:5672/%2f".into());

    smol::block_on(async {
        let conn = Connection::connect(
            &addr,
            ConnectionProperties::default().with_async_io(SmolExecutor),
        )
        .await?;

        info!("CONNECTED");

        let channel_a = conn.create_channel().await?;
        let channel_b = conn.create_channel().await?;

        let queue = channel_a
            .queue_declare(
                "hello",
                QueueDeclareOptions::default(),
                FieldTable::default(),
            )
            .await?;

        info!("Declared queue {:?}", queue);

        let consumer = channel_b
            .basic_consume(
                "hello",
                "my_consumer",
                BasicConsumeOptions::default(),
                FieldTable::default(),
            )
            .await?;

        consumer
            .set_delegate(move |delivery: DeliveryResult| async move {
                let delivery = delivery.expect("error caught in in consumer");
                if let Some((channel, delivery)) = delivery {
                    channel
                        .basic_ack(delivery.delivery_tag, BasicAckOptions::default())
                        .await
                        .expect("failed to ack");
                }
            })
            .expect("set_delegate");

        let payload = b"Hello world!";

        loop {
            let confirm = channel_a
                .basic_publish(
                    "",
                    "hello",
                    BasicPublishOptions::default(),
                    payload.to_vec(),
                    BasicProperties::default(),
                )
                .await?
                .await?;
            assert_eq!(confirm, Confirmation::NotRequested);
        }
    })
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use hyper::header::ContentType;
use hyper::mime::{Attr, Mime, SubLevel, TopLevel, Value};
use hyper_serde::Serde;
use ipc_channel::ipc;
use msg::constellation_msg::{PipelineId, ReferrerPolicy};
use net_traits::{LoadContext, LoadData, LoadOrigin, NetworkError};
use net_traits::LoadConsumer::Channel;
use net_traits::ProgressMsg::{Done, Payload};
use url::Url;

struct DataLoadTest;

impl LoadOrigin for DataLoadTest {
    fn referrer_url(&self) -> Option<Url> {
        None
    }
    fn referrer_policy(&self) -> Option<ReferrerPolicy> {
        None
    }
    fn pipeline_id(&self) -> Option<PipelineId> {
        None
    }
}

#[cfg(test)]
fn assert_parse(url:          &'static str,
                content_type: Option<ContentType>,
                charset:      Option<String>,
                data:         Option<Vec<u8>>) {
    use net::data_loader::load;
    use net::mime_classifier::MimeClassifier;
    use net::resource_thread::CancellationListener;
    use std::sync::Arc;

    let (start_chan, start_port) = ipc::channel().unwrap();
    let classifier = Arc::new(MimeClassifier::new());
    load(LoadData::new(LoadContext::Browsing, Url::parse(url).unwrap(), &DataLoadTest),
         Channel(start_chan),
         classifier, CancellationListener::new(None));

    let response = start_port.recv().unwrap();
    assert_eq!(&response.metadata.content_type.map(Serde::into_inner),
               &content_type);
    assert_eq!(&response.metadata.charset,      &charset);

    let progress = response.progress_port.recv().unwrap();

    match data {
        None => {
            assert_eq!(progress, Done(Err(NetworkError::Internal("invalid data uri".to_owned()))));
        }
        Some(dat) => {
            assert_eq!(progress, Payload(dat));
            assert_eq!(response.progress_port.recv().unwrap(), Done(Ok(())));
        }
    }
}

#[test]
fn empty_invalid() {
    assert_parse("data:", None, None, None);
}

#[test]
fn plain() {
    assert_parse(
        "data:,hello%20world",
        Some(ContentType(Mime(TopLevel::Text, SubLevel::Plain,
                              vec!((Attr::Charset, Value::Ext("us-ascii".to_owned())))))),
        Some("US-ASCII".to_owned()), Some(b"hello world".iter().map(|&x| x).collect()));
}

#[test]
fn plain_ct() {
    assert_parse(
        "data:text/plain,hello",
        Some(ContentType(Mime(TopLevel::Text, SubLevel::Plain, vec!()))),
        None,
        Some(b"hello".iter().map(|&x| x).collect()));
}

#[test]
fn plain_charset() {
    assert_parse("data:text/plain;charset=latin1,hello",
        Some(ContentType(Mime(TopLevel::Text,
                              SubLevel::Plain,
                              vec!((Attr::Charset, Value::Ext("latin1".to_owned())))))),
        Some("latin1".to_owned()), Some(b"hello".iter().map(|&x| x).collect()));
}

#[test]
fn plain_only_charset() {
    assert_parse(
        "data:;charset=utf-8,hello",
        Some(ContentType(Mime(TopLevel::Text,
                              SubLevel::Plain,
                              vec!((Attr::Charset, Value::Utf8))))),
        Some("utf-8".to_owned()), Some(b"hello".iter().map(|&x| x).collect()));
}

#[test]
fn base64() {
    assert_parse(
        "data:;base64,C62+7w==",
        Some(ContentType(Mime(TopLevel::Text,
                              SubLevel::Plain,
                              vec!((Attr::Charset, Value::Ext("us-ascii".to_owned())))))),
        Some("US-ASCII".to_owned()), Some(vec!(0x0B, 0xAD, 0xBE, 0xEF)));
}

#[test]
fn base64_ct() {
    assert_parse("data:application/octet-stream;base64,C62+7w==",
        Some(ContentType(Mime(TopLevel::Application, SubLevel::Ext("octet-stream".to_owned()), vec!()))),
        None,
        Some(vec!(0x0B, 0xAD, 0xBE, 0xEF)));
}

#[test]
fn base64_charset() {
    assert_parse("data:text/plain;charset=koi8-r;base64,8PLl9+XkIO3l5Pfl5A==",
        Some(ContentType(Mime(TopLevel::Text, SubLevel::Plain,
                              vec!((Attr::Charset, Value::Ext("koi8-r".to_owned())))))),
        Some("koi8-r".to_owned()),
        Some(vec!(0xF0, 0xF2, 0xE5, 0xF7, 0xE5, 0xE4, 0x20, 0xED, 0xE5, 0xE4, 0xF7, 0xE5, 0xE4)));
}
<commit_msg>Make the data_loader test more efficient and readable.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use hyper::header::ContentType;
use hyper::mime::{Attr, Mime, SubLevel, TopLevel, Value};
use hyper_serde::Serde;
use ipc_channel::ipc;
use msg::constellation_msg::{PipelineId, ReferrerPolicy};
use net_traits::{LoadContext, LoadData, LoadOrigin, NetworkError};
use net_traits::LoadConsumer::Channel;
use net_traits::ProgressMsg::{Done, Payload};
use std::ops::Deref;
use url::Url;

struct DataLoadTest;

impl LoadOrigin for DataLoadTest {
    fn referrer_url(&self) -> Option<Url> {
        None
    }
    fn referrer_policy(&self) -> Option<ReferrerPolicy> {
        None
    }
    fn pipeline_id(&self) -> Option<PipelineId> {
        None
    }
}

#[cfg(test)]
fn assert_parse(url:          &'static str,
                content_type: Option<ContentType>,
                charset:      Option<&str>,
                data:         Option<&[u8]>) {
    use net::data_loader::load;
    use net::mime_classifier::MimeClassifier;
    use net::resource_thread::CancellationListener;
    use std::sync::Arc;

    let (start_chan, start_port) = ipc::channel().unwrap();
    let classifier = Arc::new(MimeClassifier::new());
    load(LoadData::new(LoadContext::Browsing, Url::parse(url).unwrap(), &DataLoadTest),
         Channel(start_chan),
         classifier, CancellationListener::new(None));

    let response = start_port.recv().unwrap();
    assert_eq!(&response.metadata.content_type.map(Serde::into_inner),
               &content_type);
    assert_eq!(response.metadata.charset.as_ref().map(String::deref), charset);

    let progress = response.progress_port.recv().unwrap();

    match data {
        None => {
            assert_eq!(progress, Done(Err(NetworkError::Internal("invalid data uri".to_owned()))));
        }
        Some(dat) => {
            match progress {
                Payload(d) => assert_eq!(d, dat),
                _ => panic!(),
            }
            assert_eq!(response.progress_port.recv().unwrap(), Done(Ok(())));
        }
    }
}

#[test]
fn empty_invalid() {
    assert_parse("data:", None, None, None);
}

#[test]
fn plain() {
    assert_parse(
        "data:,hello%20world",
        Some(ContentType(Mime(TopLevel::Text, SubLevel::Plain,
                              vec!((Attr::Charset, Value::Ext("us-ascii".to_owned())))))),
        Some("US-ASCII"),
        Some(b"hello world"));
}

#[test]
fn plain_ct() {
    assert_parse(
        "data:text/plain,hello",
        Some(ContentType(Mime(TopLevel::Text, SubLevel::Plain, vec!()))),
        None,
        Some(b"hello"));
}

#[test]
fn plain_charset() {
    assert_parse(
        "data:text/plain;charset=latin1,hello",
        Some(ContentType(Mime(TopLevel::Text,
                              SubLevel::Plain,
                              vec!((Attr::Charset, Value::Ext("latin1".to_owned())))))),
        Some("latin1"),
        Some(b"hello"));
}

#[test]
fn plain_only_charset() {
    assert_parse(
        "data:;charset=utf-8,hello",
        Some(ContentType(Mime(TopLevel::Text,
                              SubLevel::Plain,
                              vec!((Attr::Charset, Value::Utf8))))),
        Some("utf-8"),
        Some(b"hello"));
}

#[test]
fn base64() {
    assert_parse(
        "data:;base64,C62+7w==",
        Some(ContentType(Mime(TopLevel::Text,
                              SubLevel::Plain,
                              vec!((Attr::Charset, Value::Ext("us-ascii".to_owned())))))),
        Some("US-ASCII"),
        Some(&[0x0B, 0xAD, 0xBE, 0xEF]));
}

#[test]
fn base64_ct() {
    assert_parse(
        "data:application/octet-stream;base64,C62+7w==",
        Some(ContentType(Mime(TopLevel::Application, SubLevel::Ext("octet-stream".to_owned()), vec!()))),
        None,
        Some(&[0x0B, 0xAD, 0xBE, 0xEF]));
}

#[test]
fn base64_charset() {
    assert_parse(
        "data:text/plain;charset=koi8-r;base64,8PLl9+XkIO3l5Pfl5A==",
        Some(ContentType(Mime(TopLevel::Text, SubLevel::Plain,
                              vec!((Attr::Charset, Value::Ext("koi8-r".to_owned())))))),
        Some("koi8-r"),
        Some(&[0xF0, 0xF2, 0xE5, 0xF7, 0xE5, 0xE4, 0x20, 0xED, 0xE5, 0xE4, 0xF7, 0xE5, 0xE4]));
}
<|endoftext|>


<commit_before>// Copyright 2015, 2016 Ethcore (UK) Ltd.
// This file is part of Parity.

// Parity is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Parity is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Parity.  If not, see <http://www.gnu.org/licenses/>.

//! jsonrpc server over win named pipes
//!
//! ```no_run
//! extern crate jsonrpc_core;
//! extern crate jsonrpc_ipc_server;
//!
//! use std::sync::Arc;
//! use jsonrpc_core::*;
//! use jsonrpc_ipc_server::Server;
//!
//! fn main() {
//! 	let mut io = IoHandler::new();
//! 	io.add_method("say_hello", |_params| {
//!			Ok(Value::String("hello".into()))
//!		});
//! 	let server = Server::new("/tmp/json-ipc-test.ipc", io).unwrap();
//!     ::std::thread::spawn(move || server.run());
//! }
//! ```

//! Named pipes library

use miow::pipe::{NamedPipe, NamedPipeBuilder};
use std;
use std::io;
use std::io::{Read, Write};
use std::sync::atomic::*;
use std::sync::{Arc, Mutex};
use jsonrpc_core::{Metadata, MetaIoHandler};
use jsonrpc_core::reactor::{RpcHandler, RpcEventLoop, RpcEventLoopHandle};
use validator;

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

const MAX_REQUEST_LEN: u32 = 65536;
const REQUEST_READ_BATCH: usize = 4096;

#[derive(Debug)]
pub enum Error {
	Io(std::io::Error),
	NotStarted,
	AlreadyStopping,
	NotStopped,
	IsStopping,
}

impl std::convert::From<std::io::Error> for Error {
	fn from(io_error: std::io::Error) -> Error {
		Error::Io(io_error)
	}
}

pub struct PipeHandler<M: Metadata = ()> {
	waiting_pipe: NamedPipe,
	io_handler: RpcHandler<M>,
}

impl<M: Metadata> PipeHandler<M> {
	/// start ipc rpc server
	pub fn start(addr: &str, io_handler: RpcHandler<M>) -> Result<PipeHandler<M>> {
		Ok(PipeHandler {
			waiting_pipe: try!(
				NamedPipeBuilder::new(addr)
					.first(true)
					.accept_remote(true)
					.max_instances(255)
					.inbound(true)
					.outbound(true)
					.out_buffer_size(MAX_REQUEST_LEN)
					.in_buffer_size(MAX_REQUEST_LEN)
					.create()
			),
			io_handler: io_handler.clone(),
		})
	}

	fn handle_incoming(&mut self, addr: &str, stop: Arc<AtomicBool>) -> io::Result<()> {
		trace!(target: "ipc", "Waiting for client: [{}]", addr);
		try!(self.waiting_pipe.connect());
		trace!(target: "ipc", "Received connection to address [{}]", addr);
		if stop.load(Ordering::Relaxed) {
			trace!(target: "ipc", "Stopped listening sequence [{}]", addr);
			return Ok(())
		}

		let mut connected_pipe = std::mem::replace::<NamedPipe>(&mut self.waiting_pipe,
			try!(NamedPipeBuilder::new(addr)
				.first(false)
				.accept_remote(true)
				.inbound(true)
				.outbound(true)
				.out_buffer_size(MAX_REQUEST_LEN)
				.in_buffer_size(MAX_REQUEST_LEN)
				.create()));

		let thread_handler = self.io_handler.clone();
		std::thread::spawn(move || {
			let mut buf = vec![0u8; MAX_REQUEST_LEN as usize];
			let mut fin = REQUEST_READ_BATCH;
			loop {
				let start = fin - REQUEST_READ_BATCH;
				trace!(target: "ipc", "Reading {} - {} of the buffer", start, fin);
				match connected_pipe.read(&mut buf[start..fin]) {
					Ok(size) => {
						let (requests, last_index) = {
							let effective = &buf[0..start + size];
							fin = fin + size;
							trace!(target: "ipc", "Received rpc data: {} bytes", effective.len());

							validator::extract_requests(effective)
						};
						if requests.len() > 0 {
							let mut response_buf = Vec::new();
							for rpc_msg in requests  {
								trace!(target: "ipc", "Request: {}", rpc_msg);
								let meta = Default::default();
								let response: Option<String> = thread_handler.handle_request_sync(&rpc_msg, meta);

								if let Some(response_str) = response {
									trace!(target: "ipc", "Response: {}", &response_str);
									response_buf.extend(response_str.into_bytes());
								}
							}

							if let Err(write_err) = connected_pipe.write_all(&response_buf[..]).and_then(|_| connected_pipe.flush()) {
								trace!(target: "ipc", "Response write error: {:?}", write_err);
							}
							else {
								trace!(target: "ipc", "Sent rpc response: {} bytes", response_buf.len());
							}

							let leftover_len = start + size - (last_index + 1);
							if leftover_len > 0 {
								let leftover = buf[last_index + 1..start + size].to_vec();
								buf[0..leftover_len].copy_from_slice(&leftover[..]);
							}
							fin = leftover_len + REQUEST_READ_BATCH;
						}
						else { continue; }
					},
					Err(e) => {
						// closed connection
						trace!(target: "ipc", "Dropped connection {:?}", e);
						break;
					}
				}
			}
		});

		Ok(())
	}
}

pub struct Server<M: Metadata = ()> {
	is_stopping: Arc<AtomicBool>,
	is_stopped: Arc<AtomicBool>,
	io_handler: RpcHandler<M>,
	rpc_event_loop: Mutex<Option<RpcEventLoopHandle>>,
	addr: String,
}

impl<M: Metadata> Server<M> {
	/// New server
	pub fn new<T>(socket_addr: &str, io_handler: T) -> Result<Server<M>> where
		T: Into<MetaIoHandler<M>>,
	{
		let rpc_loop = RpcEventLoop::spawn();
		let mut server = try!(Self::with_rpc_handler(socket_addr, rpc_loop.handler(Arc::new(io_handler.into()))));
		server.rpc_event_loop = Mutex::new(Some(rpc_loop.into()));
		Ok(server)
	}

	// New Server using RpcHandler
	pub fn with_rpc_handler(socket_addr: &str, io_handler: RpcHandler<M>) -> Result<Server<M>> {
		Ok(Server {
			io_handler: io_handler,
			is_stopping: Arc::new(AtomicBool::new(false)),
			is_stopped: Arc::new(AtomicBool::new(true)),
			rpc_event_loop: Mutex::new(None),
			addr: socket_addr.to_owned(),
		})
	}

	/// Run server (in this thread)
	pub fn run(&self) -> Result<()> {
		let mut pipe_handler = try!(PipeHandler::start(&self.addr, self.io_handler.clone()));
		loop  {
			try!(pipe_handler.handle_incoming(&self.addr, Arc::new(AtomicBool::new(false))));
		}
	}

	/// Run server (in separate thread)
	pub fn run_async(&self) -> Result<()> {
		if self.is_stopping.load(Ordering::Relaxed) { return Err(Error::IsStopping) }
		if !self.is_stopped.load(Ordering::Relaxed) { return Err(Error::NotStopped) }

		trace!(target: "ipc", "Started named pipes server [{}]", self.addr);

		let thread_stopping = self.is_stopping.clone();
		let thread_stopped = self.is_stopped.clone();
		let thread_handler = self.io_handler.clone();
		let addr = self.addr.clone();
		std::thread::spawn(move || {
			let mut pipe_handler = PipeHandler::start(&addr, thread_handler).unwrap();
			while !thread_stopping.load(Ordering::Relaxed) {
				trace!(target: "ipc", "Accepting pipe connection");
				if let Err(pipe_listener_error) = pipe_handler.handle_incoming(&addr, thread_stopping.clone()) {
					trace!(target: "ipc", "Pipe listening error: {:?}", pipe_listener_error);
				}
			}
			thread_stopped.store(true, Ordering::Relaxed);
		});

		self.is_stopped.store(false, Ordering::Relaxed);
		Ok(())
	}

	pub fn stop_async(&self) -> Result<()> {
		self.rpc_event_loop.lock().unwrap().take().map(|s| s.close());
		if self.is_stopped.load(Ordering::Relaxed) { return Err(Error::NotStarted) }
		if self.is_stopping.load(Ordering::Relaxed) { return Err(Error::AlreadyStopping)}
		self.is_stopping.store(true, Ordering::Relaxed);
		Ok(())
	}

	pub fn stop(&self) -> Result<()> {
		self.rpc_event_loop.lock().unwrap().take().map(|s| s.close());
		if self.is_stopped.load(Ordering::Relaxed) { return Err(Error::NotStarted) }
		if self.is_stopping.load(Ordering::Relaxed) { return Err(Error::AlreadyStopping)}
		self.is_stopping.store(true, Ordering::Relaxed);
		while !self.is_stopped.load(Ordering::Relaxed) { std::thread::park_timeout(std::time::Duration::new(0, 50)); }
		Ok(())
	}
}

impl<M: Metadata> Drop for Server<M> {
	fn drop(&mut self) {
		self.stop_async().unwrap_or_else(|_| {}); // ignore error - can be stopped already
		// todo : no stable logging for windows?
		trace!(target: "ipc", "IPC Server : shutdown");
	}
}
<commit_msg>add wait with timeout<commit_after>// Copyright 2015, 2016 Ethcore (UK) Ltd.
// This file is part of Parity.

// Parity is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Parity is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Parity.  If not, see <http://www.gnu.org/licenses/>.

//! jsonrpc server over win named pipes
//!
//! ```no_run
//! extern crate jsonrpc_core;
//! extern crate jsonrpc_ipc_server;
//!
//! use std::sync::Arc;
//! use jsonrpc_core::*;
//! use jsonrpc_ipc_server::Server;
//!
//! fn main() {
//! 	let mut io = IoHandler::new();
//! 	io.add_method("say_hello", |_params| {
//!			Ok(Value::String("hello".into()))
//!		});
//! 	let server = Server::new("/tmp/json-ipc-test.ipc", io).unwrap();
//!     ::std::thread::spawn(move || server.run());
//! }
//! ```

//! Named pipes library

use miow::pipe::{NamedPipe, NamedPipeBuilder};
use std;
use std::io;
use std::io::{Read, Write};
use std::sync::atomic::*;
use std::sync::{Arc, Mutex};
use jsonrpc_core::{Metadata, MetaIoHandler};
use jsonrpc_core::reactor::{RpcHandler, RpcEventLoop, RpcEventLoopHandle};
use validator;

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

const MAX_REQUEST_LEN: u32 = 65536;
const REQUEST_READ_BATCH: usize = 4096;

#[derive(Debug)]
pub enum Error {
	Io(std::io::Error),
	NotStarted,
	AlreadyStopping,
	NotStopped,
	IsStopping,
}

impl std::convert::From<std::io::Error> for Error {
	fn from(io_error: std::io::Error) -> Error {
		Error::Io(io_error)
	}
}

pub struct PipeHandler<M: Metadata = ()> {
	waiting_pipe: NamedPipe,
	io_handler: RpcHandler<M>,
}

impl<M: Metadata> PipeHandler<M> {
	/// start ipc rpc server
	pub fn start(addr: &str, io_handler: RpcHandler<M>) -> Result<PipeHandler<M>> {
		Ok(PipeHandler {
			waiting_pipe: try!(
				NamedPipeBuilder::new(addr)
					.first(true)
					.accept_remote(true)
					.max_instances(255)
					.inbound(true)
					.outbound(true)
					.out_buffer_size(MAX_REQUEST_LEN)
					.in_buffer_size(MAX_REQUEST_LEN)
					.create()
			),
			io_handler: io_handler.clone(),
		})
	}

	fn handle_incoming(&mut self, addr: &str, stop: Arc<AtomicBool>) -> io::Result<()> {
		trace!(target: "ipc", "Waiting for client: [{}]", addr);

		// blocking wait with small timeouts
		// allows check if the server is actually stopped to quit gracefully
		// (`connect` does not allow that, it will block indefinitely)
		loop {
			if let Ok(_) = NamedPipe::wait(addr, Some(200)) {
				try!(self.waiting_pipe.connect());
				trace!(target: "ipc", "Received connection to address [{}]", addr);
				break;
			}
			if stop.load(Ordering::Relaxed) {
				trace!(target: "ipc", "Stopped listening sequence [{}]", addr);
				return Ok(())
			}
		}

		let mut connected_pipe = std::mem::replace::<NamedPipe>(&mut self.waiting_pipe,
			try!(NamedPipeBuilder::new(addr)
				.first(false)
				.accept_remote(true)
				.inbound(true)
				.outbound(true)
				.out_buffer_size(MAX_REQUEST_LEN)
				.in_buffer_size(MAX_REQUEST_LEN)
				.create()));

		let thread_handler = self.io_handler.clone();
		std::thread::spawn(move || {
			let mut buf = vec![0u8; MAX_REQUEST_LEN as usize];
			let mut fin = REQUEST_READ_BATCH;
			loop {
				let start = fin - REQUEST_READ_BATCH;
				trace!(target: "ipc", "Reading {} - {} of the buffer", start, fin);
				match connected_pipe.read(&mut buf[start..fin]) {
					Ok(size) => {
						let (requests, last_index) = {
							let effective = &buf[0..start + size];
							fin = fin + size;
							trace!(target: "ipc", "Received rpc data: {} bytes", effective.len());

							validator::extract_requests(effective)
						};
						if requests.len() > 0 {
							let mut response_buf = Vec::new();
							for rpc_msg in requests  {
								trace!(target: "ipc", "Request: {}", rpc_msg);
								let meta = Default::default();
								let response: Option<String> = thread_handler.handle_request_sync(&rpc_msg, meta);

								if let Some(response_str) = response {
									trace!(target: "ipc", "Response: {}", &response_str);
									response_buf.extend(response_str.into_bytes());
								}
							}

							if let Err(write_err) = connected_pipe.write_all(&response_buf[..]).and_then(|_| connected_pipe.flush()) {
								trace!(target: "ipc", "Response write error: {:?}", write_err);
							}
							else {
								trace!(target: "ipc", "Sent rpc response: {} bytes", response_buf.len());
							}

							let leftover_len = start + size - (last_index + 1);
							if leftover_len > 0 {
								let leftover = buf[last_index + 1..start + size].to_vec();
								buf[0..leftover_len].copy_from_slice(&leftover[..]);
							}
							fin = leftover_len + REQUEST_READ_BATCH;
						}
						else { continue; }
					},
					Err(e) => {
						// closed connection
						trace!(target: "ipc", "Dropped connection {:?}", e);
						break;
					}
				}
			}
		});

		Ok(())
	}
}

pub struct Server<M: Metadata = ()> {
	is_stopping: Arc<AtomicBool>,
	is_stopped: Arc<AtomicBool>,
	io_handler: RpcHandler<M>,
	rpc_event_loop: Mutex<Option<RpcEventLoopHandle>>,
	addr: String,
}

impl<M: Metadata> Server<M> {
	/// New server
	pub fn new<T>(socket_addr: &str, io_handler: T) -> Result<Server<M>> where
		T: Into<MetaIoHandler<M>>,
	{
		let rpc_loop = RpcEventLoop::spawn();
		let mut server = try!(Self::with_rpc_handler(socket_addr, rpc_loop.handler(Arc::new(io_handler.into()))));
		server.rpc_event_loop = Mutex::new(Some(rpc_loop.into()));
		Ok(server)
	}

	// New Server using RpcHandler
	pub fn with_rpc_handler(socket_addr: &str, io_handler: RpcHandler<M>) -> Result<Server<M>> {
		Ok(Server {
			io_handler: io_handler,
			is_stopping: Arc::new(AtomicBool::new(false)),
			is_stopped: Arc::new(AtomicBool::new(true)),
			rpc_event_loop: Mutex::new(None),
			addr: socket_addr.to_owned(),
		})
	}

	/// Run server (in this thread)
	pub fn run(&self) -> Result<()> {
		let mut pipe_handler = try!(PipeHandler::start(&self.addr, self.io_handler.clone()));
		loop  {
			try!(pipe_handler.handle_incoming(&self.addr, Arc::new(AtomicBool::new(false))));
		}
	}

	/// Run server (in separate thread)
	pub fn run_async(&self) -> Result<()> {
		if self.is_stopping.load(Ordering::Relaxed) { return Err(Error::IsStopping) }
		if !self.is_stopped.load(Ordering::Relaxed) { return Err(Error::NotStopped) }

		trace!(target: "ipc", "Started named pipes server [{}]", self.addr);

		let thread_stopping = self.is_stopping.clone();
		let thread_stopped = self.is_stopped.clone();
		let thread_handler = self.io_handler.clone();
		let addr = self.addr.clone();
		std::thread::spawn(move || {
			let mut pipe_handler = PipeHandler::start(&addr, thread_handler).unwrap();
			while !thread_stopping.load(Ordering::Relaxed) {
				trace!(target: "ipc", "Accepting pipe connection");
				if let Err(pipe_listener_error) = pipe_handler.handle_incoming(&addr, thread_stopping.clone()) {
					trace!(target: "ipc", "Pipe listening error: {:?}", pipe_listener_error);
				}
			}
			thread_stopped.store(true, Ordering::Relaxed);
		});

		self.is_stopped.store(false, Ordering::Relaxed);
		Ok(())
	}

	pub fn stop_async(&self) -> Result<()> {
		self.rpc_event_loop.lock().unwrap().take().map(|s| s.close());
		if self.is_stopped.load(Ordering::Relaxed) { return Err(Error::NotStarted) }
		if self.is_stopping.load(Ordering::Relaxed) { return Err(Error::AlreadyStopping)}
		self.is_stopping.store(true, Ordering::Relaxed);
		Ok(())
	}

	pub fn stop(&self) -> Result<()> {
		self.rpc_event_loop.lock().unwrap().take().map(|s| s.close());
		if self.is_stopped.load(Ordering::Relaxed) { return Err(Error::NotStarted) }
		if self.is_stopping.load(Ordering::Relaxed) { return Err(Error::AlreadyStopping)}
		self.is_stopping.store(true, Ordering::Relaxed);
		while !self.is_stopped.load(Ordering::Relaxed) { std::thread::park_timeout(std::time::Duration::new(0, 50)); }
		Ok(())
	}
}

impl<M: Metadata> Drop for Server<M> {
	fn drop(&mut self) {
		self.stop_async().unwrap_or_else(|_| {}); // ignore error - can be stopped already
		// todo : no stable logging for windows?
		trace!(target: "ipc", "IPC Server : shutdown");
	}
}
<|endoftext|>


<commit_before>extern crate crossbeam;
extern crate walkdir;

pub mod single_threaded;

use self::crossbeam::deque::fifo;
use self::crossbeam::deque::Pop;
use self::crossbeam::deque::Steal;
use self::crossbeam::deque::Worker;
use self::walkdir::WalkDir;
use colored::*;
use error::*;
use regex::{Regex, RegexSet};
use std::ffi::OsStr;
use std::fs;
use std::path::Path;
use std::path::PathBuf;
use std::process::exit;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::thread;
use types::FileSize;
use utils::size;

pub use walk_parallel::single_threaded::*;

/// Enum for messaging between workers/stealers
pub enum Status<T> {
    Done,
    Data(T),
}

/// The 'Walk' struct contains all the information we need to traverse a directory.
#[derive(Debug)]
pub struct Walk {
    pub path: PathBuf,
    gitignore: Option<RegexSet>,
    hgignore: Option<RegexSet>,
    darcs_boring: Option<RegexSet>,
    excludes: Option<Regex>,
    max_depth: Option<u8>,
    threshold: Option<u64>,
    start_depth: usize,
    nproc: usize,
    show_files: bool,
    get_blocks: bool,
    follow_symlinks: bool,
    artifacts_only: bool,
}

impl Walk {
    /// function to make output from a 'Walk', using one thread. It also takes an 'Arc<AtomicU64>'
    /// and will add the relevant directory sizes to it.
    pub fn print_dir(w: &Walk, total: &Arc<AtomicUsize>) {
        let excludes = match w.excludes {
            Some(ref x) => Some(x),
            _ => None,
        };

        let v = if excludes.is_some() || w.artifacts_only {
            read_all(
                &w.path,
                w.start_depth as u8,
                w.max_depth,
                excludes,
                &w.gitignore,
                false,
                w.artifacts_only,
            )
        } else {
            read_all_fast(&w.path, w.start_depth as u8, w.max_depth)
        };

        let subdir_size = v.file_size.get();

        total.fetch_add(subdir_size as usize, Ordering::Relaxed);

        let mut to_print = if let Some(m) = w.threshold {
            subdir_size > m
        } else {
            true
        };

        if let Some(0) = w.max_depth {
            to_print = false;
        }

        if to_print {
            // filter by depth
            let mut v_filtered = v.filtered(w.threshold, !w.show_files, w.max_depth);

            v_filtered.display_tree(&w.path);
        }
    }

    /// set the maximum depth to display
    pub fn set_depth(&mut self, d: u8) {
        self.max_depth = Some(d);
    }

    /// set the regex for excludes
    pub fn set_regex(&mut self, r: Regex) {
        self.excludes = Some(r);
    }

    /// set the minumum file size
    pub fn set_threshold(&mut self, n: u64) {
        self.threshold = Some(n);
    }

    /// include files when printing
    pub fn with_files(&mut self) {
        self.show_files = true;
    }

    /// include files when printing
    pub fn blocks(&mut self) {
        self.get_blocks = true;
    }

    /// include files when printing
    pub fn artifacts_only(&mut self) {
        self.artifacts_only = true;
    }

    fn get_proc(&self) -> usize {
        self.nproc
    }

    /// Create a new 'Walk' from a 'PathBuf' and the number
    /// of processor cores to be used.
    pub fn new(p: PathBuf, n: usize) -> Walk {
        Walk {
            path: p,
            gitignore: None,
            hgignore: None,
            darcs_boring: None,
            excludes: None,
            max_depth: None,
            threshold: None,
            start_depth: 0,
            nproc: n,
            show_files: false,
            get_blocks: false,
            follow_symlinks: false,
            artifacts_only: false,
        }
    }

    fn bump_depth(&mut self) {
        self.start_depth += 1;
    }

    /// This takes a 'Walk' and a 'Worker<Status<Walk>>' and executes the walk *in parallel*,
    /// creating new work for each subdirectory. It's not the most efficient concurrency
    /// imaginable, but it's fast and easy-ish to use. It *also* takes in an 'Arc<AtomicU64>',
    /// which it updates with any file sizes in the directory.
    pub fn push_subdir(w: &Walk, worker: &mut Worker<Status<Walk>>, total: &Arc<AtomicUsize>) {
        let in_paths = &w.path;

        // fill up queue + print out files
        if let Ok(paths) = fs::read_dir(in_paths) {
            // iterate over all the entries in the directory
            for p in paths {
                let val = match p {
                    Ok(x) => x,
                    _ => {
                        eprintln!("{}: path error at {:?}.", "Error".red(), p);
                        exit(0x0001)
                    }
                };

                let exclude_check = if let Some(ref x) = w.excludes {
                    if let Some(r) = val.path().into_os_string().to_str() {
                        !x.is_match(r)
                    } else {
                        eprintln!(
                            "{}: ignoring invalid unicode at: {:?}",
                            "Warning".yellow(),
                            val.path().display()
                        );
                        true
                    }
                } else {
                    true
                };

                if exclude_check {
                    match val.file_type() {
                        Ok(t) => {
                            if t.is_dir() {
                                let mut new_path = w.path.to_owned();
                                new_path.push(val.file_name());
                                let mut new_walk = Walk::new(new_path, w.get_proc());
                                if w.excludes.is_some() {
                                    new_walk.set_regex(w.excludes.clone().unwrap());
                                }
                                if w.show_files {
                                    new_walk.with_files();
                                }
                                new_walk.bump_depth();
                                if let Some(d) = w.max_depth {
                                    new_walk.set_depth(d);
                                }
                                if let Some(b) = w.threshold {
                                    new_walk.set_threshold(b);
                                }
                                worker.push(Status::Data(new_walk)); // pass a vector of Arc's to do 2-level traversals?
                            } else if t.is_file() {
                                if let Ok(l) = val.metadata() {
                                    let size = size(&l, w.get_blocks);
                                    total.fetch_add(size as usize, Ordering::Relaxed);
                                    if w.show_files && size != 0 {
                                        let to_formatted = format!("{}", FileSize::new(size));
                                        println!(
                                            "{}\t {}",
                                            &to_formatted.green(),
                                            val.path().display()
                                        );
                                    }
                                } else {
                                    eprintln!(
                                        "{}: could not find filesize for file at {}.",
                                        "Warning".yellow(),
                                        val.path().display()
                                    );
                                }
                            }
                        }
                        _ => eprintln!(
                            "{}: could not determine file type for: {}",
                            "Warning".yellow(),
                            val.path().display()
                        ),
                    }
                }
            }

            // send "done" messages to all the workers
            let iter = 0..(w.get_proc());
            iter.map(|_| worker.push(Status::Done)).count();
        }
        // if we can't read the directory contents, figure out why
        // 1: check the path exists
        else if !in_paths.exists() {
            eprintln!(
                "{}: path '{}' does not exist, or you do not have permission to enter.",
                "Error".red(),
                &in_paths.display()
            );
        }
        // 2: check the path is actually a directory
        else if !in_paths.is_dir() {
            if w.artifacts_only {
                eprintln!(
                    "{}: {} is not a directory; not searching for artifacts",
                    "Warning".yellow(),
                    &in_paths.display()
                );
            }

            if let Ok(l) = in_paths.metadata() {
                let size = size(&l, w.get_blocks); // l.len();
                let to_formatted = format!("{}", FileSize::new(size));
                println!("{}\t {}", &to_formatted.green(), in_paths.display());
            } else {
                panic!("{}", Internal::IoError);
            }
        }
        // 3: otherwise, give a warning about permissions
        else {
            eprintln!(
                "{}: permission denied for directory: {}",
                "Warning".yellow(),
                &in_paths.display()
            );
        }
    }
}

fn ats_cgen(p: Option<&OsStr>) -> bool {
    lazy_static! {
        static ref DATS_C: Regex =
            Regex::new(r"(_(d|h|s)ats\.c|_lats\.dats|_sats\.c|_stub\.h)$").unwrap();
    }
    match p {
        Some(p) => DATS_C.is_match(&p.to_string_lossy().to_string()),
        None => false,
    }
}

fn latex_log<P: AsRef<Path>>(p: P) -> bool {
    lazy_static! {
        static ref LOG: Regex = Regex::new(r"\.log$").unwrap();
    }

    if LOG.is_match(&p.as_ref().to_string_lossy().to_string()) {
        let mut parent = (&p.as_ref())
            .parent()
            .unwrap()
            .to_string_lossy()
            .to_string();
        parent.push_str("/*.tex");
        glob_exists(&parent)
    } else {
        false
    }
}

// TODO figure out why the unwrap_or is failing?
// FIXME take optional reference to a regex
pub fn clean_project_dirs<P: AsRef<Path>>(p: P, exclude: &Option<Regex>, _: bool) {
    lazy_static! {
        static ref REGEX: Regex =
            Regex::new(r"\.(a|i|ii|la|lo|o|keter|bc|dyn_o|d|rlib|crate|hi|hc|chi|dyn_hi|jsexe|webapp|js\.externs|ibc|toc|aux|fdb_latexmk|bbl|blg|fls|egg-info|whl|js_a|js_hi|jld|ji|js_o|so.*|dump-.*|vmb|crx|orig|elmo|elmi|hspec-failures|pyc|mod|vo|beam|agdai|go\.(v|teak|xmldef|rewrittenast|rewrittengo|simplego|tree-(bind|eval|finish|parse))|p_hi|p_o|prof|hide-cache|ghc\.environment\..*\d.\d.\d|tix|synctex\.gz|hl|sandbox\.config|hp|eventlog|ipa|ttc|chs\.h|chi)$")
            .unwrap();
    }

    for dir in WalkDir::new(p)
        .into_iter()
        .filter_map(|e| e.ok())
        .filter(|p| {
            exclude
                .clone()
                .map(|e| e.is_match(&p.path().to_string_lossy().to_string()))
                != Some(false)
        })
        .filter(|p| {
            REGEX.is_match(&p.path().to_string_lossy().to_string())
                || is_project_dir(
                    &p.path().to_string_lossy().to_string(),
                    &p.path()
                        .file_name()
                        .map(|x| x.to_string_lossy().to_string())
                        .unwrap_or_else(|| "".to_string()),
                )
                || latex_log(&p.path())
                || ats_cgen(p.path().file_name())
                || ({
                    let x = &p.path().to_string_lossy().to_string();
                    x.ends_with("/flxg_stats.txt")
                })
        })
    {
        if dir.file_type().is_file() {
            fs::remove_file(dir.path()).unwrap_or(());
        } else if dir.file_type().is_dir() {
            fs::remove_dir_all(dir.path()).unwrap_or(());
        }
    }
}

/// Given a 'Walk' struct, traverse it concurrently and print out any relevant outputs.
/// Currently, this only works for a depth of two, which is probably bad.
pub fn print_parallel(w: Walk) {
    // initialize the total at 0 and create a reference to it
    let val = AtomicUsize::new(0);
    let arc = Arc::new(val);
    let arc_producer = arc.clone();
    let arc_child = arc.clone();
    let path_display = w.path.clone();

    // set up worker & stealer
    let (mut worker, stealer) = fifo();

    // set up our iterator for the workers
    let iter = 0..(&w.get_proc() - 1);

    let _ = (&w.path).is_file();

    // create the producer in another thread
    let child_producer = thread::spawn(move || {
        let arc_local = arc_producer.clone();

        // assign work to everyone
        Walk::push_subdir(&w, &mut worker, &arc_local);

        // start popping off values in the worker's thread
        loop {
            if let Pop::Data(p) = worker.pop() {
                match p {
                    Status::Data(d) => Walk::print_dir(&d, &arc_local),
                    _ => break,
                }
            }
        }
    });

    // create a vector of thread handles so that it doesn't execute
    // everything sequentially.
    let mut threads = Vec::new();

    // set up as many workers as we have threads
    for _ in iter {
        // create a new stealer
        let stealer_clone = stealer.clone();

        let arc_local = arc_child.clone();

        // run the stealer in a new thread
        let child_consumer = thread::spawn(move || loop {
            if let Steal::Data(p) = stealer_clone.steal() {
                match p {
                    Status::Data(d) => Walk::print_dir(&d, &arc_local),
                    _ => break,
                }
            }
        });

        threads.push(child_consumer);
    }

    // join the child producer to the main thread
    let _ = child_producer.join();

    // join the workers to the main thread
    let _ = threads
        .into_iter()
        .map(|v| {
            let result = v.join();
            if let Ok(exit) = result {
                exit
            } else if let Err(e) = result {
                panic!("{:?}", e)
            }
        })
        .count();

    // get the total size
    let m = arc.load(Ordering::SeqCst); // TODO - check if this works with Relaxed?
    let size = FileSize::new(m as u64);

    // print directory total.
    if size != FileSize::new(0) {
        let to_formatted = format!("{}", size);
        println!("{}\t {}", &to_formatted.green(), path_display.display());
    }
}
<commit_msg>Doesn't accidentally corrupt your repositories<commit_after>extern crate crossbeam;
extern crate walkdir;

pub mod single_threaded;

use self::crossbeam::deque::fifo;
use self::crossbeam::deque::Pop;
use self::crossbeam::deque::Steal;
use self::crossbeam::deque::Worker;
use self::walkdir::WalkDir;
use colored::*;
use error::*;
use regex::{Regex, RegexSet};
use std::ffi::OsStr;
use std::fs;
use std::path::Path;
use std::path::PathBuf;
use std::process::exit;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::thread;
use types::FileSize;
use utils::size;

pub use walk_parallel::single_threaded::*;

/// Enum for messaging between workers/stealers
pub enum Status<T> {
    Done,
    Data(T),
}

/// The 'Walk' struct contains all the information we need to traverse a directory.
#[derive(Debug)]
pub struct Walk {
    pub path: PathBuf,
    gitignore: Option<RegexSet>,
    hgignore: Option<RegexSet>,
    darcs_boring: Option<RegexSet>,
    excludes: Option<Regex>,
    max_depth: Option<u8>,
    threshold: Option<u64>,
    start_depth: usize,
    nproc: usize,
    show_files: bool,
    get_blocks: bool,
    follow_symlinks: bool,
    artifacts_only: bool,
}

impl Walk {
    /// function to make output from a 'Walk', using one thread. It also takes an 'Arc<AtomicU64>'
    /// and will add the relevant directory sizes to it.
    pub fn print_dir(w: &Walk, total: &Arc<AtomicUsize>) {
        let excludes = match w.excludes {
            Some(ref x) => Some(x),
            _ => None,
        };

        let v = if excludes.is_some() || w.artifacts_only {
            read_all(
                &w.path,
                w.start_depth as u8,
                w.max_depth,
                excludes,
                &w.gitignore,
                false,
                w.artifacts_only,
            )
        } else {
            read_all_fast(&w.path, w.start_depth as u8, w.max_depth)
        };

        let subdir_size = v.file_size.get();

        total.fetch_add(subdir_size as usize, Ordering::Relaxed);

        let mut to_print = if let Some(m) = w.threshold {
            subdir_size > m
        } else {
            true
        };

        if let Some(0) = w.max_depth {
            to_print = false;
        }

        if to_print {
            // filter by depth
            let mut v_filtered = v.filtered(w.threshold, !w.show_files, w.max_depth);

            v_filtered.display_tree(&w.path);
        }
    }

    /// set the maximum depth to display
    pub fn set_depth(&mut self, d: u8) {
        self.max_depth = Some(d);
    }

    /// set the regex for excludes
    pub fn set_regex(&mut self, r: Regex) {
        self.excludes = Some(r);
    }

    /// set the minumum file size
    pub fn set_threshold(&mut self, n: u64) {
        self.threshold = Some(n);
    }

    /// include files when printing
    pub fn with_files(&mut self) {
        self.show_files = true;
    }

    /// include files when printing
    pub fn blocks(&mut self) {
        self.get_blocks = true;
    }

    /// include files when printing
    pub fn artifacts_only(&mut self) {
        self.artifacts_only = true;
    }

    fn get_proc(&self) -> usize {
        self.nproc
    }

    /// Create a new 'Walk' from a 'PathBuf' and the number
    /// of processor cores to be used.
    pub fn new(p: PathBuf, n: usize) -> Walk {
        Walk {
            path: p,
            gitignore: None,
            hgignore: None,
            darcs_boring: None,
            excludes: None,
            max_depth: None,
            threshold: None,
            start_depth: 0,
            nproc: n,
            show_files: false,
            get_blocks: false,
            follow_symlinks: false,
            artifacts_only: false,
        }
    }

    fn bump_depth(&mut self) {
        self.start_depth += 1;
    }

    /// This takes a 'Walk' and a 'Worker<Status<Walk>>' and executes the walk *in parallel*,
    /// creating new work for each subdirectory. It's not the most efficient concurrency
    /// imaginable, but it's fast and easy-ish to use. It *also* takes in an 'Arc<AtomicU64>',
    /// which it updates with any file sizes in the directory.
    pub fn push_subdir(w: &Walk, worker: &mut Worker<Status<Walk>>, total: &Arc<AtomicUsize>) {
        let in_paths = &w.path;

        // fill up queue + print out files
        if let Ok(paths) = fs::read_dir(in_paths) {
            // iterate over all the entries in the directory
            for p in paths {
                let val = match p {
                    Ok(x) => x,
                    _ => {
                        eprintln!("{}: path error at {:?}.", "Error".red(), p);
                        exit(0x0001)
                    }
                };

                let exclude_check = if let Some(ref x) = w.excludes {
                    if let Some(r) = val.path().into_os_string().to_str() {
                        !x.is_match(r)
                    } else {
                        eprintln!(
                            "{}: ignoring invalid unicode at: {:?}",
                            "Warning".yellow(),
                            val.path().display()
                        );
                        true
                    }
                } else {
                    true
                };

                if exclude_check {
                    match val.file_type() {
                        Ok(t) => {
                            if t.is_dir() {
                                let mut new_path = w.path.to_owned();
                                new_path.push(val.file_name());
                                let mut new_walk = Walk::new(new_path, w.get_proc());
                                if w.excludes.is_some() {
                                    new_walk.set_regex(w.excludes.clone().unwrap());
                                }
                                if w.show_files {
                                    new_walk.with_files();
                                }
                                new_walk.bump_depth();
                                if let Some(d) = w.max_depth {
                                    new_walk.set_depth(d);
                                }
                                if let Some(b) = w.threshold {
                                    new_walk.set_threshold(b);
                                }
                                worker.push(Status::Data(new_walk)); // pass a vector of Arc's to do 2-level traversals?
                            } else if t.is_file() {
                                if let Ok(l) = val.metadata() {
                                    let size = size(&l, w.get_blocks);
                                    total.fetch_add(size as usize, Ordering::Relaxed);
                                    if w.show_files && size != 0 {
                                        let to_formatted = format!("{}", FileSize::new(size));
                                        println!(
                                            "{}\t {}",
                                            &to_formatted.green(),
                                            val.path().display()
                                        );
                                    }
                                } else {
                                    eprintln!(
                                        "{}: could not find filesize for file at {}.",
                                        "Warning".yellow(),
                                        val.path().display()
                                    );
                                }
                            }
                        }
                        _ => eprintln!(
                            "{}: could not determine file type for: {}",
                            "Warning".yellow(),
                            val.path().display()
                        ),
                    }
                }
            }

            // send "done" messages to all the workers
            let iter = 0..(w.get_proc());
            iter.map(|_| worker.push(Status::Done)).count();
        }
        // if we can't read the directory contents, figure out why
        // 1: check the path exists
        else if !in_paths.exists() {
            eprintln!(
                "{}: path '{}' does not exist, or you do not have permission to enter.",
                "Error".red(),
                &in_paths.display()
            );
        }
        // 2: check the path is actually a directory
        else if !in_paths.is_dir() {
            if w.artifacts_only {
                eprintln!(
                    "{}: {} is not a directory; not searching for artifacts",
                    "Warning".yellow(),
                    &in_paths.display()
                );
            }

            if let Ok(l) = in_paths.metadata() {
                let size = size(&l, w.get_blocks); // l.len();
                let to_formatted = format!("{}", FileSize::new(size));
                println!("{}\t {}", &to_formatted.green(), in_paths.display());
            } else {
                panic!("{}", Internal::IoError);
            }
        }
        // 3: otherwise, give a warning about permissions
        else {
            eprintln!(
                "{}: permission denied for directory: {}",
                "Warning".yellow(),
                &in_paths.display()
            );
        }
    }
}

fn ats_cgen(p: Option<&OsStr>) -> bool {
    lazy_static! {
        static ref DATS_C: Regex =
            Regex::new(r"(_(d|h|s)ats\.c|_lats\.dats|_sats\.c|_stub\.h)$").unwrap();
    }
    match p {
        Some(p) => DATS_C.is_match(&p.to_string_lossy().to_string()),
        None => false,
    }
}

fn latex_log<P: AsRef<Path>>(p: P) -> bool {
    lazy_static! {
        static ref LOG: Regex = Regex::new(r"\.log$").unwrap();
    }

    if LOG.is_match(&p.as_ref().to_string_lossy().to_string()) {
        let mut parent = (&p.as_ref())
            .parent()
            .unwrap()
            .to_string_lossy()
            .to_string();
        parent.push_str("/*.tex");
        glob_exists(&parent)
    } else {
        false
    }
}

// TODO figure out why the unwrap_or is failing?
// FIXME take optional reference to a regex
pub fn clean_project_dirs<P: AsRef<Path>>(p: P, exclude: &Option<Regex>, _: bool) {
    lazy_static! {
        static ref REGEX: Regex =
            Regex::new(r"\.(a|i|ii|la|lo|o|keter|bc|dyn_o|d|rlib|crate|hi|hc|chi|dyn_hi|jsexe|webapp|js\.externs|ibc|toc|aux|fdb_latexmk|bbl|blg|fls|egg-info|whl|js_a|js_hi|jld|ji|js_o|so.*|dump-.*|vmb|crx|orig|elmo|elmi|hspec-failures|pyc|mod|vo|beam|agdai|go\.(v|teak|xmldef|rewrittenast|rewrittengo|simplego|tree-(bind|eval|finish|parse))|p_hi|p_o|prof|hide-cache|ghc\.environment\..*\d.\d.\d|tix|synctex\.gz|hl|sandbox\.config|hp|eventlog|ipa|ttc|chs\.h|chi)$")
            .unwrap();
    }
    lazy_static! {
        static ref SRC_CONTROL: Regex = Regex::new(r"(_darcs|\.(git|hg|pijul))").unwrap();
    }

    for dir in WalkDir::new(p)
        .into_iter()
        .filter_map(|e| e.ok())
        .filter(|p| {
            exclude
                .clone()
                .map(|e| e.is_match(&p.path().to_string_lossy().to_string()))
                != Some(false)
        })
        .filter(|p| !SRC_CONTROL.is_match(&p.path().to_string_lossy().to_string()))
        .filter(|p| {
            REGEX.is_match(&p.path().to_string_lossy().to_string())
                || is_project_dir(
                    &p.path().to_string_lossy().to_string(),
                    &p.path()
                        .file_name()
                        .map(|x| x.to_string_lossy().to_string())
                        .unwrap_or_else(|| "".to_string()),
                )
                || latex_log(&p.path())
                || ats_cgen(p.path().file_name())
                || ({
                    let x = &p.path().to_string_lossy().to_string();
                    x.ends_with("/flxg_stats.txt")
                })
        })
    {
        if dir.file_type().is_file() {
            fs::remove_file(dir.path()).unwrap_or(());
        } else if dir.file_type().is_dir() {
            fs::remove_dir_all(dir.path()).unwrap_or(());
        }
    }
}

/// Given a 'Walk' struct, traverse it concurrently and print out any relevant outputs.
/// Currently, this only works for a depth of two, which is probably bad.
pub fn print_parallel(w: Walk) {
    // initialize the total at 0 and create a reference to it
    let val = AtomicUsize::new(0);
    let arc = Arc::new(val);
    let arc_producer = arc.clone();
    let arc_child = arc.clone();
    let path_display = w.path.clone();

    // set up worker & stealer
    let (mut worker, stealer) = fifo();

    // set up our iterator for the workers
    let iter = 0..(&w.get_proc() - 1);

    let _ = (&w.path).is_file();

    // create the producer in another thread
    let child_producer = thread::spawn(move || {
        let arc_local = arc_producer.clone();

        // assign work to everyone
        Walk::push_subdir(&w, &mut worker, &arc_local);

        // start popping off values in the worker's thread
        loop {
            if let Pop::Data(p) = worker.pop() {
                match p {
                    Status::Data(d) => Walk::print_dir(&d, &arc_local),
                    _ => break,
                }
            }
        }
    });

    // create a vector of thread handles so that it doesn't execute
    // everything sequentially.
    let mut threads = Vec::new();

    // set up as many workers as we have threads
    for _ in iter {
        // create a new stealer
        let stealer_clone = stealer.clone();

        let arc_local = arc_child.clone();

        // run the stealer in a new thread
        let child_consumer = thread::spawn(move || loop {
            if let Steal::Data(p) = stealer_clone.steal() {
                match p {
                    Status::Data(d) => Walk::print_dir(&d, &arc_local),
                    _ => break,
                }
            }
        });

        threads.push(child_consumer);
    }

    // join the child producer to the main thread
    let _ = child_producer.join();

    // join the workers to the main thread
    let _ = threads
        .into_iter()
        .map(|v| {
            let result = v.join();
            if let Ok(exit) = result {
                exit
            } else if let Err(e) = result {
                panic!("{:?}", e)
            }
        })
        .count();

    // get the total size
    let m = arc.load(Ordering::SeqCst); // TODO - check if this works with Relaxed?
    let size = FileSize::new(m as u64);

    // print directory total.
    if size != FileSize::new(0) {
        let to_formatted = format!("{}", size);
        println!("{}\t {}", &to_formatted.green(), path_display.display());
    }
}
<|endoftext|>


<commit_before>fn main() {
    let mut prevprev = 1;
    let mut prev = 2;
    let mut sum = 2;
    loop {
        let current = prevprev + prev;
        if (current % 2 == 0) { sum += current; }
        if (current >= 4000000) { break; }
        prevprev = prev;
        prev = current;
    }
    println!("sum == {}", sum);
}
<commit_msg>euler: fix bug (that didn't affect result) in rust solution of problem 2<commit_after>fn main() {
    let mut prevprev = 1;
    let mut prev = 2;
    let mut sum = 2;
    loop {
        let current = prevprev + prev;
        if (current >= 4000000) { break; }
        if (current % 2 == 0) { sum += current; }
        prevprev = prev;
        prev = current;
    }
    println!("sum == {}", sum);
}
<|endoftext|>


<commit_before>use std::net::{TcpStream, Shutdown};
use std::io::{Stdin, stdin, stdout, Read, Write, Error, Result, Cursor};
use std::fs::File;
use std::env;
use std::str;
use std::option::Option;
use std::process::{Command, Child, exit};
use std::path::PathBuf;
use std::thread::sleep;
use std::time::Duration;
use wait_timeout::ChildExt;
use Input::*;

extern crate wait_timeout;

const MAX_RETRIES: usize = 5;

const JGRAB_INFO: &str = "\
=============== JGrab Client ================
 - https://github.com/renatoathaydes/jgrab -
=============================================
Jgrab can execute Java code from stdin (if not given any argument),
a Java file, or a Java snippet.

This is the native JGrab Client, written in Rust!

A Java daemon is started the first time the JGrab Client is run so
that subsequent runs are much faster.";

const JGRAB_USAGE: &str = "\
Usage:
  jgrab [<option> | java_file | -e java_snippet]
Options:
  --stop -s
    Stops the JGrab daemon.
  --help -h
    Shows usage help.";

/// All possible sources of input for the JGrab Client
enum Input {
    FileInput(File),
    StdinInput(Stdin),
    TextInput(Cursor<String>)
}

impl Read for Input {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        match *self {
            FileInput(ref mut r) => r.read(buf),
            StdinInput(ref mut r) => r.read(buf),
            TextInput(ref mut r) => r.read(buf),
        }
    }
}

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

    let input: Input;

    if args.len() == 1 {
        // no args, pipe stdin
        input = StdinInput(stdin());
    } else if args.len() == 2 && !args[1].starts_with('-') {
        // there's one argument and it is not an option, so it must be a file

        match File::open(&args[1]) {
            Ok(file) => input = FileInput(file),
            Err(err) => error(&format!("unable to read file: {}", err))
        }
    } else if args.len() == 2 {
        match args[1].trim() {
            "--help" | "-h" => {
                println!("{}\n\n{}", JGRAB_INFO, JGRAB_USAGE);
                return
            }
            "--stop" | "-s" => {
                if let Err(_) = connect() {
                    log("daemon is not running");
                    return
                }
                input = TextInput(Cursor::new("--stop".to_string()))
            }
            _ => {
                usage_error(&format!("invalid option"));
            }
        }
    } else {
        // just pass on the arguments to JGrab
        input = TextInput(Cursor::new(create_message(&args)))
    }

    send_message_retrying(input);
}

fn create_message(args: &Vec<String>) -> String {
    args[1..].join(" ")
}

fn send_message_retrying<R: Read>(mut reader: R) {
    if let Some(_) = send_message(&mut reader, false) {
        // failed to connect, try to start the daemon, then retry
        let mut retries = MAX_RETRIES;

        let mut child = start_daemon();
        check_status(&mut child);

        while retries > 0 {
            if let Some(err) = send_message(&mut reader, true) {
                check_status(&mut child);

                log(&format!("unable to connect to JGrab daemon: {}", err));
                log(&format!("will re-try {} more times", &mut retries));
                retries -= 1;
            } else {
                break // success
            }
        }

        if retries == 0 {
            error("unable to start JGrab daemon. \
                   Make sure JGrab's port [5002] is not already bound");
        }
    }
}

fn connect() -> Result<TcpStream> {
    TcpStream::connect("127.0.0.1:5002")
}

fn send_message<R: Read>(reader: &mut R,
                         is_retry: bool) -> Option<Error> {
    match connect() {
        Ok(mut stream) => {
            if is_retry {
                log("Connected!");
            }

            let mut socket_message = [0u8; 1024];

            loop {
                match reader.read(&mut socket_message) {
                    Ok(n) => {
                        if n == 0 {
                            break;
                        } else {
                            stream.write(&socket_message[0..n]).unwrap();
                        }
                    }
                    Err(err) => error(&err.to_string())
                }
            }

            stream.shutdown(Shutdown::Write).unwrap();

            let mut client_buffer = socket_message;

            loop {
                match stream.read(&mut client_buffer) {
                    Ok(n) => {
                        if n == 0 {
                            break;
                        } else {
                            stdout().write(&client_buffer[0..n]).unwrap();
                        }
                    }
                    Err(err) => error(&err.to_string())
                }
            }

            Option::None
        }
        Err(err) => Option::Some(err)
    }
}

fn start_daemon() -> Child {
    log("Starting daemon");

    if let Some(user_home) = env::home_dir() {
        let mut jgrab_jar: PathBuf = user_home;
        jgrab_jar.push(".jgrab");
        jgrab_jar.push("jgrab.jar");

        if jgrab_jar.as_path().is_file() {
            let cmd = Command::new("java")
                .arg("-jar")
                .arg(jgrab_jar.into_os_string().into_string().unwrap())
                .arg("--daemon")
                .spawn();

            match cmd {
                Ok(child) => {
                    log(&format!("Daemon started, pid={}", child.id()));
                    child
                }
                Err(err) => error(&err.to_string())
            }
        } else {
            error("The JGrab jar is not installed! Please install it as explained \
                   in https://github.com/renatoathaydes/jgrab");
        }
    } else {
        error("user.home could not be found, making it impossible to start the JGrab Daemon");
    }
}

fn check_status(child: &mut Child) {
    let timeout = Duration::from_secs(1);

    match child.wait_timeout(timeout) {
        Ok(Some(status)) => error(&format!(
            "The JGrab daemon has died prematurely, {}", status)),
        Ok(None) => {
            // give time for the socket server to become active
            sleep(Duration::from_secs(1));
        }
        Err(e) => error(&format!(
            "unable to wait for JGrab daemon process status: {}", e)),
    }
}

fn log(message: &str) {
    println!("=== JGrab Client - {} ===", message)
}

fn usage_error(message: &str) -> ! {
    println!("### {} ###\n\n{}", message, JGRAB_USAGE);
    exit(2)
}

fn error(message: &str) -> ! {
    println!("### JGrab Client Error - {} ###", message);
    exit(1)
}
<commit_msg>Added JGrab Client -v --version option.<commit_after>use std::net::{TcpStream, Shutdown};
use std::io::{Stdin, stdin, stdout, Read, Write, Error, Result, Cursor};
use std::fs::File;
use std::env;
use std::str;
use std::option::Option;
use std::process::{Command, Child, exit};
use std::path::PathBuf;
use std::thread::sleep;
use std::time::Duration;
use wait_timeout::ChildExt;
use Input::*;

extern crate wait_timeout;

const MAX_RETRIES: usize = 5;

const VERSION: &str = env!("CARGO_PKG_VERSION");

const JGRAB_INFO: &str = "\
=============== JGrab Client ================
 - https://github.com/renatoathaydes/jgrab -
=============================================
Jgrab can execute Java code from stdin (if not given any argument),
a Java file, or a Java snippet.

This is the native JGrab Client, written in Rust!

A Java daemon is started the first time the JGrab Client is run so
that subsequent runs are much faster.";

const JGRAB_USAGE: &str = "\
Usage:
  jgrab [<option> | java_file | -e java_snippet]
Options:
  --stop -s
    Stops the JGrab daemon.
  --help -h
    Shows usage help.";

/// All possible sources of input for the JGrab Client
enum Input {
    FileInput(File),
    StdinInput(Stdin),
    TextInput(Cursor<String>)
}

impl Read for Input {
    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        match *self {
            FileInput(ref mut r) => r.read(buf),
            StdinInput(ref mut r) => r.read(buf),
            TextInput(ref mut r) => r.read(buf),
        }
    }
}

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

    let input: Input;

    if args.len() == 1 {
        // no args, pipe stdin
        input = StdinInput(stdin());
    } else if args.len() == 2 && !args[1].starts_with('-') {
        // there's one argument and it is not an option, so it must be a file

        match File::open(&args[1]) {
            Ok(file) => input = FileInput(file),
            Err(err) => error(&format!("unable to read file: {}", err))
        }
    } else if args.len() == 2 {
        match args[1].trim() {
            "--help" | "-h" => {
                println!("{}\n\n{}", JGRAB_INFO, JGRAB_USAGE);
                return
            }
            "--stop" | "-s" => {
                if let Err(_) = connect() {
                    log("daemon is not running");
                    return
                }
                input = TextInput(Cursor::new("--stop".to_string()))
            }
            "--version" | "-v" => {
                println!("JGrab Client Version: {}", VERSION);
                return
            }
            _ => {
                usage_error(&format!("invalid option"));
            }
        }
    } else {
        // just pass on the arguments to JGrab
        input = TextInput(Cursor::new(create_message(&args)))
    }

    send_message_retrying(input);
}

fn create_message(args: &Vec<String>) -> String {
    args[1..].join(" ")
}

fn send_message_retrying<R: Read>(mut reader: R) {
    if let Some(_) = send_message(&mut reader, false) {
        // failed to connect, try to start the daemon, then retry
        let mut retries = MAX_RETRIES;

        let mut child = start_daemon();
        check_status(&mut child);

        while retries > 0 {
            if let Some(err) = send_message(&mut reader, true) {
                check_status(&mut child);

                log(&format!("unable to connect to JGrab daemon: {}", err));
                log(&format!("will re-try {} more times", &mut retries));
                retries -= 1;
            } else {
                break // success
            }
        }

        if retries == 0 {
            error("unable to start JGrab daemon. \
                   Make sure JGrab's port [5002] is not already bound");
        }
    }
}

fn connect() -> Result<TcpStream> {
    TcpStream::connect("127.0.0.1:5002")
}

fn send_message<R: Read>(reader: &mut R,
                         is_retry: bool) -> Option<Error> {
    match connect() {
        Ok(mut stream) => {
            if is_retry {
                log("Connected!");
            }

            let mut socket_message = [0u8; 1024];

            loop {
                match reader.read(&mut socket_message) {
                    Ok(n) => {
                        if n == 0 {
                            break;
                        } else {
                            stream.write(&socket_message[0..n]).unwrap();
                        }
                    }
                    Err(err) => error(&err.to_string())
                }
            }

            stream.shutdown(Shutdown::Write).unwrap();

            let mut client_buffer = socket_message;

            loop {
                match stream.read(&mut client_buffer) {
                    Ok(n) => {
                        if n == 0 {
                            break;
                        } else {
                            stdout().write(&client_buffer[0..n]).unwrap();
                        }
                    }
                    Err(err) => error(&err.to_string())
                }
            }

            Option::None
        }
        Err(err) => Option::Some(err)
    }
}

fn start_daemon() -> Child {
    log("Starting daemon");

    if let Some(user_home) = env::home_dir() {
        let mut jgrab_jar: PathBuf = user_home;
        jgrab_jar.push(".jgrab");
        jgrab_jar.push("jgrab.jar");

        if jgrab_jar.as_path().is_file() {
            let cmd = Command::new("java")
                .arg("-jar")
                .arg(jgrab_jar.into_os_string().into_string().unwrap())
                .arg("--daemon")
                .spawn();

            match cmd {
                Ok(child) => {
                    log(&format!("Daemon started, pid={}", child.id()));
                    child
                }
                Err(err) => error(&err.to_string())
            }
        } else {
            error("The JGrab jar is not installed! Please install it as explained \
                   in https://github.com/renatoathaydes/jgrab");
        }
    } else {
        error("user.home could not be found, making it impossible to start the JGrab Daemon");
    }
}

fn check_status(child: &mut Child) {
    let timeout = Duration::from_secs(1);

    match child.wait_timeout(timeout) {
        Ok(Some(status)) => error(&format!(
            "The JGrab daemon has died prematurely, {}", status)),
        Ok(None) => {
            // give time for the socket server to become active
            sleep(Duration::from_secs(1));
        }
        Err(e) => error(&format!(
            "unable to wait for JGrab daemon process status: {}", e)),
    }
}

fn log(message: &str) {
    println!("=== JGrab Client - {} ===", message)
}

fn usage_error(message: &str) -> ! {
    println!("### {} ###\n\n{}", message, JGRAB_USAGE);
    exit(2)
}

fn error(message: &str) -> ! {
    println!("### JGrab Client Error - {} ###", message);
    exit(1)
}
<|endoftext|>