<commit_before>use std::cell::RefCell;
use std::io::{self, Write};
use std::net::Shutdown;
use std::rc::Rc;
use mio::net::TcpStream;
use mio::{Event, PollOpt, Ready, Token};

use super::binary;
use super::close_listener::CloseListener;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::ipv4_packet_buffer::IPv4PacketBuffer;
use super::router::Router;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;

const TAG: &'static str = "Client";

pub struct Client {
    id: u32,
    stream: TcpStream,
    token: Token,
    client_to_network: IPv4PacketBuffer,
    network_to_client: StreamBuffer,
    router: Router,
    close_listener: Box<CloseListener<Client>>,
    closed: bool,
    // number of remaining bytes of "id" to send to the client before relaying any data
    pending_id_bytes: usize,
}

impl Client {
    pub fn new(id: u32, selector: &mut Selector, stream: TcpStream, close_listener: Box<CloseListener<Client>>) -> io::Result<Rc<RefCell<Self>>> {
        let rc = Rc::new(RefCell::new(Self {
            id: id,
            stream: stream,
            token: Token(0), // default value, will be set afterwards
            client_to_network: IPv4PacketBuffer::new(),
            network_to_client: StreamBuffer::new(16 * MAX_PACKET_LENGTH),
            router: Router::new(),
            closed: false,
            close_listener: close_listener,
            pending_id_bytes: 4,
        }));
        // set client as router owner
        rc.borrow_mut().router.set_client(Rc::downgrade(&rc));

        {
            let rc_clone = rc.clone();
            let handler = move |selector: &mut Selector, ready| {
                let mut self_ref = rc_clone.borrow_mut();
                self_ref.on_ready(selector, ready);
            };
            let mut self_ref = rc.borrow_mut();
            // on start, we are interested only in writing (we must first send the client id)
            let token = selector.register(&self_ref.stream, handler, Ready::writable(), PollOpt::level())?;
            self_ref.token = token;
        }
        Ok(rc)
    }

    pub fn id(&self) -> u32 {
        return self.id;
    }

