<commit_before>#[cfg(test)]
mod tests;
pub mod response;

pub use message::response::*;
use notification::Notification;
use std::collections::HashMap;
use std::collections::BTreeMap;
use std::io::prelude::*;
use hyper::Client;
use hyper::header::{Authorization, ContentType};
use hyper::status::StatusCode;
use rustc_serialize::json::{self, Json, ToJson};

#[derive(PartialEq, Debug, Clone)]
pub enum Priority {
    Normal, High
}

/// Represents a FCM message. Construct the FCM message
/// using various utility methods and finally send it.
/// # Examples:
/// ```rust
/// use fcm::Message;
///
/// let message = Message::new("<registration id>").dry_run(true);
/// ```
pub struct Message<'a> {
    to: &'a str,
    registration_ids: Option<Vec<String>>,
    collapse_key: Option<&'a str>,
    priority: Option<Priority>,
    content_available: Option<bool>,
    delay_while_idle: Option<bool>,
    time_to_live: Option<i32>,
    restricted_package_name: Option<&'a str>,
    dry_run: Option<bool>,
    data: Option<HashMap<String, String>>,
    notification: Option<Notification<'a>>,
}

impl <'a> ToJson for Message<'a> {
    fn to_json(&self) -> Json {
        let mut root = BTreeMap::new();

        root.insert("to".to_string(), self.to.to_json());

        if self.registration_ids.is_some() {
            root.insert("registration_ids".to_string(),
                    self.registration_ids.clone().unwrap().to_json());
        }

        if self.collapse_key.is_some() {
            root.insert("collapse_key".to_string(), self.collapse_key.clone().unwrap().to_json());
        }

        if self.priority.is_some() {
            root.insert("priority".to_string(), match self.priority.clone().unwrap() {
                Priority::Normal => Json::String("normal".to_string()),
                Priority::High => Json::String("high".to_string()),
            });
        }

        if self.content_available.is_some() {
            root.insert("content_available".to_string(), self.content_available.unwrap().to_json());
        }

        if self.delay_while_idle.is_some() {
            root.insert("delay_while_idle".to_string(), self.delay_while_idle.unwrap().to_json());
        }

        if self.time_to_live.is_some() {
            root.insert("time_to_live".to_string(), self.time_to_live.unwrap().to_json());
        }

        if self.restricted_package_name.is_some() {
            root.insert("restricted_package_name".to_string(), self.restricted_package_name.clone().unwrap().to_json());
        }

        if self.dry_run.is_some() {
            root.insert("dry_run".to_string(), self.dry_run.unwrap().to_json());
        }

        if self.data.is_some() {
            root.insert("data".to_string(), self.data.clone().unwrap().to_json());
        }

        if self.notification.is_some() {
            root.insert("notification".to_string(), self.notification.clone().unwrap().to_json());
        }

        Json::Object(root)
    }
}

impl <'a> Message<'a> {
    /// Get a new instance of Message. You need to supply either
    /// a registration id, or a topic (/topic/...).
    pub fn new(to: &'a str) -> Message {
        Message {
            to: to,
            registration_ids: None,
            collapse_key: None,
            priority: None,
            content_available: None,
            delay_while_idle: None,
            time_to_live: None,
            restricted_package_name: None,
            dry_run: None,
            data: None,
            notification: None,
        }
    }

    /// Set various registration ids to which the message ought to be sent.
    pub fn registration_ids(mut self, ids: Vec<&'a str>) -> Message<'a> {
        self.registration_ids = Some(ids.iter().map(|s| s.to_string()).collect());
        self
    }

    /// Set this parameter to identify groups of messages that can be collapsed.
    pub fn collapse_key(mut self, collapse_key: &'a str) -> Message<'a> {
        self.collapse_key = Some(collapse_key);
        self
    }

    /// Set the priority of the message. You can set Normal or High priorities.
    /// # Examples:
    /// ```rust
    /// use fcm::{Message, Priority};
    ///
    /// let message = Message::new("<registration id>")
    ///     .priority(Priority::High);
    /// ```
    pub fn priority(mut self, priority: Priority) -> Message<'a> {
        self.priority = Some(priority);
        self
    }

    /// To set the `content-available` field on iOS
    pub fn content_available(mut self, content_available: bool) -> Message<'a> {
        self.content_available = Some(content_available);
        self
    }

    /// When set to `true`, sends the message only when the device is active.
    pub fn delay_while_idle(mut self, delay_while_idle: bool) -> Message<'a> {
        self.delay_while_idle = Some(delay_while_idle);
        self
    }

    /// How long (in seconds) to keep the message on FCM servers in case the device
    /// is offline. The maximum and default is 4 weeks.
    pub fn time_to_live(mut self, time_to_live: i32) -> Message<'a> {
        self.time_to_live = Some(time_to_live);
        self
    }

    /// Package name of the application where the registration tokens must match.
    pub fn restricted_package_name(mut self, restricted_package_name: &'a str) -> Message<'a> {
        self.restricted_package_name = Some(restricted_package_name);
        self
    }

    /// When set to `true`, allows you to test FCM without actually sending the message.
    pub fn dry_run(mut self, dry_run: bool) -> Message<'a> {
        self.dry_run = Some(dry_run);
        self
    }

    /// Use this to add custom key-value pairs to the message. This data
    /// must be handled appropriately on the client end.
    /// # Examples:
    /// ```rust
    /// use fcm::Message;
    /// use std::collections::HashMap;
    ///
    /// let mut map = HashMap::new();
    /// map.insert("message", "Howdy!");
    ///
    /// let message = Message::new("<registration id>").data(map);
    /// ```
    pub fn data(mut self, data: HashMap<&'a str, &'a str>) -> Message<'a> {
        let mut datamap: HashMap<String, String> = HashMap::new();
        for (key, val) in data.iter() {
            datamap.insert(key.to_string(), val.to_string());
        }

        self.data = Some(datamap);
        self
    }

    /// Use this to set a `Notification` for the message.
    /// # Examples:
    /// ```rust
    /// use fcm::{Message, NotificationBuilder};
    ///
    /// let notification = NotificationBuilder::new("Hey!")
    ///     .body("Do you want to catch up later?")
    ///     .finalize();
    ///
    /// let message = Message::new("<registration id>")
    ///     .notification(notification);
    /// ```
    pub fn notification(mut self, notification: Notification<'a>) -> Message<'a> {
        self.notification = Some(notification);
        self
    }

    /// Send the message using your FCM API Key.
    /// # Examples:
    /// ```no_run
    /// use fcm::Message;
    /// use std::collections::HashMap;
    ///
    /// let mut map = HashMap::new();
    /// map.insert("message", "Howdy!");
    ///
    /// let result = Message::new("<registration id>")
    ///     .data(map)
    ///     .send("<FCM API Key>");
    /// ```
    pub fn send(self, api_key: &'a str) -> Result<response::FcmResponse, response::FcmError> {
        let payload = self.to_json().to_string();
        let auth_header = "key=".to_string() + api_key;
        let client = Client::new();

        let response = client
            .post("https://fcm.googleapis.com/fcm/send")
            .body(&payload)
            .header(Authorization(auth_header))
            .header(ContentType::json())
            .send();

        match response {
            Ok(mut response) => {
                let mut body = String::new();
                response.read_to_string(&mut body).unwrap();

                Message::parse_response(response.status, &body)
            },
            Err(_) => {
                Message::parse_response(StatusCode::InternalServerError, "Server Error")
            }
        }
    }

    fn parse_response(status: StatusCode, body: &str) -> Result<response::FcmResponse, response::FcmError> {
        use hyper::status::StatusCode::*;
        use hyper::status::StatusClass::ServerError;
        match status {
            Ok =>
                Result::Ok(json::decode(body).unwrap()),
            Unauthorized =>
                Err(response::FcmError::Unauthorized),
            BadRequest =>
                Err(response::FcmError::InvalidMessage(body.to_string())),
            status if status.is_server_error() =>
                Err(response::FcmError::ServerError),
            _ =>
                Err(response::FcmError::InvalidMessage("Unknown Error".to_string())),
        }
    }
}
<commit_msg>Remove unused import<commit_after>#[cfg(test)]
mod tests;
pub mod response;

pub use message::response::*;
use notification::Notification;
use std::collections::HashMap;
use std::collections::BTreeMap;
use std::io::prelude::*;
use hyper::Client;
use hyper::header::{Authorization, ContentType};
use hyper::status::StatusCode;
use rustc_serialize::json::{self, Json, ToJson};

#[derive(PartialEq, Debug, Clone)]
pub enum Priority {
    Normal, High
}

/// Represents a FCM message. Construct the FCM message
/// using various utility methods and finally send it.
/// # Examples:
/// ```rust
/// use fcm::Message;
///
/// let message = Message::new("<registration id>").dry_run(true);
/// ```
pub struct Message<'a> {
    to: &'a str,
    registration_ids: Option<Vec<String>>,
    collapse_key: Option<&'a str>,
    priority: Option<Priority>,
    content_available: Option<bool>,
    delay_while_idle: Option<bool>,
    time_to_live: Option<i32>,
    restricted_package_name: Option<&'a str>,
    dry_run: Option<bool>,
    data: Option<HashMap<String, String>>,
    notification: Option<Notification<'a>>,
}

