<commit_before>use futures::{Async, Future, Poll, Sink};
use futures::stream::Stream;
use futures::sink::Send;
use futures_state_stream::{StateStream, StreamEvent};

use native_tls::TlsConnector;

use std::io;
use std::net::ToSocketAddrs;

use tokio_core::net::{TcpStream, TcpStreamNew};
use tokio_io::AsyncRead;
use tokio_core::reactor::Handle;
use tokio_tls::{ConnectAsync, TlsConnectorExt};

use proto::*;

pub struct Client {
    transport: ImapTransport,
    state: ClientState,
}

impl Client {
    pub fn connect(server: &str, handle: &Handle) -> ConnectFuture {
        let addr = format!("{}:993", server);
        let addr = addr.to_socket_addrs().unwrap().next().unwrap();
        let stream = TcpStream::connect(&addr, handle);
        ConnectFuture::TcpConnecting(stream, server.to_string())
    }

    pub fn call(self, cmd: Command) -> CommandFuture {
        let Client { transport, mut state } = self;
        let request_id = state.request_ids.next().unwrap();
        let (cmd_bytes, next_state) = cmd.to_parts();
        let future = transport.send(Request(request_id.clone(), cmd_bytes));
        CommandFuture::new(future, state, request_id, next_state)
    }
}

pub struct CommandFuture {
    future: Option<Send<ImapTransport>>,
    transport: Option<ImapTransport>,
    state: Option<ClientState>,
    request_id: RequestId,
    next_state: Option<State>,
    done: bool,
}

impl CommandFuture {
    pub fn new(future: Send<ImapTransport>, state: ClientState,
               request_id: RequestId, next_state: Option<State>) -> CommandFuture {
        CommandFuture {
            future: Some(future),
            transport: None,
            state: Some(state),
            request_id: request_id,
            next_state: next_state,
            done: false,
        }
    }
}

impl StateStream for CommandFuture {
    type Item = ResponseData;
    type State = Client;
    type Error = io::Error;
    fn poll(&mut self) -> Poll<StreamEvent<Self::Item, Self::State>, Self::Error> {
        if self.future.is_some() {
            let mut future = self.future.take().unwrap();
            match future.poll() {
                Ok(Async::Ready(transport)) => {
                    self.transport = Some(transport);
                },
                Ok(Async::NotReady) => {
                    self.future = Some(future);
                    return Ok(Async::NotReady);
                },
                Err(e) => {
                    return Err(e);
                },
            }
        }
        if !self.transport.is_some() {
            return Ok(Async::NotReady);
        }
        let mut transport = self.transport.take().unwrap();
        if self.done {
            let mut state = self.state.take().unwrap();
            if self.next_state.is_some() {
                state.state = self.next_state.take().unwrap();
            }
            let client = Client { transport, state };
            return Ok(Async::Ready(StreamEvent::Done(client)));
        }
        loop {
            match transport.poll() {
                Ok(Async::Ready(Some(rsp))) => {
                    if let Some(req_id) = rsp.request_id() {
                        self.done = *req_id == self.request_id;
                    };
                    self.transport = Some(transport);
                    return Ok(Async::Ready(StreamEvent::Next(rsp)));
                },
                Ok(Async::Ready(None)) | Ok(Async::NotReady) => {
                    break;
                },
                Err(e) => {
                    return Err(e);
                },
            }
        }
        self.transport = Some(transport);
        Ok(Async::NotReady)
    }
}

pub enum ConnectFuture {
    #[doc(hidden)]
    TcpConnecting(TcpStreamNew, String),
    #[doc(hidden)]
    TlsHandshake(ConnectAsync<TcpStream>),
    #[doc(hidden)]
    ServerGreeting(Option<ImapTransport>),
}

impl Future for ConnectFuture {
    type Item = (Client, ResponseData);
    type Error = io::Error;
    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut new = None;
        if let ConnectFuture::TcpConnecting(ref mut future, ref domain) = *self {
            let stream = try_ready!(future.poll());
            let ctx = TlsConnector::builder().unwrap().build().unwrap();
            let future = ctx.connect_async(domain, stream);
            new = Some(ConnectFuture::TlsHandshake(future));
        }
        if new.is_some() {
            *self = new.take().unwrap();
        }
        if let ConnectFuture::TlsHandshake(ref mut future) = *self {
            let transport = try_ready!(future.map_err(|e| {
                io::Error::new(io::ErrorKind::Other, e)
            }).poll()).framed(ImapCodec);
            new = Some(ConnectFuture::ServerGreeting(Some(transport)));
        }
        if new.is_some() {
            *self = new.take().unwrap();
        }
        if let ConnectFuture::ServerGreeting(ref mut wrapped) = *self {
            let msg = try_ready!(wrapped.as_mut().unwrap().poll()).unwrap();
            return Ok(Async::Ready((Client {
                transport: wrapped.take().unwrap(),
                state: ClientState::new(),
            }, msg)));
        }
        Ok(Async::NotReady)
    }
}

pub struct ClientState {
    state: State,
    request_ids: IdGenerator,
}

impl ClientState {
    pub fn new() -> ClientState {
        ClientState {
            state: State::NotAuthenticated,
            request_ids: IdGenerator::new(),
        }
    }
}
<commit_msg>Rename CommandFuture to ResponseStream<commit_after>use futures::{Async, Future, Poll, Sink};
use futures::stream::Stream;
use futures::sink::Send;
use futures_state_stream::{StateStream, StreamEvent};

use native_tls::TlsConnector;

use std::io;
use std::net::ToSocketAddrs;

use tokio_core::net::{TcpStream, TcpStreamNew};
use tokio_io::AsyncRead;
use tokio_core::reactor::Handle;
use tokio_tls::{ConnectAsync, TlsConnectorExt};

use proto::*;

pub struct Client {
    transport: ImapTransport,
    state: ClientState,
}

impl Client {
    pub fn connect(server: &str, handle: &Handle) -> ConnectFuture {
        let addr = format!("{}:993", server);
        let addr = addr.to_socket_addrs().unwrap().next().unwrap();
        let stream = TcpStream::connect(&addr, handle);
        ConnectFuture::TcpConnecting(stream, server.to_string())
    }

    pub fn call(self, cmd: Command) -> ResponseStream {
        let Client { transport, mut state } = self;
        let request_id = state.request_ids.next().unwrap();
        let (cmd_bytes, next_state) = cmd.to_parts();
        let future = transport.send(Request(request_id.clone(), cmd_bytes));
        ResponseStream::new(future, state, request_id, next_state)
    }
}

pub struct ResponseStream {
    future: Option<Send<ImapTransport>>,
    transport: Option<ImapTransport>,
    state: Option<ClientState>,
    request_id: RequestId,
    next_state: Option<State>,
    done: bool,
}

impl ResponseStream {
    pub fn new(future: Send<ImapTransport>, state: ClientState,
               request_id: RequestId, next_state: Option<State>) -> ResponseStream {
        ResponseStream {
            future: Some(future),
            transport: None,
            state: Some(state),
            request_id: request_id,
            next_state: next_state,
            done: false,
        }
    }
}

impl StateStream for ResponseStream {
    type Item = ResponseData;
    type State = Client;
    type Error = io::Error;
    fn poll(&mut self) -> Poll<StreamEvent<Self::Item, Self::State>, Self::Error> {
        if self.future.is_some() {
            let mut future = self.future.take().unwrap();
            match future.poll() {
                Ok(Async::Ready(transport)) => {
                    self.transport = Some(transport);
                },
                Ok(Async::NotReady) => {
                    self.future = Some(future);
                    return Ok(Async::NotReady);
                },
                Err(e) => {
                    return Err(e);
                },
            }
        }
        if !self.transport.is_some() {
            return Ok(Async::NotReady);
        }
        let mut transport = self.transport.take().unwrap();
        if self.done {
            let mut state = self.state.take().unwrap();
            if self.next_state.is_some() {
                state.state = self.next_state.take().unwrap();
            }
            let client = Client { transport, state };
            return Ok(Async::Ready(StreamEvent::Done(client)));
        }
        loop {
            match transport.poll() {
                Ok(Async::Ready(Some(rsp))) => {
                    if let Some(req_id) = rsp.request_id() {
                        self.done = *req_id == self.request_id;
                    };
                    self.transport = Some(transport);
                    return Ok(Async::Ready(StreamEvent::Next(rsp)));
                },
                Ok(Async::Ready(None)) | Ok(Async::NotReady) => {
                    break;
                },
                Err(e) => {
                    return Err(e);
                },
            }
        }
        self.transport = Some(transport);
        Ok(Async::NotReady)
    }
}

pub enum ConnectFuture {
    #[doc(hidden)]
    TcpConnecting(TcpStreamNew, String),
    #[doc(hidden)]
    TlsHandshake(ConnectAsync<TcpStream>),
    #[doc(hidden)]
    ServerGreeting(Option<ImapTransport>),
}

impl Future for ConnectFuture {
    type Item = (Client, ResponseData);
    type Error = io::Error;
    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut new = None;
        if let ConnectFuture::TcpConnecting(ref mut future, ref domain) = *self {
            let stream = try_ready!(future.poll());
            let ctx = TlsConnector::builder().unwrap().build().unwrap();
            let future = ctx.connect_async(domain, stream);
            new = Some(ConnectFuture::TlsHandshake(future));
        }
        if new.is_some() {
            *self = new.take().unwrap();
        }
        if let ConnectFuture::TlsHandshake(ref mut future) = *self {
            let transport = try_ready!(future.map_err(|e| {
                io::Error::new(io::ErrorKind::Other, e)
            }).poll()).framed(ImapCodec);
            new = Some(ConnectFuture::ServerGreeting(Some(transport)));
        }
        if new.is_some() {
            *self = new.take().unwrap();
        }
        if let ConnectFuture::ServerGreeting(ref mut wrapped) = *self {
            let msg = try_ready!(wrapped.as_mut().unwrap().poll()).unwrap();
            return Ok(Async::Ready((Client {
                transport: wrapped.take().unwrap(),
                state: ClientState::new(),
            }, msg)));
        }
        Ok(Async::NotReady)
    }
}

pub struct ClientState {
    state: State,
    request_ids: IdGenerator,
}

impl ClientState {
    pub fn new() -> ClientState {
        ClientState {
            state: State::NotAuthenticated,
            request_ids: IdGenerator::new(),
        }
    }
}
<|endoftext|>


<commit_before>use std::io::net::udp::UdpSocket;
use std::io::net::ip::SocketAddr;
use std::io::{IoResult, IoError, OtherIoError, TimedOut};
use std::io::Timer;
use std::comm::{Disconnected, Empty, Select};
use std::time::duration::Duration;
use packet::{Packet, PacketAccept, PacketReject, PacketDisconnect, PacketMessage, TaskCommand, Disconnect};
use shared::ConnectionConfig;
use time::now;

pub enum ConnectionState {
    CommsDisconnected,
    CommsConnecting,
    CommsConnected
}