    pub fn router(&mut self) -> &mut Router {
        &mut self.router
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        selector.deregister(&self.stream, self.token).unwrap();
        // shutdown only (there is no close), the socket will be closed on drop
        if let Err(_) = self.stream.shutdown(Shutdown::Both) {
            warn!(target: TAG, "Cannot shutdown client socket");
        }
        self.router.clear(selector);
        self.close_listener.on_closed(self);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        if self.must_send_id() {
            match self.send_id() {
                Ok(_) => {
                    if self.pending_id_bytes == 0 {
                        debug!(target:TAG, "Client id #{} sent to client", self.id);
                    }
                }
                Err(_) => {
                    error!(target: TAG, "Cannot write client id #{}", self.id);
                    self.close(selector);
                }
            }
        } else {
            match self.write() {
                Ok(_) => self.process_pending(),
                Err(_) => {
                    error!(target: TAG, "Cannot write");
                    self.close(selector);
                }
            }
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        match self.read() {
            Ok(true) => self.push_to_network(selector),
            Ok(false) => {
                debug!(target: TAG, "EOF reached");
                self.close(selector);
            }
            Err(_) => {
                error!(target: TAG, "Cannot read");
                self.close(selector);
            }
        }
    }

    pub fn send_to_client(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        if ipv4_packet.length() as usize <= self.network_to_client.remaining() {
            self.network_to_client.read_from(ipv4_packet.raw())?;
            self.update_interests(selector);
            Ok(())
        } else {
            warn!(target: TAG, "Client buffer full");
            Err(io::Error::new(io::ErrorKind::WouldBlock, "Client buffer full"))
        }
    }

    fn send_id(&mut self) -> io::Result<()> {
        assert!(self.must_send_id());
        let raw_id = binary::to_byte_array(self.id);
        let w = self.stream.write(&raw_id[4 - self.pending_id_bytes..])?;
        self.pending_id_bytes -= w;
        Ok(())
    }

    fn update_interests(&mut self, selector: &mut Selector) {
        let ready = if self.network_to_client.is_empty() {
            Ready::readable()
        } else {
            Ready::readable() | Ready::writable()
        };
        selector.reregister(&self.stream, self.token, ready, PollOpt::level()).expect("Cannot register on poll");
    }

    fn read(&mut self) -> io::Result<(bool)> {
        self.client_to_network.read_from(&mut self.stream)
    }

    fn write(&mut self) -> io::Result<()> {
        self.network_to_client.write_to(&mut self.stream)?;
        Ok(())
    }

    fn push_to_network(&mut self, selector: &mut Selector) {
        while self.push_one_packet_to_network(selector) {
            self.client_to_network.next();
        }
    }

    fn push_one_packet_to_network(&mut self, selector: &mut Selector) -> bool {
        match self.client_to_network.as_ipv4_packet() {
            Some(ref packet) => {
                self.router.send_to_network(selector, packet);
                true
            }
            None => false
        }
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn process_pending(&mut self) {
        // TODO
    }

    pub fn clean_expired_connections(&mut self, selector: &mut Selector) {
        self.router.clean_expired_connections(selector);
    }

    fn must_send_id(&self) -> bool{
        self.pending_id_bytes > 0
    }
}
<commit_msg>Expose a method to send raw data to client<commit_after>use std::cell::RefCell;
use std::io::{self, Write};
use std::net::Shutdown;
use std::rc::Rc;
use mio::net::TcpStream;
use mio::{Event, PollOpt, Ready, Token};

use super::binary;
use super::close_listener::CloseListener;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::ipv4_packet_buffer::IPv4PacketBuffer;
use super::router::Router;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;

const TAG: &'static str = "Client";

pub struct Client {
    id: u32,
    stream: TcpStream,
    token: Token,
    client_to_network: IPv4PacketBuffer,
    network_to_client: StreamBuffer,
    router: Router,
    close_listener: Box<CloseListener<Client>>,
    closed: bool,
    // number of remaining bytes of "id" to send to the client before relaying any data
    pending_id_bytes: usize,
}

impl Client {
    pub fn new(id: u32, selector: &mut Selector, stream: TcpStream, close_listener: Box<CloseListener<Client>>) -> io::Result<Rc<RefCell<Self>>> {
        let rc = Rc::new(RefCell::new(Self {
            id: id,
            stream: stream,
            token: Token(0), // default value, will be set afterwards
            client_to_network: IPv4PacketBuffer::new(),
            network_to_client: StreamBuffer::new(16 * MAX_PACKET_LENGTH),
            router: Router::new(),
            closed: false,
            close_listener: close_listener,
            pending_id_bytes: 4,
        }));
        // set client as router owner
        rc.borrow_mut().router.set_client(Rc::downgrade(&rc));

        {
            let rc_clone = rc.clone();
            let handler = move |selector: &mut Selector, ready| {
                let mut self_ref = rc_clone.borrow_mut();
                self_ref.on_ready(selector, ready);
            };
            let mut self_ref = rc.borrow_mut();
            // on start, we are interested only in writing (we must first send the client id)
            let token = selector.register(&self_ref.stream, handler, Ready::writable(), PollOpt::level())?;
            self_ref.token = token;
        }
        Ok(rc)
    }

    pub fn id(&self) -> u32 {
        return self.id;
    }

    pub fn router(&mut self) -> &mut Router {
        &mut self.router
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        selector.deregister(&self.stream, self.token).unwrap();
        // shutdown only (there is no close), the socket will be closed on drop
        if let Err(_) = self.stream.shutdown(Shutdown::Both) {
            warn!(target: TAG, "Cannot shutdown client socket");
        }
        self.router.clear(selector);
        self.close_listener.on_closed(self);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        if self.must_send_id() {
            match self.send_id() {
                Ok(_) => {
                    if self.pending_id_bytes == 0 {
                        debug!(target:TAG, "Client id #{} sent to client", self.id);
                    }
                }
                Err(_) => {
                    error!(target: TAG, "Cannot write client id #{}", self.id);
                    self.close(selector);
                }
            }
        } else {
            match self.write() {
                Ok(_) => self.process_pending(),
                Err(_) => {
                    error!(target: TAG, "Cannot write");
                    self.close(selector);
                }
            }
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        match self.read() {
            Ok(true) => self.push_to_network(selector),
            Ok(false) => {
                debug!(target: TAG, "EOF reached");
                self.close(selector);
            }
            Err(_) => {
                error!(target: TAG, "Cannot read");
                self.close(selector);
            }
        }
    }

    pub fn send_to_client(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        self.send_raw_to_client(selector, ipv4_packet.raw())
    }

    pub fn send_raw_to_client(&mut self, selector: &mut Selector, raw: &[u8]) -> io::Result<()> {
        if raw.len() as usize <= self.network_to_client.remaining() {
            self.network_to_client.read_from(raw)?;
            self.update_interests(selector);
            Ok(())
        } else {
            warn!(target: TAG, "Client buffer full");
            Err(io::Error::new(io::ErrorKind::WouldBlock, "Client buffer full"))
        }
    }

    fn send_id(&mut self) -> io::Result<()> {
        assert!(self.must_send_id());
        let raw_id = binary::to_byte_array(self.id);
        let w = self.stream.write(&raw_id[4 - self.pending_id_bytes..])?;
        self.pending_id_bytes -= w;
        Ok(())
    }

    fn update_interests(&mut self, selector: &mut Selector) {
        let ready = if self.network_to_client.is_empty() {
            Ready::readable()
        } else {
            Ready::readable() | Ready::writable()
        };
        selector.reregister(&self.stream, self.token, ready, PollOpt::level()).expect("Cannot register on poll");
    }

    fn read(&mut self) -> io::Result<(bool)> {
        self.client_to_network.read_from(&mut self.stream)
    }

    fn write(&mut self) -> io::Result<()> {
        self.network_to_client.write_to(&mut self.stream)?;
        Ok(())
    }

    fn push_to_network(&mut self, selector: &mut Selector) {
        while self.push_one_packet_to_network(selector) {
            self.client_to_network.next();
        }
    }

    fn push_one_packet_to_network(&mut self, selector: &mut Selector) -> bool {
        match self.client_to_network.as_ipv4_packet() {
            Some(ref packet) => {
                self.router.send_to_network(selector, packet);
                true
            }
            None => false
        }
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn process_pending(&mut self) {
        // TODO
    }

    pub fn clean_expired_connections(&mut self, selector: &mut Selector) {
        self.router.clean_expired_connections(selector);
    }

    fn must_send_id(&self) -> bool{
        self.pending_id_bytes > 0
    }
}
<|endoftext|>


<commit_before>#![feature(try_from)]
extern crate bytes;
extern crate clap;
extern crate cpuprofiler;
extern crate dotenv;
extern crate error;
extern crate futures;
extern crate http;
extern crate httparse;
extern crate hyper;
extern crate jsonrpc_types;
extern crate libproto;
#[macro_use]
extern crate log;
extern crate logger;
extern crate net2;
extern crate num_cpus;
extern crate protobuf;
extern crate pubsub;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate serde_json;
extern crate threadpool;
extern crate time;
extern crate tokio_core;
extern crate tokio_io;
extern crate toml;
extern crate unicase;
#[macro_use]
extern crate util;
extern crate uuid;
extern crate ws;

mod config;
mod helper;
mod ws_handler;
mod mq_handler;
mod http_server;
mod response;

use clap::App;
use config::{NewTxFlowConfig, ProfileConfig};
use cpuprofiler::PROFILER;
use http_server::Server;
use libproto::Message;
use libproto::request::{self as reqlib, BatchRequest};
use protobuf::RepeatedField;
use pubsub::start_pubsub;
use std::collections::HashMap;
use std::convert::TryInto;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender};
use std::thread;
use std::time::{Duration, SystemTime};
use tokio_core::reactor::Core;
use util::{set_panic_handler, Mutex};
use uuid::Uuid;
use ws_handler::WsFactory;

pub const TOPIC_NEW_TX: &str = "jsonrpc.new_tx";
pub const TOPIC_NEW_TX_BATCH: &str = "jsonrpc.new_tx_batch";

fn main() {
    micro_service_init!("cita-jsonrpc", "CITA:jsonrpc");

    // todo load config
    let matches = App::new("JsonRpc")
        .version("0.1")
        .author("Cryptape")
        .about("CITA JSON-RPC by Rust")
        .args_from_usage("-c, --config=[FILE] 'Sets a custom config file'")
        .get_matches();

    let mut config_path = "./jsonrpc.toml";
    if let Some(c) = matches.value_of("config") {
        info!("Value for config: {}", c);
        config_path = c;
    }

    let config = config::Config::new(config_path);
    info!(
        "CITA:jsonrpc config \n {:?}",
        toml::to_string_pretty(&config).unwrap()
    );

    //enable HTTP or WebSocket server!
    if !config.ws_config.enable && !config.http_config.enable {
        error!("enable HTTP or WebSocket server!");
        std::process::exit(-1);
    }

    start_profile(&config.profile_config);

    // init pubsub
    let (tx_sub, rx_sub) = channel();
    let (tx_pub, rx_pub) = channel();
    //used for buffer message
    let (tx_relay, rx_relay) = channel();
    start_pubsub("jsonrpc", vec!["auth.rpc", "chain.rpc"], tx_sub, rx_pub);

    let backlog_capacity = config.backlog_capacity;

    let responses = Arc::new(Mutex::new(HashMap::with_capacity(backlog_capacity)));
    let http_responses = Arc::clone(&responses);
    let ws_responses = Arc::clone(&responses);
    let mut mq_handle = mq_handler::MqHandler::new(responses);

    //dispatch
    let tx_flow_config = config.new_tx_flow_config;
    thread::spawn(move || {
        let mut new_tx_request_buffer = Vec::new();
        let mut time_stamp = SystemTime::now();
        loop {
            if let Ok(res) = rx_relay.try_recv() {
                let (topic, req): (String, reqlib::Request) = res;
                forward_service(
                    topic,
                    req,
                    &mut new_tx_request_buffer,
                    &mut time_stamp,
                    &tx_pub,
                    &tx_flow_config,
                );
            } else {
                if !new_tx_request_buffer.is_empty() {
                    batch_forward_new_tx(&mut new_tx_request_buffer, &mut time_stamp, &tx_pub);
                }
                thread::sleep(Duration::new(0, tx_flow_config.buffer_duration));
            }
        }
    });

    //ws
    if config.ws_config.enable {
        let ws_config = config.ws_config.clone();
        let tx = tx_relay.clone();
        thread::spawn(move || {
            let url = ws_config.listen_ip.clone() + ":" + &ws_config.listen_port.clone().to_string();
            //let factory = WsFactory::new(ws_responses, tx_pub, 0);
            let factory = WsFactory::new(ws_responses, tx, 0);
            info!("WebSocket Listening on {}", url);
            let mut ws_build = ws::Builder::new();
            ws_build.with_settings(ws_config.into());
            let ws_server = ws_build.build(factory).unwrap();
            let _ = ws_server.listen(url);
        });
    }

    if config.http_config.enable {
        let http_config = config.http_config.clone();
        let addr = http_config.listen_ip.clone() + ":" + &http_config.listen_port.clone().to_string();
        info!("Http Listening on {}", &addr);

        let threads: usize = config
            .http_config
            .thread_number
            .unwrap_or_else(num_cpus::get);

        for i in 0..threads {
            let addr = addr.clone().parse().unwrap();
            let tx = tx_relay.clone();
            let timeout = http_config.timeout;
            let http_responses = Arc::clone(&http_responses);
            let allow_origin = http_config.allow_origin.clone();
            let _ = thread::Builder::new()
                .name(format!("worker{}", i))
                .spawn(move || {
                    let core = Core::new().unwrap();
                    let handle = core.handle();
                    let timeout = Duration::from_secs(timeout);
                    let listener = http_server::listener(&addr, &handle).unwrap();
                    Server::start(core, listener, tx, http_responses, timeout, &allow_origin);
                })
                .unwrap();
        }
    }

    loop {
        let (key, msg) = rx_sub.recv().unwrap();
        mq_handle.handle(&key, &msg);
    }
}

fn batch_forward_new_tx(
    new_tx_request_buffer: &mut Vec<reqlib::Request>,
    time_stamp: &mut SystemTime,
    tx_pub: &Sender<(String, Vec<u8>)>,
) {
    trace!(
        "Going to send new tx batch to auth with {} new tx and buffer time cost is {:?} ",
        new_tx_request_buffer.len(),
        time_stamp.elapsed().unwrap()
    );
    let mut batch_request = BatchRequest::new();
    batch_request.set_new_tx_requests(RepeatedField::from_slice(&new_tx_request_buffer[..]));

    let request_id = Uuid::new_v4().as_bytes().to_vec();
    let mut request = reqlib::Request::new();
    request.set_batch_req(batch_request);
    request.set_request_id(request_id);

    let data: Message = request.into();
    tx_pub
        .send((String::from(TOPIC_NEW_TX_BATCH), data.try_into().unwrap()))
        .unwrap();
    *time_stamp = SystemTime::now();
    new_tx_request_buffer.clear();
}

fn forward_service(
    topic: String,
    req: reqlib::Request,
    new_tx_request_buffer: &mut Vec<reqlib::Request>,
    time_stamp: &mut SystemTime,
    tx_pub: &Sender<(String, Vec<u8>)>,
    config: &NewTxFlowConfig,
) {
    if topic.as_str() != TOPIC_NEW_TX {
        let data: Message = req.into();
        tx_pub.send((topic, data.try_into().unwrap())).unwrap();
    } else {
        new_tx_request_buffer.push(req);
        trace!(
            "New tx is pushed and has {} new tx and buffer time cost is {:?}",
            new_tx_request_buffer.len(),
            time_stamp.elapsed().unwrap()
        );
        if new_tx_request_buffer.len() > config.count_per_batch
            || time_stamp.elapsed().unwrap().subsec_nanos() > config.buffer_duration
        {
            batch_forward_new_tx(new_tx_request_buffer, time_stamp, tx_pub);
        }
    }
}

fn start_profile(config: &ProfileConfig) {
    if config.enable && config.flag_prof_start != 0 && config.flag_prof_duration != 0 {
        let start = config.flag_prof_start;
        let duration = config.flag_prof_duration;
        thread::spawn(move || {
            thread::sleep(Duration::new(start, 0));
            PROFILER
                .lock()
                .unwrap()
                .start("./jsonrpc.profile")
                .expect("Couldn't start");
            thread::sleep(Duration::new(duration, 0));
            PROFILER.lock().unwrap().stop().unwrap();
        });
    }
}
<commit_msg>Remove unused imports.<commit_after>#![feature(try_from)]
extern crate bytes;
extern crate clap;
extern crate cpuprofiler;
extern crate dotenv;
extern crate error;
extern crate futures;
extern crate http;
extern crate httparse;
extern crate hyper;
extern crate jsonrpc_types;
extern crate libproto;
#[macro_use]
extern crate log;
extern crate logger;
extern crate net2;
extern crate num_cpus;
extern crate protobuf;
extern crate pubsub;
extern crate serde;
#[macro_use]
extern crate serde_derive;
#[cfg(not(test))]
extern crate serde_json;
#[cfg(test)]
#[macro_use]
extern crate serde_json;
extern crate threadpool;
extern crate time;
extern crate tokio_core;
extern crate tokio_io;
extern crate toml;
extern crate unicase;
#[macro_use]
extern crate util;
extern crate uuid;
extern crate ws;

mod config;
mod helper;
mod ws_handler;
mod mq_handler;
mod http_server;
mod response;

use clap::App;
use config::{NewTxFlowConfig, ProfileConfig};
use cpuprofiler::PROFILER;
use http_server::Server;
use libproto::Message;
use libproto::request::{self as reqlib, BatchRequest};
use protobuf::RepeatedField;
use pubsub::start_pubsub;
use std::collections::HashMap;
use std::convert::TryInto;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender};
use std::thread;
use std::time::{Duration, SystemTime};
use tokio_core::reactor::Core;
use util::{set_panic_handler, Mutex};
use uuid::Uuid;
use ws_handler::WsFactory;

pub const TOPIC_NEW_TX: &str = "jsonrpc.new_tx";
pub const TOPIC_NEW_TX_BATCH: &str = "jsonrpc.new_tx_batch";

fn main() {
    micro_service_init!("cita-jsonrpc", "CITA:jsonrpc");

    // todo load config
    let matches = App::new("JsonRpc")
        .version("0.1")
        .author("Cryptape")
        .about("CITA JSON-RPC by Rust")
        .args_from_usage("-c, --config=[FILE] 'Sets a custom config file'")
        .get_matches();

    let mut config_path = "./jsonrpc.toml";
    if let Some(c) = matches.value_of("config") {
        info!("Value for config: {}", c);
        config_path = c;
    }

    let config = config::Config::new(config_path);
    info!(
        "CITA:jsonrpc config \n {:?}",
        toml::to_string_pretty(&config).unwrap()
    );

    //enable HTTP or WebSocket server!
    if !config.ws_config.enable && !config.http_config.enable {
        error!("enable HTTP or WebSocket server!");
        std::process::exit(-1);
    }

    start_profile(&config.profile_config);

    // init pubsub
    let (tx_sub, rx_sub) = channel();
    let (tx_pub, rx_pub) = channel();
    //used for buffer message
    let (tx_relay, rx_relay) = channel();
    start_pubsub("jsonrpc", vec!["auth.rpc", "chain.rpc"], tx_sub, rx_pub);

    let backlog_capacity = config.backlog_capacity;

    let responses = Arc::new(Mutex::new(HashMap::with_capacity(backlog_capacity)));
    let http_responses = Arc::clone(&responses);
    let ws_responses = Arc::clone(&responses);
    let mut mq_handle = mq_handler::MqHandler::new(responses);

    //dispatch
    let tx_flow_config = config.new_tx_flow_config;
    thread::spawn(move || {
        let mut new_tx_request_buffer = Vec::new();
        let mut time_stamp = SystemTime::now();
        loop {
            if let Ok(res) = rx_relay.try_recv() {
                let (topic, req): (String, reqlib::Request) = res;
                forward_service(
                    topic,
                    req,
                    &mut new_tx_request_buffer,
                    &mut time_stamp,
                    &tx_pub,
                    &tx_flow_config,
                );
            } else {
                if !new_tx_request_buffer.is_empty() {
                    batch_forward_new_tx(&mut new_tx_request_buffer, &mut time_stamp, &tx_pub);
                }
                thread::sleep(Duration::new(0, tx_flow_config.buffer_duration));
            }
        }
    });

    //ws
    if config.ws_config.enable {
        let ws_config = config.ws_config.clone();
        let tx = tx_relay.clone();
        thread::spawn(move || {
            let url = ws_config.listen_ip.clone() + ":" + &ws_config.listen_port.clone().to_string();
            //let factory = WsFactory::new(ws_responses, tx_pub, 0);
            let factory = WsFactory::new(ws_responses, tx, 0);
            info!("WebSocket Listening on {}", url);
            let mut ws_build = ws::Builder::new();
            ws_build.with_settings(ws_config.into());
            let ws_server = ws_build.build(factory).unwrap();
            let _ = ws_server.listen(url);
        });
    }

    if config.http_config.enable {
        let http_config = config.http_config.clone();
        let addr = http_config.listen_ip.clone() + ":" + &http_config.listen_port.clone().to_string();
        info!("Http Listening on {}", &addr);

        let threads: usize = config
            .http_config
            .thread_number
            .unwrap_or_else(num_cpus::get);

        for i in 0..threads {
            let addr = addr.clone().parse().unwrap();
            let tx = tx_relay.clone();
            let timeout = http_config.timeout;
            let http_responses = Arc::clone(&http_responses);
            let allow_origin = http_config.allow_origin.clone();
            let _ = thread::Builder::new()
                .name(format!("worker{}", i))
                .spawn(move || {
                    let core = Core::new().unwrap();
                    let handle = core.handle();
                    let timeout = Duration::from_secs(timeout);
                    let listener = http_server::listener(&addr, &handle).unwrap();
                    Server::start(core, listener, tx, http_responses, timeout, &allow_origin);
                })
                .unwrap();
        }
    }

    loop {
        let (key, msg) = rx_sub.recv().unwrap();
        mq_handle.handle(&key, &msg);
    }
}

fn batch_forward_new_tx(
    new_tx_request_buffer: &mut Vec<reqlib::Request>,
    time_stamp: &mut SystemTime,
    tx_pub: &Sender<(String, Vec<u8>)>,
) {
    trace!(
        "Going to send new tx batch to auth with {} new tx and buffer time cost is {:?} ",
        new_tx_request_buffer.len(),
        time_stamp.elapsed().unwrap()
    );
    let mut batch_request = BatchRequest::new();
    batch_request.set_new_tx_requests(RepeatedField::from_slice(&new_tx_request_buffer[..]));

    let request_id = Uuid::new_v4().as_bytes().to_vec();
    let mut request = reqlib::Request::new();
    request.set_batch_req(batch_request);
    request.set_request_id(request_id);

    let data: Message = request.into();
    tx_pub
        .send((String::from(TOPIC_NEW_TX_BATCH), data.try_into().unwrap()))
        .unwrap();
    *time_stamp = SystemTime::now();
    new_tx_request_buffer.clear();
}

fn forward_service(
    topic: String,
    req: reqlib::Request,
    new_tx_request_buffer: &mut Vec<reqlib::Request>,
    time_stamp: &mut SystemTime,
    tx_pub: &Sender<(String, Vec<u8>)>,
    config: &NewTxFlowConfig,
) {
    if topic.as_str() != TOPIC_NEW_TX {
        let data: Message = req.into();
        tx_pub.send((topic, data.try_into().unwrap())).unwrap();
    } else {
        new_tx_request_buffer.push(req);
        trace!(
            "New tx is pushed and has {} new tx and buffer time cost is {:?}",
            new_tx_request_buffer.len(),
            time_stamp.elapsed().unwrap()
        );
        if new_tx_request_buffer.len() > config.count_per_batch
            || time_stamp.elapsed().unwrap().subsec_nanos() > config.buffer_duration
        {
            batch_forward_new_tx(new_tx_request_buffer, time_stamp, tx_pub);
        }
    }
}

fn start_profile(config: &ProfileConfig) {
    if config.enable && config.flag_prof_start != 0 && config.flag_prof_duration != 0 {
        let start = config.flag_prof_start;
        let duration = config.flag_prof_duration;
        thread::spawn(move || {
            thread::sleep(Duration::new(start, 0));
            PROFILER
                .lock()
                .unwrap()
                .start("./jsonrpc.profile")
                .expect("Couldn't start");
            thread::sleep(Duration::new(duration, 0));
            PROFILER.lock().unwrap().stop().unwrap();
        });
    }
}
<|endoftext|>


<commit_before>//! Detecting diagnostic items.
//!
//! Diagnostic items are items that are not language-inherent, but can reasonably be expected to
//! exist for diagnostic purposes. This allows diagnostic authors to refer to specific items
//! directly, without having to guess module paths and crates.
//! Examples are:
//!
//! * Traits like `Debug`, that have no bearing on language semantics
//!
//! * Compiler internal types like `Ty` and `TyCtxt`

use rustc_ast as ast;
use rustc_data_structures::fx::FxHashMap;
use rustc_hir as hir;
use rustc_hir::itemlikevisit::ItemLikeVisitor;
use rustc_middle::ty::query::Providers;
use rustc_middle::ty::TyCtxt;
use rustc_session::Session;
use rustc_span::def_id::{DefId, LOCAL_CRATE};
use rustc_span::symbol::{sym, Symbol};

struct DiagnosticItemCollector<'tcx> {
    // items from this crate
    items: FxHashMap<Symbol, DefId>,
    tcx: TyCtxt<'tcx>,
}

impl<'v, 'tcx> ItemLikeVisitor<'v> for DiagnosticItemCollector<'tcx> {
    fn visit_item(&mut self, item: &hir::Item<'_>) {
        self.observe_item(&item.attrs, item.hir_id);
    }