impl <'a> ToJson for Message<'a> {
    fn to_json(&self) -> Json {
        let mut root = BTreeMap::new();

        root.insert("to".to_string(), self.to.to_json());

        if self.registration_ids.is_some() {
            root.insert("registration_ids".to_string(),
                    self.registration_ids.clone().unwrap().to_json());
        }

        if self.collapse_key.is_some() {
            root.insert("collapse_key".to_string(), self.collapse_key.clone().unwrap().to_json());
        }

        if self.priority.is_some() {
            root.insert("priority".to_string(), match self.priority.clone().unwrap() {
                Priority::Normal => Json::String("normal".to_string()),
                Priority::High => Json::String("high".to_string()),
            });
        }

        if self.content_available.is_some() {
            root.insert("content_available".to_string(), self.content_available.unwrap().to_json());
        }

        if self.delay_while_idle.is_some() {
            root.insert("delay_while_idle".to_string(), self.delay_while_idle.unwrap().to_json());
        }

        if self.time_to_live.is_some() {
            root.insert("time_to_live".to_string(), self.time_to_live.unwrap().to_json());
        }

        if self.restricted_package_name.is_some() {
            root.insert("restricted_package_name".to_string(), self.restricted_package_name.clone().unwrap().to_json());
        }

        if self.dry_run.is_some() {
            root.insert("dry_run".to_string(), self.dry_run.unwrap().to_json());
        }

        if self.data.is_some() {
            root.insert("data".to_string(), self.data.clone().unwrap().to_json());
        }

        if self.notification.is_some() {
            root.insert("notification".to_string(), self.notification.clone().unwrap().to_json());
        }

        Json::Object(root)
    }
}

impl <'a> Message<'a> {
    /// Get a new instance of Message. You need to supply either
    /// a registration id, or a topic (/topic/...).
    pub fn new(to: &'a str) -> Message {
        Message {
            to: to,
            registration_ids: None,
            collapse_key: None,
            priority: None,
            content_available: None,
            delay_while_idle: None,
            time_to_live: None,
            restricted_package_name: None,
            dry_run: None,
            data: None,
            notification: None,
        }
    }

    /// Set various registration ids to which the message ought to be sent.
    pub fn registration_ids(mut self, ids: Vec<&'a str>) -> Message<'a> {
        self.registration_ids = Some(ids.iter().map(|s| s.to_string()).collect());
        self
    }

    /// Set this parameter to identify groups of messages that can be collapsed.
    pub fn collapse_key(mut self, collapse_key: &'a str) -> Message<'a> {
        self.collapse_key = Some(collapse_key);
        self
    }

    /// Set the priority of the message. You can set Normal or High priorities.
    /// # Examples:
    /// ```rust
    /// use fcm::{Message, Priority};
    ///
    /// let message = Message::new("<registration id>")
    ///     .priority(Priority::High);
    /// ```
    pub fn priority(mut self, priority: Priority) -> Message<'a> {
        self.priority = Some(priority);
        self
    }

    /// To set the `content-available` field on iOS
    pub fn content_available(mut self, content_available: bool) -> Message<'a> {
        self.content_available = Some(content_available);
        self
    }

    /// When set to `true`, sends the message only when the device is active.
    pub fn delay_while_idle(mut self, delay_while_idle: bool) -> Message<'a> {
        self.delay_while_idle = Some(delay_while_idle);
        self
    }

    /// How long (in seconds) to keep the message on FCM servers in case the device
    /// is offline. The maximum and default is 4 weeks.
    pub fn time_to_live(mut self, time_to_live: i32) -> Message<'a> {
        self.time_to_live = Some(time_to_live);
        self
    }

    /// Package name of the application where the registration tokens must match.
    pub fn restricted_package_name(mut self, restricted_package_name: &'a str) -> Message<'a> {
        self.restricted_package_name = Some(restricted_package_name);
        self
    }

    /// When set to `true`, allows you to test FCM without actually sending the message.
    pub fn dry_run(mut self, dry_run: bool) -> Message<'a> {
        self.dry_run = Some(dry_run);
        self
    }

    /// Use this to add custom key-value pairs to the message. This data
    /// must be handled appropriately on the client end.
    /// # Examples:
    /// ```rust
    /// use fcm::Message;
    /// use std::collections::HashMap;
    ///
    /// let mut map = HashMap::new();
    /// map.insert("message", "Howdy!");
    ///
    /// let message = Message::new("<registration id>").data(map);
    /// ```
    pub fn data(mut self, data: HashMap<&'a str, &'a str>) -> Message<'a> {
        let mut datamap: HashMap<String, String> = HashMap::new();
        for (key, val) in data.iter() {
            datamap.insert(key.to_string(), val.to_string());
        }

        self.data = Some(datamap);
        self
    }

    /// Use this to set a `Notification` for the message.
    /// # Examples:
    /// ```rust
    /// use fcm::{Message, NotificationBuilder};
    ///
    /// let notification = NotificationBuilder::new("Hey!")
    ///     .body("Do you want to catch up later?")
    ///     .finalize();
    ///
    /// let message = Message::new("<registration id>")
    ///     .notification(notification);
    /// ```
    pub fn notification(mut self, notification: Notification<'a>) -> Message<'a> {
        self.notification = Some(notification);
        self
    }

    /// Send the message using your FCM API Key.
    /// # Examples:
    /// ```no_run
    /// use fcm::Message;
    /// use std::collections::HashMap;
    ///
    /// let mut map = HashMap::new();
    /// map.insert("message", "Howdy!");
    ///
    /// let result = Message::new("<registration id>")
    ///     .data(map)
    ///     .send("<FCM API Key>");
    /// ```
    pub fn send(self, api_key: &'a str) -> Result<response::FcmResponse, response::FcmError> {
        let payload = self.to_json().to_string();
        let auth_header = "key=".to_string() + api_key;
        let client = Client::new();

        let response = client
            .post("https://fcm.googleapis.com/fcm/send")
            .body(&payload)
            .header(Authorization(auth_header))
            .header(ContentType::json())
            .send();

        match response {
            Ok(mut response) => {
                let mut body = String::new();
                response.read_to_string(&mut body).unwrap();

                Message::parse_response(response.status, &body)
            },
            Err(_) => {
                Message::parse_response(StatusCode::InternalServerError, "Server Error")
            }
        }
    }

    fn parse_response(status: StatusCode, body: &str) -> Result<response::FcmResponse, response::FcmError> {
        use hyper::status::StatusCode::*;
        match status {
            Ok =>
                Result::Ok(json::decode(body).unwrap()),
            Unauthorized =>
                Err(response::FcmError::Unauthorized),
            BadRequest =>
                Err(response::FcmError::InvalidMessage(body.to_string())),
            status if status.is_server_error() =>
                Err(response::FcmError::ServerError),
            _ =>
                Err(response::FcmError::InvalidMessage("Unknown Error".to_string())),
        }
    }
}
<|endoftext|>


<commit_before>use std::cmp::min;
use std::net::{SocketAddr,IpAddr};
use std::io::Write;
use mio::*;
use mio::tcp::TcpStream;
use pool::{Pool,Checkout,Reset};
use time::{Duration, precise_time_s, precise_time_ns};
use uuid::Uuid;
use network::{ClientResult,Protocol};
use network::buffer_queue::BufferQueue;
use network::proxy::Readiness;
use network::socket::{SocketHandler,SocketResult};

type BackendToken = Token;

#[derive(PartialEq)]
pub enum ClientStatus {
  Normal,
  DefaultAnswer,
}

pub struct Pipe<Front:SocketHandler> {
  pub frontend:       Front,
  backend:            TcpStream,
  token:              Option<Token>,
  backend_token:      Option<Token>,
  pub front_buf:      Checkout<BufferQueue>,
  back_buf:           Checkout<BufferQueue>,
  front_buf_position: usize,
  back_buf_position:  usize,
  pub app_id:         Option<String>,
  pub request_id:     String,
  pub server_context: String,
  pub readiness:      Readiness,
  pub log_ctx:        String,
  public_address:     Option<IpAddr>,
}

impl<Front:SocketHandler> Pipe<Front> {
  pub fn new(server_context: &str, frontend: Front, backend: TcpStream, front_buf: Checkout<BufferQueue>, back_buf: Checkout<BufferQueue>, public_address: Option<IpAddr>) -> Option<Pipe<Front>> {
    let request_id = Uuid::new_v4().hyphenated().to_string();
    let log_ctx    = format!("{}\t{}\tunknown\t", server_context, &request_id);
    let client = Pipe {
      frontend:           frontend,
      backend:            backend,
      token:              None,
      backend_token:      None,
      front_buf:          front_buf,
      back_buf:           back_buf,
      front_buf_position: 0,
      back_buf_position:  0,
      app_id:             None,
      request_id:         request_id,
      server_context:     String::from(server_context),
      readiness:          Readiness {
                            front_interest:  Ready::readable() | Ready::hup() | Ready::error(),
                            back_interest:   Ready::readable() | Ready::hup() | Ready::error(),
                            front_readiness: Ready::none(),
                            back_readiness:  Ready::none(),
      },
      log_ctx:            log_ctx,
      public_address:     public_address,
    };

    info!("created pipe");
    Some(client)
  }