pub enum PollFailResult {
    PollEmpty,
    PollDisconnected
}

fn reader_process(mut reader: UdpSocket, send: Sender<Packet>, recv: Receiver<TaskCommand>, target_addr: SocketAddr, protocol_id: u32, timeout_period: u32) {
    let mut buf = [0, ..1023];
    reader.set_timeout(Some(1000));

    let mut expires = now().to_timespec().sec + timeout_period as i64;

    loop {
        match reader.recv_from(buf) {
            Ok((amt, src)) => {
                if src == target_addr {
                    match Packet::deserialize(buf.slice_to(amt)) {
                        Ok(packet) => {
                            if packet.protocol_id == protocol_id {
                                send.send(packet);
                                expires = now().to_timespec().sec + timeout_period as i64;
                            }
                        },
                        Err(_) => ()
                    }
                }
            },
            Err(e) => {
                match e.kind {
                    TimedOut => {
                        match recv.try_recv() {
                            Ok(Disconnect) => {
                                break;
                            },
                            Err(Disconnected) => {
                                break;
                            },
                            Err(Empty) => {
                                //Keep going
                            }
                        }
                    },
                    _ => ()
                }
            }
        };
        if now().to_timespec().sec > expires {
            send.send(Packet::disconnect(protocol_id))
        }
    }
}

fn writer_process(mut writer: UdpSocket, _send: Sender<TaskCommand>, recv: Receiver<Packet>, target_addr: SocketAddr) {
    for msg in recv.iter() {
        match msg.serialize() {
            Ok(msg) => {
                match writer.send_to(msg.as_slice(), target_addr) {
                    Ok(()) => (),
                    Err(e) => println!("Error sending data - {}", e)
                }
            },
            Err(_) => ()
        }
    }
}

/**
 * Clientside implementation of UDP networking
 */
pub struct Client <T> {
    pub addr: SocketAddr,
    pub target_addr: SocketAddr,
    pub config: ConnectionConfig<T>,

    connection_state: ConnectionState,

    reader_send: Sender<TaskCommand>,
    reader_receive: Receiver<Packet>,
    writer_send: Sender<Packet>,
    writer_receive: Receiver<TaskCommand>,
}

/**
 * Additional configuration options for a Client connection
 */
pub struct ClientConnectionConfig {
    pub max_connect_retries: uint,
    pub connect_attempt_timeout: Duration
}

impl ClientConnectionConfig {

    pub fn new(max_connect_retries: uint, connect_attempt_timeout: Duration) -> ClientConnectionConfig {
        ClientConnectionConfig {
            max_connect_retries: max_connect_retries,
            connect_attempt_timeout: connect_attempt_timeout
        }
    }
}

impl <T> Client <T> {
    /**
     * Connect our Client to a target Server
     */
    pub fn connect(addr: SocketAddr, target_addr: SocketAddr, config: ConnectionConfig<T>, client_connection_config: ClientConnectionConfig) -> IoResult<Client<T>> {
         match UdpSocket::bind(addr) {
            Ok(reader) => {
                let writer = reader.clone();

                let (reader_send, reader_task_receive) = channel();
                let (reader_task_send, reader_receive) = channel();

                let protocol_id = config.protocol_id;
                let timeout_period = config.timeout_period;

                spawn(proc() {
                    reader_process(reader, reader_task_send, reader_task_receive, target_addr, protocol_id, timeout_period);
                });

                let (writer_send, writer_task_receive) = channel();
                let (writer_task_send, writer_receive) = channel();

                spawn(proc() {
                    writer_process(writer, writer_task_send, writer_task_receive, target_addr);
                });

                let mut client = Client {
                    addr: addr,
                    target_addr: target_addr,
                    reader_send: reader_send,
                    reader_receive: reader_receive,
                    writer_send: writer_send,
                    writer_receive: writer_receive,
                    connection_state: CommsDisconnected,
                    config: config
                };

                if client.connection_dance(client_connection_config.max_connect_retries, client_connection_config.connect_attempt_timeout) {
                    Ok(client)
                } else {
                    Err(IoError {
                        kind: OtherIoError,
                        desc: "Failed to connect",
                        detail: None
                    })
                }
            }
            Err(e) => Err(e)
        }
    }

    /**
     * A blocking connection request
     */
    fn connection_dance(&mut self, max_attempts: uint, timeout: Duration) -> bool {
        self.connection_state = CommsConnecting;
        let mut timer = Timer::new().unwrap();
        let mut attempts = 0u;

        while attempts < max_attempts && match self.connection_state { CommsConnecting => true, _ => false } {
            self.writer_send.send(Packet::connect(self.config.protocol_id));

            let timeout = timer.oneshot(timeout);

            //FIXME: Replace with the select! macro when it starts working
            let sel = Select::new();
            let mut reader = sel.handle(&self.reader_receive);
            let mut timeout = sel.handle(&timeout);
            unsafe { reader.add(); timeout.add(); }
            let ret = sel.wait();
            if ret == reader.id() {
                let packet = self.reader_receive.recv();
                match packet.packet_type {
                    PacketAccept => {
                        self.connection_state = CommsConnected;
                    }
                    PacketReject => {
                        self.connection_state = CommsDisconnected;
                    }
                    PacketDisconnect => {
                        self.connection_state = CommsDisconnected;
                    }
                    _ => (),
                }
            } else if ret == timeout.id() {
                let () = timeout.recv();
                attempts += 1;
            } else {
                fail!("What");
            }
        }

        match self.connection_state {
            CommsConnecting => {
                self.connection_state = CommsDisconnected;
                false
            },
            CommsDisconnected => false,
            CommsConnected => {
                true
            }
        }
    }

    /**
     * Pop the last event off of our comms queue, if any
     */
    pub fn poll(&mut self) -> Result<T, PollFailResult> {
        match self.connection_state {
            CommsConnected => {
                let mut result = Err(PollEmpty);
                loop {
                    match self.reader_receive.try_recv() {
                        Ok(value) => {
                            match value.packet_type {
                                PacketDisconnect => {
                                    self.connection_state = CommsDisconnected;
                                    result = Err(PollDisconnected);
                                    break;
                                },
                                PacketMessage => {
                                    match (self.config.packet_deserializer)(&value.packet_content.unwrap()) {
                                        Some(deserialized) => {
                                            result = Ok(deserialized);
                                            break;
                                        },
                                        None => ()
                                    }
                                },
                                _ => ()
                            }
                        },
                        _ => break
                    };
                }
                result
            },
            _ => Err(PollDisconnected)
        }
    }

    /**
     * Send a packet to the server
     */
    pub fn send(&mut self, packet: &T) {
        self.writer_send.send(Packet::message(self.config.protocol_id, (self.config.packet_serializer)(packet)));
    }
}

#[unsafe_destructor]
impl<T> Drop for Client<T> {

    fn drop(&mut self) {
        self.reader_send.send(Disconnect);
    }
}
<commit_msg>Swapped bad fail for more idiomatic 'unreachable'<commit_after>use std::io::net::udp::UdpSocket;
use std::io::net::ip::SocketAddr;
use std::io::{IoResult, IoError, OtherIoError, TimedOut};
use std::io::Timer;
use std::comm::{Disconnected, Empty, Select};
use std::time::duration::Duration;
use packet::{Packet, PacketAccept, PacketReject, PacketDisconnect, PacketMessage, TaskCommand, Disconnect};
use shared::ConnectionConfig;
use time::now;

pub enum ConnectionState {
    CommsDisconnected,
    CommsConnecting,
    CommsConnected
}

pub enum PollFailResult {
    PollEmpty,
    PollDisconnected
}

fn reader_process(mut reader: UdpSocket, send: Sender<Packet>, recv: Receiver<TaskCommand>, target_addr: SocketAddr, protocol_id: u32, timeout_period: u32) {
    let mut buf = [0, ..1023];
    reader.set_timeout(Some(1000));

    let mut expires = now().to_timespec().sec + timeout_period as i64;

    loop {
        match reader.recv_from(buf) {
            Ok((amt, src)) => {
                if src == target_addr {
                    match Packet::deserialize(buf.slice_to(amt)) {
                        Ok(packet) => {
                            if packet.protocol_id == protocol_id {
                                send.send(packet);
                                expires = now().to_timespec().sec + timeout_period as i64;
                            }
                        },
                        Err(_) => ()
                    }
                }
            },
            Err(e) => {
                match e.kind {
                    TimedOut => {
                        match recv.try_recv() {
                            Ok(Disconnect) => {
                                break;
                            },
                            Err(Disconnected) => {
                                break;
                            },
                            Err(Empty) => {
                                //Keep going
                            }
                        }
                    },
                    _ => ()
                }
            }
        };
        if now().to_timespec().sec > expires {
            send.send(Packet::disconnect(protocol_id))
        }
    }
}

fn writer_process(mut writer: UdpSocket, _send: Sender<TaskCommand>, recv: Receiver<Packet>, target_addr: SocketAddr) {
    for msg in recv.iter() {
        match msg.serialize() {
            Ok(msg) => {
                match writer.send_to(msg.as_slice(), target_addr) {
                    Ok(()) => (),
                    Err(e) => println!("Error sending data - {}", e)
                }
            },
            Err(_) => ()
        }
    }
}

/**
 * Clientside implementation of UDP networking
 */
pub struct Client <T> {
    pub addr: SocketAddr,
    pub target_addr: SocketAddr,
    pub config: ConnectionConfig<T>,

    connection_state: ConnectionState,

    reader_send: Sender<TaskCommand>,
    reader_receive: Receiver<Packet>,
    writer_send: Sender<Packet>,
    writer_receive: Receiver<TaskCommand>,
}

/**
 * Additional configuration options for a Client connection
 */
pub struct ClientConnectionConfig {
    pub max_connect_retries: uint,
    pub connect_attempt_timeout: Duration
}

impl ClientConnectionConfig {

    pub fn new(max_connect_retries: uint, connect_attempt_timeout: Duration) -> ClientConnectionConfig {
        ClientConnectionConfig {
            max_connect_retries: max_connect_retries,
            connect_attempt_timeout: connect_attempt_timeout
        }
    }
}