    fn visit_trait_item(&mut self, trait_item: &hir::TraitItem<'_>) {
        self.observe_item(&trait_item.attrs, trait_item.hir_id);
    }

    fn visit_impl_item(&mut self, impl_item: &hir::ImplItem<'_>) {
        self.observe_item(&impl_item.attrs, impl_item.hir_id);
    }

    fn visit_foreign_item(&mut self, _foreign_item: &hir::ForeignItem<'_>) {}
}

impl<'tcx> DiagnosticItemCollector<'tcx> {
    fn new(tcx: TyCtxt<'tcx>) -> DiagnosticItemCollector<'tcx> {
        DiagnosticItemCollector { tcx, items: Default::default() }
    }

    fn observe_item(&mut self, attrs: &[ast::Attribute], hir_id: hir::HirId) {
        if let Some(name) = extract(&self.tcx.sess, attrs) {
            let def_id = self.tcx.hir().local_def_id(hir_id);
            // insert into our table
            collect_item(self.tcx, &mut self.items, name, def_id.to_def_id());
        }
    }
}

fn collect_item(
    tcx: TyCtxt<'_>,
    items: &mut FxHashMap<Symbol, DefId>,
    name: Symbol,
    item_def_id: DefId,
) {
    // Check for duplicates.
    if let Some(original_def_id) = items.insert(name, item_def_id) {
        if original_def_id != item_def_id {
            let mut err = match tcx.hir().span_if_local(item_def_id) {
                Some(span) => tcx.sess.struct_span_err(
                    span,
                    &format!("duplicate diagnostic item found: `{}`.", name),
                ),
                None => tcx.sess.struct_err(&format!(
                    "duplicate diagnostic item in crate `{}`: `{}`.",
                    tcx.crate_name(item_def_id.krate),
                    name
                )),
            };
            if let Some(span) = tcx.hir().span_if_local(original_def_id) {
                err.span_note(span, "the diagnostic item is first defined here");
            } else {
                err.note(&format!(
                    "the diagnostic item is first defined in crate `{}`.",
                    tcx.crate_name(original_def_id.krate)
                ));
            }
            err.emit();
        }
    }
}

/// Extract the first `rustc_diagnostic_item = "$name"` out of a list of attributes.
fn extract(sess: &Session, attrs: &[ast::Attribute]) -> Option<Symbol> {
    attrs.iter().find_map(|attr| {
        if sess.check_name(attr, sym::rustc_diagnostic_item) { attr.value_str() } else { None }
    })
}

/// Traverse and collect the diagnostic items in the current
fn collect<'tcx>(tcx: TyCtxt<'tcx>) -> FxHashMap<Symbol, DefId> {
    // Initialize the collector.
    let mut collector = DiagnosticItemCollector::new(tcx);

    // Collect diagnostic items in this crate.
    tcx.hir().krate().visit_all_item_likes(&mut collector);
    // FIXME(visit_all_item_likes): Foreign items are not visited
    // here, so we have to manually look at them for now.
    for (_, foreign_module) in tcx.foreign_modules(LOCAL_CRATE).iter() {
        for &foreign_item in foreign_module.foreign_items.iter() {
            match tcx.hir().get(tcx.hir().local_def_id_to_hir_id(foreign_item.expect_local())) {
                hir::Node::ForeignItem(item) => {
                    collector.observe_item(item.attrs, item.hir_id);
                }
                item => bug!("unexpected foreign item {:?}", item),
            }
        }
    }

    for m in tcx.hir().krate().exported_macros {
        collector.observe_item(m.attrs, m.hir_id);
    }

    collector.items
}

/// Traverse and collect all the diagnostic items in all crates.
fn collect_all<'tcx>(tcx: TyCtxt<'tcx>) -> FxHashMap<Symbol, DefId> {
    // Initialize the collector.
    let mut collector = FxHashMap::default();

    // Collect diagnostic items in other crates.
    for &cnum in tcx.crates().iter().chain(std::iter::once(&LOCAL_CRATE)) {
        for (&name, &def_id) in tcx.diagnostic_items(cnum).iter() {
            collect_item(tcx, &mut collector, name, def_id);
        }
    }

    collector
}

pub fn provide(providers: &mut Providers) {
    providers.diagnostic_items = |tcx, id| {
        assert_eq!(id, LOCAL_CRATE);
        collect(tcx)
    };
    providers.all_diagnostic_items = |tcx, id| {
        assert_eq!(id, LOCAL_CRATE);
        collect_all(tcx)
    };
}
<commit_msg>Use visitor for its very purpose.<commit_after>//! Detecting diagnostic items.
//!
//! Diagnostic items are items that are not language-inherent, but can reasonably be expected to
//! exist for diagnostic purposes. This allows diagnostic authors to refer to specific items
//! directly, without having to guess module paths and crates.
//! Examples are:
//!
//! * Traits like `Debug`, that have no bearing on language semantics
//!
//! * Compiler internal types like `Ty` and `TyCtxt`

use rustc_ast as ast;
use rustc_data_structures::fx::FxHashMap;
use rustc_hir as hir;
use rustc_hir::itemlikevisit::ItemLikeVisitor;
use rustc_middle::ty::query::Providers;
use rustc_middle::ty::TyCtxt;
use rustc_session::Session;
use rustc_span::def_id::{DefId, LOCAL_CRATE};
use rustc_span::symbol::{sym, Symbol};

struct DiagnosticItemCollector<'tcx> {
    // items from this crate
    items: FxHashMap<Symbol, DefId>,
    tcx: TyCtxt<'tcx>,
}