  fn tokens(&self) -> Option<(Token,Token)> {
    if let Some(front) = self.token {
      if let Some(back) = self.backend_token {
        return Some((front, back))
      }
    }
    None
  }

  pub fn front_socket(&self) -> &TcpStream {
    self.frontend.socket_ref()
  }

  pub fn back_socket(&self)  -> Option<&TcpStream> {
    Some(&self.backend)
  }

  pub fn front_token(&self)  -> Option<Token> {
    self.token
  }

  pub fn back_token(&self)   -> Option<Token> {
    self.backend_token
  }

  pub fn close(&mut self) {
  }

  pub fn log_context(&self) -> String {
    if let Some(ref app_id) = self.app_id {
      format!("{}\t{}\t{}\t", self.server_context, self.request_id, app_id)
    } else {
      format!("{}\t{}\tunknown\t", self.server_context, self.request_id)
    }
  }

  pub fn set_front_token(&mut self, token: Token) {
    self.token         = Some(token);
  }

  pub fn set_back_token(&mut self, token: Token) {
    self.backend_token = Some(token);
  }

  pub fn readiness(&mut self) -> &mut Readiness {
    &mut self.readiness
  }

  pub fn front_hup(&mut self) -> ClientResult {
    if self.backend_token == None {
      ClientResult::CloseClient
    } else {
      ClientResult::Continue
    }
  }

  pub fn back_hup(&mut self) -> ClientResult {
    if self.token == None {
      ClientResult::CloseClient
    } else {
      ClientResult::Continue
    }
  }

  // Read content from the client
  pub fn readable(&mut self) -> ClientResult {
    info!("pipe readable");
    if self.front_buf.buffer.available_space() == 0 {
      self.readiness.front_interest.remove(Ready::readable());
      self.readiness.back_interest.insert(Ready::writable());
      return ClientResult::Continue;
    }

    let (sz, res) = self.frontend.socket_read(self.front_buf.buffer.space());
    debug!("{}\tFRONT [{:?}]: read {} bytes", self.log_ctx, self.token, sz);

    if sz > 0 {
      self.front_buf.buffer.fill(sz);
      self.front_buf.sliced_input(sz);
      self.front_buf.consume_parsed_data(sz);

      if self.front_buf.buffer.available_space() == 0 {
        self.readiness.front_interest.remove(Ready::readable());
        self.readiness.back_interest.insert(Ready::writable());
      }
    } else {
      self.readiness.front_readiness.remove(Ready::readable());
    }

    match res {
      SocketResult::Error => {
        error!("{}\t[{:?}] front socket error, closing the connection", self.log_ctx, self.token);
        self.readiness.reset();
        return ClientResult::CloseClient;
      },
      SocketResult::WouldBlock => {
        self.readiness.front_readiness.remove(Ready::readable());
      },
      SocketResult::Continue => {}
    };

    self.readiness.back_interest.insert(Ready::writable());
    ClientResult::Continue
  }

  // Forward content to client
  pub fn writable(&mut self) -> ClientResult {
    info!("pipe writable");
    if self.back_buf.output_data_size() == 0 || self.back_buf.next_output_data().len() == 0 {
      self.readiness.back_interest.insert(Ready::readable());
      self.readiness.front_interest.remove(Ready::writable());
      return ClientResult::Continue;
    }

    let mut sz = 0usize;
    let mut res = SocketResult::Continue;
    while res == SocketResult::Continue && self.back_buf.output_data_size() > 0 {
      // no more data in buffer, stop here
      if self.back_buf.next_output_data().len() == 0 {
        self.readiness.back_interest.insert(Ready::readable());
        self.readiness.front_interest.remove(Ready::writable());
        return ClientResult::Continue;
      }
      let (current_sz, current_res) = self.frontend.socket_write(self.back_buf.next_output_data());
      res = current_res;
      self.back_buf.consume_output_data(current_sz);
      self.back_buf_position += current_sz;
      sz += current_sz;
    }

    if sz > 0 {
      self.readiness.back_interest.insert(Ready::readable());
    }

    if let Some((front,back)) = self.tokens() {
      debug!("{}\tFRONT [{}<-{}]: wrote {} bytes of {}, buffer position {} restart position {}",
        self.log_ctx, front.0, back.0, sz, self.back_buf.output_data_size(),
        self.back_buf.buffer_position, self.back_buf.start_parsing_position);
    }

    match res {
      SocketResult::Error => {
        error!("{}\t[{:?}] error writing to front socket, closing", self.log_ctx, self.token);
        self.readiness.reset();
        return ClientResult::CloseClient;
      },
      SocketResult::WouldBlock => {
        self.readiness.front_readiness.remove(Ready::writable());
      },
      SocketResult::Continue => {},
    }

    ClientResult::Continue
  }

  // Forward content to application
  pub fn back_writable(&mut self) -> ClientResult {
    info!("pipe back_writable");
    if self.front_buf.output_data_size() == 0 || self.front_buf.next_output_data().len() == 0 {
      self.readiness.front_interest.insert(Ready::readable());
      self.readiness.back_interest.remove(Ready::writable());
      return ClientResult::Continue;
    }

    let tokens = self.tokens().clone();
    let output_size = self.front_buf.output_data_size();

    let mut sz = 0usize;
    let mut socket_res = SocketResult::Continue;

    while socket_res == SocketResult::Continue && self.front_buf.output_data_size() > 0 {
      // no more data in buffer, stop here
      if self.front_buf.next_output_data().len() == 0 {
        self.readiness.front_interest.insert(Ready::readable());
        self.readiness.back_interest.remove(Ready::writable());
        return ClientResult::Continue;
      }
      let (current_sz, current_res) = self.backend.socket_write(self.front_buf.next_output_data());
      socket_res = current_res;
      self.front_buf.consume_output_data(current_sz);
      self.front_buf_position += current_sz;
      sz += current_sz;
    }

    if let Some((front,back)) = tokens {
      debug!("{}\tBACK [{}->{}]: wrote {} bytes of {}", self.log_ctx, front.0, back.0, sz, output_size);
    }
    match socket_res {
      SocketResult::Error => {
        error!("{}\tback socket write error, closing connection", self.log_ctx);
        self.readiness.reset();
        return ClientResult::CloseBothFailure;
      },
      SocketResult::WouldBlock => {
        self.readiness.back_readiness.remove(Ready::writable());

      },
      SocketResult::Continue => {}
    }
    ClientResult::Continue
  }

  // Read content from application
  pub fn back_readable(&mut self) -> ClientResult {
    info!("pipe back_readable");
    if self.back_buf.buffer.available_space() == 0 {
      self.readiness.back_interest.remove(Ready::readable());
      return ClientResult::Continue;
    }

    let tokens     = self.tokens().clone();
    let (sz, r) = self.backend.socket_read(&mut self.back_buf.buffer.space());
    self.back_buf.buffer.fill(sz);
    self.back_buf.sliced_input(sz);
    self.back_buf.consume_parsed_data(sz);

    if let Some((front,back)) = tokens {
      debug!("{}\tBACK  [{}<-{}]: read {} bytes", self.log_ctx, front.0, back.0, sz);
    }

    if r != SocketResult::Continue || sz == 0 {
      self.readiness.back_readiness.remove(Ready::readable());
    }

    if r == SocketResult::Error {
      error!("{}\tback socket read error, closing connection", self.log_ctx);
      self.readiness.reset();
      return ClientResult::CloseBothFailure;
    }

    ClientResult::Continue
  }
}

<commit_msg>websocket support for HTTP<commit_after>use std::cmp::min;
use std::net::{SocketAddr,IpAddr};
use std::io::Write;
use mio::*;
use mio::tcp::TcpStream;
use pool::{Pool,Checkout,Reset};
use time::{Duration, precise_time_s, precise_time_ns};
use uuid::Uuid;
use network::{ClientResult,Protocol};
use network::buffer_queue::BufferQueue;
use network::proxy::Readiness;
use network::socket::{SocketHandler,SocketResult};
use nom::HexDisplay;

type BackendToken = Token;

#[derive(PartialEq)]
pub enum ClientStatus {
  Normal,
  DefaultAnswer,
}

pub struct Pipe<Front:SocketHandler> {
  pub frontend:       Front,
  backend:            TcpStream,
  token:              Option<Token>,
  backend_token:      Option<Token>,
  pub front_buf:      Checkout<BufferQueue>,
  back_buf:           Checkout<BufferQueue>,
  front_buf_position: usize,
  back_buf_position:  usize,
  pub app_id:         Option<String>,
  pub request_id:     String,
  pub server_context: String,
  pub readiness:      Readiness,
  pub log_ctx:        String,
  public_address:     Option<IpAddr>,
}