impl <T> Client <T> {
    /**
     * Connect our Client to a target Server
     */
    pub fn connect(addr: SocketAddr, target_addr: SocketAddr, config: ConnectionConfig<T>, client_connection_config: ClientConnectionConfig) -> IoResult<Client<T>> {
         match UdpSocket::bind(addr) {
            Ok(reader) => {
                let writer = reader.clone();

                let (reader_send, reader_task_receive) = channel();
                let (reader_task_send, reader_receive) = channel();

                let protocol_id = config.protocol_id;
                let timeout_period = config.timeout_period;

                spawn(proc() {
                    reader_process(reader, reader_task_send, reader_task_receive, target_addr, protocol_id, timeout_period);
                });

                let (writer_send, writer_task_receive) = channel();
                let (writer_task_send, writer_receive) = channel();

                spawn(proc() {
                    writer_process(writer, writer_task_send, writer_task_receive, target_addr);
                });

                let mut client = Client {
                    addr: addr,
                    target_addr: target_addr,
                    reader_send: reader_send,
                    reader_receive: reader_receive,
                    writer_send: writer_send,
                    writer_receive: writer_receive,
                    connection_state: CommsDisconnected,
                    config: config
                };

                if client.connection_dance(client_connection_config.max_connect_retries, client_connection_config.connect_attempt_timeout) {
                    Ok(client)
                } else {
                    Err(IoError {
                        kind: OtherIoError,
                        desc: "Failed to connect",
                        detail: None
                    })
                }
            }
            Err(e) => Err(e)
        }
    }

    /**
     * A blocking connection request
     */
    fn connection_dance(&mut self, max_attempts: uint, timeout: Duration) -> bool {
        self.connection_state = CommsConnecting;
        let mut timer = Timer::new().unwrap();
        let mut attempts = 0u;

        while attempts < max_attempts && match self.connection_state { CommsConnecting => true, _ => false } {
            self.writer_send.send(Packet::connect(self.config.protocol_id));

            let timeout = timer.oneshot(timeout);

            //FIXME: Replace with the select! macro when it starts working
            let sel = Select::new();
            let mut reader = sel.handle(&self.reader_receive);
            let mut timeout = sel.handle(&timeout);
            unsafe { reader.add(); timeout.add(); }
            let ret = sel.wait();
            if ret == reader.id() {
                let packet = self.reader_receive.recv();
                match packet.packet_type {
                    PacketAccept => {
                        self.connection_state = CommsConnected;
                    }
                    PacketReject => {
                        self.connection_state = CommsDisconnected;
                    }
                    PacketDisconnect => {
                        self.connection_state = CommsDisconnected;
                    }
                    _ => (),
                }
            } else if ret == timeout.id() {
                let () = timeout.recv();
                attempts += 1;
            } else {
                unreachable!();
            }
        }

        match self.connection_state {
            CommsConnecting => {
                self.connection_state = CommsDisconnected;
                false
            },
            CommsDisconnected => false,
            CommsConnected => {
                true
            }
        }
    }

    /**
     * Pop the last event off of our comms queue, if any
     */
    pub fn poll(&mut self) -> Result<T, PollFailResult> {
        match self.connection_state {
            CommsConnected => {
                let mut result = Err(PollEmpty);
                loop {
                    match self.reader_receive.try_recv() {
                        Ok(value) => {
                            match value.packet_type {
                                PacketDisconnect => {
                                    self.connection_state = CommsDisconnected;
                                    result = Err(PollDisconnected);
                                    break;
                                },
                                PacketMessage => {
                                    match (self.config.packet_deserializer)(&value.packet_content.unwrap()) {
                                        Some(deserialized) => {
                                            result = Ok(deserialized);
                                            break;
                                        },
                                        None => ()
                                    }
                                },
                                _ => ()
                            }
                        },
                        _ => break
                    };
                }
                result
            },
            _ => Err(PollDisconnected)
        }
    }

    /**
     * Send a packet to the server
     */
    pub fn send(&mut self, packet: &T) {
        self.writer_send.send(Packet::message(self.config.protocol_id, (self.config.packet_serializer)(packet)));
    }
}

#[unsafe_destructor]
impl<T> Drop for Client<T> {

    fn drop(&mut self) {
        self.reader_send.send(Disconnect);
    }
}
<|endoftext|>


<commit_before>//! Some common code used by both the event and reply modules.
use reply;
use serde_json as json;
use std::collections::HashMap;

/// Recursively build the tree of containers from the given json value.
pub fn build_tree(val: &json::Value) -> reply::Node {
    reply::Node {
        focus: match val.get("focus") {
            Some(xs) => xs
                .as_array()
                .unwrap()
                .iter()
                .map(|x| x.as_i64().unwrap())
                .collect(),
            None => vec![],
        },
        nodes: match val.get("nodes") {
            Some(nds) => nds
                .as_array()
                .unwrap()
                .iter()
                .map(|n| build_tree(n))
                .collect(),
            None => vec![],
        },
        floating_nodes: match val.get("floating_nodes") {
            Some(nds) => nds
                .as_array()
                .unwrap()
                .iter()
                .map(|n| build_tree(n))
                .collect(),
            None => vec![],
        },
        id: val.get("id").unwrap().as_i64().unwrap(),
        name: match val.get("name") {
            Some(n) => match n.as_str() {
                Some(s) => Some(s.to_owned()),
                None => None,
            },
            None => None,
        },
        nodetype: match val.get("type").unwrap().as_str().unwrap() {
            "root" => reply::NodeType::Root,
            "output" => reply::NodeType::Output,
            "con" => reply::NodeType::Con,
            "floating_con" => reply::NodeType::FloatingCon,
            "workspace" => reply::NodeType::Workspace,
            "dockarea" => reply::NodeType::DockArea,
            other => {
                warn!(target: "i3ipc", "Unknown NodeType {}", other);
                reply::NodeType::Unknown
            }
        },
        border: match val.get("border").unwrap().as_str().unwrap() {
            "normal" => reply::NodeBorder::Normal,
            "none" => reply::NodeBorder::None,
            "pixel" => reply::NodeBorder::Pixel,
            other => {
                warn!(target: "i3ipc", "Unknown NodeBorder {}", other);
                reply::NodeBorder::Unknown
            }
        },
        current_border_width: val.get("current_border_width").unwrap().as_i64().unwrap() as i32,
        layout: match val.get("layout").unwrap().as_str().unwrap() {
            "splith" => reply::NodeLayout::SplitH,
            "splitv" => reply::NodeLayout::SplitV,
            "stacked" => reply::NodeLayout::Stacked,
            "tabbed" => reply::NodeLayout::Tabbed,
            "dockarea" => reply::NodeLayout::DockArea,
            "output" => reply::NodeLayout::Output,
            other => {
                warn!(target: "i3ipc", "Unknown NodeLayout {}", other);
                reply::NodeLayout::Unknown
            }
        },
        percent: match *val.get("percent").unwrap() {
            json::Value::Number(ref f) => Some(f.as_f64().unwrap()),
            json::Value::Null => None,
            _ => unreachable!(),
        },
        rect: build_rect(val.get("rect").unwrap()),
        window_rect: build_rect(val.get("window_rect").unwrap()),
        deco_rect: build_rect(val.get("deco_rect").unwrap()),
        geometry: build_rect(val.get("geometry").unwrap()),
        window: match val.get("window").unwrap().clone() {
            json::Value::Number(i) => Some(i.as_i64().unwrap() as i32),
            json::Value::Null => None,
            _ => unreachable!(),
        },
        window_properties: build_window_properties(val.get("window_properties")),
        urgent: val.get("urgent").unwrap().as_bool().unwrap(),
        focused: val.get("focused").unwrap().as_bool().unwrap(),
    }
}

pub fn build_window_properties(j: Option<&json::Value>) -> Option<HashMap<reply::WindowProperty, String>> {
    match j {
        None => None,
        Some(props) => {
            let properties = props.as_object().unwrap();
            let mut map = HashMap::new();
            for (key, val) in properties {
                let window_property = match key.as_ref() {
                    "class" => Some(reply::WindowProperty::Class),
                    "instance" => Some(reply::WindowProperty::Instance),
                    "window_role" => Some(reply::WindowProperty::WindowRole),
                    "title" => Some(reply::WindowProperty::Title),
                    "transient_for" => Some(reply::WindowProperty::TransientFor),
                    other => {
                        warn!(target: "i3ipc", "Unknown WindowProperty {}", other);
                        return None;
                    }
                };
                if let Some(window_property) = window_property {
                    map.insert(window_property, val.as_str().unwrap_or_default().to_string());
                }
            }
            Some(map)
        }
    }
}

pub fn build_rect(jrect: &json::Value) -> (i32, i32, i32, i32) {
    let x = jrect.get("x").unwrap().as_i64().unwrap() as i32;
    let y = jrect.get("y").unwrap().as_i64().unwrap() as i32;
    let width = jrect.get("width").unwrap().as_i64().unwrap() as i32;
    let height = jrect.get("height").unwrap().as_i64().unwrap() as i32;
    (x, y, width, height)
}

pub fn build_bar_config(j: &json::Value) -> reply::BarConfig {
    reply::BarConfig {
        id: j.get("id").unwrap().as_str().unwrap().to_owned(),
        mode: j.get("mode").unwrap().as_str().unwrap().to_owned(),
        position: j.get("position").unwrap().as_str().unwrap().to_owned(),
        status_command: j
            .get("status_command")
            .unwrap()
            .as_str()
            .unwrap()
            .to_owned(),
        font: j.get("font").unwrap().as_str().unwrap().to_owned(),
        workspace_buttons: j.get("workspace_buttons").unwrap().as_bool().unwrap(),
        binding_mode_indicator: j.get("binding_mode_indicator").unwrap().as_bool().unwrap(),
        verbose: j.get("verbose").unwrap().as_bool().unwrap(),
        colors: {
            let colors = j.get("colors").unwrap().as_object().unwrap();
            let mut map = HashMap::new();
            for c in colors.keys() {
                let enum_key = match c.as_ref() {
                    "background" => reply::ColorableBarPart::Background,
                    "statusline" => reply::ColorableBarPart::Statusline,
                    "separator" => reply::ColorableBarPart::Separator,

                    #[cfg(feature = "i3-4-12")]
                    "focused_background" => reply::ColorableBarPart::FocusedBackground,

                    #[cfg(feature = "i3-4-12")]
                    "focused_statusline" => reply::ColorableBarPart::FocusedStatusline,

                    #[cfg(feature = "i3-4-12")]
                    "focused_separator" => reply::ColorableBarPart::FocusedSeparator,

                    "focused_workspace_text" => reply::ColorableBarPart::FocusedWorkspaceText,
                    "focused_workspace_bg" => reply::ColorableBarPart::FocusedWorkspaceBg,
                    "focused_workspace_border" => reply::ColorableBarPart::FocusedWorkspaceBorder,
                    "active_workspace_text" => reply::ColorableBarPart::ActiveWorkspaceText,
                    "active_workspace_bg" => reply::ColorableBarPart::ActiveWorkspaceBg,
                    "active_workspace_border" => reply::ColorableBarPart::ActiveWorkspaceBorder,
                    "inactive_workspace_text" => reply::ColorableBarPart::InactiveWorkspaceText,
                    "inactive_workspace_bg" => reply::ColorableBarPart::InactiveWorkspaceBg,
                    "inactive_workspace_border" => reply::ColorableBarPart::InactiveWorkspaceBorder,
                    "urgent_workspace_text" => reply::ColorableBarPart::UrgentWorkspaceText,
                    "urgent_workspace_bg" => reply::ColorableBarPart::UrgentWorkspaceBg,
                    "urgent_workspace_border" => reply::ColorableBarPart::UrgentWorkspaceBorder,
                    "binding_mode_text" => reply::ColorableBarPart::BindingModeText,
                    "binding_mode_bg" => reply::ColorableBarPart::BindingModeBg,
                    "binding_mode_border" => reply::ColorableBarPart::BindingModeBorder,
                    other => {
                        warn!(target: "i3ipc", "Unknown ColorableBarPart {}", other);
                        reply::ColorableBarPart::Unknown
                    }
                };
                let hex = colors.get(c).unwrap().as_str().unwrap().to_owned();
                map.insert(enum_key, hex);
            }
            map
        },
    }
}
<commit_msg>apply cargo fmt<commit_after>//! Some common code used by both the event and reply modules.
use reply;
use serde_json as json;
use std::collections::HashMap;