impl<'v, 'tcx> ItemLikeVisitor<'v> for DiagnosticItemCollector<'tcx> {
    fn visit_item(&mut self, item: &hir::Item<'_>) {
        self.observe_item(&item.attrs, item.hir_id);
    }

    fn visit_trait_item(&mut self, trait_item: &hir::TraitItem<'_>) {
        self.observe_item(&trait_item.attrs, trait_item.hir_id);
    }

    fn visit_impl_item(&mut self, impl_item: &hir::ImplItem<'_>) {
        self.observe_item(&impl_item.attrs, impl_item.hir_id);
    }

    fn visit_foreign_item(&mut self, foreign_item: &hir::ForeignItem<'_>) {
        self.observe_item(foreign_item.attrs, foreign_item.hir_id);
    }
}

impl<'tcx> DiagnosticItemCollector<'tcx> {
    fn new(tcx: TyCtxt<'tcx>) -> DiagnosticItemCollector<'tcx> {
        DiagnosticItemCollector { tcx, items: Default::default() }
    }

    fn observe_item(&mut self, attrs: &[ast::Attribute], hir_id: hir::HirId) {
        if let Some(name) = extract(&self.tcx.sess, attrs) {
            let def_id = self.tcx.hir().local_def_id(hir_id);
            // insert into our table
            collect_item(self.tcx, &mut self.items, name, def_id.to_def_id());
        }
    }
}

fn collect_item(
    tcx: TyCtxt<'_>,
    items: &mut FxHashMap<Symbol, DefId>,
    name: Symbol,
    item_def_id: DefId,
) {
    // Check for duplicates.
    if let Some(original_def_id) = items.insert(name, item_def_id) {
        if original_def_id != item_def_id {
            let mut err = match tcx.hir().span_if_local(item_def_id) {
                Some(span) => tcx.sess.struct_span_err(
                    span,
                    &format!("duplicate diagnostic item found: `{}`.", name),
                ),
                None => tcx.sess.struct_err(&format!(
                    "duplicate diagnostic item in crate `{}`: `{}`.",
                    tcx.crate_name(item_def_id.krate),
                    name
                )),
            };
            if let Some(span) = tcx.hir().span_if_local(original_def_id) {
                err.span_note(span, "the diagnostic item is first defined here");
            } else {
                err.note(&format!(
                    "the diagnostic item is first defined in crate `{}`.",
                    tcx.crate_name(original_def_id.krate)
                ));
            }
            err.emit();
        }
    }
}

/// Extract the first `rustc_diagnostic_item = "$name"` out of a list of attributes.
fn extract(sess: &Session, attrs: &[ast::Attribute]) -> Option<Symbol> {
    attrs.iter().find_map(|attr| {
        if sess.check_name(attr, sym::rustc_diagnostic_item) { attr.value_str() } else { None }
    })
}

/// Traverse and collect the diagnostic items in the current
fn collect<'tcx>(tcx: TyCtxt<'tcx>) -> FxHashMap<Symbol, DefId> {
    // Initialize the collector.
    let mut collector = DiagnosticItemCollector::new(tcx);

    // Collect diagnostic items in this crate.
    tcx.hir().krate().visit_all_item_likes(&mut collector);

    for m in tcx.hir().krate().exported_macros {
        collector.observe_item(m.attrs, m.hir_id);
    }

    collector.items
}

/// Traverse and collect all the diagnostic items in all crates.
fn collect_all<'tcx>(tcx: TyCtxt<'tcx>) -> FxHashMap<Symbol, DefId> {
    // Initialize the collector.
    let mut collector = FxHashMap::default();

    // Collect diagnostic items in other crates.
    for &cnum in tcx.crates().iter().chain(std::iter::once(&LOCAL_CRATE)) {
        for (&name, &def_id) in tcx.diagnostic_items(cnum).iter() {
            collect_item(tcx, &mut collector, name, def_id);
        }
    }

    collector
}

pub fn provide(providers: &mut Providers) {
    providers.diagnostic_items = |tcx, id| {
        assert_eq!(id, LOCAL_CRATE);
        collect(tcx)
    };
    providers.all_diagnostic_items = |tcx, id| {
        assert_eq!(id, LOCAL_CRATE);
        collect_all(tcx)
    };
}
<|endoftext|>


<commit_before>use crate::{api_client::{self,
                         Client},
            common::ui::{Status,
                         UIWriter,
                         UI},
            error::{Error,
                    Result},
            PRODUCT,
            VERSION};
use reqwest::StatusCode;

pub fn start(ui: &mut UI, bldr_url: &str, token: &str, origin: &str, account: &str) -> Result<()> {
    let api_client = Client::new(bldr_url, PRODUCT, VERSION, None).map_err(Error::APIClient)?;

    ui.status(Status::Transferring,
              format!("ownership of origin {} to {}.", origin, account))?;

    match api_client.transfer_origin(origin, token, account) {
        Ok(_) => {
            ui.status(Status::Transferred, format!("ownership successfully!"))
              .map_err(Error::from)
        }
        Err(api_client::Error::APIError(StatusCode::FORBIDDEN, msg)) => {
            ui.fatal("Failed to transfer origin ownership!")?;
            ui.fatal("Either you are not the current owner, the operation requested has already \
                      been completed, or the new owner is not yet a member.")?;
            Err(Error::APIClient(api_client::Error::APIError(StatusCode::FORBIDDEN, msg)))
        }
        Err(api_client::Error::APIError(StatusCode::UNAUTHORIZED, msg)) => {
            ui.fatal("Failed to transfer origin ownership!")?;
            Err(Error::APIClient(api_client::Error::APIError(StatusCode::UNAUTHORIZED, msg)))
        }
        Err(api_client::Error::APIError(StatusCode::UNPROCESSABLE_ENTITY, msg)) => {
            ui.fatal("Failed to transfer origin ownership!")?;
            ui.fatal("Did you try to transfer ownership to yourself?")?;
            Err(Error::APIClient(api_client::Error::APIError(StatusCode::UNPROCESSABLE_ENTITY,
                                                             msg)))
        }
        Err(api_client::Error::APIError(StatusCode::NOT_FOUND, msg)) => {
            ui.fatal("Failed to transfer origin ownership!")?;
            ui.fatal("The origin or the origin owner account (or both) does not exist.")?;
            Err(Error::APIClient(api_client::Error::APIError(StatusCode::NOT_FOUND, msg)))
        }
        Err(e) => {
            ui.fatal(format!("Failed to transfer origin {} ownership to {}, {:?}",
                             origin, account, e))?;
            Err(Error::from(e))
        }
    }
}
<commit_msg>removing useless format<commit_after>use crate::{api_client::{self,
                         Client},
            common::ui::{Status,
                         UIWriter,
                         UI},
            error::{Error,
                    Result},
            PRODUCT,
            VERSION};
use reqwest::StatusCode;

pub fn start(ui: &mut UI, bldr_url: &str, token: &str, origin: &str, account: &str) -> Result<()> {
    let api_client = Client::new(bldr_url, PRODUCT, VERSION, None).map_err(Error::APIClient)?;

    ui.status(Status::Transferring,
              format!("ownership of origin {} to {}.", origin, account))?;

    match api_client.transfer_origin(origin, token, account) {
        Ok(_) => {
            ui.status(Status::Transferred, "ownership successfully!".to_string())
              .map_err(Error::from)
        }
        Err(api_client::Error::APIError(StatusCode::FORBIDDEN, msg)) => {
            ui.fatal("Failed to transfer origin ownership!")?;
            ui.fatal("Either you are not the current owner, the operation requested has already \
                      been completed, or the new owner is not yet a member.")?;
            Err(Error::APIClient(api_client::Error::APIError(StatusCode::FORBIDDEN, msg)))
        }
        Err(api_client::Error::APIError(StatusCode::UNAUTHORIZED, msg)) => {
            ui.fatal("Failed to transfer origin ownership!")?;
            Err(Error::APIClient(api_client::Error::APIError(StatusCode::UNAUTHORIZED, msg)))
        }
        Err(api_client::Error::APIError(StatusCode::UNPROCESSABLE_ENTITY, msg)) => {
            ui.fatal("Failed to transfer origin ownership!")?;
            ui.fatal("Did you try to transfer ownership to yourself?")?;
            Err(Error::APIClient(api_client::Error::APIError(StatusCode::UNPROCESSABLE_ENTITY,
                                                             msg)))
        }
        Err(api_client::Error::APIError(StatusCode::NOT_FOUND, msg)) => {
            ui.fatal("Failed to transfer origin ownership!")?;
            ui.fatal("The origin or the origin owner account (or both) does not exist.")?;
            Err(Error::APIClient(api_client::Error::APIError(StatusCode::NOT_FOUND, msg)))
        }
        Err(e) => {
            ui.fatal(format!("Failed to transfer origin {} ownership to {}, {:?}",
                             origin, account, e))?;
            Err(Error::from(e))
        }
    }
}
<|endoftext|>


<commit_before>//! Abstractions for dealing with the main behaviors we need when
//! dealing with container images, while remaining unconcerned about
//! which underlying tool is actually performing those tasks.
//!
//! This allows us to swap out the `docker` CLI for `buildah` if we
//! want to create containers as a non-root user, for instance.

use crate::error::Result;
use clap::{Arg,
           ArgMatches};