impl<Front:SocketHandler> Pipe<Front> {
  pub fn new(server_context: &str, frontend: Front, backend: TcpStream, front_buf: Checkout<BufferQueue>, back_buf: Checkout<BufferQueue>, public_address: Option<IpAddr>) -> Option<Pipe<Front>> {
    let request_id = Uuid::new_v4().hyphenated().to_string();
    let log_ctx    = format!("{}\t{}\tunknown\t", server_context, &request_id);
    let client = Pipe {
      frontend:           frontend,
      backend:            backend,
      token:              None,
      backend_token:      None,
      front_buf:          front_buf,
      back_buf:           back_buf,
      front_buf_position: 0,
      back_buf_position:  0,
      app_id:             None,
      request_id:         request_id,
      server_context:     String::from(server_context),
      readiness:          Readiness {
                            front_interest:  Ready::readable() | Ready::writable() |Ready::hup() | Ready::error(),
                            back_interest:   Ready::readable() | Ready::writable() |Ready::hup() | Ready::error(),
                            front_readiness: Ready::none(),
                            back_readiness:  Ready::none(),
      },
      log_ctx:            log_ctx,
      public_address:     public_address,
    };

    info!("created pipe");
    Some(client)
  }

  fn tokens(&self) -> Option<(Token,Token)> {
    if let Some(front) = self.token {
      if let Some(back) = self.backend_token {
        return Some((front, back))
      }
    }
    None
  }

  pub fn front_socket(&self) -> &TcpStream {
    self.frontend.socket_ref()
  }

  pub fn back_socket(&self)  -> Option<&TcpStream> {
    Some(&self.backend)
  }

  pub fn front_token(&self)  -> Option<Token> {
    self.token
  }

  pub fn back_token(&self)   -> Option<Token> {
    self.backend_token
  }

  pub fn close(&mut self) {
  }

  pub fn log_context(&self) -> String {
    if let Some(ref app_id) = self.app_id {
      format!("{}\t{}\t{}\t", self.server_context, self.request_id, app_id)
    } else {
      format!("{}\t{}\tunknown\t", self.server_context, self.request_id)
    }
  }

  pub fn set_front_token(&mut self, token: Token) {
    self.token         = Some(token);
  }

  pub fn set_back_token(&mut self, token: Token) {
    self.backend_token = Some(token);
  }

  pub fn readiness(&mut self) -> &mut Readiness {
    &mut self.readiness
  }

  pub fn front_hup(&mut self) -> ClientResult {
    if self.backend_token == None {
      ClientResult::CloseClient
    } else {
      ClientResult::Continue
    }
  }

  pub fn back_hup(&mut self) -> ClientResult {
    if self.token == None {
      ClientResult::CloseClient
    } else {
      ClientResult::Continue
    }
  }

  // Read content from the client
  pub fn readable(&mut self) -> ClientResult {
    info!("pipe readable");
    if self.front_buf.buffer.available_space() == 0 {
      self.readiness.front_interest.remove(Ready::readable());
      self.readiness.back_interest.insert(Ready::writable());
      return ClientResult::Continue;
    }

    let (sz, res) = self.frontend.socket_read(self.front_buf.buffer.space());
    debug!("{}\tFRONT [{:?}]: read {} bytes", self.log_ctx, self.token, sz);

    if sz > 0 {
      //FIXME: replace with copy()
      self.front_buf.buffer.fill(sz);
      self.front_buf.sliced_input(sz);
      self.front_buf.consume_parsed_data(sz);
      self.front_buf.slice_output(sz);

      if self.front_buf.buffer.available_space() == 0 {
        self.readiness.front_interest.remove(Ready::readable());
      }
      self.readiness.back_interest.insert(Ready::writable());
    } else {
      self.readiness.front_readiness.remove(Ready::readable());
    }

    match res {
      SocketResult::Error => {
        error!("{}\t[{:?}] front socket error, closing the connection", self.log_ctx, self.token);
        self.readiness.reset();
        return ClientResult::CloseClient;
      },
      SocketResult::WouldBlock => {
        self.readiness.front_readiness.remove(Ready::readable());
      },
      SocketResult::Continue => {}
    };

    self.readiness.back_interest.insert(Ready::writable());
    ClientResult::Continue
  }

  // Forward content to client
  pub fn writable(&mut self) -> ClientResult {
    info!("pipe writable");
    if self.back_buf.output_data_size() == 0 || self.back_buf.next_output_data().len() == 0 {
      self.readiness.back_interest.insert(Ready::readable());
      self.readiness.front_interest.remove(Ready::writable());
      return ClientResult::Continue;
    }

    let mut sz = 0usize;
    let mut res = SocketResult::Continue;
    while res == SocketResult::Continue && self.back_buf.output_data_size() > 0 {
      // no more data in buffer, stop here
      if self.back_buf.next_output_data().len() == 0 {
        self.readiness.back_interest.insert(Ready::readable());
        self.readiness.front_interest.remove(Ready::writable());
        return ClientResult::Continue;
      }
      let (current_sz, current_res) = self.frontend.socket_write(self.back_buf.next_output_data());
      res = current_res;
      self.back_buf.consume_output_data(current_sz);
      self.back_buf_position += current_sz;
      sz += current_sz;
    }

    if sz > 0 {
      self.readiness.back_interest.insert(Ready::readable());
    }

    if let Some((front,back)) = self.tokens() {
      debug!("{}\tFRONT [{}<-{}]: wrote {} bytes of {}, buffer position {} restart position {}",
        self.log_ctx, front.0, back.0, sz, self.back_buf.output_data_size(),
        self.back_buf.buffer_position, self.back_buf.start_parsing_position);
    }

    match res {
      SocketResult::Error => {
        error!("{}\t[{:?}] error writing to front socket, closing", self.log_ctx, self.token);
        self.readiness.reset();
        return ClientResult::CloseClient;
      },
      SocketResult::WouldBlock => {
        self.readiness.front_readiness.remove(Ready::writable());
      },
      SocketResult::Continue => {},
    }

    ClientResult::Continue
  }

  // Forward content to application
  pub fn back_writable(&mut self) -> ClientResult {
    info!("pipe back_writable");
    if self.front_buf.output_data_size() == 0 || self.front_buf.next_output_data().len() == 0 {
      self.readiness.front_interest.insert(Ready::readable());
      self.readiness.back_interest.remove(Ready::writable());
      return ClientResult::Continue;
    }

    let tokens = self.tokens().clone();
    let output_size = self.front_buf.output_data_size();

    let mut sz = 0usize;
    let mut socket_res = SocketResult::Continue;

    while socket_res == SocketResult::Continue && self.front_buf.output_data_size() > 0 {
      // no more data in buffer, stop here
      if self.front_buf.next_output_data().len() == 0 {
        self.readiness.front_interest.insert(Ready::readable());
        self.readiness.back_interest.remove(Ready::writable());
        return ClientResult::Continue;
      }
      let (current_sz, current_res) = self.backend.socket_write(self.front_buf.next_output_data());
      socket_res = current_res;
      self.front_buf.consume_output_data(current_sz);
      self.front_buf_position += current_sz;
      sz += current_sz;
    }

    if let Some((front,back)) = tokens {
      debug!("{}\tBACK [{}->{}]: wrote {} bytes of {}", self.log_ctx, front.0, back.0, sz, output_size);
    }
    match socket_res {
      SocketResult::Error => {
        error!("{}\tback socket write error, closing connection", self.log_ctx);
        self.readiness.reset();
        return ClientResult::CloseBothFailure;
      },
      SocketResult::WouldBlock => {
        self.readiness.back_readiness.remove(Ready::writable());

      },
      SocketResult::Continue => {}
    }
    ClientResult::Continue
  }

  // Read content from application
  pub fn back_readable(&mut self) -> ClientResult {
    info!("pipe back_readable");
    if self.back_buf.buffer.available_space() == 0 {
      self.readiness.back_interest.remove(Ready::readable());
      return ClientResult::Continue;
    }

    let tokens     = self.tokens().clone();
    let (sz, r) = self.backend.socket_read(&mut self.back_buf.buffer.space());
    self.back_buf.buffer.fill(sz);
    self.back_buf.sliced_input(sz);
    self.back_buf.consume_parsed_data(sz);
    self.back_buf.slice_output(sz);

    if let Some((front,back)) = tokens {
      debug!("{}\tBACK  [{}<-{}]: read {} bytes", self.log_ctx, front.0, back.0, sz);
    }

    if r != SocketResult::Continue || sz == 0 {
      self.readiness.back_readiness.remove(Ready::readable());
    }
    if sz > 0 {
      self.readiness.front_interest.insert(Ready::writable());
    }

    if r == SocketResult::Error {
      error!("{}\tback socket read error, closing connection", self.log_ctx);
      self.readiness.reset();
      return ClientResult::CloseBothFailure;
    }

    ClientResult::Continue
  }
}

<|endoftext|>


<commit_before>//! The JIT driver uses [`cranelift_jit`] to JIT execute programs without writing any object
//! files.

use std::cell::RefCell;
use std::ffi::CString;
use std::os::raw::{c_char, c_int};

use cranelift_codegen::binemit::{NullStackMapSink, NullTrapSink};
use rustc_codegen_ssa::CrateInfo;
use rustc_middle::mir::mono::MonoItem;

use cranelift_jit::{JITBuilder, JITModule};

use crate::{prelude::*, BackendConfig};
use crate::{CodegenCx, CodegenMode};