/// Recursively build the tree of containers from the given json value.
pub fn build_tree(val: &json::Value) -> reply::Node {
    reply::Node {
        focus: match val.get("focus") {
            Some(xs) => xs
                .as_array()
                .unwrap()
                .iter()
                .map(|x| x.as_i64().unwrap())
                .collect(),
            None => vec![],
        },
        nodes: match val.get("nodes") {
            Some(nds) => nds
                .as_array()
                .unwrap()
                .iter()
                .map(|n| build_tree(n))
                .collect(),
            None => vec![],
        },
        floating_nodes: match val.get("floating_nodes") {
            Some(nds) => nds
                .as_array()
                .unwrap()
                .iter()
                .map(|n| build_tree(n))
                .collect(),
            None => vec![],
        },
        id: val.get("id").unwrap().as_i64().unwrap(),
        name: match val.get("name") {
            Some(n) => match n.as_str() {
                Some(s) => Some(s.to_owned()),
                None => None,
            },
            None => None,
        },
        nodetype: match val.get("type").unwrap().as_str().unwrap() {
            "root" => reply::NodeType::Root,
            "output" => reply::NodeType::Output,
            "con" => reply::NodeType::Con,
            "floating_con" => reply::NodeType::FloatingCon,
            "workspace" => reply::NodeType::Workspace,
            "dockarea" => reply::NodeType::DockArea,
            other => {
                warn!(target: "i3ipc", "Unknown NodeType {}", other);
                reply::NodeType::Unknown
            }
        },
        border: match val.get("border").unwrap().as_str().unwrap() {
            "normal" => reply::NodeBorder::Normal,
            "none" => reply::NodeBorder::None,
            "pixel" => reply::NodeBorder::Pixel,
            other => {
                warn!(target: "i3ipc", "Unknown NodeBorder {}", other);
                reply::NodeBorder::Unknown
            }
        },
        current_border_width: val.get("current_border_width").unwrap().as_i64().unwrap() as i32,
        layout: match val.get("layout").unwrap().as_str().unwrap() {
            "splith" => reply::NodeLayout::SplitH,
            "splitv" => reply::NodeLayout::SplitV,
            "stacked" => reply::NodeLayout::Stacked,
            "tabbed" => reply::NodeLayout::Tabbed,
            "dockarea" => reply::NodeLayout::DockArea,
            "output" => reply::NodeLayout::Output,
            other => {
                warn!(target: "i3ipc", "Unknown NodeLayout {}", other);
                reply::NodeLayout::Unknown
            }
        },
        percent: match *val.get("percent").unwrap() {
            json::Value::Number(ref f) => Some(f.as_f64().unwrap()),
            json::Value::Null => None,
            _ => unreachable!(),
        },
        rect: build_rect(val.get("rect").unwrap()),
        window_rect: build_rect(val.get("window_rect").unwrap()),
        deco_rect: build_rect(val.get("deco_rect").unwrap()),
        geometry: build_rect(val.get("geometry").unwrap()),
        window: match val.get("window").unwrap().clone() {
            json::Value::Number(i) => Some(i.as_i64().unwrap() as i32),
            json::Value::Null => None,
            _ => unreachable!(),
        },
        window_properties: build_window_properties(val.get("window_properties")),
        urgent: val.get("urgent").unwrap().as_bool().unwrap(),
        focused: val.get("focused").unwrap().as_bool().unwrap(),
    }
}

pub fn build_window_properties(
    j: Option<&json::Value>,
) -> Option<HashMap<reply::WindowProperty, String>> {
    match j {
        None => None,
        Some(props) => {
            let properties = props.as_object().unwrap();
            let mut map = HashMap::new();
            for (key, val) in properties {
                let window_property = match key.as_ref() {
                    "class" => Some(reply::WindowProperty::Class),
                    "instance" => Some(reply::WindowProperty::Instance),
                    "window_role" => Some(reply::WindowProperty::WindowRole),
                    "title" => Some(reply::WindowProperty::Title),
                    "transient_for" => Some(reply::WindowProperty::TransientFor),
                    other => {
                        warn!(target: "i3ipc", "Unknown WindowProperty {}", other);
                        return None;
                    }
                };
                if let Some(window_property) = window_property {
                    map.insert(
                        window_property,
                        val.as_str().unwrap_or_default().to_string(),
                    );
                }
            }
            Some(map)
        }
    }
}

pub fn build_rect(jrect: &json::Value) -> (i32, i32, i32, i32) {
    let x = jrect.get("x").unwrap().as_i64().unwrap() as i32;
    let y = jrect.get("y").unwrap().as_i64().unwrap() as i32;
    let width = jrect.get("width").unwrap().as_i64().unwrap() as i32;
    let height = jrect.get("height").unwrap().as_i64().unwrap() as i32;
    (x, y, width, height)
}

pub fn build_bar_config(j: &json::Value) -> reply::BarConfig {
    reply::BarConfig {
        id: j.get("id").unwrap().as_str().unwrap().to_owned(),
        mode: j.get("mode").unwrap().as_str().unwrap().to_owned(),
        position: j.get("position").unwrap().as_str().unwrap().to_owned(),
        status_command: j
            .get("status_command")
            .unwrap()
            .as_str()
            .unwrap()
            .to_owned(),
        font: j.get("font").unwrap().as_str().unwrap().to_owned(),
        workspace_buttons: j.get("workspace_buttons").unwrap().as_bool().unwrap(),
        binding_mode_indicator: j.get("binding_mode_indicator").unwrap().as_bool().unwrap(),
        verbose: j.get("verbose").unwrap().as_bool().unwrap(),
        colors: {
            let colors = j.get("colors").unwrap().as_object().unwrap();
            let mut map = HashMap::new();
            for c in colors.keys() {
                let enum_key = match c.as_ref() {
                    "background" => reply::ColorableBarPart::Background,
                    "statusline" => reply::ColorableBarPart::Statusline,
                    "separator" => reply::ColorableBarPart::Separator,

                    #[cfg(feature = "i3-4-12")]
                    "focused_background" => reply::ColorableBarPart::FocusedBackground,

                    #[cfg(feature = "i3-4-12")]
                    "focused_statusline" => reply::ColorableBarPart::FocusedStatusline,

                    #[cfg(feature = "i3-4-12")]
                    "focused_separator" => reply::ColorableBarPart::FocusedSeparator,

                    "focused_workspace_text" => reply::ColorableBarPart::FocusedWorkspaceText,
                    "focused_workspace_bg" => reply::ColorableBarPart::FocusedWorkspaceBg,
                    "focused_workspace_border" => reply::ColorableBarPart::FocusedWorkspaceBorder,
                    "active_workspace_text" => reply::ColorableBarPart::ActiveWorkspaceText,
                    "active_workspace_bg" => reply::ColorableBarPart::ActiveWorkspaceBg,
                    "active_workspace_border" => reply::ColorableBarPart::ActiveWorkspaceBorder,
                    "inactive_workspace_text" => reply::ColorableBarPart::InactiveWorkspaceText,
                    "inactive_workspace_bg" => reply::ColorableBarPart::InactiveWorkspaceBg,
                    "inactive_workspace_border" => reply::ColorableBarPart::InactiveWorkspaceBorder,
                    "urgent_workspace_text" => reply::ColorableBarPart::UrgentWorkspaceText,
                    "urgent_workspace_bg" => reply::ColorableBarPart::UrgentWorkspaceBg,
                    "urgent_workspace_border" => reply::ColorableBarPart::UrgentWorkspaceBorder,
                    "binding_mode_text" => reply::ColorableBarPart::BindingModeText,
                    "binding_mode_bg" => reply::ColorableBarPart::BindingModeBg,
                    "binding_mode_border" => reply::ColorableBarPart::BindingModeBorder,
                    other => {
                        warn!(target: "i3ipc", "Unknown ColorableBarPart {}", other);
                        reply::ColorableBarPart::Unknown
                    }
                };
                let hex = colors.get(c).unwrap().as_str().unwrap().to_owned();
                map.insert(enum_key, hex);
            }
            map
        },
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use private::capability::ClientHook;
use private::units::*;
use message;
use message::{Allocator, ReaderSegments};
use {Error, OutputSegments, Result, Word};
use std::collections::HashMap;

pub type SegmentId = u32;

pub struct SegmentReader {
    pub arena: ArenaPtr,
    pub ptr: *const Word,
    pub size: WordCount32,
    pub read_limiter: ::std::rc::Rc<ReadLimiter>,
}

impl SegmentReader {
    #[inline]
    pub unsafe fn get_start_ptr(&self) -> *const Word {
        self.ptr.offset(0)
    }

    #[inline]
    pub fn contains_interval(&self, from : *const Word, to : *const Word) -> bool {
        let this_begin : usize = self.ptr as usize;
        let this_end : usize = unsafe { self.ptr.offset(self.size as isize) as usize };
        from as usize >= this_begin && to as usize <= this_end && from as usize <= to as usize &&
            self.read_limiter.can_read((to as usize - from as usize) as u64 / BYTES_PER_WORD as u64)
    }

    #[inline]
    pub fn amplified_read(&self, virtual_amount : u64) -> bool {
        self.read_limiter.can_read(virtual_amount)
    }
}

pub struct SegmentBuilder {
    pub reader: SegmentReader,
    pub id: SegmentId,
    pos: *mut Word,
}

unsafe impl Send for SegmentBuilder {}

impl SegmentBuilder {
    pub fn new(arena : *mut BuilderArena,
               limiter : ::std::rc::Rc<ReadLimiter>,
               id : SegmentId,
               ptr : *mut Word,
               size : WordCount32) -> SegmentBuilder {
        SegmentBuilder {
            reader : SegmentReader {
                arena : ArenaPtr::Builder(arena),
                ptr : unsafe {::std::mem::transmute(ptr)},
                size : size,
                read_limiter: limiter,
            },
            id : id,
            pos : ptr,
        }
    }

    pub fn get_word_offset_to(&mut self, ptr : *mut Word) -> WordCount32 {
        let this_addr : usize = self.reader.ptr as usize;
        let ptr_addr : usize = ptr as usize;
        assert!(ptr_addr >= this_addr);
        let result = (ptr_addr - this_addr) / BYTES_PER_WORD;
        result as u32
    }

    #[inline]
    pub fn current_size(&self) -> WordCount32 {
        ((self.pos as usize - self.reader.ptr as usize) / BYTES_PER_WORD) as u32
    }

    #[inline]
    pub fn allocate(&mut self, amount : WordCount32) -> Option<*mut Word> {
        if amount > self.reader.size - self.current_size() {
            None
        } else {
            let result = self.pos;
            self.pos = unsafe { self.pos.offset(amount as isize) };
            Some(result)
        }
    }

    #[inline]
    pub fn get_ptr_unchecked(&self, offset : WordCount32) -> *mut Word {
        unsafe {
            ::std::mem::transmute(self.reader.ptr.offset(offset as isize))
        }
    }

    #[inline]
    pub fn get_segment_id(&self) -> SegmentId { self.id }

    #[inline]
    pub fn get_arena<'a> (&'a mut self) -> *mut BuilderArena {
        match self.reader.arena {
            ArenaPtr::Builder(b) => b,
            _ => unreachable!()
        }
    }

    pub fn currently_allocated<'a>(&'a self) -> &'a [Word] {
        unsafe { ::std::slice::from_raw_parts(self.get_ptr_unchecked(0), self.current_size() as usize) }
    }
}