use habitat_core::fs::find_command;
use std::{convert::TryFrom,
          path::{Path,
                 PathBuf},
          process::{Command,
                    ExitStatus},
          result::Result as StdResult,
          str::FromStr};

#[cfg(not(windows))]
mod buildah;
mod docker;

#[derive(Debug, Fail)]
enum EngineError {
    #[fail(display = "Container image build failed with exit code: {}", _0)]
    BuildFailed(ExitStatus),
    #[fail(display = "Could not find the container engine executable '{}' on the PATH",
           _0)]
    ExecutableNotFound(String),
    #[fail(display = "Could not determine container image ID for: {}", _0)]
    ImageIdNotFound(String),
    #[fail(display = "Removing local container images failed with exit code: {}",
           _0)]
    RemoveFailed(ExitStatus),
    #[fail(display = "Container image push failed with exit code: {}", _0)]
    PushFailed(ExitStatus),
    #[fail(display = "Unknown Container Engine '{}' was specified.", _0)]
    UnknownEngine(String),
    #[fail(display = "Cannot use `--engine=buildah` with `--multi-layer` due to https://github.com/containers/buildah/issues/2215. Please use `--engine=docker` or remove `--multi-layer`.")]
    BuildahIncompatibleWithMultiLayer,
}

/// Due to a bug in Buildah, any layers that we create in a
/// multi-layer build won't get reused, which eliminates any benefit
/// we might get from them.
///
/// Until that bug is fixed, we'll prevent using Buildah to create
/// multi-layer images, as the confusion arising from generating
/// multiple layers but not being able to reuse any of them is
/// something that's better to avoid.
///
/// When https://github.com/containers/buildah/issues/2215 is fixed,
/// we can update our Buildah dependency and remove this check.
pub fn fail_if_buildah_and_multilayer(matches: &ArgMatches) -> Result<()> {
    if matches.value_of("ENGINE") == Some("buildah") && matches.is_present("MULTI_LAYER") {
        return Err(EngineError::BuildahIncompatibleWithMultiLayer.into());
    }
    Ok(())
}

/// Things that can build containers!
#[derive(Clone, Copy, Debug)]
enum EngineKind {
    Docker,
    #[cfg(not(windows))]
    Buildah,
}

impl FromStr for EngineKind {
    type Err = EngineError;

    fn from_str(s: &str) -> StdResult<Self, Self::Err> {
        match s {
            "docker" => Ok(EngineKind::Docker),
            #[cfg(not(windows))]
            "buildah" => Ok(EngineKind::Buildah),
            _ => Err(EngineError::UnknownEngine(s.to_string())),
        }
    }
}

/// Define the CLAP CLI argument for specifying a container build
/// engine to use.
#[rustfmt::skip] // otherwise the long_help formatting goes crazy
pub fn cli_arg<'a, 'b>() -> Arg<'a, 'b> {
    let arg =
        Arg::with_name("ENGINE").value_name("ENGINE")
        .long("engine")
        .required(true)
        .takes_value(true)
        .multiple(false)
        .default_value("docker")
        .help("The name of the container creation engine to use.");

    // TODO (CM): Find a way to tie this more closely to the
    // Engine enum values.
    if cfg!(windows) {
        // Since there is effectively no choice of engine for
        // Windows, we hide the CLI option and don't document it
        // any further.
        arg.possible_values(&["docker"]).hidden(true)
    } else {
        arg.long_help(
"Using the `docker` engine allows you to use Docker to create
your container images. You must ensure that a Docker daemon
is running on the host where this command is executed, and
that the user executing the command has permission to access
the Docker socket.

Using the `buildah` engine allows you to create container images
as an unprivileged user, and without having to use a Docker
daemon. This is the recommended engine for use in CI systems and
other environments where security is of particular concern.
Please see https://buildah.io for more details.

Both engines create equivalent container images.
",
        )
            .possible_values(&["docker", "buildah"])
    }
}

impl TryFrom<&ArgMatches<'_>> for Box<dyn Engine> {
    type Error = failure::Error;

    fn try_from(value: &ArgMatches) -> StdResult<Self, Self::Error> {
        let engine_kind =
            clap::value_t!(value, "ENGINE", EngineKind).expect("ENGINE is a required option");
        match engine_kind {
            EngineKind::Docker => Ok(Box::new(docker::DockerEngine::new()?)),
            #[cfg(not(windows))]
            EngineKind::Buildah => Ok(Box::new(buildah::BuildahEngine::new()?)),
        }
    }
}

pub trait Engine {
    /// A command that takes a container image reference and returns
    /// the ID of that image on the first line of standard output.
    fn image_id_command(&self, image_reference: &str) -> Command;

    /// A command that removes the specified local container image;
    fn image_removal_command(&self, image_reference: &str) -> Command;

    /// A command that pushes the specified container image, using
    /// configuration stored in `config_dir`.
    // TODO (CM): accept repository URL information
    // TODO (CM): worth taking credential / repo information and
    // handling the config directory stuff internally?
    fn image_push_command(&self, image_reference: &str, config_dir: &Path) -> Command;

    fn build_command(&self, build_context: &Path, tags: &[String], memory: Option<&str>)
                     -> Command;

    /// Retrieve the ID of the given image, which is expected to exist.
    fn image_id(&self, image_reference: &str) -> Result<String> {
        let mut cmd = self.image_id_command(image_reference);
        debug!("Running: {:?}", cmd);
        let output = cmd.output()?;
        let stdout = String::from_utf8_lossy(&output.stdout);
        match stdout.lines().next() {
            Some(id) => Ok(id.to_string()),
            None => Err(EngineError::ImageIdNotFound(image_reference.to_string()).into()),
        }
    }

    /// Delete the referenced image in the local image store.
    fn remove_image(&self, image_reference: &str) -> Result<()> {
        run(self.image_removal_command(image_reference),
            EngineError::RemoveFailed)
    }

    /// Pushes the specified container image to a remote repository, using
    /// configuration stored in `config_dir`.
    // TODO (CM): accept repository URL information
    // TODO (CM): worth taking credential / repo information and
    // handling the config directory stuff internally?
    fn push_image(&self, image_reference: &str, config_dir: &Path) -> Result<()> {
        run(self.image_push_command(image_reference, config_dir),
            EngineError::PushFailed)
    }

    /// Actually create the image.
    ///
    /// `build_context` will serve as the build context directory, and
    /// a suitable `Dockerfile` is expected to be present in it. The
    /// image will be tagged with each of `tags`.
    ///
    /// `memory` governs how much memory is provided to the build
    /// process.
    ///
    /// Returns the ID of the image that was built.
    fn build(&self, build_context: &Path, tags: &[String], memory: Option<&str>) -> Result<String> {
        run(self.build_command(build_context, tags, memory),
            EngineError::BuildFailed)?;

        let identifier = tags.first()
                             .expect("There should always be at least one tag");
        self.image_id(identifier)
    }
}

/// General helper function for actually executing all these commands.
///
/// Not part of the trait because nobody need to be calling this from
/// outside.
fn run<F>(mut cmd: Command, err_fn: F) -> Result<()>
    where F: Fn(ExitStatus) -> EngineError
{
    debug!("Running: {:?}", &cmd);
    let exit_status = cmd.spawn()?.wait()?;
    if !exit_status.success() {
        return Err(err_fn(exit_status).into());
    }
    Ok(())
}

fn resolve_engine_binary(binary_name: &str) -> StdResult<PathBuf, EngineError> {
    find_command(binary_name).ok_or_else(|| {
                                 EngineError::ExecutableNotFound(binary_name.to_string())
                             })
}
<commit_msg>Allow container export engine to be specified in the environment<commit_after>//! Abstractions for dealing with the main behaviors we need when
//! dealing with container images, while remaining unconcerned about
//! which underlying tool is actually performing those tasks.
//!
//! This allows us to swap out the `docker` CLI for `buildah` if we
//! want to create containers as a non-root user, for instance.

use crate::error::Result;
use clap::{Arg,
           ArgMatches};
use habitat_core::fs::find_command;
use std::{convert::TryFrom,
          path::{Path,
                 PathBuf},
          process::{Command,
                    ExitStatus},
          result::Result as StdResult,
          str::FromStr};

#[cfg(not(windows))]
mod buildah;
mod docker;

#[derive(Debug, Fail)]
enum EngineError {
    #[fail(display = "Container image build failed with exit code: {}", _0)]
    BuildFailed(ExitStatus),
    #[fail(display = "Could not find the container engine executable '{}' on the PATH",
           _0)]
    ExecutableNotFound(String),
    #[fail(display = "Could not determine container image ID for: {}", _0)]
    ImageIdNotFound(String),
    #[fail(display = "Removing local container images failed with exit code: {}",
           _0)]
    RemoveFailed(ExitStatus),
    #[fail(display = "Container image push failed with exit code: {}", _0)]
    PushFailed(ExitStatus),
    #[fail(display = "Unknown Container Engine '{}' was specified.", _0)]
    UnknownEngine(String),
    #[fail(display = "Cannot use `--engine=buildah` with `--multi-layer` due to https://github.com/containers/buildah/issues/2215. Please use `--engine=docker` or remove `--multi-layer`.")]
    BuildahIncompatibleWithMultiLayer,
}

/// Due to a bug in Buildah, any layers that we create in a
/// multi-layer build won't get reused, which eliminates any benefit
/// we might get from them.
///
/// Until that bug is fixed, we'll prevent using Buildah to create
/// multi-layer images, as the confusion arising from generating
/// multiple layers but not being able to reuse any of them is
/// something that's better to avoid.
///
/// When https://github.com/containers/buildah/issues/2215 is fixed,
/// we can update our Buildah dependency and remove this check.
pub fn fail_if_buildah_and_multilayer(matches: &ArgMatches) -> Result<()> {
    if matches.value_of("ENGINE") == Some("buildah") && matches.is_present("MULTI_LAYER") {
        return Err(EngineError::BuildahIncompatibleWithMultiLayer.into());
    }
    Ok(())
}

/// Things that can build containers!
#[derive(Clone, Copy, Debug)]
enum EngineKind {
    Docker,
    #[cfg(not(windows))]
    Buildah,
}

impl FromStr for EngineKind {
    type Err = EngineError;

    fn from_str(s: &str) -> StdResult<Self, Self::Err> {
        match s {
            "docker" => Ok(EngineKind::Docker),
            #[cfg(not(windows))]
            "buildah" => Ok(EngineKind::Buildah),
            _ => Err(EngineError::UnknownEngine(s.to_string())),
        }
    }
}