struct JitState {
    backend_config: BackendConfig,
    jit_module: JITModule,
}

thread_local! {
    static LAZY_JIT_STATE: RefCell<Option<JitState>> = RefCell::new(None);
}

fn create_jit_module<'tcx>(
    tcx: TyCtxt<'tcx>,
    backend_config: &BackendConfig,
    hotswap: bool,
) -> (JITModule, CodegenCx<'tcx>) {
    let imported_symbols = load_imported_symbols_for_jit(tcx);

    let isa = crate::build_isa(tcx.sess, backend_config);
    let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
    jit_builder.hotswap(hotswap);
    crate::compiler_builtins::register_functions_for_jit(&mut jit_builder);
    jit_builder.symbols(imported_symbols);
    let mut jit_module = JITModule::new(jit_builder);

    let mut cx = crate::CodegenCx::new(tcx, backend_config.clone(), jit_module.isa(), false);

    crate::allocator::codegen(tcx, &mut jit_module, &mut cx.unwind_context);
    crate::main_shim::maybe_create_entry_wrapper(
        tcx,
        &mut jit_module,
        &mut cx.unwind_context,
        true,
    );

    (jit_module, cx)
}

pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
    if !tcx.sess.opts.output_types.should_codegen() {
        tcx.sess.fatal("JIT mode doesn't work with `cargo check`");
    }

    if !tcx.sess.crate_types().contains(&rustc_session::config::CrateType::Executable) {
        tcx.sess.fatal("can't jit non-executable crate");
    }

    let (mut jit_module, mut cx) = create_jit_module(
        tcx,
        &backend_config,
        matches!(backend_config.codegen_mode, CodegenMode::JitLazy),
    );

    let (_, cgus) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);
    let mono_items = cgus
        .iter()
        .map(|cgu| cgu.items_in_deterministic_order(tcx).into_iter())
        .flatten()
        .collect::<FxHashMap<_, (_, _)>>()
        .into_iter()
        .collect::<Vec<(_, (_, _))>>();

    super::time(tcx, backend_config.display_cg_time, "codegen mono items", || {
        super::predefine_mono_items(tcx, &mut jit_module, &mono_items);
        for (mono_item, _) in mono_items {
            match mono_item {
                MonoItem::Fn(inst) => match backend_config.codegen_mode {
                    CodegenMode::Aot => unreachable!(),
                    CodegenMode::Jit => {
                        cx.tcx.sess.time("codegen fn", || {
                            crate::base::codegen_fn(&mut cx, &mut jit_module, inst)
                        });
                    }
                    CodegenMode::JitLazy => codegen_shim(&mut cx, &mut jit_module, inst),
                },
                MonoItem::Static(def_id) => {
                    crate::constant::codegen_static(tcx, &mut jit_module, def_id);
                }
                MonoItem::GlobalAsm(item_id) => {
                    let item = tcx.hir().item(item_id);
                    tcx.sess.span_fatal(item.span, "Global asm is not supported in JIT mode");
                }
            }
        }
    });

    if !cx.global_asm.is_empty() {
        tcx.sess.fatal("Inline asm is not supported in JIT mode");
    }

    tcx.sess.abort_if_errors();

    jit_module.finalize_definitions();
    unsafe { cx.unwind_context.register_jit(&jit_module) };

    println!(
        "Rustc codegen cranelift will JIT run the executable, because -Cllvm-args=mode=jit was passed"
    );

    let args = std::iter::once(&*tcx.crate_name(LOCAL_CRATE).as_str().to_string())
        .chain(backend_config.jit_args.iter().map(|arg| &**arg))
        .map(|arg| CString::new(arg).unwrap())
        .collect::<Vec<_>>();
    let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();

    // Push a null pointer as a terminating argument. This is required by POSIX and
    // useful as some dynamic linkers use it as a marker to jump over.
    argv.push(std::ptr::null());

    let start_sig = Signature {
        params: vec![
            AbiParam::new(jit_module.target_config().pointer_type()),
            AbiParam::new(jit_module.target_config().pointer_type()),
        ],
        returns: vec![AbiParam::new(jit_module.target_config().pointer_type() /*isize*/)],
        call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
    };
    let start_func_id = jit_module.declare_function("main", Linkage::Import, &start_sig).unwrap();
    let finalized_start: *const u8 = jit_module.get_finalized_function(start_func_id);

    LAZY_JIT_STATE.with(|lazy_jit_state| {
        let mut lazy_jit_state = lazy_jit_state.borrow_mut();
        assert!(lazy_jit_state.is_none());
        *lazy_jit_state = Some(JitState { backend_config, jit_module });
    });

    let (main_def_id, entry_ty) = tcx.entry_fn(LOCAL_CRATE).unwrap();
    let instance = Instance::mono(tcx, main_def_id).polymorphize(tcx);

    match entry_ty {
        EntryFnType::Main => {
            // FIXME set program arguments somehow

            let main_sig = Signature {
                params: vec![],
                returns: vec![],
                call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
            };
            let main_func_id = jit_module
                .declare_function(tcx.symbol_name(instance).name, Linkage::Import, &main_sig)
                .unwrap();
            let finalized_main: *const u8 = jit_module.get_finalized_function(main_func_id);

            CURRENT_MODULE.with(|current_module| {
                assert!(current_module.borrow_mut().replace(jit_module).is_none())
            });

            let f: extern "C" fn() = unsafe { ::std::mem::transmute(finalized_main) };
            f();
            std::process::exit(0);
        }
        EntryFnType::Start => {
            let start_sig = Signature {
                params: vec![
                    AbiParam::new(jit_module.target_config().pointer_type()),
                    AbiParam::new(jit_module.target_config().pointer_type()),
                ],
                returns: vec![AbiParam::new(
                    jit_module.target_config().pointer_type(), /*isize*/
                )],
                call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
            };
            let start_func_id = jit_module
                .declare_function(tcx.symbol_name(instance).name, Linkage::Import, &start_sig)
                .unwrap();
            let finalized_start: *const u8 = jit_module.get_finalized_function(start_func_id);

            CURRENT_MODULE.with(|current_module| {
                assert!(current_module.borrow_mut().replace(jit_module).is_none())
            });

            let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
                unsafe { ::std::mem::transmute(finalized_start) };
            let ret = f(args.len() as c_int, argv.as_ptr());
            std::process::exit(ret);
        }
    }
}

#[no_mangle]
extern "C" fn __clif_jit_fn(instance_ptr: *const Instance<'static>) -> *const u8 {
    rustc_middle::ty::tls::with(|tcx| {
        // lift is used to ensure the correct lifetime for instance.
        let instance = tcx.lift(unsafe { *instance_ptr }).unwrap();

        LAZY_JIT_STATE.with(|lazy_jit_state| {
            let mut lazy_jit_state = lazy_jit_state.borrow_mut();
            let lazy_jit_state = lazy_jit_state.as_mut().unwrap();
            let jit_module = &mut lazy_jit_state.jit_module;
            let backend_config = lazy_jit_state.backend_config.clone();

            let name = tcx.symbol_name(instance).name;
            let sig = crate::abi::get_function_sig(tcx, jit_module.isa().triple(), instance);
            let func_id = jit_module.declare_function(name, Linkage::Export, &sig).unwrap();
            jit_module.prepare_for_function_redefine(func_id).unwrap();

            let mut cx = crate::CodegenCx::new(tcx, backend_config, jit_module.isa(), false);
            tcx.sess.time("codegen fn", || crate::base::codegen_fn(&mut cx, jit_module, instance));

            assert!(cx.global_asm.is_empty());
            jit_module.finalize_definitions();
            unsafe { cx.unwind_context.register_jit(&jit_module) };
            jit_module.get_finalized_function(func_id)
        })
    })
}

fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
    use rustc_middle::middle::dependency_format::Linkage;

    let mut dylib_paths = Vec::new();

    let crate_info = CrateInfo::new(tcx);
    let formats = tcx.dependency_formats(LOCAL_CRATE);
    let data = &formats
        .iter()
        .find(|(crate_type, _data)| *crate_type == rustc_session::config::CrateType::Executable)
        .unwrap()
        .1;
    for &(cnum, _) in &crate_info.used_crates_dynamic {
        let src = &crate_info.used_crate_source[&cnum];
        match data[cnum.as_usize() - 1] {
            Linkage::NotLinked | Linkage::IncludedFromDylib => {}
            Linkage::Static => {
                let name = tcx.crate_name(cnum);
                let mut err =
                    tcx.sess.struct_err(&format!("Can't load static lib {}", name.as_str()));
                err.note("rustc_codegen_cranelift can only load dylibs in JIT mode.");
                err.emit();
            }
            Linkage::Dynamic => {
                dylib_paths.push(src.dylib.as_ref().unwrap().0.clone());
            }
        }
    }

    let mut imported_symbols = Vec::new();
    for path in dylib_paths {
        use object::{Object, ObjectSymbol};
        let lib = libloading::Library::new(&path).unwrap();
        let obj = std::fs::read(path).unwrap();
        let obj = object::File::parse(&obj).unwrap();
        imported_symbols.extend(obj.dynamic_symbols().filter_map(|symbol| {
            let name = symbol.name().unwrap().to_string();
            if name.is_empty() || !symbol.is_global() || symbol.is_undefined() {
                return None;
            }
            if name.starts_with("rust_metadata_") {
                // The metadata is part of a section that is not loaded by the dynamic linker in
                // case of cg_llvm.
                return None;
            }
            let dlsym_name = if cfg!(target_os = "macos") {
                // On macOS `dlsym` expects the name without leading `_`.
                assert!(name.starts_with('_'), "{:?}", name);
                &name[1..]
            } else {
                &name
            };
            let symbol: libloading::Symbol<'_, *const u8> =
                unsafe { lib.get(dlsym_name.as_bytes()) }.unwrap();
            Some((name, *symbol))
        }));
        std::mem::forget(lib)
    }

    tcx.sess.abort_if_errors();

    imported_symbols
}

fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx>, module: &mut JITModule, inst: Instance<'tcx>) {
    let tcx = cx.tcx;

    let pointer_type = module.target_config().pointer_type();

    let name = tcx.symbol_name(inst).name;
    let sig = crate::abi::get_function_sig(tcx, module.isa().triple(), inst);
    let func_id = module.declare_function(name, Linkage::Export, &sig).unwrap();

    let instance_ptr = Box::into_raw(Box::new(inst));

    let jit_fn = module
        .declare_function(
            "__clif_jit_fn",
            Linkage::Import,
            &Signature {
                call_conv: module.target_config().default_call_conv,
                params: vec![AbiParam::new(pointer_type)],
                returns: vec![AbiParam::new(pointer_type)],
            },
        )
        .unwrap();

    cx.cached_context.clear();
    let trampoline = &mut cx.cached_context.func;
    trampoline.signature = sig.clone();

    let mut builder_ctx = FunctionBuilderContext::new();
    let mut trampoline_builder = FunctionBuilder::new(trampoline, &mut builder_ctx);

    let jit_fn = module.declare_func_in_func(jit_fn, trampoline_builder.func);
    let sig_ref = trampoline_builder.func.import_signature(sig);

    let entry_block = trampoline_builder.create_block();
    trampoline_builder.append_block_params_for_function_params(entry_block);
    let fn_args = trampoline_builder.func.dfg.block_params(entry_block).to_vec();

    trampoline_builder.switch_to_block(entry_block);
    let instance_ptr = trampoline_builder.ins().iconst(pointer_type, instance_ptr as u64 as i64);
    let jitted_fn = trampoline_builder.ins().call(jit_fn, &[instance_ptr]);
    let jitted_fn = trampoline_builder.func.dfg.inst_results(jitted_fn)[0];
    let call_inst = trampoline_builder.ins().call_indirect(sig_ref, jitted_fn, &fn_args);
    let ret_vals = trampoline_builder.func.dfg.inst_results(call_inst).to_vec();
    trampoline_builder.ins().return_(&ret_vals);

    module
        .define_function(
            func_id,
            &mut cx.cached_context,
            &mut NullTrapSink {},
            &mut NullStackMapSink {},
        )
        .unwrap();
}
<commit_msg>[cg_clif] Fix run_jit from sync<commit_after>//! The JIT driver uses [`cranelift_jit`] to JIT execute programs without writing any object
//! files.

use std::cell::RefCell;
use std::ffi::CString;
use std::os::raw::{c_char, c_int};

use cranelift_codegen::binemit::{NullStackMapSink, NullTrapSink};
use rustc_codegen_ssa::CrateInfo;
use rustc_middle::mir::mono::MonoItem;
use rustc_session::config::EntryFnType;

use cranelift_jit::{JITBuilder, JITModule};

use crate::{prelude::*, BackendConfig};
use crate::{CodegenCx, CodegenMode};

struct JitState {
    backend_config: BackendConfig,
    jit_module: JITModule,
}

thread_local! {
    static LAZY_JIT_STATE: RefCell<Option<JitState>> = RefCell::new(None);
}

fn create_jit_module<'tcx>(
    tcx: TyCtxt<'tcx>,
    backend_config: &BackendConfig,
    hotswap: bool,
) -> (JITModule, CodegenCx<'tcx>) {
    let imported_symbols = load_imported_symbols_for_jit(tcx);

    let isa = crate::build_isa(tcx.sess, backend_config);
    let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
    jit_builder.hotswap(hotswap);
    crate::compiler_builtins::register_functions_for_jit(&mut jit_builder);
    jit_builder.symbols(imported_symbols);
    let mut jit_module = JITModule::new(jit_builder);

    let mut cx = crate::CodegenCx::new(tcx, backend_config.clone(), jit_module.isa(), false);

    crate::allocator::codegen(tcx, &mut jit_module, &mut cx.unwind_context);
    crate::main_shim::maybe_create_entry_wrapper(
        tcx,
        &mut jit_module,
        &mut cx.unwind_context,
        true,
    );

    (jit_module, cx)
}

pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
    if !tcx.sess.opts.output_types.should_codegen() {
        tcx.sess.fatal("JIT mode doesn't work with `cargo check`");
    }

    if !tcx.sess.crate_types().contains(&rustc_session::config::CrateType::Executable) {
        tcx.sess.fatal("can't jit non-executable crate");
    }

    let (mut jit_module, mut cx) = create_jit_module(
        tcx,
        &backend_config,
        matches!(backend_config.codegen_mode, CodegenMode::JitLazy),
    );

    let (_, cgus) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);
    let mono_items = cgus
        .iter()
        .map(|cgu| cgu.items_in_deterministic_order(tcx).into_iter())
        .flatten()
        .collect::<FxHashMap<_, (_, _)>>()
        .into_iter()
        .collect::<Vec<(_, (_, _))>>();

    super::time(tcx, backend_config.display_cg_time, "codegen mono items", || {
        super::predefine_mono_items(tcx, &mut jit_module, &mono_items);
        for (mono_item, _) in mono_items {
            match mono_item {
                MonoItem::Fn(inst) => match backend_config.codegen_mode {
                    CodegenMode::Aot => unreachable!(),
                    CodegenMode::Jit => {
                        cx.tcx.sess.time("codegen fn", || {
                            crate::base::codegen_fn(&mut cx, &mut jit_module, inst)
                        });
                    }
                    CodegenMode::JitLazy => codegen_shim(&mut cx, &mut jit_module, inst),
                },
                MonoItem::Static(def_id) => {
                    crate::constant::codegen_static(tcx, &mut jit_module, def_id);
                }
                MonoItem::GlobalAsm(item_id) => {
                    let item = tcx.hir().item(item_id);
                    tcx.sess.span_fatal(item.span, "Global asm is not supported in JIT mode");
                }
            }
        }
    });

    if !cx.global_asm.is_empty() {
        tcx.sess.fatal("Inline asm is not supported in JIT mode");
    }

    tcx.sess.abort_if_errors();

    jit_module.finalize_definitions();
    unsafe { cx.unwind_context.register_jit(&jit_module) };

    println!(
        "Rustc codegen cranelift will JIT run the executable, because -Cllvm-args=mode=jit was passed"
    );

    let args = std::iter::once(&*tcx.crate_name(LOCAL_CRATE).as_str().to_string())
        .chain(backend_config.jit_args.iter().map(|arg| &**arg))
        .map(|arg| CString::new(arg).unwrap())
        .collect::<Vec<_>>();
    let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();

    // Push a null pointer as a terminating argument. This is required by POSIX and
    // useful as some dynamic linkers use it as a marker to jump over.
    argv.push(std::ptr::null());

    let start_sig = Signature {
        params: vec![
            AbiParam::new(jit_module.target_config().pointer_type()),
            AbiParam::new(jit_module.target_config().pointer_type()),
        ],
        returns: vec![AbiParam::new(jit_module.target_config().pointer_type() /*isize*/)],
        call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
    };
    let start_func_id = jit_module.declare_function("main", Linkage::Import, &start_sig).unwrap();
    let finalized_start: *const u8 = jit_module.get_finalized_function(start_func_id);

    LAZY_JIT_STATE.with(|lazy_jit_state| {
        let mut lazy_jit_state = lazy_jit_state.borrow_mut();
        assert!(lazy_jit_state.is_none());
        *lazy_jit_state = Some(JitState { backend_config, jit_module });
    });

    let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
        unsafe { ::std::mem::transmute(finalized_start) };
    let ret = f(args.len() as c_int, argv.as_ptr());
    std::process::exit(ret);
}