pub struct ReadLimiter {
    pub limit : ::std::cell::Cell<u64>,
}

impl ReadLimiter {
    pub fn new(limit : u64) -> ReadLimiter {
        ReadLimiter { limit : ::std::cell::Cell::new(limit) }
    }

    #[inline]
    pub fn can_read(&self, amount : u64) -> bool {
        let current = self.limit.get();
        if amount > current {
            // TODO arena->reportReadLimitReached()
            false
        } else {
            self.limit.set(current - amount);
            true
        }
    }
}

pub struct ReaderArena {
    raw_segments: &'static ReaderSegments,
    pub segment0: SegmentReader,
    pub more_segments: HashMap<SegmentId, Box<SegmentReader>>,
    pub cap_table : Vec<Option<Box<ClientHook+Send>>>,
    pub read_limiter : ::std::rc::Rc<ReadLimiter>,
}

impl ReaderArena {
    pub fn new(segments: &'static ReaderSegments, options: message::ReaderOptions) -> Box<ReaderArena> {

        let limiter = ::std::rc::Rc::new(ReadLimiter::new(options.traversal_limit_in_words));

        let segment0 = segments.get_segment(0).expect("segment zero does not exist");

        let segment0_reader =  SegmentReader {
            arena: ArenaPtr::Null,
            ptr: unsafe { segment0.get_unchecked(0) },
            size: segment0.len() as u32,
            read_limiter: limiter.clone(),
        };

        let mut arena = Box::new(ReaderArena {
            raw_segments: segments,
            segment0: segment0_reader,
            more_segments : HashMap::new(),
            cap_table : Vec::new(),
            read_limiter : limiter.clone(),
        });

        let arena_ptr = ArenaPtr::Reader(&mut *arena);
        arena.segment0.arena = arena_ptr;
        arena
    }

    fn try_get_segment(&mut self, id: SegmentId) -> Result<*const SegmentReader> {
        if id == 0 {
            Ok(&self.segment0)
        } else if self.more_segments.contains_key(&id) {
            Ok(&*self.more_segments[&id])
        } else {
            let cloned_limiter = self.read_limiter.clone();
            let new_segment = match self.raw_segments.get_segment(id) {
                Some(s) => s,
                None => {
                    return Err(Error::new_decode_error("Invalid segment id.", Some(format!("{}", id))));
                }
            };
            let new_segment_reader = SegmentReader {
                arena: ArenaPtr::Reader(&mut *self),
                ptr: unsafe { new_segment.get_unchecked(0) },
                size: new_segment.len() as u32,
                read_limiter: cloned_limiter
            };
            self.more_segments.insert(id, Box::new(new_segment_reader));
            Ok(&*self.more_segments[&id])
        }
    }

    #[inline]
    pub fn init_cap_table(&mut self, cap_table : Vec<Option<Box<ClientHook+Send>>>) {
        self.cap_table = cap_table;
    }
}

pub struct BuilderArena {
    allocator: &'static mut Allocator,
    pub segment0 : SegmentBuilder,
    pub more_segments : Vec<Box<SegmentBuilder>>,
    pub cap_table : Vec<Option<Box<ClientHook+Send>>>,
    pub dummy_limiter : ::std::rc::Rc<ReadLimiter>,
}

impl BuilderArena  {
    pub fn new(allocator: &'static mut Allocator) -> Box<BuilderArena> {
        let limiter = ::std::rc::Rc::new(ReadLimiter::new(::std::u64::MAX));
        let (first_segment, num_words) = allocator.allocate_segment(2);

        let mut result = Box::new(BuilderArena {
            allocator: allocator,
            segment0: SegmentBuilder {
                reader: SegmentReader {
                    ptr: first_segment,
                    size: num_words,
                    arena: ArenaPtr::Null,
                    read_limiter: limiter.clone()},
                id: 0,
                pos: first_segment,
            },
            more_segments: Vec::new(),
            cap_table: Vec::new(),
            dummy_limiter: limiter,
        });

        let arena_ptr = ArenaPtr::Builder(&mut *result);
        result.segment0.reader.arena = arena_ptr;
        result
    }

    pub fn try_get_segment(&self, id: SegmentId) -> Result<*const SegmentReader> {
        if id == 0 {
            Ok(&self.segment0.reader)
        } else if ((id - 1) as usize) < self.more_segments.len() {
            Ok(&self.more_segments[(id - 1) as usize].reader)
        } else {
            Err(Error::new_decode_error("Invalid segment id.", Some(format!("{}", id))))
        }
    }

    #[inline]
    pub fn allocate(&mut self, amount: WordCount32) -> (*mut SegmentBuilder, *mut Word) {
        unsafe {
            match self.segment0.allocate(amount) {
                Some(result) => { return (&mut self.segment0, result) }
                None => {}
            }

            //# Need to fall back to additional segments.

            let id = {
                let len = self.more_segments.len();
                if len == 0 { 1 }
                else {
                    let result_ptr : *mut SegmentBuilder = &mut *self.more_segments[len-1];
                    match self.more_segments[len - 1].allocate(amount) {
                        Some(result) => { return (result_ptr, result) }
                        None => { len + 1 }
                    }
                }};

            let (words, size) = self.allocator.allocate_segment(amount);
            let mut new_builder = Box::new(SegmentBuilder::new(self, self.dummy_limiter.clone(),
                                                               id as u32, words, size));
            let builder_ptr : *mut SegmentBuilder = &mut *new_builder;

            self.more_segments.push(new_builder);

            (builder_ptr, (*builder_ptr).allocate(amount).unwrap() )
        }
    }

    pub fn get_segment(&mut self, id : SegmentId) -> Result<*mut SegmentBuilder> {
        if id == 0 {
            Ok(&mut self.segment0)
        } else if ((id - 1) as usize) < self.more_segments.len() {
            Ok(&mut *self.more_segments[(id - 1) as usize])
        } else {
            Err(Error::new_decode_error("Invalid segment id.", Some(format!("{}", id))))
        }
    }

    pub fn get_segments_for_output<'a>(&'a self) -> OutputSegments<'a> {
        if self.more_segments.len() == 0 {
            OutputSegments::SingleSegment([self.segment0.currently_allocated()])
        } else {
            let mut v = Vec::new();
            v.push(self.segment0.currently_allocated());
            for seg in self.more_segments.iter() {
                v.push(seg.currently_allocated());
            }
            OutputSegments::MultiSegment(v)
        }
    }

    pub fn get_cap_table<'a>(&'a self) -> &'a [Option<Box<ClientHook+Send>>] {
        &self.cap_table
    }

    pub fn inject_cap(&mut self, cap : Box<ClientHook+Send>) -> u32 {
        self.cap_table.push(Some(cap));
        self.cap_table.len() as u32 - 1
    }
}

#[derive(Clone, Copy)]
pub enum ArenaPtr {
    Reader(*mut ReaderArena),
    Builder(*mut BuilderArena),
    Null
}

impl ArenaPtr {
    pub fn try_get_segment(&self, id: SegmentId) -> Result<*const SegmentReader> {
        unsafe {
            match self {
                &ArenaPtr::Reader(reader) => {
                    (&mut *reader).try_get_segment(id)
                }
                &ArenaPtr::Builder(builder) => {
                    (&*builder).try_get_segment(id)
                }
                &ArenaPtr::Null => {
                    Err(Error::new_decode_error("Null arena.", None))
                }
            }
        }
    }

    pub fn extract_cap(&self, index: usize) -> Option<Box<ClientHook+Send>> {
        unsafe {
            match self {
                &ArenaPtr::Reader(reader) => {
                    if index < (*reader).cap_table.len() {
                        match (*reader).cap_table[index] {
                            Some( ref hook ) => { Some(hook.copy()) }
                            None => {
                                None
                            }
                        }
                    } else {
                        None
                    }
                }
                &ArenaPtr::Builder(builder) => {
                    if index < (*builder).cap_table.len() {
                        match (*builder).cap_table[index] {
                            Some( ref hook ) => { Some(hook.copy()) }
                            None => {
                                None
                            }
                        }
                    } else {
                        None
                    }
                }
                &ArenaPtr::Null => {
                    panic!();
                }
            }
        }
    }
}
<commit_msg>arena.rs style<commit_after>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

use std::cell::Cell;
use std::collections::HashMap;
use std::mem;
use std::rc::Rc;
use std::slice;
use std::u64;

use private::capability::ClientHook;
use private::units::*;
use message;
use message::{Allocator, ReaderSegments};
use {Error, OutputSegments, Result, Word};

pub type SegmentId = u32;

pub struct SegmentReader {
    pub arena: ArenaPtr,
    pub ptr: *const Word,
    pub size: WordCount32,
    pub read_limiter: Rc<ReadLimiter>,
}

impl SegmentReader {
    #[inline]
    pub unsafe fn get_start_ptr(&self) -> *const Word {
        self.ptr.offset(0)
    }

    #[inline]
    pub fn contains_interval(&self, from: *const Word, to: *const Word) -> bool {
        let this_begin: usize = self.ptr as usize;
        let this_end: usize = unsafe { self.ptr.offset(self.size as isize) as usize };
        from as usize >= this_begin && to as usize <= this_end && from as usize <= to as usize &&
            self.read_limiter.can_read((to as usize - from as usize) as u64 / BYTES_PER_WORD as u64)
    }

    #[inline]
    pub fn amplified_read(&self, virtual_amount: u64) -> bool {
        self.read_limiter.can_read(virtual_amount)
    }
}

pub struct SegmentBuilder {
    pub reader: SegmentReader,
    pub id: SegmentId,
    pos: *mut Word,
}