/// Define the CLAP CLI argument for specifying a container build
/// engine to use.
#[rustfmt::skip] // otherwise the long_help formatting goes crazy
pub fn cli_arg<'a, 'b>() -> Arg<'a, 'b> {
    let arg =
        Arg::with_name("ENGINE").value_name("ENGINE")
        .long("engine")
        .required(true)
        .env("HAB_PKG_EXPORT_CONTAINER_ENGINE")
        .takes_value(true)
        .multiple(false)
        .default_value("docker")
        .help("The name of the container creation engine to use.");

    // TODO (CM): Find a way to tie this more closely to the
    // Engine enum values.
    if cfg!(windows) {
        // Since there is effectively no choice of engine for
        // Windows, we hide the CLI option and don't document it
        // any further.
        arg.possible_values(&["docker"]).hidden(true)
    } else {
        arg.long_help(
"Using the `docker` engine allows you to use Docker to create
your container images. You must ensure that a Docker daemon
is running on the host where this command is executed, and
that the user executing the command has permission to access
the Docker socket.

Using the `buildah` engine allows you to create container images
as an unprivileged user, and without having to use a Docker
daemon. This is the recommended engine for use in CI systems and
other environments where security is of particular concern.
Please see https://buildah.io for more details.

Both engines create equivalent container images.
",
        )
            .possible_values(&["docker", "buildah"])
    }
}

impl TryFrom<&ArgMatches<'_>> for Box<dyn Engine> {
    type Error = failure::Error;

    fn try_from(value: &ArgMatches) -> StdResult<Self, Self::Error> {
        let engine_kind =
            clap::value_t!(value, "ENGINE", EngineKind).expect("ENGINE is a required option");
        match engine_kind {
            EngineKind::Docker => Ok(Box::new(docker::DockerEngine::new()?)),
            #[cfg(not(windows))]
            EngineKind::Buildah => Ok(Box::new(buildah::BuildahEngine::new()?)),
        }
    }
}

pub trait Engine {
    /// A command that takes a container image reference and returns
    /// the ID of that image on the first line of standard output.
    fn image_id_command(&self, image_reference: &str) -> Command;

    /// A command that removes the specified local container image;
    fn image_removal_command(&self, image_reference: &str) -> Command;

    /// A command that pushes the specified container image, using
    /// configuration stored in `config_dir`.
    // TODO (CM): accept repository URL information
    // TODO (CM): worth taking credential / repo information and
    // handling the config directory stuff internally?
    fn image_push_command(&self, image_reference: &str, config_dir: &Path) -> Command;

    fn build_command(&self, build_context: &Path, tags: &[String], memory: Option<&str>)
                     -> Command;

    /// Retrieve the ID of the given image, which is expected to exist.
    fn image_id(&self, image_reference: &str) -> Result<String> {
        let mut cmd = self.image_id_command(image_reference);
        debug!("Running: {:?}", cmd);
        let output = cmd.output()?;
        let stdout = String::from_utf8_lossy(&output.stdout);
        match stdout.lines().next() {
            Some(id) => Ok(id.to_string()),
            None => Err(EngineError::ImageIdNotFound(image_reference.to_string()).into()),
        }
    }

    /// Delete the referenced image in the local image store.
    fn remove_image(&self, image_reference: &str) -> Result<()> {
        run(self.image_removal_command(image_reference),
            EngineError::RemoveFailed)
    }

    /// Pushes the specified container image to a remote repository, using
    /// configuration stored in `config_dir`.
    // TODO (CM): accept repository URL information
    // TODO (CM): worth taking credential / repo information and
    // handling the config directory stuff internally?
    fn push_image(&self, image_reference: &str, config_dir: &Path) -> Result<()> {
        run(self.image_push_command(image_reference, config_dir),
            EngineError::PushFailed)
    }

    /// Actually create the image.
    ///
    /// `build_context` will serve as the build context directory, and
    /// a suitable `Dockerfile` is expected to be present in it. The
    /// image will be tagged with each of `tags`.
    ///
    /// `memory` governs how much memory is provided to the build
    /// process.
    ///
    /// Returns the ID of the image that was built.
    fn build(&self, build_context: &Path, tags: &[String], memory: Option<&str>) -> Result<String> {
        run(self.build_command(build_context, tags, memory),
            EngineError::BuildFailed)?;

        let identifier = tags.first()
                             .expect("There should always be at least one tag");
        self.image_id(identifier)
    }
}

/// General helper function for actually executing all these commands.
///
/// Not part of the trait because nobody need to be calling this from
/// outside.
fn run<F>(mut cmd: Command, err_fn: F) -> Result<()>
    where F: Fn(ExitStatus) -> EngineError
{
    debug!("Running: {:?}", &cmd);
    let exit_status = cmd.spawn()?.wait()?;
    if !exit_status.success() {
        return Err(err_fn(exit_status).into());
    }
    Ok(())
}

fn resolve_engine_binary(binary_name: &str) -> StdResult<PathBuf, EngineError> {
    find_command(binary_name).ok_or_else(|| {
                                 EngineError::ExecutableNotFound(binary_name.to_string())
                             })
}
<|endoftext|>


<commit_before>// external crates
extern crate test;

// local crates
extern crate pretty;

// local mod imports
use trees::{
    Tree,
};

// custom mod imports
#[path="../examples/trees.rs"]
mod trees;

#[bench]
fn bench(b:&mut test::Bencher) -> () {
    let bbbbbbs =
        [ Tree::new("ccc", [])
        , Tree::new("dd", [])
        ];
    let ffffs =
        [ Tree::new("gg", [])
        , Tree::new("hhh", [])
        , Tree::new("ii", [])
        ];
    let aaas =
        [ Tree::new("bbbbbb", bbbbbbs)
        , Tree::new("eee", [])
        , Tree::new("ffff", ffffs)
        ];
    let example = Tree::new("aaa", aaas);
    let task = || {
        example.pretty().render(70)
    };
    b.iter(task);
}
<commit_msg>Run more iterations during benchmark<commit_after>// external crates
extern crate test;

// local crates
extern crate pretty;

// local mod imports
use trees::{
    Tree,
};

// custom mod imports
#[path="../examples/trees.rs"]
mod trees;

#[bench]
fn bench(b:&mut test::Bencher) -> () {
    let bbbbbbs =
        [ Tree::new("ccc", [])
        , Tree::new("dd", [])
        ];
    let ffffs =
        [ Tree::new("gg", [])
        , Tree::new("hhh", [])
        , Tree::new("ii", [])
        ];
    let aaas =
        [ Tree::new("bbbbbb", bbbbbbs)
        , Tree::new("eee", [])
        , Tree::new("ffff", ffffs)
        ];
    let example = Tree::new("aaa", aaas);
    let task = || {
        for _ in range (0u, 100u) {
            example.pretty().render(70);
        }
    };
    b.iter(task);
}
<|endoftext|>


<commit_before>#![cfg_attr(feature = "clippy", plugin(clippy))]
#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", allow(too_many_arguments))]
#![cfg_attr(feature = "clippy", allow(used_underscore_binding))]
#![cfg_attr(not(feature = "with-syntex"), feature(rustc_private, plugin))]

// The `quote!` macro requires deep recursion.
#![recursion_limit = "192"]

extern crate serde_codegen_internals as internals;

#[cfg(feature = "with-syntex")]
extern crate syntex;

#[cfg(feature = "with-syntex")]
#[macro_use]
extern crate syntex_syntax as syntax;

#[cfg(not(feature = "with-syntex"))]
#[macro_use]
extern crate syntax;

#[cfg(not(feature = "with-syntex"))]
extern crate rustc_plugin;

extern crate syn;
#[macro_use]
extern crate quote;

#[cfg(feature = "with-syntex")]
use std::path::Path;

#[cfg(not(feature = "with-syntex"))]
use syntax::feature_gate::AttributeType;

mod bound;
mod de;
mod ser;

#[cfg(feature = "with-syntex")]
fn syntex_registry() -> syntex::Registry {
    use syntax::{ast, fold};

    /// Strip the serde attributes from the crate.
    #[cfg(feature = "with-syntex")]
    fn strip_attributes(krate: ast::Crate) -> ast::Crate {
        /// Helper folder that strips the serde attributes after the extensions have been expanded.
        struct StripAttributeFolder;

        impl fold::Folder for StripAttributeFolder {
            fn fold_attribute(&mut self, attr: ast::Attribute) -> Option<ast::Attribute> {
                match attr.node.value.node {
                    ast::MetaItemKind::List(ref n, _) if n == &"serde" => { return None; }
                    _ => {}
                }

                Some(attr)
            }

            fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
                fold::noop_fold_mac(mac, self)
            }
        }

        fold::Folder::fold_crate(&mut StripAttributeFolder, krate)
    }

    let mut reg = syntex::Registry::new();

    reg.add_attr("feature(custom_derive)");
    reg.add_attr("feature(custom_attribute)");

    reg.add_decorator("derive_Serialize", shim::expand_derive_serialize);
    reg.add_decorator("derive_Deserialize", shim::expand_derive_deserialize);

    reg.add_post_expansion_pass(strip_attributes);

    reg
}

#[cfg(feature = "with-syntex")]
pub fn expand_str(src: &str) -> Result<String, syntex::Error> {
    let src = src.to_owned();

    let expand_thread = move || {
        syntex_registry().expand_str("", "", &src)
    };

    syntex::with_extra_stack(expand_thread)
}

#[cfg(feature = "with-syntex")]
pub fn expand<S, D>(src: S, dst: D) -> Result<(), syntex::Error>
    where S: AsRef<Path>,
          D: AsRef<Path>,
{
    let src = src.as_ref().to_owned();
    let dst = dst.as_ref().to_owned();

    let expand_thread = move || {
        syntex_registry().expand("", src, dst)
    };

    syntex::with_extra_stack(expand_thread)
}

#[cfg(not(feature = "with-syntex"))]
pub fn register(reg: &mut rustc_plugin::Registry) {
    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Serialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(shim::expand_derive_serialize)));

    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Deserialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(shim::expand_derive_deserialize)));

    reg.register_attribute("serde".to_owned(), AttributeType::Normal);
}