#[no_mangle]
extern "C" fn __clif_jit_fn(instance_ptr: *const Instance<'static>) -> *const u8 {
    rustc_middle::ty::tls::with(|tcx| {
        // lift is used to ensure the correct lifetime for instance.
        let instance = tcx.lift(unsafe { *instance_ptr }).unwrap();

        LAZY_JIT_STATE.with(|lazy_jit_state| {
            let mut lazy_jit_state = lazy_jit_state.borrow_mut();
            let lazy_jit_state = lazy_jit_state.as_mut().unwrap();
            let jit_module = &mut lazy_jit_state.jit_module;
            let backend_config = lazy_jit_state.backend_config.clone();

            let name = tcx.symbol_name(instance).name;
            let sig = crate::abi::get_function_sig(tcx, jit_module.isa().triple(), instance);
            let func_id = jit_module.declare_function(name, Linkage::Export, &sig).unwrap();
            jit_module.prepare_for_function_redefine(func_id).unwrap();

            let mut cx = crate::CodegenCx::new(tcx, backend_config, jit_module.isa(), false);
            tcx.sess.time("codegen fn", || crate::base::codegen_fn(&mut cx, jit_module, instance));

            assert!(cx.global_asm.is_empty());
            jit_module.finalize_definitions();
            unsafe { cx.unwind_context.register_jit(&jit_module) };
            jit_module.get_finalized_function(func_id)
        })
    })
}

fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
    use rustc_middle::middle::dependency_format::Linkage;

    let mut dylib_paths = Vec::new();

    let crate_info = CrateInfo::new(tcx);
    let formats = tcx.dependency_formats(LOCAL_CRATE);
    let data = &formats
        .iter()
        .find(|(crate_type, _data)| *crate_type == rustc_session::config::CrateType::Executable)
        .unwrap()
        .1;
    for &(cnum, _) in &crate_info.used_crates_dynamic {
        let src = &crate_info.used_crate_source[&cnum];
        match data[cnum.as_usize() - 1] {
            Linkage::NotLinked | Linkage::IncludedFromDylib => {}
            Linkage::Static => {
                let name = tcx.crate_name(cnum);
                let mut err =
                    tcx.sess.struct_err(&format!("Can't load static lib {}", name.as_str()));
                err.note("rustc_codegen_cranelift can only load dylibs in JIT mode.");
                err.emit();
            }
            Linkage::Dynamic => {
                dylib_paths.push(src.dylib.as_ref().unwrap().0.clone());
            }
        }
    }

    let mut imported_symbols = Vec::new();
    for path in dylib_paths {
        use object::{Object, ObjectSymbol};
        let lib = libloading::Library::new(&path).unwrap();
        let obj = std::fs::read(path).unwrap();
        let obj = object::File::parse(&obj).unwrap();
        imported_symbols.extend(obj.dynamic_symbols().filter_map(|symbol| {
            let name = symbol.name().unwrap().to_string();
            if name.is_empty() || !symbol.is_global() || symbol.is_undefined() {
                return None;
            }
            if name.starts_with("rust_metadata_") {
                // The metadata is part of a section that is not loaded by the dynamic linker in
                // case of cg_llvm.
                return None;
            }
            let dlsym_name = if cfg!(target_os = "macos") {
                // On macOS `dlsym` expects the name without leading `_`.
                assert!(name.starts_with('_'), "{:?}", name);
                &name[1..]
            } else {
                &name
            };
            let symbol: libloading::Symbol<'_, *const u8> =
                unsafe { lib.get(dlsym_name.as_bytes()) }.unwrap();
            Some((name, *symbol))
        }));
        std::mem::forget(lib)
    }

    tcx.sess.abort_if_errors();

    imported_symbols
}

fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx>, module: &mut JITModule, inst: Instance<'tcx>) {
    let tcx = cx.tcx;

    let pointer_type = module.target_config().pointer_type();

    let name = tcx.symbol_name(inst).name;
    let sig = crate::abi::get_function_sig(tcx, module.isa().triple(), inst);
    let func_id = module.declare_function(name, Linkage::Export, &sig).unwrap();

    let instance_ptr = Box::into_raw(Box::new(inst));

    let jit_fn = module
        .declare_function(
            "__clif_jit_fn",
            Linkage::Import,
            &Signature {
                call_conv: module.target_config().default_call_conv,
                params: vec![AbiParam::new(pointer_type)],
                returns: vec![AbiParam::new(pointer_type)],
            },
        )
        .unwrap();

    cx.cached_context.clear();
    let trampoline = &mut cx.cached_context.func;
    trampoline.signature = sig.clone();

    let mut builder_ctx = FunctionBuilderContext::new();
    let mut trampoline_builder = FunctionBuilder::new(trampoline, &mut builder_ctx);

    let jit_fn = module.declare_func_in_func(jit_fn, trampoline_builder.func);
    let sig_ref = trampoline_builder.func.import_signature(sig);

    let entry_block = trampoline_builder.create_block();
    trampoline_builder.append_block_params_for_function_params(entry_block);
    let fn_args = trampoline_builder.func.dfg.block_params(entry_block).to_vec();

    trampoline_builder.switch_to_block(entry_block);
    let instance_ptr = trampoline_builder.ins().iconst(pointer_type, instance_ptr as u64 as i64);
    let jitted_fn = trampoline_builder.ins().call(jit_fn, &[instance_ptr]);
    let jitted_fn = trampoline_builder.func.dfg.inst_results(jitted_fn)[0];
    let call_inst = trampoline_builder.ins().call_indirect(sig_ref, jitted_fn, &fn_args);
    let ret_vals = trampoline_builder.func.dfg.inst_results(call_inst).to_vec();
    trampoline_builder.ins().return_(&ret_vals);

    module
        .define_function(
            func_id,
            &mut cx.cached_context,
            &mut NullTrapSink {},
            &mut NullStackMapSink {},
        )
        .unwrap();
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2011-2017, UDI Contributors
// All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
#![deny(warnings)]

extern crate goblin;

use std::collections::HashMap;
use std::fs::File;
use std::io::Write;
use std::io::Read;
use std::path::PathBuf;
use std::process::Command;

const SIMPLE_BINARY: &'static str =
    "simple-debug-noopt-dynamic.cd52194667df0781720ff834a56df134fef7fb51";

const WAITTHREAD_BINARY: &'static str =
    "waitthread-debug-noopt-dynamic.9bb7ace83d00e653b003e05b15ee7f9944c7f261";

pub fn setup(out_path: &PathBuf, zip_path: &PathBuf) {

    let dest_dir = out_path.join("native-file-tests");
    Command::new("unzip").arg("-j")
                         .arg(zip_path.to_str().unwrap())
                         .arg("-d")
                         .arg(dest_dir.to_str().unwrap())
                         .spawn()
                         .expect("Failed to start extract of native file tests zip")
                         .wait()
                         .expect("Failed to extract native file tests zip");

    let simple_path = dest_dir.join(SIMPLE_BINARY);
    let waitthread_path = dest_dir.join(WAITTHREAD_BINARY);

    let mod_file_path = out_path.join("native_file_tests.rs");
    let mut mod_file = File::create(&mod_file_path).unwrap();

    let symbols = build_symbols(&dest_dir);

    let sym_defs = symbols.iter()
                          .map(|e| format!("pub const {}: u64 = {};\n", e.0, e.1))
                          .fold("".to_owned(), |acc, elem| acc + &elem);

    let mod_file_content = format!("
pub const SIMPLE_EXEC_PATH: &'static str = \"{}\";
pub const WAITTHREAD_EXEC_PATH: &'static str = \"{}\";

{}
",
    simple_path.to_str().unwrap(),
    waitthread_path.to_str().unwrap(),
    sym_defs);

    mod_file.write_all(&mod_file_content.into_bytes())
            .expect("Failed to write native file tests module");
}

fn build_symbols(dest_dir: &PathBuf) -> HashMap<&str, String> {
    let mut symbols = HashMap::new();

    add_simple_binary_symbols(dest_dir, &mut symbols);
    add_waitthread_binary_symbols(dest_dir, &mut symbols);

    symbols
}

fn add_simple_binary_symbols(dest_dir: &PathBuf, symbols: &mut HashMap<&str, String>) {
    let path = dest_dir.join(SIMPLE_BINARY);

    for_each_symbols(&path, |name, value, size| {
        if name == "function1" {
            symbols.insert("SIMPLE_FUNCTION1", format!("0x{:x}", value));
        } else if name == "function2" {
            symbols.insert("SIMPLE_FUNCTION2", format!("0x{:x}", value));
            symbols.insert("SIMPLE_FUNCTION2_LENGTH", size.to_string());
        }
    });
}

fn add_waitthread_binary_symbols(dest_dir: &PathBuf, symbols: &mut HashMap<&str, String>) {
    let path = dest_dir.join(WAITTHREAD_BINARY);

    for_each_symbols(&path, |name, value, _| {
        if name == "breakpoint_func" {
            symbols.insert("THREAD_BREAK_FUNC", format!("0x{:x}", value));
        }
    });
}

fn for_each_symbols<F>(path: &PathBuf, mut func: F)
    where F: FnMut(&String, u64, u64) {

    let mut binary_file = File::open(&path).expect(&format!("Symbol file {:?} error",
                                                           path));
    let mut buffer = vec![];
    binary_file.read_to_end(&mut buffer).unwrap();

    match goblin::parse(&buffer).unwrap() {
        goblin::Object::Elf(elf) => {
            let strtab = &(elf.strtab);
            for sym in &(elf.syms) {
                func(&(strtab.get(sym.st_name).unwrap().to_owned()),
                     sym.st_value,
                     sym.st_size);
            }
        },
        _ => {
            panic!(format!("Unsupported file type"));
        }
    }
}
<commit_msg>Update rust binding to avoid unzipping the archive unconditionally.<commit_after>//
// Copyright (c) 2011-2017, UDI Contributors
// All rights reserved.
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
#![deny(warnings)]