unsafe impl Send for SegmentBuilder {}

impl SegmentBuilder {
    pub fn new(arena: *mut BuilderArena,
               limiter: Rc<ReadLimiter>,
               id: SegmentId,
               ptr: *mut Word,
               size: WordCount32) -> SegmentBuilder {
        SegmentBuilder {
            reader: SegmentReader {
                arena: ArenaPtr::Builder(arena),
                ptr: unsafe {mem::transmute(ptr)},
                size: size,
                read_limiter: limiter,
            },
            id: id,
            pos: ptr,
        }
    }

    pub fn get_word_offset_to(&mut self, ptr: *mut Word) -> WordCount32 {
        let this_addr: usize = self.reader.ptr as usize;
        let ptr_addr: usize = ptr as usize;
        assert!(ptr_addr >= this_addr);
        let result = (ptr_addr - this_addr) / BYTES_PER_WORD;
        result as u32
    }

    #[inline]
    pub fn current_size(&self) -> WordCount32 {
        ((self.pos as usize - self.reader.ptr as usize) / BYTES_PER_WORD) as u32
    }

    #[inline]
    pub fn allocate(&mut self, amount: WordCount32) -> Option<*mut Word> {
        if amount > self.reader.size - self.current_size() {
            None
        } else {
            let result = self.pos;
            self.pos = unsafe { self.pos.offset(amount as isize) };
            Some(result)
        }
    }

    #[inline]
    pub fn get_ptr_unchecked(&self, offset: WordCount32) -> *mut Word {
        unsafe {
            mem::transmute(self.reader.ptr.offset(offset as isize))
        }
    }

    #[inline]
    pub fn get_segment_id(&self) -> SegmentId { self.id }

    #[inline]
    pub fn get_arena<'a> (&'a mut self) -> *mut BuilderArena {
        match self.reader.arena {
            ArenaPtr::Builder(b) => b,
            _ => unreachable!()
        }
    }

    pub fn currently_allocated<'a>(&'a self) -> &'a [Word] {
        unsafe { slice::from_raw_parts(self.get_ptr_unchecked(0), self.current_size() as usize) }
    }
}

pub struct ReadLimiter {
    pub limit: Cell<u64>,
}

impl ReadLimiter {
    pub fn new(limit: u64) -> ReadLimiter {
        ReadLimiter { limit: Cell::new(limit) }
    }

    #[inline]
    pub fn can_read(&self, amount: u64) -> bool {
        let current = self.limit.get();
        if amount > current {
            // TODO arena->reportReadLimitReached()
            false
        } else {
            self.limit.set(current - amount);
            true
        }
    }
}

pub struct ReaderArena {
    raw_segments: &'static ReaderSegments,
    pub segment0: SegmentReader,
    pub more_segments: HashMap<SegmentId, Box<SegmentReader>>,
    pub cap_table: Vec<Option<Box<ClientHook+Send>>>,
    pub read_limiter: Rc<ReadLimiter>,
}

impl ReaderArena {
    pub fn new(segments: &'static ReaderSegments,
               options: message::ReaderOptions)
               -> Box<ReaderArena> {

        let limiter = Rc::new(ReadLimiter::new(options.traversal_limit_in_words));

        let segment0 = segments.get_segment(0).expect("segment zero does not exist");

        let segment0_reader =  SegmentReader {
            arena: ArenaPtr::Null,
            ptr: unsafe { segment0.get_unchecked(0) },
            size: segment0.len() as u32,
            read_limiter: limiter.clone(),
        };

        let mut arena = Box::new(ReaderArena {
            raw_segments: segments,
            segment0: segment0_reader,
            more_segments: HashMap::new(),
            cap_table: Vec::new(),
            read_limiter: limiter.clone(),
        });

        let arena_ptr = ArenaPtr::Reader(&mut *arena);
        arena.segment0.arena = arena_ptr;
        arena
    }

    fn try_get_segment(&mut self, id: SegmentId) -> Result<*const SegmentReader> {
        if id == 0 {
            Ok(&self.segment0)
        } else if self.more_segments.contains_key(&id) {
            Ok(&*self.more_segments[&id])
        } else {
            let cloned_limiter = self.read_limiter.clone();
            let new_segment = match self.raw_segments.get_segment(id) {
                Some(s) => s,
                None => {
                    return Err(Error::new_decode_error("Invalid segment id.",
                                                       Some(format!("{}", id))));
                }
            };
            let new_segment_reader = SegmentReader {
                arena: ArenaPtr::Reader(&mut *self),
                ptr: unsafe { new_segment.get_unchecked(0) },
                size: new_segment.len() as u32,
                read_limiter: cloned_limiter
            };
            self.more_segments.insert(id, Box::new(new_segment_reader));
            Ok(&*self.more_segments[&id])
        }
    }

    #[inline]
    pub fn init_cap_table(&mut self, cap_table: Vec<Option<Box<ClientHook+Send>>>) {
        self.cap_table = cap_table;
    }
}

pub struct BuilderArena {
    allocator: &'static mut Allocator,
    pub segment0: SegmentBuilder,
    pub more_segments: Vec<Box<SegmentBuilder>>,
    pub cap_table: Vec<Option<Box<ClientHook+Send>>>,
    pub dummy_limiter: Rc<ReadLimiter>,
}

impl BuilderArena  {
    pub fn new(allocator: &'static mut Allocator) -> Box<BuilderArena> {
        let limiter = Rc::new(ReadLimiter::new(u64::MAX));
        let (first_segment, num_words) = allocator.allocate_segment(2);

        let mut result = Box::new(BuilderArena {
            allocator: allocator,
            segment0: SegmentBuilder {
                reader: SegmentReader {
                    ptr: first_segment,
                    size: num_words,
                    arena: ArenaPtr::Null,
                    read_limiter: limiter.clone()},
                id: 0,
                pos: first_segment,
            },
            more_segments: Vec::new(),
            cap_table: Vec::new(),
            dummy_limiter: limiter,
        });

        let arena_ptr = ArenaPtr::Builder(&mut *result);
        result.segment0.reader.arena = arena_ptr;
        result
    }

    pub fn try_get_segment(&self, id: SegmentId) -> Result<*const SegmentReader> {
        if id == 0 {
            Ok(&self.segment0.reader)
        } else if ((id - 1) as usize) < self.more_segments.len() {
            Ok(&self.more_segments[(id - 1) as usize].reader)
        } else {
            Err(Error::new_decode_error("Invalid segment id.", Some(format!("{}", id))))
        }
    }

    #[inline]
    pub fn allocate(&mut self, amount: WordCount32) -> (*mut SegmentBuilder, *mut Word) {
        unsafe {
            match self.segment0.allocate(amount) {
                Some(result) => { return (&mut self.segment0, result) }
                None => {}
            }

            //# Need to fall back to additional segments.

            let id = {
                let len = self.more_segments.len();
                if len == 0 { 1 }
                else {
                    let result_ptr: *mut SegmentBuilder = &mut *self.more_segments[len-1];
                    match self.more_segments[len - 1].allocate(amount) {
                        Some(result) => { return (result_ptr, result) }
                        None => { len + 1 }
                    }
                }};

            let (words, size) = self.allocator.allocate_segment(amount);
            let mut new_builder = Box::new(SegmentBuilder::new(self, self.dummy_limiter.clone(),
                                                               id as u32, words, size));
            let builder_ptr: *mut SegmentBuilder = &mut *new_builder;

            self.more_segments.push(new_builder);

            (builder_ptr, (*builder_ptr).allocate(amount).unwrap() )
        }
    }

    pub fn get_segment(&mut self, id: SegmentId) -> Result<*mut SegmentBuilder> {
        if id == 0 {
            Ok(&mut self.segment0)
        } else if ((id - 1) as usize) < self.more_segments.len() {
            Ok(&mut *self.more_segments[(id - 1) as usize])
        } else {
            Err(Error::new_decode_error("Invalid segment id.", Some(format!("{}", id))))
        }
    }

    pub fn get_segments_for_output<'a>(&'a self) -> OutputSegments<'a> {
        if self.more_segments.len() == 0 {
            OutputSegments::SingleSegment([self.segment0.currently_allocated()])
        } else {
            let mut v = Vec::new();
            v.push(self.segment0.currently_allocated());
            for seg in self.more_segments.iter() {
                v.push(seg.currently_allocated());
            }
            OutputSegments::MultiSegment(v)
        }
    }

    pub fn get_cap_table<'a>(&'a self) -> &'a [Option<Box<ClientHook+Send>>] {
        &self.cap_table
    }

    pub fn inject_cap(&mut self, cap: Box<ClientHook+Send>) -> u32 {
        self.cap_table.push(Some(cap));
        self.cap_table.len() as u32 - 1
    }
}

#[derive(Clone, Copy)]
pub enum ArenaPtr {
    Reader(*mut ReaderArena),
    Builder(*mut BuilderArena),
    Null
}

impl ArenaPtr {
    pub fn try_get_segment(&self, id: SegmentId) -> Result<*const SegmentReader> {
        unsafe {
            match self {
                &ArenaPtr::Reader(reader) => {
                    (&mut *reader).try_get_segment(id)
                }
                &ArenaPtr::Builder(builder) => {
                    (&*builder).try_get_segment(id)
                }
                &ArenaPtr::Null => {
                    Err(Error::new_decode_error("Null arena.", None))
                }
            }
        }
    }

    pub fn extract_cap(&self, index: usize) -> Option<Box<ClientHook+Send>> {
        unsafe {
            match self {
                &ArenaPtr::Reader(reader) => {
                    if index < (*reader).cap_table.len() {
                        match (*reader).cap_table[index] {
                            Some( ref hook ) => { Some(hook.copy()) }
                            None => {
                                None
                            }
                        }
                    } else {
                        None
                    }
                }
                &ArenaPtr::Builder(builder) => {
                    if index < (*builder).cap_table.len() {
                        match (*builder).cap_table[index] {
                            Some( ref hook ) => { Some(hook.copy()) }
                            None => {
                                None
                            }
                        }
                    } else {
                        None
                    }
                }
                &ArenaPtr::Null => {
                    panic!();
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// Deny warnings, except in `dev` mode

// To update this list
// 1. Run `rustc -W help`
// 2. Grab all `default=warn` warnings
// 3. Paste them here, deleting `warnings`, and any with `deprecated` in the name
#![deny(const_err,
        dead_code,
        illegal_floating_point_literal_pattern,
        improper_ctypes,
        non_camel_case_types,
        non_shorthand_field_patterns,
        non_snake_case,
        non_upper_case_globals,
        no_mangle_generic_items,
        overflowing_literals,
        path_statements,
        patterns_in_fns_without_body,
        plugin_as_library,
        private_in_public,
        private_no_mangle_fns,
        private_no_mangle_statics,
        renamed_and_removed_lints,
        stable_features,
        unconditional_recursion,
        unions_with_drop_fields,
        unknown_lints,
        unreachable_code,
        unreachable_patterns,
        unused_allocation,
        unused_assignments,
        unused_attributes,
        unused_comparisons,
        unused_features,
        unused_imports,
        unused_macros,
        unused_must_use,
        unused_mut,
        unused_parens,
        unused_unsafe,
        unused_variables,
        while_true)]
// This list is select `allow` warnings
#![deny(trivial_casts,
       trivial_numeric_casts,
       unused_extern_crates,
       unused_import_braces)]
#![cfg_attr(feature="cargo-clippy", allow(
        cyclomatic_complexity,
        needless_pass_by_value))]