macro_rules! shim {
    ($name:ident $pkg:ident :: $func:ident) => {
        pub fn $func(
            cx: &mut ::syntax::ext::base::ExtCtxt,
            span: ::syntax::codemap::Span,
            meta_item: &::syntax::ast::MetaItem,
            annotatable: &::syntax::ext::base::Annotatable,
            push: &mut FnMut(::syntax::ext::base::Annotatable)
        ) {
            let item = match *annotatable {
                ::syntax::ext::base::Annotatable::Item(ref item) => item,
                _ => {
                    cx.span_err(
                        meta_item.span,
                        concat!("`#[derive(",
                                stringify!($name),
                                ")]` may only be applied to structs and enums"));
                    return;
                }
            };

            use syntax::{attr, ast, visit};
            struct MarkSerdeAttributesUsed;
            impl visit::Visitor for MarkSerdeAttributesUsed {
                fn visit_attribute(&mut self, attr: &ast::Attribute) {
                    match attr.node.value.node {
                        ast::MetaItemKind::List(ref name, _) if name == "serde" => {
                            attr::mark_used(attr);
                        }
                        _ => {}
                    }
                }
            }
            visit::walk_item(&mut MarkSerdeAttributesUsed, item);

            use syntax::print::pprust;
            let s = pprust::item_to_string(item);

            use {syn, $pkg};
            let syn_item = syn::parse_macro_input(&s).unwrap();
            let expanded = match $pkg::$func(&syn_item) {
                Ok(expanded) => expanded.to_string(),
                Err(msg) => {
                    cx.span_err(span, &msg);
                    return;
                }
            };

            use syntax::parse;
            let name = stringify!($name).to_string();
            let sess = cx.parse_sess;
            let impl_item = parse::parse_item_from_source_str(name, expanded, sess);
            push(::syntax::ext::base::Annotatable::Item(impl_item.unwrap().unwrap()));
        }
    };
}

mod shim {
    shim!(Serialize ser::expand_derive_serialize);
    shim!(Deserialize de::expand_derive_deserialize);
}

#[cfg(feature = "with-syn")]
pub fn expand_derive_serialize(item: &str) -> Result<String, String> {
    let syn_item = syn::parse_macro_input(item).unwrap();
    ser::expand_derive_serialize(&syn_item).map(|derive| derive.to_string())
}

#[cfg(feature = "with-syn")]
pub fn expand_derive_deserialize(item: &str) -> Result<String, String> {
    let syn_item = syn::parse_macro_input(item).unwrap();
    de::expand_derive_deserialize(&syn_item).map(|derive| derive.to_string())
}
<commit_msg>Hide the serde_codegen functions from docs<commit_after>#![cfg_attr(feature = "clippy", plugin(clippy))]
#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", allow(too_many_arguments))]
#![cfg_attr(feature = "clippy", allow(used_underscore_binding))]
#![cfg_attr(not(feature = "with-syntex"), feature(rustc_private, plugin))]

// The `quote!` macro requires deep recursion.
#![recursion_limit = "192"]

extern crate serde_codegen_internals as internals;

#[cfg(feature = "with-syntex")]
extern crate syntex;

#[cfg(feature = "with-syntex")]
#[macro_use]
extern crate syntex_syntax as syntax;

#[cfg(not(feature = "with-syntex"))]
#[macro_use]
extern crate syntax;

#[cfg(not(feature = "with-syntex"))]
extern crate rustc_plugin;

extern crate syn;
#[macro_use]
extern crate quote;

#[cfg(feature = "with-syntex")]
use std::path::Path;

#[cfg(not(feature = "with-syntex"))]
use syntax::feature_gate::AttributeType;

mod bound;
mod de;
mod ser;

#[cfg(feature = "with-syntex")]
fn syntex_registry() -> syntex::Registry {
    use syntax::{ast, fold};

    /// Strip the serde attributes from the crate.
    #[cfg(feature = "with-syntex")]
    fn strip_attributes(krate: ast::Crate) -> ast::Crate {
        /// Helper folder that strips the serde attributes after the extensions have been expanded.
        struct StripAttributeFolder;

        impl fold::Folder for StripAttributeFolder {
            fn fold_attribute(&mut self, attr: ast::Attribute) -> Option<ast::Attribute> {
                match attr.node.value.node {
                    ast::MetaItemKind::List(ref n, _) if n == &"serde" => { return None; }
                    _ => {}
                }

                Some(attr)
            }

            fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac {
                fold::noop_fold_mac(mac, self)
            }
        }

        fold::Folder::fold_crate(&mut StripAttributeFolder, krate)
    }

    let mut reg = syntex::Registry::new();

    reg.add_attr("feature(custom_derive)");
    reg.add_attr("feature(custom_attribute)");

    reg.add_decorator("derive_Serialize", shim::expand_derive_serialize);
    reg.add_decorator("derive_Deserialize", shim::expand_derive_deserialize);

    reg.add_post_expansion_pass(strip_attributes);

    reg
}

#[cfg(feature = "with-syntex")]
pub fn expand_str(src: &str) -> Result<String, syntex::Error> {
    let src = src.to_owned();

    let expand_thread = move || {
        syntex_registry().expand_str("", "", &src)
    };

    syntex::with_extra_stack(expand_thread)
}

#[cfg(feature = "with-syntex")]
pub fn expand<S, D>(src: S, dst: D) -> Result<(), syntex::Error>
    where S: AsRef<Path>,
          D: AsRef<Path>,
{
    let src = src.as_ref().to_owned();
    let dst = dst.as_ref().to_owned();

    let expand_thread = move || {
        syntex_registry().expand("", src, dst)
    };

    syntex::with_extra_stack(expand_thread)
}

#[cfg(not(feature = "with-syntex"))]
pub fn register(reg: &mut rustc_plugin::Registry) {
    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Serialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(shim::expand_derive_serialize)));

    reg.register_syntax_extension(
        syntax::parse::token::intern("derive_Deserialize"),
        syntax::ext::base::MultiDecorator(
            Box::new(shim::expand_derive_deserialize)));

    reg.register_attribute("serde".to_owned(), AttributeType::Normal);
}

macro_rules! shim {
    ($name:ident $pkg:ident :: $func:ident) => {
        pub fn $func(
            cx: &mut ::syntax::ext::base::ExtCtxt,
            span: ::syntax::codemap::Span,
            meta_item: &::syntax::ast::MetaItem,
            annotatable: &::syntax::ext::base::Annotatable,
            push: &mut FnMut(::syntax::ext::base::Annotatable)
        ) {
            let item = match *annotatable {
                ::syntax::ext::base::Annotatable::Item(ref item) => item,
                _ => {
                    cx.span_err(
                        meta_item.span,
                        concat!("`#[derive(",
                                stringify!($name),
                                ")]` may only be applied to structs and enums"));
                    return;
                }
            };

            use syntax::{attr, ast, visit};
            struct MarkSerdeAttributesUsed;
            impl visit::Visitor for MarkSerdeAttributesUsed {
                fn visit_attribute(&mut self, attr: &ast::Attribute) {
                    match attr.node.value.node {
                        ast::MetaItemKind::List(ref name, _) if name == "serde" => {
                            attr::mark_used(attr);
                        }
                        _ => {}
                    }
                }
            }
            visit::walk_item(&mut MarkSerdeAttributesUsed, item);

            use syntax::print::pprust;
            let s = pprust::item_to_string(item);

            use {syn, $pkg};
            let syn_item = syn::parse_macro_input(&s).unwrap();
            let expanded = match $pkg::$func(&syn_item) {
                Ok(expanded) => expanded.to_string(),
                Err(msg) => {
                    cx.span_err(span, &msg);
                    return;
                }
            };

            use syntax::parse;
            let name = stringify!($name).to_string();
            let sess = cx.parse_sess;
            let impl_item = parse::parse_item_from_source_str(name, expanded, sess);
            push(::syntax::ext::base::Annotatable::Item(impl_item.unwrap().unwrap()));
        }
    };
}

mod shim {
    shim!(Serialize ser::expand_derive_serialize);
    shim!(Deserialize de::expand_derive_deserialize);
}

#[cfg(feature = "with-syn")]
#[doc(hidden)]
/// Not public API. Use the serde_derive crate.
pub fn expand_derive_serialize(item: &str) -> Result<String, String> {
    let syn_item = syn::parse_macro_input(item).unwrap();
    ser::expand_derive_serialize(&syn_item).map(|derive| derive.to_string())
}

#[cfg(feature = "with-syn")]
#[doc(hidden)]
/// Not public API. Use the serde_derive crate.
pub fn expand_derive_deserialize(item: &str) -> Result<String, String> {
    let syn_item = syn::parse_macro_input(item).unwrap();
    de::expand_derive_deserialize(&syn_item).map(|derive| derive.to_string())
}
<|endoftext|>


<commit_before>use cairo::Status;
use gdk_pixbuf::{PixbufLoader, PixbufLoaderExt};
use gio;
use glib::translate::*;
use std::cell::RefCell;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::rc::Rc;

use allowed_url::{AllowedUrl, Fragment};
use error::LoadingError;
use handle::LoadOptions;
use io;
use node::RsvgNode;
use properties::ComputedValues;
use surface_utils::shared_surface::SharedImageSurface;
use xml::XmlState;
use xml2_load::xml_state_load_from_possibly_compressed_stream;

/// A loaded SVG file and its derived data
///
/// This contains the tree of nodes (SVG elements), the mapping
/// of id to node, and the CSS styles defined for this SVG.
pub struct Svg {
    tree: RsvgNode,

    ids: HashMap<String, RsvgNode>,

    // These require interior mutability because we load the extern
    // resources all over the place.  Eventually we'll be able to do this
    // once, at loading time, and keep this immutable.
    externs: RefCell<Resources>,
    images: RefCell<Images>,

    // Once we do not need to load externs, we can drop this as well
    load_options: LoadOptions,
}

impl Svg {
    pub fn new(tree: RsvgNode, ids: HashMap<String, RsvgNode>, load_options: LoadOptions) -> Svg {
        let values = ComputedValues::default();
        tree.cascade(&values);

        Svg {
            tree,
            ids,
            externs: RefCell::new(Resources::new()),
            images: RefCell::new(Images::new()),
            load_options,
        }
    }

    pub fn load_from_stream(
        load_options: &LoadOptions,
        stream: &gio::InputStream,
        cancellable: Option<&gio::Cancellable>,
    ) -> Result<Svg, LoadingError> {
        let load_flags = load_options.flags;
        let mut xml = XmlState::new(load_options);

        xml_state_load_from_possibly_compressed_stream(&mut xml, load_flags, stream, cancellable)?;

        xml.steal_result()
    }

    pub fn root(&self) -> RsvgNode {
        self.tree.clone()
    }

    pub fn lookup(&self, fragment: &Fragment) -> Result<RsvgNode, LoadingError> {
        if fragment.uri().is_some() {
            self.externs
                .borrow_mut()
                .lookup(&self.load_options, fragment)
        } else {
            self.lookup_node_by_id(fragment.fragment())
                .ok_or(LoadingError::BadUrl)
        }
    }

    pub fn lookup_node_by_id(&self, id: &str) -> Option<RsvgNode> {
        self.ids.get(id).map(|n| (*n).clone())
    }

    pub fn lookup_image(&self, href: &str) -> Result<SharedImageSurface, LoadingError> {
        self.images.borrow_mut().lookup(&self.load_options, href)
    }
}

struct Resources {
    resources: HashMap<AllowedUrl, Result<Rc<Svg>, LoadingError>>,
}