extern crate goblin;

use std::collections::HashMap;
use std::fs::File;
use std::io::Write;
use std::io::Read;
use std::path::PathBuf;
use std::process::Command;

// TODO dynamically select correct binaries based on platform
const SIMPLE_BINARY: &'static str =
    "simple-debug-noopt-dynamic.e537d89753d335c4819963172276d34dcd5932d2";

const WAITTHREAD_BINARY: &'static str =
    "waitthread-debug-noopt-dynamic.0e8e1118053d6ea41cf58968d1e56e9708ff6758";

pub fn setup(manifest_path: &PathBuf, out_path: &PathBuf, zip_path: &PathBuf) {

    let dest_dir = manifest_path.join("native-file-tests");

    if !dest_dir.exists() {
        Command::new("unzip").arg("-j")
                             .arg(zip_path.to_str().unwrap())
                             .arg("-d")
                             .arg(dest_dir.to_str().unwrap())
                             .spawn()
                             .expect("Failed to start extract of native file tests zip")
                             .wait()
                             .expect("Failed to extract native file tests zip");
    }

    let simple_path = dest_dir.join(SIMPLE_BINARY);
    let waitthread_path = dest_dir.join(WAITTHREAD_BINARY);

    let mod_file_path = out_path.join("native_file_tests.rs");
    let mut mod_file = File::create(&mod_file_path).unwrap();

    let symbols = build_symbols(&dest_dir);

    let sym_defs = symbols.iter()
                          .map(|e| format!("pub const {}: u64 = {};\n", e.0, e.1))
                          .fold("".to_owned(), |acc, elem| acc + &elem);

    let mod_file_content = format!("
pub const SIMPLE_EXEC_PATH: &'static str = \"{}\";
pub const WAITTHREAD_EXEC_PATH: &'static str = \"{}\";

{}
",
    simple_path.to_str().unwrap(),
    waitthread_path.to_str().unwrap(),
    sym_defs);

    mod_file.write_all(&mod_file_content.into_bytes())
            .expect("Failed to write native file tests module");
}

fn build_symbols(dest_dir: &PathBuf) -> HashMap<&str, String> {
    let mut symbols = HashMap::new();

    add_simple_binary_symbols(dest_dir, &mut symbols);
    add_waitthread_binary_symbols(dest_dir, &mut symbols);

    symbols
}

fn add_simple_binary_symbols(dest_dir: &PathBuf, symbols: &mut HashMap<&str, String>) {
    let path = dest_dir.join(SIMPLE_BINARY);

    for_each_symbols(&path, |name, value, size| {
        if name == "function1" {
            symbols.insert("SIMPLE_FUNCTION1", format!("0x{:x}", value));
        } else if name == "function2" {
            symbols.insert("SIMPLE_FUNCTION2", format!("0x{:x}", value));
            symbols.insert("SIMPLE_FUNCTION2_LENGTH", size.to_string());
        }
    });
}

fn add_waitthread_binary_symbols(dest_dir: &PathBuf, symbols: &mut HashMap<&str, String>) {
    let path = dest_dir.join(WAITTHREAD_BINARY);

    for_each_symbols(&path, |name, value, _| {
        match name.as_str() {
            "breakpoint_thr_func" => {
                symbols.insert("THREAD_BREAK_FUNC", format!("0x{:x}", value));
            },
            "start_notification" => {
                symbols.insert("START_NOTIFICATION_FUNC", format!("0x{:x}", value));
            },
            "term_notification" => {
                symbols.insert("TERM_NOTIFICATION_FUNC", format!("0x{:x}", value));
            },
            _ => {}
        };
    });
}

fn for_each_symbols<F>(path: &PathBuf, mut func: F)
    where F: FnMut(&String, u64, u64) {

    let mut binary_file = File::open(&path).expect(&format!("Symbol file {:?} error",
                                                           path));
    let mut buffer = vec![];
    binary_file.read_to_end(&mut buffer).unwrap();

    match goblin::parse(&buffer).unwrap() {
        goblin::Object::Elf(elf) => {
            let strtab = &(elf.strtab);
            for sym in &(elf.syms) {
                func(&(strtab.get(sym.st_name).unwrap().to_owned()),
                     sym.st_value,
                     sym.st_size);
            }
        },
        _ => {
            panic!(format!("Unsupported file type"));
        }
    }
}
<|endoftext|>


<commit_before>// Implements http://rosettacode.org/wiki/Parallel_calculations
// See http://static.rust-lang.org/doc/master/guide-tasks.html for information
// about tasks, channels, future, etc.

#[cfg(test)]
extern crate test;
extern crate sync;

use sync::Future;
use prime_decomposition::factor;

mod prime_decomposition;

// Returns the minimal prime factor of a number
fn min_factor(x: uint) -> uint {
    // factor returns a sorted vector, so we just take the first element
    *factor(x).get(0)
}

// Returns the largest minimal factor of the numbers in a slice
// The function is implemented using a channel
#[cfg(test)]
fn largest_min_factor_chan(numbers: &[uint]) -> uint {
    let (sender, receiver) = channel();

    // Send all the minimal factors
    for &x in numbers.iter() {
        let child_sender = sender.clone();
        spawn(proc() { child_sender.send(min_factor(x)) });
    }

    // Receive them and keep the largest one
    numbers.iter().fold(0u, |max, _| {
        std::cmp::max(receiver.recv(), max)
    })
}

// Returns the largest minimal factor of the numbers in a slice
// The function is implemented using the Future struct
fn largest_min_factor_fut(numbers: &[uint]) -> uint {
    // We will save the future values of the minimal factor in the results vec
    let mut results = Vec::from_fn(numbers.len(), |i| {
        let number = numbers[i];
        Future::spawn(proc() { min_factor(number) })
    });

    // Get the largest minimal factor of all results
    results.mut_iter().map(|r| r.get()).max().unwrap()
}

#[cfg(not(test))]
fn main() {
    // Numbers to be factorized
    let numbers = &[1122725u,
                   1125827,
                   1122725,
                   1152800,
                   1157978,
                   1099726];

    let max = largest_min_factor_fut(numbers);
    println!("The largest minimal factor is {}", max);
}

// We dont have benchmarks because the Bencher doesn't work good with tasks
#[test]
fn test_basic() {
    let numbers = &[25, 80, 256, 55, 18, 19];
    assert_eq!(largest_min_factor_fut(numbers), 19);
}

#[test]
fn test_equivalence() {
    let numbers = &[1122725u,
                   1125827,
                   1122725,
                   1152800,
                   1157978,
                   1099726];
    assert_eq!(largest_min_factor_chan(numbers),
                largest_min_factor_fut(numbers));
}
<commit_msg>Future is now in std::sync<commit_after>// Implements http://rosettacode.org/wiki/Parallel_calculations
// See http://static.rust-lang.org/doc/master/guide-tasks.html for information
// about tasks, channels, future, etc.

#[cfg(test)]
extern crate test;
extern crate sync;

use std::sync::Future;
use prime_decomposition::factor;

mod prime_decomposition;

// Returns the minimal prime factor of a number
fn min_factor(x: uint) -> uint {
    // factor returns a sorted vector, so we just take the first element
    *factor(x).get(0)
}

// Returns the largest minimal factor of the numbers in a slice
// The function is implemented using a channel
#[cfg(test)]
fn largest_min_factor_chan(numbers: &[uint]) -> uint {
    let (sender, receiver) = channel();

    // Send all the minimal factors
    for &x in numbers.iter() {
        let child_sender = sender.clone();
        spawn(proc() { child_sender.send(min_factor(x)) });
    }

    // Receive them and keep the largest one
    numbers.iter().fold(0u, |max, _| {
        std::cmp::max(receiver.recv(), max)
    })
}

// Returns the largest minimal factor of the numbers in a slice
// The function is implemented using the Future struct
fn largest_min_factor_fut(numbers: &[uint]) -> uint {
    // We will save the future values of the minimal factor in the results vec
    let mut results = Vec::from_fn(numbers.len(), |i| {
        let number = numbers[i];
        Future::spawn(proc() { min_factor(number) })
    });

    // Get the largest minimal factor of all results
    results.mut_iter().map(|r| r.get()).max().unwrap()
}

#[cfg(not(test))]
fn main() {
    // Numbers to be factorized
    let numbers = &[1122725u,
                   1125827,
                   1122725,
                   1152800,
                   1157978,
                   1099726];

    let max = largest_min_factor_fut(numbers);
    println!("The largest minimal factor is {}", max);
}

// We dont have benchmarks because the Bencher doesn't work good with tasks
#[test]
fn test_basic() {
    let numbers = &[25, 80, 256, 55, 18, 19];
    assert_eq!(largest_min_factor_fut(numbers), 19);
}

#[test]
fn test_equivalence() {
    let numbers = &[1122725u,
                   1125827,
                   1122725,
                   1152800,
                   1157978,
                   1099726];
    assert_eq!(largest_min_factor_chan(numbers),
                largest_min_factor_fut(numbers));
}
<|endoftext|>