#![cfg_attr(feature="dev", warn(warnings))]

extern crate cobalt;
extern crate env_logger;
extern crate notify;
extern crate ghp;

extern crate hyper;

#[macro_use]
extern crate error_chain;

#[macro_use]
extern crate clap;

#[macro_use]
extern crate log;

mod args;
mod build;
mod error;
mod new;
mod serve;

use clap::{Arg, App, SubCommand, AppSettings};
use cobalt::jekyll;
use cobalt::{list_syntaxes, list_syntax_themes};

use error::*;

quick_main!(run);

fn run() -> Result<()> {
    let app_cli = App::new("Cobalt")
        .version(crate_version!())
        .author("Benny Klotz <r3qnbenni@gmail.com>, Johann Hofmann")
        .about("A static site generator written in Rust.")
        .setting(AppSettings::SubcommandRequired)
        .setting(AppSettings::GlobalVersion)
        .args(&args::get_logging_args())
        .subcommand(new::init_command_args())
        .subcommand(new::new_command_args())
        .subcommand(build::build_command_args())
        .subcommand(build::clean_command_args())
        .subcommand(serve::serve_command_args())
        .subcommand(serve::watch_command_args())
        .subcommand(build::import_command_args())
        .subcommand(SubCommand::with_name("list-syntax-themes").about("list available themes"))
        .subcommand(SubCommand::with_name("list-syntaxes").about("list supported syntaxes"))
        .subcommand(SubCommand::with_name("convert-jekyll")
                        .about("convert jekyll website to cobalt")
                        .arg(Arg::with_name("jksrc")
                                 .long("jksrc")
                                 .value_name("JEKYLL-FILE-OR-DIR")
                                 .help("Jekyll posts' directory")
                                 .required(true)
                                 .takes_value(true))
                        .arg(Arg::with_name("jkdst")
                                 .long("jkdst")
                                 .value_name("DIR")
                                 .help("Output dir of converted posts")
                                 .takes_value(true)
                                 .default_value("./posts")));

    let global_matches = app_cli.get_matches();

    let (command, matches) = match global_matches.subcommand() {
        (command, Some(matches)) => (command, matches),
        (_, None) => unreachable!(),
    };

    let mut builder = args::get_logging(&global_matches, matches)?;
    builder.init().unwrap();

    match command {
        "init" => new::init_command(matches),
        "new" => new::new_command(matches),
        "build" => build::build_command(matches),
        "clean" => build::clean_command(matches),
        "serve" => serve::serve_command(matches),
        "watch" => serve::watch_command(matches),
        "import" => build::import_command(matches),
        "list-syntax-themes" => {
            for name in list_syntax_themes() {
                println!("{}", name);
            }
            Ok(())
        }
        "list-syntaxes" => {
            for name in list_syntaxes() {
                println!("{}", name);
            }
            Ok(())
        }
        "convert-jekyll" => {
            let source = matches.value_of("jksrc").unwrap().to_string();
            let dest = matches.value_of("jkdst").unwrap().to_string();
            jekyll::jk_document::convert_from_jk(std::path::Path::new(&source),
                                                 std::path::Path::new(&dest))
                .chain_err(|| "Jekyll conversion failed.")
        }
        _ => {
            bail!(global_matches.usage());
        }
    }.chain_err(|| format!("{} command failed", command))?;

    Ok(())
}
<commit_msg>fix(jekyll): Clean up flag names<commit_after>// Deny warnings, except in `dev` mode

// To update this list
// 1. Run `rustc -W help`
// 2. Grab all `default=warn` warnings
// 3. Paste them here, deleting `warnings`, and any with `deprecated` in the name
#![deny(const_err,
        dead_code,
        illegal_floating_point_literal_pattern,
        improper_ctypes,
        non_camel_case_types,
        non_shorthand_field_patterns,
        non_snake_case,
        non_upper_case_globals,
        no_mangle_generic_items,
        overflowing_literals,
        path_statements,
        patterns_in_fns_without_body,
        plugin_as_library,
        private_in_public,
        private_no_mangle_fns,
        private_no_mangle_statics,
        renamed_and_removed_lints,
        stable_features,
        unconditional_recursion,
        unions_with_drop_fields,
        unknown_lints,
        unreachable_code,
        unreachable_patterns,
        unused_allocation,
        unused_assignments,
        unused_attributes,
        unused_comparisons,
        unused_features,
        unused_imports,
        unused_macros,
        unused_must_use,
        unused_mut,
        unused_parens,
        unused_unsafe,
        unused_variables,
        while_true)]
// This list is select `allow` warnings
#![deny(trivial_casts,
       trivial_numeric_casts,
       unused_extern_crates,
       unused_import_braces)]
#![cfg_attr(feature="cargo-clippy", allow(
        cyclomatic_complexity,
        needless_pass_by_value))]
#![cfg_attr(feature="dev", warn(warnings))]

extern crate cobalt;
extern crate env_logger;
extern crate notify;
extern crate ghp;

extern crate hyper;

#[macro_use]
extern crate error_chain;

#[macro_use]
extern crate clap;

#[macro_use]
extern crate log;

mod args;
mod build;
mod error;
mod new;
mod serve;

use clap::{Arg, App, SubCommand, AppSettings};
use cobalt::jekyll;
use cobalt::{list_syntaxes, list_syntax_themes};

use error::*;

quick_main!(run);

fn run() -> Result<()> {
    let app_cli = App::new("Cobalt")
        .version(crate_version!())
        .author("Benny Klotz <r3qnbenni@gmail.com>, Johann Hofmann")
        .about("A static site generator written in Rust.")
        .setting(AppSettings::SubcommandRequired)
        .setting(AppSettings::GlobalVersion)
        .args(&args::get_logging_args())
        .subcommand(new::init_command_args())
        .subcommand(new::new_command_args())
        .subcommand(build::build_command_args())
        .subcommand(build::clean_command_args())
        .subcommand(serve::serve_command_args())
        .subcommand(serve::watch_command_args())
        .subcommand(build::import_command_args())
        .subcommand(SubCommand::with_name("list-syntax-themes").about("list available themes"))
        .subcommand(SubCommand::with_name("list-syntaxes").about("list supported syntaxes"))
        .subcommand(SubCommand::with_name("convert-jekyll")
                        .about("convert jekyll website to cobalt")
                        .arg(Arg::with_name("source")
                                 .long("source")
                                 .value_name("JEKYLL-FILE-OR-DIR")
                                 .help("Jekyll posts' directory")
                                 .required(true)
                                 .takes_value(true))
                        .arg(Arg::with_name("destination")
                                 .long("destination")
                                 .value_name("DIR")
                                 .help("Output dir of converted posts")
                                 .takes_value(true)
                                 .default_value("./posts")));

    let global_matches = app_cli.get_matches();

    let (command, matches) = match global_matches.subcommand() {
        (command, Some(matches)) => (command, matches),
        (_, None) => unreachable!(),
    };

    let mut builder = args::get_logging(&global_matches, matches)?;
    builder.init().unwrap();

    match command {
        "init" => new::init_command(matches),
        "new" => new::new_command(matches),
        "build" => build::build_command(matches),
        "clean" => build::clean_command(matches),
        "serve" => serve::serve_command(matches),
        "watch" => serve::watch_command(matches),
        "import" => build::import_command(matches),
        "list-syntax-themes" => {
            for name in list_syntax_themes() {
                println!("{}", name);
            }
            Ok(())
        }
        "list-syntaxes" => {
            for name in list_syntaxes() {
                println!("{}", name);
            }
            Ok(())
        }
        "convert-jekyll" => {
            let source = matches.value_of("source").unwrap().to_string();
            let dest = matches.value_of("destination").unwrap().to_string();
            jekyll::jk_document::convert_from_jk(std::path::Path::new(&source),
                                                 std::path::Path::new(&dest))
                .chain_err(|| "Jekyll conversion failed.")
        }
        _ => {
            bail!(global_matches.usage());
        }
    }.chain_err(|| format!("{} command failed", command))?;

    Ok(())
}
<|endoftext|>


<commit_before>// Copyright Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::net::SocketAddr;

use hyper::service::{make_service_fn, service_fn};
use hyper::{Body, Method, Request, Response, Server, StatusCode};
use tokio::net::{TcpListener, TcpStream};
use tokio_stream::StreamExt;
use tracing::{error, info, warn};

use crate::config::Config;
use crate::identity::CertificateProvider;
use crate::proxy::inbound::InboundConnect::Hbone;
use crate::socket::relay;
use crate::tls::TlsError;
use crate::workload::WorkloadInformation;

use super::Error;

pub struct Inbound {
    cfg: Config,
    listener: TcpListener,
    cert_manager: Box<dyn CertificateProvider>,
    workloads: WorkloadInformation,
}

impl Inbound {
    pub async fn new(
        cfg: Config,
        workloads: WorkloadInformation,
        cert_manager: Box<dyn CertificateProvider>,
    ) -> Result<Inbound, Error> {
        let listener: TcpListener = TcpListener::bind(cfg.inbound_addr)
            .await
            .map_err(|e| Error::Bind(cfg.inbound_addr, e))?;
        match crate::socket::set_transparent(&listener) {
            Err(_e) => info!("running without transparent mode"),
            _ => info!("running with transparent mode"),
        };
        Ok(Inbound {
            cfg,
            workloads,
            listener,
            cert_manager,
        })
    }

    pub(super) fn address(&self) -> SocketAddr {
        self.listener.local_addr().unwrap()
    }

    pub(super) async fn run(self) {
        let addr = self.listener.local_addr().unwrap();
        let service =
            make_service_fn(|_| async { Ok::<_, hyper::Error>(service_fn(Self::serve_connect)) });
        let boring_acceptor = crate::tls::BoringTlsAcceptor {
            acceptor: InboundCertProvider {
                workloads: self.workloads.clone(),
                cert_manager: self.cert_manager.clone(),
            },
        };
        let mut listener =
            hyper::server::conn::AddrIncoming::from_listener(self.listener).expect("hbone bind");
        listener.set_nodelay(true);
        let incoming = hyper::server::accept::from_stream(
            tls_listener::builder(boring_acceptor)
                .listen(listener)
                .filter(|conn| {
                    // Avoid 'By default, if a client fails the TLS handshake, that is treated as an error, and the TlsListener will return an Err'
                    if let Err(err) = conn {
                        warn!("TLS handshake error: {}", err);
                        false
                    } else {
                        info!("TLS handshake succeeded");
                        true
                    }
                }),
        );

        let server = Server::builder(incoming)
            .http2_only(true)
            .http2_initial_stream_window_size(self.cfg.window_size)
            .http2_initial_connection_window_size(self.cfg.connection_window_size)
            .http2_max_frame_size(self.cfg.frame_size)
            .serve(service);

        info!("HBONE listener established {}", addr);

        if let Err(e) = server.await {
            error!("server error: {}", e);
        }
    }