impl Resources {
    pub fn new() -> Resources {
        Resources {
            resources: Default::default(),
        }
    }

    pub fn lookup(
        &mut self,
        load_options: &LoadOptions,
        fragment: &Fragment,
    ) -> Result<RsvgNode, LoadingError> {
        if let Some(ref href) = fragment.uri() {
            self.get_extern_svg(load_options, href).and_then(|svg| {
                svg.lookup_node_by_id(fragment.fragment())
                    .ok_or(LoadingError::BadUrl)
            })
        } else {
            unreachable!();
        }
    }

    fn get_extern_svg(
        &mut self,
        load_options: &LoadOptions,
        href: &str,
    ) -> Result<Rc<Svg>, LoadingError> {
        let aurl = AllowedUrl::from_href(href, load_options.base_url.as_ref())
            .map_err(|_| LoadingError::BadUrl)?;

        match self.resources.entry(aurl) {
            Entry::Occupied(e) => e.get().clone(),
            Entry::Vacant(e) => {
                let svg = load_svg(load_options, e.key()).map(|s| Rc::new(s));
                let res = e.insert(svg);
                res.clone()
            }
        }
    }
}

struct Images {
    images: HashMap<AllowedUrl, Result<SharedImageSurface, LoadingError>>,
}

impl Images {
    pub fn new() -> Images {
        Images {
            images: Default::default(),
        }
    }

    pub fn lookup(
        &mut self,
        load_options: &LoadOptions,
        href: &str,
    ) -> Result<SharedImageSurface, LoadingError> {
        let aurl = AllowedUrl::from_href(href, load_options.base_url.as_ref())
            .map_err(|_| LoadingError::BadUrl)?;

        match self.images.entry(aurl) {
            Entry::Occupied(e) => e.get().clone(),
            Entry::Vacant(e) => {
                let surface = load_image(load_options, e.key());
                let res = e.insert(surface);
                res.clone()
            }
        }
    }
}

fn load_svg(load_options: &LoadOptions, aurl: &AllowedUrl) -> Result<Svg, LoadingError> {
    // FIXME: pass a cancellable to these
    io::acquire_stream(aurl, None).and_then(|stream| {
        Svg::load_from_stream(&load_options.copy_with_base_url(aurl), &stream, None)
    })
}

fn load_image(
    load_options: &LoadOptions,
    aurl: &AllowedUrl,
) -> Result<SharedImageSurface, LoadingError> {
    let data = io::acquire_data(&aurl, None)?;

    if data.data.len() == 0 {
        return Err(LoadingError::EmptyData);
    }

    let loader = if let Some(ref content_type) = data.content_type {
        PixbufLoader::new_with_mime_type(content_type)?
    } else {
        PixbufLoader::new()
    };

    loader.write(&data.data)?;
    loader.close()?;

    let pixbuf = loader.get_pixbuf().ok_or(LoadingError::Unknown)?;

    let surface = SharedImageSurface::from_pixbuf(&pixbuf)?;

    if load_options.flags.keep_image_data {
        if let Some(mime_type) = data.content_type {
            extern "C" {
                fn cairo_surface_set_mime_data(
                    surface: *mut cairo_sys::cairo_surface_t,
                    mime_type: *const libc::c_char,
                    data: *mut libc::c_char,
                    length: libc::c_ulong,
                    destroy: cairo_sys::cairo_destroy_func_t,
                    closure: *mut libc::c_void,
                ) -> Status;
            }

            let data_ptr = ToGlibContainerFromSlice::to_glib_full_from_slice(&data.data);

            unsafe {
                let status = cairo_surface_set_mime_data(
                    surface.to_glib_none().0,
                    mime_type.to_glib_none().0,
                    data_ptr as *mut _,
                    data.data.len() as libc::c_ulong,
                    Some(glib_sys::g_free),
                    data_ptr as *mut _,
                );

                if status != Status::Success {
                    return Err(LoadingError::Cairo(status));
                }
            }
        }
    }

    Ok(surface)
}
<commit_msg>Adjust for changes in cairo::Status vs. cairo_sys::cairo_status_t<commit_after>use gdk_pixbuf::{PixbufLoader, PixbufLoaderExt};
use gio;
use glib::translate::*;
use std::cell::RefCell;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::rc::Rc;

use allowed_url::{AllowedUrl, Fragment};
use error::LoadingError;
use handle::LoadOptions;
use io;
use node::RsvgNode;
use properties::ComputedValues;
use surface_utils::shared_surface::SharedImageSurface;
use xml::XmlState;
use xml2_load::xml_state_load_from_possibly_compressed_stream;

/// A loaded SVG file and its derived data
///
/// This contains the tree of nodes (SVG elements), the mapping
/// of id to node, and the CSS styles defined for this SVG.
pub struct Svg {
    tree: RsvgNode,

    ids: HashMap<String, RsvgNode>,

    // These require interior mutability because we load the extern
    // resources all over the place.  Eventually we'll be able to do this
    // once, at loading time, and keep this immutable.
    externs: RefCell<Resources>,
    images: RefCell<Images>,

    // Once we do not need to load externs, we can drop this as well
    load_options: LoadOptions,
}

impl Svg {
    pub fn new(tree: RsvgNode, ids: HashMap<String, RsvgNode>, load_options: LoadOptions) -> Svg {
        let values = ComputedValues::default();
        tree.cascade(&values);

        Svg {
            tree,
            ids,
            externs: RefCell::new(Resources::new()),
            images: RefCell::new(Images::new()),
            load_options,
        }
    }

    pub fn load_from_stream(
        load_options: &LoadOptions,
        stream: &gio::InputStream,
        cancellable: Option<&gio::Cancellable>,
    ) -> Result<Svg, LoadingError> {
        let load_flags = load_options.flags;
        let mut xml = XmlState::new(load_options);

        xml_state_load_from_possibly_compressed_stream(&mut xml, load_flags, stream, cancellable)?;

        xml.steal_result()
    }

    pub fn root(&self) -> RsvgNode {
        self.tree.clone()
    }

    pub fn lookup(&self, fragment: &Fragment) -> Result<RsvgNode, LoadingError> {
        if fragment.uri().is_some() {
            self.externs
                .borrow_mut()
                .lookup(&self.load_options, fragment)
        } else {
            self.lookup_node_by_id(fragment.fragment())
                .ok_or(LoadingError::BadUrl)
        }
    }

    pub fn lookup_node_by_id(&self, id: &str) -> Option<RsvgNode> {
        self.ids.get(id).map(|n| (*n).clone())
    }

    pub fn lookup_image(&self, href: &str) -> Result<SharedImageSurface, LoadingError> {
        self.images.borrow_mut().lookup(&self.load_options, href)
    }
}

struct Resources {
    resources: HashMap<AllowedUrl, Result<Rc<Svg>, LoadingError>>,
}

impl Resources {
    pub fn new() -> Resources {
        Resources {
            resources: Default::default(),
        }
    }

    pub fn lookup(
        &mut self,
        load_options: &LoadOptions,
        fragment: &Fragment,
    ) -> Result<RsvgNode, LoadingError> {
        if let Some(ref href) = fragment.uri() {
            self.get_extern_svg(load_options, href).and_then(|svg| {
                svg.lookup_node_by_id(fragment.fragment())
                    .ok_or(LoadingError::BadUrl)
            })
        } else {
            unreachable!();
        }
    }

    fn get_extern_svg(
        &mut self,
        load_options: &LoadOptions,
        href: &str,
    ) -> Result<Rc<Svg>, LoadingError> {
        let aurl = AllowedUrl::from_href(href, load_options.base_url.as_ref())
            .map_err(|_| LoadingError::BadUrl)?;

        match self.resources.entry(aurl) {
            Entry::Occupied(e) => e.get().clone(),
            Entry::Vacant(e) => {
                let svg = load_svg(load_options, e.key()).map(|s| Rc::new(s));
                let res = e.insert(svg);
                res.clone()
            }
        }
    }
}

struct Images {
    images: HashMap<AllowedUrl, Result<SharedImageSurface, LoadingError>>,
}

impl Images {
    pub fn new() -> Images {
        Images {
            images: Default::default(),
        }
    }

    pub fn lookup(
        &mut self,
        load_options: &LoadOptions,
        href: &str,
    ) -> Result<SharedImageSurface, LoadingError> {
        let aurl = AllowedUrl::from_href(href, load_options.base_url.as_ref())
            .map_err(|_| LoadingError::BadUrl)?;

        match self.images.entry(aurl) {
            Entry::Occupied(e) => e.get().clone(),
            Entry::Vacant(e) => {
                let surface = load_image(load_options, e.key());
                let res = e.insert(surface);
                res.clone()
            }
        }
    }
}

fn load_svg(load_options: &LoadOptions, aurl: &AllowedUrl) -> Result<Svg, LoadingError> {
    // FIXME: pass a cancellable to these
    io::acquire_stream(aurl, None).and_then(|stream| {
        Svg::load_from_stream(&load_options.copy_with_base_url(aurl), &stream, None)
    })
}

fn load_image(
    load_options: &LoadOptions,
    aurl: &AllowedUrl,
) -> Result<SharedImageSurface, LoadingError> {
    let data = io::acquire_data(&aurl, None)?;

    if data.data.len() == 0 {
        return Err(LoadingError::EmptyData);
    }

    let loader = if let Some(ref content_type) = data.content_type {
        PixbufLoader::new_with_mime_type(content_type)?
    } else {
        PixbufLoader::new()
    };

    loader.write(&data.data)?;
    loader.close()?;

    let pixbuf = loader.get_pixbuf().ok_or(LoadingError::Unknown)?;

    let surface = SharedImageSurface::from_pixbuf(&pixbuf)?;

    if load_options.flags.keep_image_data {
        if let Some(mime_type) = data.content_type {
            extern "C" {
                fn cairo_surface_set_mime_data(
                    surface: *mut cairo_sys::cairo_surface_t,
                    mime_type: *const libc::c_char,
                    data: *mut libc::c_char,
                    length: libc::c_ulong,
                    destroy: cairo_sys::cairo_destroy_func_t,
                    closure: *mut libc::c_void,
                ) -> cairo_sys::cairo_status_t;
            }

            let data_ptr = ToGlibContainerFromSlice::to_glib_full_from_slice(&data.data);

            unsafe {
                let status = cairo_surface_set_mime_data(
                    surface.to_glib_none().0,
                    mime_type.to_glib_none().0,
                    data_ptr as *mut _,
                    data.data.len() as libc::c_ulong,
                    Some(glib_sys::g_free),
                    data_ptr as *mut _,
                );

                if status != cairo_sys::STATUS_SUCCESS {
                    return Err(LoadingError::Cairo(status.into()));
                }
            }
        }
    }

    Ok(surface)
}
<|endoftext|>