    /// handle_inbound serves an inbound connection with a target address `addr`.
    pub(super) async fn handle_inbound(
        request_type: InboundConnect,
        addr: SocketAddr,
    ) -> Result<(), std::io::Error> {
        let stream = TcpStream::connect(addr).await;
        match stream {
            Err(err) => {
                warn!("connect to {} failed: {}", addr, err);
                Err(err)
            }
            Ok(stream) => {
                let mut stream = stream;
                stream.set_nodelay(true)?;
                tokio::task::spawn(async move {
                    match request_type {
                        InboundConnect::DirectPath(mut incoming) => {
                            let res = relay(&mut incoming, &mut stream).await;
                            if res.is_err() {
                                error!("internal server copy {:?}", res);
                            }
                        }
                        Hbone(req) => match hyper::upgrade::on(req).await {
                            Ok(mut upgraded) => {
                                super::copy_hbone("hbone server", &mut upgraded, &mut stream)
                                    .await
                                    .expect("hbone server copy");
                            }
                            Err(e) => {
                                // Not sure if this can even happen
                                error!("No upgrade {e}");
                            }
                        },
                    }
                });
                // Send back our 200. We do this regardless of if our spawned task copies the data;
                // we need to respond with headers immediately once connection is established for the
                // stream of bytes to begin.
                Ok(())
            }
        }
    }

    async fn serve_connect(req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
        let mut res = Response::new(Body::empty());
        match req.method() {
            &Method::CONNECT => {
                let uri = req.uri();
                info!("Got {} request to {}", req.method(), uri);
                let addr: Result<SocketAddr, _> = uri.to_string().as_str().parse();
                if addr.is_err() {
                    info!("Sending 400, {:?}", addr.err());
                    let mut bad_request = Response::default();
                    *bad_request.status_mut() = StatusCode::BAD_REQUEST;
                    return Ok(bad_request);
                }

                let addr: SocketAddr = addr.unwrap();
                *res.status_mut() =
                    match Self::handle_inbound(InboundConnect::Hbone(req), addr).await {
                        Ok(_) => StatusCode::OK,
                        Err(_) => StatusCode::SERVICE_UNAVAILABLE,
                    };
                Ok(res)
            }
            // Return the 404 Not Found for other routes.
            method => {
                info!("Sending 404, got {method}");
                let mut not_found = Response::default();
                *not_found.status_mut() = StatusCode::NOT_FOUND;
                Ok(not_found)
            }
        }
    }
}

pub(super) enum InboundConnect {
    /// DirectPath is an optimization when we are connecting to an endpoint on the same node.
    /// Rather than doing a full HBONE connection over the localhost network, we just pass the outbound
    /// context directly to the inbound handling in memory.
    DirectPath(TcpStream),
    /// Hbone is a standard HBONE request coming from the network.
    Hbone(Request<Body>),
}

#[derive(Clone)]
struct InboundCertProvider {
    cert_manager: Box<dyn CertificateProvider>,
    workloads: WorkloadInformation,
}

#[async_trait::async_trait]
impl crate::tls::CertProvider for InboundCertProvider {
    async fn fetch_cert(&mut self, fd: &TcpStream) -> Result<boring::ssl::SslAcceptor, TlsError> {
        let orig = crate::socket::orig_dst_addr_or_default(fd);
        let identity = {
            let remote_addr = super::to_canonical_ip(orig);
            self.workloads
                .find_workload(&remote_addr)
                .ok_or(TlsError::CertificateLookup(remote_addr))?
                .identity()
        };
        info!("tls: accepting connection to {:?} ({})", orig, identity);
        let cert = self.cert_manager.fetch_certificate(&identity).await?;
        let acc = cert.acceptor()?;
        Ok(acc)
    }
}
<commit_msg>minior improve serve_connect (#192)<commit_after>// Copyright Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::net::SocketAddr;

use hyper::service::{make_service_fn, service_fn};
use hyper::{Body, Method, Request, Response, Server, StatusCode};
use tokio::net::{TcpListener, TcpStream};
use tokio_stream::StreamExt;
use tracing::{error, info, warn};

use crate::config::Config;
use crate::identity::CertificateProvider;
use crate::proxy::inbound::InboundConnect::Hbone;
use crate::socket::relay;
use crate::tls::TlsError;
use crate::workload::WorkloadInformation;

use super::Error;

pub struct Inbound {
    cfg: Config,
    listener: TcpListener,
    cert_manager: Box<dyn CertificateProvider>,
    workloads: WorkloadInformation,
}

impl Inbound {
    pub async fn new(
        cfg: Config,
        workloads: WorkloadInformation,
        cert_manager: Box<dyn CertificateProvider>,
    ) -> Result<Inbound, Error> {
        let listener: TcpListener = TcpListener::bind(cfg.inbound_addr)
            .await
            .map_err(|e| Error::Bind(cfg.inbound_addr, e))?;
        match crate::socket::set_transparent(&listener) {
            Err(_e) => info!("running without transparent mode"),
            _ => info!("running with transparent mode"),
        };
        Ok(Inbound {
            cfg,
            workloads,
            listener,
            cert_manager,
        })
    }

    pub(super) fn address(&self) -> SocketAddr {
        self.listener.local_addr().unwrap()
    }

    pub(super) async fn run(self) {
        let addr = self.listener.local_addr().unwrap();
        let service =
            make_service_fn(|_| async { Ok::<_, hyper::Error>(service_fn(Self::serve_connect)) });
        let boring_acceptor = crate::tls::BoringTlsAcceptor {
            acceptor: InboundCertProvider {
                workloads: self.workloads.clone(),
                cert_manager: self.cert_manager.clone(),
            },
        };
        let mut listener =
            hyper::server::conn::AddrIncoming::from_listener(self.listener).expect("hbone bind");
        listener.set_nodelay(true);
        let incoming = hyper::server::accept::from_stream(
            tls_listener::builder(boring_acceptor)
                .listen(listener)
                .filter(|conn| {
                    // Avoid 'By default, if a client fails the TLS handshake, that is treated as an error, and the TlsListener will return an Err'
                    if let Err(err) = conn {
                        warn!("TLS handshake error: {}", err);
                        false
                    } else {
                        info!("TLS handshake succeeded");
                        true
                    }
                }),
        );

        let server = Server::builder(incoming)
            .http2_only(true)
            .http2_initial_stream_window_size(self.cfg.window_size)
            .http2_initial_connection_window_size(self.cfg.connection_window_size)
            .http2_max_frame_size(self.cfg.frame_size)
            .serve(service);

        info!("HBONE listener established {}", addr);

        if let Err(e) = server.await {
            error!("server error: {}", e);
        }
    }

    /// handle_inbound serves an inbound connection with a target address `addr`.
    pub(super) async fn handle_inbound(
        request_type: InboundConnect,
        addr: SocketAddr,
    ) -> Result<(), std::io::Error> {
        let stream = TcpStream::connect(addr).await;
        match stream {
            Err(err) => {
                warn!("connect to {} failed: {}", addr, err);
                Err(err)
            }
            Ok(stream) => {
                let mut stream = stream;
                stream.set_nodelay(true)?;
                tokio::task::spawn(async move {
                    match request_type {
                        InboundConnect::DirectPath(mut incoming) => {
                            let res = relay(&mut incoming, &mut stream).await;
                            if res.is_err() {
                                error!("internal server copy {:?}", res);
                            }
                        }
                        Hbone(req) => match hyper::upgrade::on(req).await {
                            Ok(mut upgraded) => {
                                super::copy_hbone("hbone server", &mut upgraded, &mut stream)
                                    .await
                                    .expect("hbone server copy");
                            }
                            Err(e) => {
                                // Not sure if this can even happen
                                error!("No upgrade {e}");
                            }
                        },
                    }
                });
                // Send back our 200. We do this regardless of if our spawned task copies the data;
                // we need to respond with headers immediately once connection is established for the
                // stream of bytes to begin.
                Ok(())
            }
        }
    }

    async fn serve_connect(req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
        match req.method() {
            &Method::CONNECT => {
                let uri = req.uri();
                info!("Got {} request to {}", req.method(), uri);
                let addr: Result<SocketAddr, _> = uri.to_string().as_str().parse();
                if addr.is_err() {
                    info!("Sending 400, {:?}", addr.err());
                    return Ok(Response::builder()
                        .status(hyper::StatusCode::BAD_REQUEST)
                        .body(Body::empty())
                        .unwrap());
                }

                let addr: SocketAddr = addr.unwrap();
                let status_code = match Self::handle_inbound(InboundConnect::Hbone(req), addr).await
                {
                    Ok(_) => StatusCode::OK,
                    Err(_) => StatusCode::SERVICE_UNAVAILABLE,
                };

                Ok(Response::builder()
                    .status(status_code)
                    .body(Body::empty())
                    .unwrap())
            }
            // Return the 404 Not Found for other routes.
            method => {
                info!("Sending 404, got {method}");
                Ok(Response::builder()
                    .status(hyper::StatusCode::NOT_FOUND)
                    .body(Body::empty())
                    .unwrap())
            }
        }
    }
}

pub(super) enum InboundConnect {
    /// DirectPath is an optimization when we are connecting to an endpoint on the same node.
    /// Rather than doing a full HBONE connection over the localhost network, we just pass the outbound
    /// context directly to the inbound handling in memory.
    DirectPath(TcpStream),
    /// Hbone is a standard HBONE request coming from the network.
    Hbone(Request<Body>),
}

#[derive(Clone)]
struct InboundCertProvider {
    cert_manager: Box<dyn CertificateProvider>,
    workloads: WorkloadInformation,
}

#[async_trait::async_trait]
impl crate::tls::CertProvider for InboundCertProvider {
    async fn fetch_cert(&mut self, fd: &TcpStream) -> Result<boring::ssl::SslAcceptor, TlsError> {
        let orig = crate::socket::orig_dst_addr_or_default(fd);
        let identity = {
            let remote_addr = super::to_canonical_ip(orig);
            self.workloads
                .find_workload(&remote_addr)
                .ok_or(TlsError::CertificateLookup(remote_addr))?
                .identity()
        };
        info!("tls: accepting connection to {:?} ({})", orig, identity);
        let cert = self.cert_manager.fetch_certificate(&identity).await?;
        let acc = cert.acceptor()?;
        Ok(acc)
    }
}
<|endoftext|>