<commit_before>extern crate dbus;

use std;
use general::*;
use device::DeviceType;
use manager::NetworkManager;

/// Get a list of Network Manager connections sorted by path.
///
/// # Examples
///
/// ```no_run
/// use network_manager::connection;
/// let manager = network_manager::manager::NetworkManager::new();
/// let connections = connection::list(&manager).unwrap();
/// println!("{:?}", connections);
/// ```
pub fn list(manager: &NetworkManager) -> Result<Vec<Connection>, String> {
    let paths = try!(manager.list_connections());

    let mut connections = Vec::new();

    for path in &paths {
        connections.push(try!(get_connection(&manager, path)))
    }

    connections.sort();

    Ok(connections)
}

#[test]
fn test_list_function() {
    let manager = NetworkManager::new();

    let connections = list(&manager).unwrap();
    assert!(connections.len() > 0);

    for (index, val) in connections.iter().enumerate() {
        assert_ne!(Connection { ..Default::default() }, val.clone());
        assert_eq!(index as i32, i32::from(val));
    }
}

/// Creates a Network Manager connection.
///
/// # Examples
///
/// ```
/// let connection = network_manager::connection::create(
///     "resin_io",
///     network_manager::device::DeviceType::WiFi,
///     network_manager::general::Security::WPA2,
///     "super_secret_passphase"
///     ).unwrap();
/// println!("{:?}", connection);
/// ```
pub fn create(s: &str, dt: DeviceType, sc: Security, p: &str) -> Result<Connection, String> {
    // Create a connection
    // Get the connection
    // Return the connection

    let settings = ConnectionSettings {
        id: "resin_io".to_string(),
        uuid: "3c8e6e8b-b895-4b07-97a5-bbc192c3b436".to_string(),
        ssid: "resin_io".to_string(),
    };

    let connection1 = Connection {
        path: "/org/freedesktop/NetworkManager/ActiveConnection/187".to_string(),
        active_path: "test".to_string(),
        settings: settings,
        state: ConnectionState::Deactivated, /* device: "wlp4s0".to_string(),
                                              * interface: DeviceType::WiFi,
                                              * security: Security::WPA2,
                                              * state: ConnectionState::Activated, */
    };

    Ok(connection1)
}

/// Deletes a Network Manager connection.
///
/// # Examples
///
/// ```
/// let manager = network_manager::manager::NetworkManager::new();
/// let mut connections = network_manager::connection::list(&manager).unwrap();
/// network_manager::connection::delete(connections.pop().unwrap()).unwrap();
/// ```
pub fn delete(connection: Connection) -> Result<(), String> {
    let message = dbus_message!(NM_SERVICE_MANAGER,
                                connection.path,
                                NM_CONNECTION_INTERFACE,
                                "Delete");
    dbus_connect!(message);

    Ok(())
}

/// Enables a Network Manager connection.
///
/// # Examples
///
/// ```no_run
/// use network_manager::connection;
/// let manager = network_manager::manager::NetworkManager::new();
/// let connections = connection::list(&manager).unwrap();
/// let mut connection = connections[0].clone();
/// connection::enable(&manager, &mut connection, 10).unwrap();
/// println!("{:?}", connection.state);
/// ```
pub fn enable(manager: &NetworkManager,
              connection: &mut Connection,
              time_out: i32)
              -> Result<(), String> {
    match connection.state {
        ConnectionState::Activated => Ok(()),
        ConnectionState::Activating => {
            wait(manager, connection, time_out, ConnectionState::Activated)
        }
        ConnectionState::Unknown => Err("Unable to get connection state".to_string()),
        _ => {
            let mut message = dbus_message!(NM_SERVICE_MANAGER,
                                            NM_SERVICE_PATH,
                                            NM_SERVICE_INTERFACE,
                                            "ActivateConnection");
            message.append_items(&[dbus::MessageItem::ObjectPath(connection
                                                                     .path
                                                                     .to_string()
                                                                     .into()),
                                   dbus::MessageItem::ObjectPath("/".into()),
                                   dbus::MessageItem::ObjectPath("/".into())]);
            dbus_connect!(message);

            wait(manager, connection, time_out, ConnectionState::Activated)
        }
    }
}

/// Disables a Network Manager connection.
///
/// # Examples
///
/// ```no_run
/// use network_manager::connection;
/// let manager = network_manager::manager::NetworkManager::new();
/// let connections = connection::list(&manager).unwrap();
/// let mut connection = connections[0].clone();
/// connection::disable(&manager, &mut connection, 10).unwrap();
/// println!("{:?}", connection.state);
/// ```
pub fn disable(manager: &NetworkManager,
               connection: &mut Connection,
               time_out: i32)
               -> Result<(), String> {
    match connection.state {
        ConnectionState::Deactivated => Ok(()),
        ConnectionState::Deactivating => {
            wait(manager, connection, time_out, ConnectionState::Deactivated)
        }
        ConnectionState::Unknown => Err("Unable to get connection state".to_string()),
        _ => {
            let mut message = dbus_message!(NM_SERVICE_MANAGER,
                                            NM_SERVICE_PATH,
                                            NM_SERVICE_INTERFACE,
                                            "DeactivateConnection");
            message.append_items(&[dbus::MessageItem::ObjectPath(connection
                                                                     .active_path
                                                                     .to_string()
                                                                     .into())]);
            dbus_connect!(message);

            wait(manager, connection, time_out, ConnectionState::Deactivated)
        }
    }
}

#[test]
fn test_enable_disable_functions() {
    let manager = NetworkManager::new();

    let connections = list(&manager).unwrap();
    let mut connection;

    // set enviorment variable $TEST_WIFI_SSID with the wifi's SSID that you want to test
    // e.g.  export TEST_WIFI_SSID="Resin.io Wifi"
    let wifiEnvVar = "TEST_WIFI_SSID";
    match env::var(wifiEnvVar) {
        Ok(ssid) => {
            connection = connections
                .iter()
                .filter(|c| c.settings.ssid == ssid)
                .nth(0)
                .unwrap()
                .clone()
        }
        Err(e) => {
            panic!("couldn't retrieve enviorment variable {}: {}",
                   wifiEnvVar,
                   e)
        }
    };

    if connection.state == ConnectionState::Activated {
        disable(&manager, &mut connection, 10).unwrap();
        assert_eq!(ConnectionState::Deactivated, connection.state);

        enable(&manager, &mut connection, 10).unwrap();
        assert_eq!(ConnectionState::Activated, connection.state);
    } else {
        enable(&manager, &mut connection, 10).unwrap();
        assert_eq!(ConnectionState::Activated, connection.state);

        disable(&manager, &mut connection, 10).unwrap();
        assert_eq!(ConnectionState::Deactivated, connection.state);
    }
}

fn get_connection(manager: &NetworkManager, path: &String) -> Result<Connection, String> {
    let mut connection = Connection::default();

    connection.path = path.clone();

    connection.settings = try!(manager.get_connection_settings(path));

    try!(update_state(manager, &mut connection));

    Ok(connection)
}

fn update_state(manager: &NetworkManager, connection: &mut Connection) -> Result<(), String> {
    let active_paths = try!(manager.get_active_connections());

    let mut settings_paths = Vec::new();

    for active_path in &active_paths {
        if let Some(settings_path) = manager.get_active_connection_path(&active_path) {
            settings_paths.push(settings_path)
        }
    }

    // TODO: Consider using Option<String> instead for deactivated connections
    connection.active_path = "".to_string();
    connection.state = ConnectionState::Deactivated;

    for (active_path, settings_path) in active_paths.iter().zip(settings_paths.iter()) {
        if connection.path == *settings_path {

            connection.active_path = active_path.clone();
            connection.state = try!(manager.get_connection_state(&active_path));

            break;
        }
    }

    Ok(())
}

fn wait(manager: &NetworkManager,
        connection: &mut Connection,
        time_out: i32,
        target_state: ConnectionState)
        -> Result<(), String> {
    if time_out == 0 {
        return Ok(());
    }

    let mut total_time = 0;

    while total_time < time_out {
        try!(update_state(manager, connection));

        if connection.state == target_state {
            return Ok(());
        }

        std::thread::sleep(std::time::Duration::from_secs(1));

        total_time += 1;
    }

    Err("service timed out".to_string())
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Connection {
    pub path: String,
    pub active_path: String,
    pub settings: ConnectionSettings,
    pub state: ConnectionState, /* device: String,
                                 * device_type: DeviceType,
                                 * security: Security, */
}

impl Default for Connection {
    fn default() -> Connection {
        Connection {
            path: "".to_string(),
            active_path: "".to_string(),
            settings: ConnectionSettings::default(),
            state: ConnectionState::Unknown,
        }
    }
}

impl Ord for Connection {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        i32::from(self).cmp(&i32::from(other))
    }
}

impl PartialOrd for Connection {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl<'a> From<&'a Connection> for i32 {
    fn from(val: &Connection) -> i32 {
        val.clone()
            .path
            .rsplit('/')
            .nth(0)
            .unwrap()
            .parse::<i32>()
            .unwrap()
    }
}

#[derive(Default, Debug, Clone, Eq, PartialEq)]
pub struct ConnectionSettings {
    pub id: String,
    pub uuid: String,
    pub ssid: String,
}
<commit_msg>Consistent and more flat doctests for connection<commit_after>extern crate dbus;

use std;
use std::env;
use general::*;
use device::DeviceType;
use manager::NetworkManager;

/// Get a list of Network Manager connections sorted by path.
///
/// # Examples
///
/// ```no_run
/// use network_manager::connection;
/// use network_manager::manager::NetworkManager;
/// let manager = NetworkManager::new();
/// let connections = connection::list(&manager).unwrap();
/// println!("{:?}", connections);
/// ```
pub fn list(manager: &NetworkManager) -> Result<Vec<Connection>, String> {
    let paths = try!(manager.list_connections());

    let mut connections = Vec::new();

    for path in &paths {
        connections.push(try!(get_connection(&manager, path)))
    }

    connections.sort();

    Ok(connections)
}

#[test]
fn test_list_function() {
    let manager = NetworkManager::new();

    let connections = list(&manager).unwrap();
    assert!(connections.len() > 0);

    for (index, val) in connections.iter().enumerate() {
        assert_ne!(Connection { ..Default::default() }, val.clone());
        assert_eq!(index as i32, i32::from(val));
    }
}

/// Creates a Network Manager connection.
///
/// # Examples
///
/// ```
/// let connection = network_manager::connection::create(
///     "resin_io",
///     network_manager::device::DeviceType::WiFi,
///     network_manager::general::Security::WPA2,
///     "super_secret_passphase"
///     ).unwrap();
/// println!("{:?}", connection);
/// ```
pub fn create(s: &str, dt: DeviceType, sc: Security, p: &str) -> Result<Connection, String> {
    // Create a connection
    // Get the connection
    // Return the connection

    let settings = ConnectionSettings {
        id: "resin_io".to_string(),
        uuid: "3c8e6e8b-b895-4b07-97a5-bbc192c3b436".to_string(),
        ssid: "resin_io".to_string(),
    };

    let connection1 = Connection {
        path: "/org/freedesktop/NetworkManager/ActiveConnection/187".to_string(),
        active_path: "test".to_string(),
        settings: settings,
        state: ConnectionState::Deactivated, /* device: "wlp4s0".to_string(),
                                              * interface: DeviceType::WiFi,
                                              * security: Security::WPA2,
                                              * state: ConnectionState::Activated, */
    };

    Ok(connection1)
}

/// Deletes a Network Manager connection.
///
/// # Examples
///
/// ```
/// use network_manager::connection;
/// use network_manager::manager::NetworkManager;
/// let manager = NetworkManager::new();
/// let mut connections = connection::list(&manager).unwrap();
/// connection::delete(connections.pop().unwrap()).unwrap();
/// ```
pub fn delete(connection: Connection) -> Result<(), String> {
    let message = dbus_message!(NM_SERVICE_MANAGER,
                                connection.path,
                                NM_CONNECTION_INTERFACE,
                                "Delete");
    dbus_connect!(message);

    Ok(())
}

/// Enables a Network Manager connection.
///
/// # Examples
///
/// ```no_run
/// use network_manager::connection;
/// use network_manager::manager::NetworkManager;
/// let manager = NetworkManager::new();
/// let connections = connection::list(&manager).unwrap();
/// let mut connection = connections[0].clone();
/// connection::enable(&manager, &mut connection, 10).unwrap();
/// println!("{:?}", connection.state);
/// ```
pub fn enable(manager: &NetworkManager,
              connection: &mut Connection,
              time_out: i32)
              -> Result<(), String> {
    match connection.state {
        ConnectionState::Activated => Ok(()),
        ConnectionState::Activating => {
            wait(manager, connection, time_out, ConnectionState::Activated)
        }
        ConnectionState::Unknown => Err("Unable to get connection state".to_string()),
        _ => {
            let mut message = dbus_message!(NM_SERVICE_MANAGER,
                                            NM_SERVICE_PATH,
                                            NM_SERVICE_INTERFACE,
                                            "ActivateConnection");
            message.append_items(&[dbus::MessageItem::ObjectPath(connection
                                                                     .path
                                                                     .to_string()
                                                                     .into()),
                                   dbus::MessageItem::ObjectPath("/".into()),
                                   dbus::MessageItem::ObjectPath("/".into())]);
            dbus_connect!(message);

            wait(manager, connection, time_out, ConnectionState::Activated)
        }
    }
}

/// Disables a Network Manager connection.
///
/// # Examples
///
/// ```no_run
/// use network_manager::connection;
/// use network_manager::manager::NetworkManager;
/// let manager = NetworkManager::new();
/// let connections = connection::list(&manager).unwrap();
/// let mut connection = connections[0].clone();
/// connection::disable(&manager, &mut connection, 10).unwrap();
/// println!("{:?}", connection.state);
/// ```
pub fn disable(manager: &NetworkManager,
               connection: &mut Connection,
               time_out: i32)
               -> Result<(), String> {
    match connection.state {
        ConnectionState::Deactivated => Ok(()),
        ConnectionState::Deactivating => {
            wait(manager, connection, time_out, ConnectionState::Deactivated)
        }
        ConnectionState::Unknown => Err("Unable to get connection state".to_string()),
        _ => {
            let mut message = dbus_message!(NM_SERVICE_MANAGER,
                                            NM_SERVICE_PATH,
                                            NM_SERVICE_INTERFACE,
                                            "DeactivateConnection");
            message.append_items(&[dbus::MessageItem::ObjectPath(connection
                                                                     .active_path
                                                                     .to_string()
                                                                     .into())]);
            dbus_connect!(message);

            wait(manager, connection, time_out, ConnectionState::Deactivated)
        }
    }
}

#[test]
fn test_enable_disable_functions() {
    let manager = NetworkManager::new();

    let connections = list(&manager).unwrap();
    let mut connection;

    // set enviorment variable $TEST_WIFI_SSID with the wifi's SSID that you want to test
    // e.g.  export TEST_WIFI_SSID="Resin.io Wifi"
    let wifiEnvVar = "TEST_WIFI_SSID";
    match env::var(wifiEnvVar) {
        Ok(ssid) => {
            connection = connections
                .iter()
                .filter(|c| c.settings.ssid == ssid)
                .nth(0)
                .unwrap()
                .clone()
        }
        Err(e) => {
            panic!("couldn't retrieve enviorment variable {}: {}",
                   wifiEnvVar,
                   e)
        }
    };

    if connection.state == ConnectionState::Activated {
        disable(&manager, &mut connection, 10).unwrap();
        assert_eq!(ConnectionState::Deactivated, connection.state);

        enable(&manager, &mut connection, 10).unwrap();
        assert_eq!(ConnectionState::Activated, connection.state);
    } else {
        enable(&manager, &mut connection, 10).unwrap();
        assert_eq!(ConnectionState::Activated, connection.state);

        disable(&manager, &mut connection, 10).unwrap();
        assert_eq!(ConnectionState::Deactivated, connection.state);
    }
}

fn get_connection(manager: &NetworkManager, path: &String) -> Result<Connection, String> {
    let mut connection = Connection::default();

    connection.path = path.clone();

    connection.settings = try!(manager.get_connection_settings(path));

    try!(update_state(manager, &mut connection));

    Ok(connection)
}

fn update_state(manager: &NetworkManager, connection: &mut Connection) -> Result<(), String> {
    let active_paths = try!(manager.get_active_connections());

    let mut settings_paths = Vec::new();

    for active_path in &active_paths {
        if let Some(settings_path) = manager.get_active_connection_path(&active_path) {
            settings_paths.push(settings_path)
        }
    }

    // TODO: Consider using Option<String> instead for deactivated connections
    connection.active_path = "".to_string();
    connection.state = ConnectionState::Deactivated;

    for (active_path, settings_path) in active_paths.iter().zip(settings_paths.iter()) {
        if connection.path == *settings_path {

            connection.active_path = active_path.clone();
            connection.state = try!(manager.get_connection_state(&active_path));

            break;
        }
    }

    Ok(())
}

fn wait(manager: &NetworkManager,
        connection: &mut Connection,
        time_out: i32,
        target_state: ConnectionState)
        -> Result<(), String> {
    if time_out == 0 {
        return Ok(());
    }

    let mut total_time = 0;

    while total_time < time_out {
        try!(update_state(manager, connection));

        if connection.state == target_state {
            return Ok(());
        }

        std::thread::sleep(std::time::Duration::from_secs(1));

        total_time += 1;
    }

    Err("service timed out".to_string())
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Connection {
    pub path: String,
    pub active_path: String,
    pub settings: ConnectionSettings,
    pub state: ConnectionState, /* device: String,
                                 * device_type: DeviceType,
                                 * security: Security, */
}

impl Default for Connection {
    fn default() -> Connection {
        Connection {
            path: "".to_string(),
            active_path: "".to_string(),
            settings: ConnectionSettings::default(),
            state: ConnectionState::Unknown,
        }
    }
}

impl Ord for Connection {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        i32::from(self).cmp(&i32::from(other))
    }
}

impl PartialOrd for Connection {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl<'a> From<&'a Connection> for i32 {
    fn from(val: &Connection) -> i32 {
        val.clone()
            .path
            .rsplit('/')
            .nth(0)
            .unwrap()
            .parse::<i32>()
            .unwrap()
    }
}

#[derive(Default, Debug, Clone, Eq, PartialEq)]
pub struct ConnectionSettings {
    pub id: String,
    pub uuid: String,
    pub ssid: String,
}
<|endoftext|>


<commit_before>use num::bigint::{BigUint, RandBigInt};
use num::One;

use rand::Rng;

use utils::modular::Inverse;
use utils::modular::Power;
use utils::primes::generate_prime;

pub struct SecretKeyExtra {
    p: BigUint,
    q: BigUint,
    dmp1: BigUint,
    dmq1: BigUint,
    qinv: BigUint
}

impl SecretKeyExtra {
    fn from_primes(p: BigUint, q: BigUint, d: &BigUint) -> Self {
        SecretKeyExtra {
            dmp1: d % (&p - BigUint::one()),
            dmq1: d % (&q - BigUint::one()),
            qinv: q.inverse(&p).unwrap(),
            p: p,
            q: q,
        }
    }
}

pub enum RSA {
    Public {
        /// Modulus
        n: BigUint,
        /// Public Exponent
        e: BigUint,
    },
    Private {
        /// Modulus
        n: BigUint,
        /// Private Exponent
        d: BigUint,
        extra: Option<SecretKeyExtra>,
    }
}

pub type KeyPair = (RSA, RSA);

impl RSA {
    pub fn keypair_from_primes<P, Q, E>(p: P, q: Q, e: E) -> KeyPair
        where P: Into<BigUint>, Q: Into<BigUint>, E: Into<BigUint> {
            let (p, q, e) = (p.into(), q.into(), e.into());

            let n = &p * &q;
            let phi_n = &n - (&p + &q - BigUint::one());

            let d = e.inverse(&phi_n).expect("e is irreversible in ring phi(pq) - error");

            let public = RSA::Public {
                n: n.clone(),
                e: e
            };
            let private = RSA::Private {
                n: n,
                extra: Some(SecretKeyExtra::from_primes(p, q, &d)),
                d: d,
            };

            (public, private)
        }

    pub fn generate_keypair<G, T>(mut rng: G, e: T, bits: usize) -> KeyPair
        where G: Rng + RandBigInt, T: Into<BigUint> {
            let e = e.into();

            let mut p = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            while (&p - BigUint::one()) % &e != BigUint::one() {
                p = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            }
            let mut q = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            while (&q - BigUint::one()) % &e != BigUint::one() {
                q = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            }

            Self::keypair_from_primes(p, q, e)
        }

    pub fn is_public(&self) -> bool {
        match *self {
            RSA::Public { .. } => true,
            _ => false
        }
    }

    pub fn is_private(&self) -> bool {
        match *self {
            RSA::Private { .. } => true,
            _ => false
        }
    }

    pub fn crypt(&self, msg: &BigUint) -> BigUint {
        match *self {
            RSA::Private { ref n, ref d, ref extra } => crypt(msg, n, d, extra.as_ref()),
            RSA::Public { ref n, ref e } => crypt(msg, n, e, None),
        }
    }
}

fn crypt(msg: &BigUint,
         modulus: &BigUint,
         exp: &BigUint,
         extra: Option<&SecretKeyExtra>) -> BigUint {
    if let Some(ref extra) = extra {
        chinese_reminders_power(msg, extra)
    } else {
        msg.pow_mod(exp, modulus)
    }
}

fn chinese_reminders_power(c: &BigUint, extra: &SecretKeyExtra) -> BigUint {
    let mut m1 = c.pow_mod(&extra.dmp1, &extra.p);
    let m2 = c.pow_mod(&extra.dmq1, &extra.q);

    while m1 < m2 {
        m1 = m1 + &extra.p;
    }

    let h = (&extra.qinv * (m1 - &m2)) % &extra.p;

    m2 + h * &extra.q
}

#[cfg(test)]
mod tests {
    use super::RSA;

    use num::bigint::{ToBigUint};

    fn keys() -> (RSA, RSA) {
        RSA::keypair_from_primes(
            61.to_biguint().unwrap(),
            53.to_biguint().unwrap(),
            17.to_biguint().unwrap())
    }

    #[test]
    fn test_encryption() {
        let (public, _) = keys();
        let c = public.crypt(&65.to_biguint().unwrap());

        assert_eq!(c, 2790.to_biguint().unwrap())
    }

    #[test]
    fn test_decryption() {
        let (_, private) = keys();
        let m = private.crypt(&2790.to_biguint().unwrap());

        assert_eq!(m, 65.to_biguint().unwrap())
    }
}
<commit_msg>fix(): fix typos<commit_after>use num::bigint::{BigUint, RandBigInt};
use num::One;

use rand::Rng;

use utils::modular::Inverse;
use utils::modular::Power;
use utils::primes::generate_prime;

pub struct SecretKeyExtra {
    p: BigUint,
    q: BigUint,
    dmp1: BigUint,
    dmq1: BigUint,
    qinv: BigUint
}

impl SecretKeyExtra {
    fn from_primes(p: BigUint, q: BigUint, d: &BigUint) -> Self {
        SecretKeyExtra {
            dmp1: d % (&p - BigUint::one()),
            dmq1: d % (&q - BigUint::one()),
            qinv: q.inverse(&p).unwrap(),
            p: p,
            q: q,
        }
    }
}

pub enum RSA {
    Public {
        /// Modulus
        n: BigUint,
        /// Public Exponent
        e: BigUint,
    },
    Private {
        /// Modulus
        n: BigUint,
        /// Private Exponent
        d: BigUint,
        extra: Option<SecretKeyExtra>,
    }
}

pub type KeyPair = (RSA, RSA);

impl RSA {
    pub fn keypair_from_primes<P, Q, E>(p: P, q: Q, e: E) -> KeyPair
        where P: Into<BigUint>, Q: Into<BigUint>, E: Into<BigUint> {
            let (p, q, e) = (p.into(), q.into(), e.into());

            let n = &p * &q;
            let phi_n = &n - (&p + &q - BigUint::one());

            let d = e.inverse(&phi_n).expect("e is irreversible in ring phi(pq) - error");

            let public = RSA::Public {
                n: n.clone(),
                e: e
            };
            let private = RSA::Private {
                n: n,
                extra: Some(SecretKeyExtra::from_primes(p, q, &d)),
                d: d,
            };

            (public, private)
        }

    pub fn generate_keypair<G, T>(mut rng: G, e: T, bits: usize) -> KeyPair
        where G: Rng + RandBigInt, T: Into<BigUint> {
            let e = e.into();

            let mut p = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            while (&p - BigUint::one()) % &e != BigUint::one() {
                p = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            }
            let mut q = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            while (&q - BigUint::one()) % &e != BigUint::one() {
                q = generate_prime(&mut rng, bits).expect("Cannot generate safe prime");
            }

            Self::keypair_from_primes(p, q, e)
        }

    pub fn is_public(&self) -> bool {
        match *self {
            RSA::Public { .. } => true,
            _ => false
        }
    }

    pub fn is_private(&self) -> bool {
        match *self {
            RSA::Private { .. } => true,
            _ => false
        }
    }

    pub fn crypt(&self, msg: &BigUint) -> BigUint {
        match *self {
            RSA::Private { ref n, ref d, ref extra } => crypt(msg, n, d, extra.as_ref()),
            RSA::Public { ref n, ref e } => crypt(msg, n, e, None),
        }
    }
}

fn crypt(msg: &BigUint,
         modulus: &BigUint,
         exp: &BigUint,
         extra: Option<&SecretKeyExtra>) -> BigUint {
    if let Some(ref extra) = extra {
        chinese_reminders_power(msg, extra)
    } else {
        msg.pow_mod(exp, modulus)
    }
}

fn chinese_remainders_power(c: &BigUint, extra: &SecretKeyExtra) -> BigUint {
    let mut m1 = c.pow_mod(&extra.dmp1, &extra.p);
    let m2 = c.pow_mod(&extra.dmq1, &extra.q);

    while m1 < m2 {
        m1 = m1 + &extra.p;
    }

    let h = (&extra.qinv * (m1 - &m2)) % &extra.p;

    m2 + h * &extra.q
}

#[cfg(test)]
mod tests {
    use super::RSA;

    use num::bigint::{ToBigUint};

    fn keys() -> (RSA, RSA) {
        RSA::keypair_from_primes(
            61.to_biguint().unwrap(),
            53.to_biguint().unwrap(),
            17.to_biguint().unwrap())
    }

    #[test]
    fn test_encryption() {
        let (public, _) = keys();
        let c = public.crypt(&65.to_biguint().unwrap());

        assert_eq!(c, 2790.to_biguint().unwrap())
    }

    #[test]
    fn test_decryption() {
        let (_, private) = keys();
        let m = private.crypt(&2790.to_biguint().unwrap());

        assert_eq!(m, 65.to_biguint().unwrap())
    }
}
<|endoftext|>


<commit_before>use super::{
    child_properties, function, general,
    general::{cfg_deprecated_string, version_condition_no_doc, version_condition_string},
    properties, signal, trait_impls,
};
use crate::{
    analysis::{self, ref_mode::RefMode, rust_type::RustType, special_functions::Type},
    case::CaseExt,
    env::Env,
    library::{self, Nullable},
    nameutil,
    traits::IntoString,
};
use std::io::{Result, Write};

pub fn generate(
    w: &mut dyn Write,
    env: &Env,
    analysis: &analysis::object::Info,
    generate_display_trait: bool,
) -> Result<()> {
    general::start_comments(w, &env.config)?;
    general::uses(w, env, &analysis.imports, analysis.version)?;

    general::define_object_type(
        w,
        env,
        &analysis.name,
        &analysis.c_type,
        analysis.c_class_type.as_deref(),
        &analysis.get_type,
        analysis.is_interface,
        &analysis.supertypes,
    )?;

    if need_generate_inherent(analysis) {
        writeln!(w)?;
        write!(w, "impl {} {{", analysis.name)?;
        for func_analysis in &analysis.constructors() {
            function::generate(
                w,
                env,
                func_analysis,
                Some(&analysis.specials),
                analysis.version,
                false,
                false,
                1,
            )?;
        }

        if !analysis.builder_properties.is_empty() {
            // generate builder method that returns the corresponding builder
            let builder_name = format!("{}Builder", analysis.name);
            writeln!(
                w,
                "
            // rustdoc-stripper-ignore-next
            /// Creates a new builder-style object to construct a [`{name}`].
            ///
            /// This method returns an instance of [`{builder_name}`] which can be used to create a [`{name}`].
            pub fn builder() -> {builder_name} {{
                {builder_name}::default()
            }}
        ",
                name = analysis.name,
                builder_name = builder_name
            )?;
        }

        if !need_generate_trait(analysis) {
            for func_analysis in &analysis.methods() {
                function::generate(
                    w,
                    env,
                    func_analysis,
                    Some(&analysis.specials),
                    analysis.version,
                    false,
                    false,
                    1,
                )?;
            }

            for property in &analysis.properties {
                properties::generate(w, env, property, false, false, 1)?;
            }

            for child_property in &analysis.child_properties {
                child_properties::generate(w, env, child_property, false, false, 1)?;
            }
        }

        for func_analysis in &analysis.functions() {
            function::generate(
                w,
                env,
                func_analysis,
                Some(&analysis.specials),
                analysis.version,
                false,
                false,
                1,
            )?;
        }

        if !need_generate_trait(analysis) {
            for signal_analysis in analysis
                .signals
                .iter()
                .chain(analysis.notify_signals.iter())
            {
                signal::generate(w, env, signal_analysis, false, false, 1)?;
            }
        }

        writeln!(w, "}}")?;

        general::declare_default_from_new(w, env, &analysis.name, &analysis.functions)?;
    }

    trait_impls::generate(
        w,
        env,
        &analysis.name,
        &analysis.functions,
        &analysis.specials,
        if need_generate_trait(analysis) {
            Some(&analysis.trait_name)
        } else {
            None
        },
        analysis.version,
        &None, // There is no need for #[cfg()] since it's applied on the whole file.
    )?;

    if !analysis.builder_properties.is_empty() {
        writeln!(w)?;
        generate_builder(w, env, analysis)?;
    }

    if analysis.concurrency != library::Concurrency::None {
        writeln!(w)?;
    }

    match analysis.concurrency {
        library::Concurrency::Send | library::Concurrency::SendSync => {
            writeln!(w, "unsafe impl Send for {} {{}}", analysis.name)?;
        }
        library::Concurrency::SendUnique => {
            if env.namespaces.is_glib_crate {
                writeln!(w, "unsafe impl ::SendUnique for {} {{", analysis.name)?;
            } else {
                writeln!(w, "unsafe impl glib::SendUnique for {} {{", analysis.name)?;
            }

            writeln!(w, "    fn is_unique(&self) -> bool {{")?;
            writeln!(w, "        self.ref_count() == 1")?;
            writeln!(w, "    }}")?;

            writeln!(w, "}}")?;
        }
        _ => (),
    }

    if let library::Concurrency::SendSync = analysis.concurrency {
        writeln!(w, "unsafe impl Sync for {} {{}}", analysis.name)?;
    }

    if !analysis.final_type {
        writeln!(w)?;
        writeln!(
            w,
            "pub const NONE_{}: Option<&{}> = None;",
            analysis.name.to_snake().to_uppercase(),
            analysis.name
        )?;
    }

    if need_generate_trait(analysis) {
        writeln!(w)?;
        generate_trait(w, env, analysis)?;
    }

    if generate_display_trait && !analysis.specials.has_trait(Type::Display) {
        writeln!(w, "\nimpl fmt::Display for {} {{", analysis.name,)?;
        // Generate Display trait implementation.
        writeln!(
            w,
            "\tfn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {{\n\
             \t\tf.write_str(\"{}\")\n\
             \t}}\n\
             }}",
            analysis.name
        )?;
    }

    Ok(())
}

// TODO: instead create a Vec<> inside the Builder instead of Options.
fn generate_builder(w: &mut dyn Write, env: &Env, analysis: &analysis::object::Info) -> Result<()> {
    let mut methods = vec![];
    let mut properties = vec![];
    writeln!(w, "#[derive(Clone, Default)]")?;
    writeln!(
        w,
        "// rustdoc-stripper-ignore-next
        /// A builder for generating a [`{}`].",
        analysis.name,
    )?;
    writeln!(w, "pub struct {}Builder {{", analysis.name)?;
    for property in &analysis.builder_properties {
        match RustType::try_new(env, property.typ) {
            Ok(type_string) => {
                let type_string = match type_string.as_str() {
                    s if nameutil::is_gstring(s) => "String",
                    "Vec<GString>" | "Vec<glib::GString>" | "Vec<crate::GString>" => "Vec<String>",
                    typ => typ,
                };
                let direction = if property.is_get {
                    library::ParameterDirection::In
                } else {
                    library::ParameterDirection::Out
                };
                let mut param_type = RustType::builder(env, property.typ)
                    .direction(direction)
                    .ref_mode(property.set_in_ref_mode)
                    .try_build()
                    .into_string();
                let (param_type_override, bounds, conversion) = match &param_type[..] {
                    "&str" => (None, String::new(), ".to_string()"),
                    "&[&str]" => (Some("Vec<String>".to_string()), String::new(), ""),
                    _ if !property.bounds.is_empty() => {
                        let (bounds, _) = function::bounds(&property.bounds, &[], false, false);
                        let alias =
                            property
                                .bounds
                                .get_parameter_bound(&property.name)
                                .map(|bound| {
                                    bound.full_type_parameter_reference(
                                        RefMode::ByRef,
                                        Nullable(false),
                                    )
                                });
                        (alias, bounds, ".clone().upcast()")
                    }
                    typ if typ.starts_with('&') => (None, String::new(), ".clone()"),
                    _ => (None, String::new(), ""),
                };
                if let Some(param_type_override) = param_type_override {
                    param_type = param_type_override.to_string();
                }
                let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
                let version_condition_string =
                    version_condition_string(env, property.version, false, 1);
                let deprecated_string =
                    cfg_deprecated_string(env, property.deprecated_version, false, 1);
                if let Some(ref version_condition_string) = version_condition_string {
                    writeln!(w, "{}", version_condition_string)?;
                }
                if let Some(ref deprecated_string) = deprecated_string {
                    writeln!(w, "{}", deprecated_string)?;
                }
                writeln!(w, "    {}: Option<{}>,", name, type_string)?;
                let version_prefix = version_condition_string
                    .map(|version| format!("{}\n", version))
                    .unwrap_or_default();

                let deprecation_prefix = deprecated_string
                    .map(|version| format!("{}\n", version))
                    .unwrap_or_default();

                methods.push(format!(
                    "\n{version_prefix}{deprecation_prefix}    pub fn {name}{bounds}(mut self, {name}: {param_type}) -> Self {{
        self.{name} = Some({name}{conversion});
        self
    }}",
                    version_prefix = version_prefix,
                    deprecation_prefix = deprecation_prefix,
                    param_type = param_type,
                    name = name,
                    conversion = conversion,
                    bounds = bounds
                ));
                properties.push(property);
            }
            Err(_) => writeln!(w, "    //{}: /*Unknown type*/,", property.name)?,
        }
    }
    writeln!(
        w,
        "}}

impl {name}Builder {{
    // rustdoc-stripper-ignore-next
    /// Create a new [`{name}Builder`].
    pub fn new() -> Self {{
        Self::default()
    }}
",
        name = analysis.name
    )?;

    writeln!(
        w,
        "
    // rustdoc-stripper-ignore-next
    /// Build the [`{name}`].
    pub fn build(self) -> {name} {{
        let mut properties: Vec<(&str, &dyn ToValue)> = vec![];",
        name = analysis.name
    )?;
    for property in &properties {
        let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
        version_condition_no_doc(w, env, property.version, false, 2)?;
        writeln!(
            w,
            "\
            if let Some(ref {field}) = self.{field} {{
                properties.push((\"{name}\", {field}));
            }}",
            name = property.name,
            field = name
        )?;
    }
    let glib_crate_name = if env.namespaces.is_glib_crate {
        "crate"
    } else {
        "glib"
    };

    // The split allows us to not have clippy::let_and_return lint disabled
    if let Some(code) = analysis.builder_postprocess.as_ref() {
        writeln!(
            w,
            r#"        let ret = {}::Object::new::<{}>(&properties)
                .expect("Failed to create an instance of {}");"#,
            glib_crate_name, analysis.name, analysis.name,
        )?;
        writeln!(w, "        {{\n            {}\n        }}", code)?;
        writeln!(w, "    ret\n    }}")?;
    } else {
        writeln!(
            w,
            r#"        {}::Object::new::<{}>(&properties)
                .expect("Failed to create an instance of {}")"#,
            glib_crate_name, analysis.name, analysis.name,
        )?;
        writeln!(w, "\n    }}")?;
    }

    for method in methods {
        writeln!(w, "{}", method)?;
    }
    writeln!(w, "}}")
}

fn generate_trait(w: &mut dyn Write, env: &Env, analysis: &analysis::object::Info) -> Result<()> {
    write!(w, "pub trait {}: 'static {{", analysis.trait_name)?;

    for func_analysis in &analysis.methods() {
        function::generate(
            w,
            env,
            func_analysis,
            Some(&analysis.specials),
            analysis.version,
            true,
            true,
            1,
        )?;
    }
    for property in &analysis.properties {
        properties::generate(w, env, property, true, true, 1)?;
    }
    for child_property in &analysis.child_properties {
        child_properties::generate(w, env, child_property, true, true, 1)?;
    }
    for signal_analysis in analysis
        .signals
        .iter()
        .chain(analysis.notify_signals.iter())
    {
        signal::generate(w, env, signal_analysis, true, true, 1)?;
    }
    writeln!(w, "}}")?;

    writeln!(w)?;
    write!(
        w,
        "impl<O: IsA<{}>> {} for O {{",
        analysis.name, analysis.trait_name,
    )?;

    for func_analysis in &analysis.methods() {
        function::generate(
            w,
            env,
            func_analysis,
            Some(&analysis.specials),
            analysis.version,
            true,
            false,
            1,
        )?;
    }
    for property in &analysis.properties {
        properties::generate(w, env, property, true, false, 1)?;
    }
    for child_property in &analysis.child_properties {
        child_properties::generate(w, env, child_property, true, false, 1)?;
    }
    for signal_analysis in analysis
        .signals
        .iter()
        .chain(analysis.notify_signals.iter())
    {
        signal::generate(w, env, signal_analysis, true, false, 1)?;
    }
    writeln!(w, "}}")?;

    Ok(())
}

fn need_generate_inherent(analysis: &analysis::object::Info) -> bool {
    analysis.has_constructors
        || analysis.has_functions
        || !need_generate_trait(analysis)
        || !analysis.builder_properties.is_empty()
}

fn need_generate_trait(analysis: &analysis::object::Info) -> bool {
    analysis.generate_trait
}

pub fn generate_reexports(
    env: &Env,
    analysis: &analysis::object::Info,
    module_name: &str,
    contents: &mut Vec<String>,
    traits: &mut Vec<String>,
) {
    let mut cfgs: Vec<String> = Vec::new();
    if let Some(cfg) = general::cfg_condition_string(&analysis.cfg_condition, false, 0) {
        cfgs.push(cfg);
    }
    if let Some(cfg) = general::version_condition_string(env, analysis.version, false, 0) {
        cfgs.push(cfg);
    }
    if let Some(cfg) = general::cfg_deprecated_string(env, analysis.deprecated_version, false, 0) {
        cfgs.push(cfg);
    }

    contents.push("".to_owned());
    contents.extend_from_slice(&cfgs);
    contents.push(format!("mod {};", module_name));
    contents.extend_from_slice(&cfgs);

    let none_type = if !analysis.final_type {
        format!(", NONE_{}", analysis.name.to_snake().to_uppercase())
    } else {
        String::new()
    };

    contents.push(format!(
        "pub use self::{}::{{{}{}}};",
        module_name, analysis.name, none_type
    ));
    if need_generate_trait(analysis) {
        for cfg in &cfgs {
            traits.push(format!("\t{}", cfg));
        }
        traits.push(format!(
            "\tpub use super::{}::{};",
            module_name, analysis.trait_name
        ));
    }

    if !analysis.builder_properties.is_empty() {
        contents.extend_from_slice(&cfgs);
        contents.push(format!(
            "pub use self::{}::{}Builder;",
            module_name, analysis.name
        ));
    }
}
<commit_msg>Fix a/an grammar by changing the sentences.<commit_after>use super::{
    child_properties, function, general,
    general::{cfg_deprecated_string, version_condition_no_doc, version_condition_string},
    properties, signal, trait_impls,
};
use crate::{
    analysis::{self, ref_mode::RefMode, rust_type::RustType, special_functions::Type},
    case::CaseExt,
    env::Env,
    library::{self, Nullable},
    nameutil,
    traits::IntoString,
};
use std::io::{Result, Write};

pub fn generate(
    w: &mut dyn Write,
    env: &Env,
    analysis: &analysis::object::Info,
    generate_display_trait: bool,
) -> Result<()> {
    general::start_comments(w, &env.config)?;
    general::uses(w, env, &analysis.imports, analysis.version)?;

    general::define_object_type(
        w,
        env,
        &analysis.name,
        &analysis.c_type,
        analysis.c_class_type.as_deref(),
        &analysis.get_type,
        analysis.is_interface,
        &analysis.supertypes,
    )?;

    if need_generate_inherent(analysis) {
        writeln!(w)?;
        write!(w, "impl {} {{", analysis.name)?;
        for func_analysis in &analysis.constructors() {
            function::generate(
                w,
                env,
                func_analysis,
                Some(&analysis.specials),
                analysis.version,
                false,
                false,
                1,
            )?;
        }

        if !analysis.builder_properties.is_empty() {
            // generate builder method that returns the corresponding builder
            let builder_name = format!("{}Builder", analysis.name);
            writeln!(
                w,
                "
            // rustdoc-stripper-ignore-next
            /// Creates a new builder-pattern struct instance to construct [`{name}`] objects.
            ///
            /// This method returns an instance of [`{builder_name}`] which can be used to create [`{name}`] objects.
            pub fn builder() -> {builder_name} {{
                {builder_name}::default()
            }}
        ",
                name = analysis.name,
                builder_name = builder_name
            )?;
        }

        if !need_generate_trait(analysis) {
            for func_analysis in &analysis.methods() {
                function::generate(
                    w,
                    env,
                    func_analysis,
                    Some(&analysis.specials),
                    analysis.version,
                    false,
                    false,
                    1,
                )?;
            }

            for property in &analysis.properties {
                properties::generate(w, env, property, false, false, 1)?;
            }

            for child_property in &analysis.child_properties {
                child_properties::generate(w, env, child_property, false, false, 1)?;
            }
        }

        for func_analysis in &analysis.functions() {
            function::generate(
                w,
                env,
                func_analysis,
                Some(&analysis.specials),
                analysis.version,
                false,
                false,
                1,
            )?;
        }

        if !need_generate_trait(analysis) {
            for signal_analysis in analysis
                .signals
                .iter()
                .chain(analysis.notify_signals.iter())
            {
                signal::generate(w, env, signal_analysis, false, false, 1)?;
            }
        }

        writeln!(w, "}}")?;

        general::declare_default_from_new(w, env, &analysis.name, &analysis.functions)?;
    }

    trait_impls::generate(
        w,
        env,
        &analysis.name,
        &analysis.functions,
        &analysis.specials,
        if need_generate_trait(analysis) {
            Some(&analysis.trait_name)
        } else {
            None
        },
        analysis.version,
        &None, // There is no need for #[cfg()] since it's applied on the whole file.
    )?;

    if !analysis.builder_properties.is_empty() {
        writeln!(w)?;
        generate_builder(w, env, analysis)?;
    }

    if analysis.concurrency != library::Concurrency::None {
        writeln!(w)?;
    }

    match analysis.concurrency {
        library::Concurrency::Send | library::Concurrency::SendSync => {
            writeln!(w, "unsafe impl Send for {} {{}}", analysis.name)?;
        }
        library::Concurrency::SendUnique => {
            if env.namespaces.is_glib_crate {
                writeln!(w, "unsafe impl ::SendUnique for {} {{", analysis.name)?;
            } else {
                writeln!(w, "unsafe impl glib::SendUnique for {} {{", analysis.name)?;
            }

            writeln!(w, "    fn is_unique(&self) -> bool {{")?;
            writeln!(w, "        self.ref_count() == 1")?;
            writeln!(w, "    }}")?;

            writeln!(w, "}}")?;
        }
        _ => (),
    }

    if let library::Concurrency::SendSync = analysis.concurrency {
        writeln!(w, "unsafe impl Sync for {} {{}}", analysis.name)?;
    }

    if !analysis.final_type {
        writeln!(w)?;
        writeln!(
            w,
            "pub const NONE_{}: Option<&{}> = None;",
            analysis.name.to_snake().to_uppercase(),
            analysis.name
        )?;
    }

    if need_generate_trait(analysis) {
        writeln!(w)?;
        generate_trait(w, env, analysis)?;
    }

    if generate_display_trait && !analysis.specials.has_trait(Type::Display) {
        writeln!(w, "\nimpl fmt::Display for {} {{", analysis.name,)?;
        // Generate Display trait implementation.
        writeln!(
            w,
            "\tfn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {{\n\
             \t\tf.write_str(\"{}\")\n\
             \t}}\n\
             }}",
            analysis.name
        )?;
    }

    Ok(())
}

// TODO: instead create a Vec<> inside the Builder instead of Options.
fn generate_builder(w: &mut dyn Write, env: &Env, analysis: &analysis::object::Info) -> Result<()> {
    let mut methods = vec![];
    let mut properties = vec![];
    writeln!(w, "#[derive(Clone, Default)]")?;
    writeln!(
        w,
        "// rustdoc-stripper-ignore-next
        /// A [builder-pattern] type to construct [`{}`] objects.
        /// [builder-pattern]: https://doc.rust-lang.org/1.0.0/style/ownership/builders.html",
        analysis.name,
    )?;
    writeln!(w, "pub struct {}Builder {{", analysis.name)?;
    for property in &analysis.builder_properties {
        match RustType::try_new(env, property.typ) {
            Ok(type_string) => {
                let type_string = match type_string.as_str() {
                    s if nameutil::is_gstring(s) => "String",
                    "Vec<GString>" | "Vec<glib::GString>" | "Vec<crate::GString>" => "Vec<String>",
                    typ => typ,
                };
                let direction = if property.is_get {
                    library::ParameterDirection::In
                } else {
                    library::ParameterDirection::Out
                };
                let mut param_type = RustType::builder(env, property.typ)
                    .direction(direction)
                    .ref_mode(property.set_in_ref_mode)
                    .try_build()
                    .into_string();
                let (param_type_override, bounds, conversion) = match &param_type[..] {
                    "&str" => (None, String::new(), ".to_string()"),
                    "&[&str]" => (Some("Vec<String>".to_string()), String::new(), ""),
                    _ if !property.bounds.is_empty() => {
                        let (bounds, _) = function::bounds(&property.bounds, &[], false, false);
                        let alias =
                            property
                                .bounds
                                .get_parameter_bound(&property.name)
                                .map(|bound| {
                                    bound.full_type_parameter_reference(
                                        RefMode::ByRef,
                                        Nullable(false),
                                    )
                                });
                        (alias, bounds, ".clone().upcast()")
                    }
                    typ if typ.starts_with('&') => (None, String::new(), ".clone()"),
                    _ => (None, String::new(), ""),
                };
                if let Some(param_type_override) = param_type_override {
                    param_type = param_type_override.to_string();
                }
                let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
                let version_condition_string =
                    version_condition_string(env, property.version, false, 1);
                let deprecated_string =
                    cfg_deprecated_string(env, property.deprecated_version, false, 1);
                if let Some(ref version_condition_string) = version_condition_string {
                    writeln!(w, "{}", version_condition_string)?;
                }
                if let Some(ref deprecated_string) = deprecated_string {
                    writeln!(w, "{}", deprecated_string)?;
                }
                writeln!(w, "    {}: Option<{}>,", name, type_string)?;
                let version_prefix = version_condition_string
                    .map(|version| format!("{}\n", version))
                    .unwrap_or_default();

                let deprecation_prefix = deprecated_string
                    .map(|version| format!("{}\n", version))
                    .unwrap_or_default();

                methods.push(format!(
                    "\n{version_prefix}{deprecation_prefix}    pub fn {name}{bounds}(mut self, {name}: {param_type}) -> Self {{
        self.{name} = Some({name}{conversion});
        self
    }}",
                    version_prefix = version_prefix,
                    deprecation_prefix = deprecation_prefix,
                    param_type = param_type,
                    name = name,
                    conversion = conversion,
                    bounds = bounds
                ));
                properties.push(property);
            }
            Err(_) => writeln!(w, "    //{}: /*Unknown type*/,", property.name)?,
        }
    }
    writeln!(
        w,
        "}}

impl {name}Builder {{
    // rustdoc-stripper-ignore-next
    /// Create a new [`{name}Builder`].
    pub fn new() -> Self {{
        Self::default()
    }}
",
        name = analysis.name
    )?;

    writeln!(
        w,
        "
    // rustdoc-stripper-ignore-next
    /// Build the [`{name}`].
    pub fn build(self) -> {name} {{
        let mut properties: Vec<(&str, &dyn ToValue)> = vec![];",
        name = analysis.name
    )?;
    for property in &properties {
        let name = nameutil::mangle_keywords(nameutil::signal_to_snake(&property.name));
        version_condition_no_doc(w, env, property.version, false, 2)?;
        writeln!(
            w,
            "\
            if let Some(ref {field}) = self.{field} {{
                properties.push((\"{name}\", {field}));
            }}",
            name = property.name,
            field = name
        )?;
    }
    let glib_crate_name = if env.namespaces.is_glib_crate {
        "crate"
    } else {
        "glib"
    };

    // The split allows us to not have clippy::let_and_return lint disabled
    if let Some(code) = analysis.builder_postprocess.as_ref() {
        writeln!(
            w,
            r#"        let ret = {}::Object::new::<{}>(&properties)
                .expect("Failed to create an instance of {}");"#,
            glib_crate_name, analysis.name, analysis.name,
        )?;
        writeln!(w, "        {{\n            {}\n        }}", code)?;
        writeln!(w, "    ret\n    }}")?;
    } else {
        writeln!(
            w,
            r#"        {}::Object::new::<{}>(&properties)
                .expect("Failed to create an instance of {}")"#,
            glib_crate_name, analysis.name, analysis.name,
        )?;
        writeln!(w, "\n    }}")?;
    }

    for method in methods {
        writeln!(w, "{}", method)?;
    }
    writeln!(w, "}}")
}

fn generate_trait(w: &mut dyn Write, env: &Env, analysis: &analysis::object::Info) -> Result<()> {
    write!(w, "pub trait {}: 'static {{", analysis.trait_name)?;

    for func_analysis in &analysis.methods() {
        function::generate(
            w,
            env,
            func_analysis,
            Some(&analysis.specials),
            analysis.version,
            true,
            true,
            1,
        )?;
    }
    for property in &analysis.properties {
        properties::generate(w, env, property, true, true, 1)?;
    }
    for child_property in &analysis.child_properties {
        child_properties::generate(w, env, child_property, true, true, 1)?;
    }
    for signal_analysis in analysis
        .signals
        .iter()
        .chain(analysis.notify_signals.iter())
    {
        signal::generate(w, env, signal_analysis, true, true, 1)?;
    }
    writeln!(w, "}}")?;

    writeln!(w)?;
    write!(
        w,
        "impl<O: IsA<{}>> {} for O {{",
        analysis.name, analysis.trait_name,
    )?;

    for func_analysis in &analysis.methods() {
        function::generate(
            w,
            env,
            func_analysis,
            Some(&analysis.specials),
            analysis.version,
            true,
            false,
            1,
        )?;
    }
    for property in &analysis.properties {
        properties::generate(w, env, property, true, false, 1)?;
    }
    for child_property in &analysis.child_properties {
        child_properties::generate(w, env, child_property, true, false, 1)?;
    }
    for signal_analysis in analysis
        .signals
        .iter()
        .chain(analysis.notify_signals.iter())
    {
        signal::generate(w, env, signal_analysis, true, false, 1)?;
    }
    writeln!(w, "}}")?;

    Ok(())
}

fn need_generate_inherent(analysis: &analysis::object::Info) -> bool {
    analysis.has_constructors
        || analysis.has_functions
        || !need_generate_trait(analysis)
        || !analysis.builder_properties.is_empty()
}

fn need_generate_trait(analysis: &analysis::object::Info) -> bool {
    analysis.generate_trait
}

pub fn generate_reexports(
    env: &Env,
    analysis: &analysis::object::Info,
    module_name: &str,
    contents: &mut Vec<String>,
    traits: &mut Vec<String>,
) {
    let mut cfgs: Vec<String> = Vec::new();
    if let Some(cfg) = general::cfg_condition_string(&analysis.cfg_condition, false, 0) {
        cfgs.push(cfg);
    }
    if let Some(cfg) = general::version_condition_string(env, analysis.version, false, 0) {
        cfgs.push(cfg);
    }
    if let Some(cfg) = general::cfg_deprecated_string(env, analysis.deprecated_version, false, 0) {
        cfgs.push(cfg);
    }

    contents.push("".to_owned());
    contents.extend_from_slice(&cfgs);
    contents.push(format!("mod {};", module_name));
    contents.extend_from_slice(&cfgs);

    let none_type = if !analysis.final_type {
        format!(", NONE_{}", analysis.name.to_snake().to_uppercase())
    } else {
        String::new()
    };

    contents.push(format!(
        "pub use self::{}::{{{}{}}};",
        module_name, analysis.name, none_type
    ));
    if need_generate_trait(analysis) {
        for cfg in &cfgs {
            traits.push(format!("\t{}", cfg));
        }
        traits.push(format!(
            "\tpub use super::{}::{};",
            module_name, analysis.trait_name
        ));
    }

    if !analysis.builder_properties.is_empty() {
        contents.extend_from_slice(&cfgs);
        contents.push(format!(
            "pub use self::{}::{}Builder;",
            module_name, analysis.name
        ));
    }
}
<|endoftext|>


<commit_before>use std::sync::{Arc, RwLock};
use std::option::Option;
use std::collections::{HashMap, HashSet};

use rocksdb::DB;
use mio::{self, EventLoop};
use protobuf;
use uuid::Uuid;

use proto::raft_serverpb::{RaftMessage, StoreIdent};
use proto::raft_cmdpb::{self as cmd, RaftCommandRequest, RaftCommandResponse};
use raftserver::{Result, other};
use proto::metapb;
use super::{Sender, Msg};
use super::keys;
use super::engine::Retriever;
use super::config::Config;
use super::peer::{Peer, PendingCmd};
use super::msg::Callback;
use super::cmd_resp::{self, bind_uuid};
use super::transport::Transport;

pub struct Store<T: Transport> {
    cfg: Config,
    ident: StoreIdent,
    engine: Arc<DB>,
    sender: Sender,
    event_loop: Option<EventLoop<Store<T>>>,

    peers: HashMap<u64, Peer>,
    pending_raft_groups: HashSet<u64>,

    trans: Arc<RwLock<T>>,
}

impl<T: Transport> Store<T> {
    pub fn new(cfg: Config, engine: Arc<DB>, trans: Arc<RwLock<T>>) -> Result<Store<T>> {
        let ident: StoreIdent = try!(load_store_ident(engine.as_ref()).and_then(|res| {
            match res {
                None => Err(other("store must be bootstrapped first")),
                Some(ident) => Ok(ident),
            }
        }));

        let event_loop = try!(EventLoop::new());
        let sender = Sender::new(event_loop.channel());

        Ok(Store {
            cfg: cfg,
            ident: ident,
            engine: engine,
            sender: sender,
            event_loop: Some(event_loop),
            peers: HashMap::new(),
            pending_raft_groups: HashSet::new(),
            trans: trans,
        })
    }

    // Do something before store runs.
    fn prepare(&mut self) -> Result<()> {
        // Scan region meta to get saved regions.
        let start_key = keys::REGION_META_MIN_KEY;
        let end_key = keys::REGION_META_MAX_KEY;
        let engine = self.engine.clone();
        try!(engine.scan(start_key,
                         end_key,
                         &mut |key, value| -> Result<bool> {
                             let (_, suffix) = try!(keys::decode_region_meta_key(key));
                             if suffix != keys::REGION_INFO_SUFFIX {
                                 return Ok(true);
                             }

                             let region = try!(protobuf::parse_from_bytes::<metapb::Region>(value));
                             let peer = try!(Peer::create(self, region));
                             // TODO: check duplicated peer id later?
                             self.peers.insert(peer.get_peer_id(), peer);
                             Ok(true)
                         }));

        Ok(())
    }

    pub fn run(&mut self) -> Result<()> {
        try!(self.prepare());

        let mut event_loop = self.event_loop.take().unwrap();
        self.register_raft_base_tick(&mut event_loop);
        try!(event_loop.run(self));
        Ok(())
    }

    pub fn get_sender(&self) -> Sender {
        self.sender.clone()
    }

    pub fn get_engine(&self) -> Arc<DB> {
        self.engine.clone()
    }

    pub fn get_ident(&self) -> &StoreIdent {
        &self.ident
    }

    pub fn get_store_id(&self) -> u64 {
        self.ident.get_store_id()
    }

    pub fn get_config(&self) -> &Config {
        &self.cfg
    }

    fn register_raft_base_tick(&self, event_loop: &mut EventLoop<Self>) {
        // If we register raft base tick failed, the whole raft can't run correctly,
        // TODO: shutdown the store?
        if let Err(e) = register_timer(event_loop,
                                       Msg::RaftBaseTick,
                                       self.cfg.raft_base_tick_interval) {
            error!("register raft base tick err: {:?}", e);
        };
    }

    fn handle_raft_base_tick(&mut self, event_loop: &mut EventLoop<Self>) {
        for (region_id, peer) in &mut self.peers {
            peer.raft_group.tick();
            self.pending_raft_groups.insert(*region_id);
        }

        self.register_raft_base_tick(event_loop);
    }

    // Clippy doesn't allow hash_map contains_key followed by insert, and suggests
    // using entry().or_insert() instead, but we can't use this because creating peer
    // may fail, so we allow map_entry.
    #[allow(map_entry)]
    fn handle_raft_message(&mut self, msg: RaftMessage) -> Result<()> {
        let region_id = msg.get_region_id();
        let from_peer = msg.get_from_peer();
        let to_peer = msg.get_to_peer();
        debug!("handle raft message for region {}, from {} to {}",
               region_id,
               from_peer.get_peer_id(),
               to_peer.get_peer_id());

        {
            let mut trans = self.trans.write().unwrap();
            trans.cache_peer(from_peer.get_peer_id(), from_peer.clone());
            trans.cache_peer(to_peer.get_peer_id(), to_peer.clone());
        }

        if !self.peers.contains_key(&region_id) {
            let peer = try!(Peer::replicate(self, region_id, to_peer.get_peer_id()));
            self.peers.insert(region_id, peer);
        }

        let mut peer = self.peers.get_mut(&region_id).unwrap();

        try!(peer.raft_group.step(msg.get_message().clone()));

        // Add in pending raft group for later handing ready.
        self.pending_raft_groups.insert(region_id);

        Ok(())
    }

    fn handle_raft_ready(&mut self) -> Result<()> {
        let ids = self.pending_raft_groups.drain();

        for region_id in ids {
            if let Some(peer) = self.peers.get_mut(&region_id) {
                if let Err(e) = peer.handle_raft_ready(&self.trans) {
                    // TODO: should we panic here or shutdown the store?
                    error!("handle raft ready err: {:?}", e);
                    return Err(e);
                }
            }
        }

        Ok(())
    }

    fn propose_raft_command(&mut self, msg: RaftCommandRequest, cb: Callback) -> Result<()> {
        let mut resp: RaftCommandResponse;
        let uuid: Uuid = match Uuid::from_bytes(msg.get_header().get_uuid()) {
            None => {
                resp = cmd_resp::message_error("missing request uuid");
                return cb.call_box((resp,));
            }
            Some(uuid) => uuid,
        };

        if msg.has_status_request() {
            // For status commands, we handle it here directly.
            resp = match self.execute_status_command(msg) {
                Err(e) => cmd_resp::message_error(format!("{:?}", e)),
                Ok(resp) => resp,
            };
            bind_uuid(&mut resp, uuid);
            return cb.call_box((resp,));
        }

        let region_id = msg.get_header().get_region_id();
        let mut peer = match self.peers.get_mut(&region_id) {
            None => {
                resp = cmd_resp::region_not_found_error(region_id);
                bind_uuid(&mut resp, uuid);
                return cb.call_box((resp,));
            }
            Some(peer) => peer,
        };

        if !peer.is_leader() {
            let trans = self.trans.read().unwrap();
            resp = cmd_resp::not_leader_error(region_id, trans.get_peer(peer.get_leader()));
            bind_uuid(&mut resp, uuid);
            return cb.call_box((resp,));
        }

        let peer_id = msg.get_header().get_peer().get_peer_id();
        if peer.get_peer_id() != peer_id {
            resp = cmd_resp::message_error(format!("mismatch peer id {} != {}",
                                                   peer.get_peer_id(),
                                                   peer_id));
            bind_uuid(&mut resp, uuid);
            return cb.call_box((resp,));
        }

        if peer.pending_cmds.contains_key(&uuid) {
            resp = cmd_resp::message_error(format!("duplicated uuid {:?}", uuid));
            bind_uuid(&mut resp, uuid);
            return cb.call_box((resp,));
        }

        // Notice:
        // Here means the peer is leader, it can still step down to follower later,
        // but it doesn't matter, if the peer is not leader, the proposing command
        // log entry can't be committed.

        // TODO: support handing read-only commands later.
        // for read-only, if we don't care stale read, we can
        // execute these commands immediately in leader.

        let mut pending_cmd = PendingCmd {
            uuid: uuid,
            cb: None,
            cmd: Some(msg),
        };

        match peer.propose_pending_cmd(&mut pending_cmd) {
            Err(e) => {
                resp = cmd_resp::message_error(format!("{:?}", e));
                bind_uuid(&mut resp, uuid);
                return cb.call_box((resp,));
            }
            Ok(()) => (), // nothing to do.
        };

        // Keep the callback in pending_cmd so that we can call it later
        // after command applied.
        pending_cmd.cb = Some(cb);
        peer.pending_cmds.insert(uuid, pending_cmd);

        self.pending_raft_groups.insert(region_id);

        // TODO: add timeout, if the command is not applied after timeout,
        // we will call the callback with timeout error.

        Ok(())
    }
}

fn load_store_ident<T: Retriever>(r: &T) -> Result<Option<StoreIdent>> {
    let ident = try!(r.get_msg::<StoreIdent>(&keys::store_ident_key()));

    Ok(ident)
}

fn register_timer<T: Transport>(event_loop: &mut EventLoop<Store<T>>,
                                msg: Msg,
                                delay: u64)
                                -> Result<mio::Timeout> {
    // TODO: now mio TimerError doesn't implement Error trait,
    // so we can't use `try!` directly.
    event_loop.timeout_ms(msg, delay).map_err(|e| other(format!("register timer err: {:?}", e)))
}

impl<T: Transport> mio::Handler for Store<T> {
    type Timeout = Msg;
    type Message = Msg;

    fn notify(&mut self, event_loop: &mut EventLoop<Self>, msg: Msg) {
        match msg {
            Msg::RaftMessage(data) => {
                if let Err(e) = self.handle_raft_message(data) {
                    error!("handle raft message err: {:?}", e);
                }
            }
            Msg::RaftCommand{request, callback} => {
                if let Err(e) = self.propose_raft_command(request, callback) {
                    error!("propose raft command err: {:?}", e);
                }
            }
            Msg::Quit => {
                info!("receive quit message");
                event_loop.shutdown();
            }
            _ => panic!("invalid notify msg type {:?}", msg),
        }
    }

    fn timeout(&mut self, event_loop: &mut EventLoop<Self>, timeout: Msg) {
        match timeout {
            Msg::RaftBaseTick => self.handle_raft_base_tick(event_loop),
            _ => panic!("invalid timeout msg type {:?}", timeout),
        }
    }

    fn tick(&mut self, event_loop: &mut EventLoop<Self>) {
        if !event_loop.is_running() {
            // TODO: do some close here.
            return;
        }

        // We handle raft ready in event loop.
        if let Err(e) = self.handle_raft_ready() {
            // TODO: should we panic here or shutdown the store?
            error!("handle raft ready err: {:?}", e);
        }
    }
}

impl<T: Transport> Store<T> {
    // Handle status commands here, separate the logic, maybe we can move it
    // to another file later.
    // Unlike other commands (write or admin), status commands only show current
    // store status, so no need to handle it in raft group.
    fn execute_status_command(&mut self,
                              request: RaftCommandRequest)
                              -> Result<RaftCommandResponse> {
        let cmd_type = request.get_status_request().get_cmd_type();
        let mut response = try!(match cmd_type {
            cmd::StatusCommandType::RegionLeader => self.execute_region_leader(request),
            e => Err(other(format!("unsupported status command type {:?}", e))),
        });
        response.set_cmd_type(cmd_type);

        let mut resp = RaftCommandResponse::new();
        resp.set_status_response(response);
        Ok(resp)
    }

    fn execute_region_leader(&mut self,
                             request: RaftCommandRequest)
                             -> Result<cmd::StatusResponse> {
        let region_id = request.get_header().get_region_id();
        let peer = match self.peers.get_mut(&region_id) {
            None => return Err(other(format!("region {} not found", region_id))),
            Some(peer) => peer,
        };

        let mut resp = cmd::StatusResponse::new();
        let trans = self.trans.read().unwrap();
        if let Some(leader) = trans.get_peer(peer.get_leader()) {
            resp.mut_region_leader().set_leader(leader);
        }

        Ok(resp)
    }
}
<commit_msg>raftserver: update event loop config and fix bug.<commit_after>use std::sync::{Arc, RwLock};
use std::option::Option;
use std::collections::{HashMap, HashSet};

use rocksdb::DB;
use mio::{self, EventLoop, EventLoopConfig};
use protobuf;
use uuid::Uuid;

use proto::raft_serverpb::{RaftMessage, StoreIdent};
use proto::raft_cmdpb::{self as cmd, RaftCommandRequest, RaftCommandResponse};
use raftserver::{Result, other};
use proto::metapb;
use super::{Sender, Msg};
use super::keys;
use super::engine::Retriever;
use super::config::Config;
use super::peer::{Peer, PendingCmd};
use super::msg::Callback;
use super::cmd_resp::{self, bind_uuid};
use super::transport::Transport;

pub struct Store<T: Transport> {
    cfg: Config,
    ident: StoreIdent,
    engine: Arc<DB>,
    sender: Sender,
    event_loop: Option<EventLoop<Store<T>>>,

    peers: HashMap<u64, Peer>,
    pending_raft_groups: HashSet<u64>,

    trans: Arc<RwLock<T>>,
}

impl<T: Transport> Store<T> {
    pub fn new(cfg: Config, engine: Arc<DB>, trans: Arc<RwLock<T>>) -> Result<Store<T>> {
        let ident: StoreIdent = try!(load_store_ident(engine.as_ref()).and_then(|res| {
            match res {
                None => Err(other("store must be bootstrapped first")),
                Some(ident) => Ok(ident),
            }
        }));

        // We use base raft tick as the event loop timer tick.
        let mut event_cfg = EventLoopConfig::new();
        event_cfg.timer_tick_ms(cfg.raft_base_tick_interval);
        let event_loop = try!(EventLoop::configured(event_cfg));

        let sender = Sender::new(event_loop.channel());

        Ok(Store {
            cfg: cfg,
            ident: ident,
            engine: engine,
            sender: sender,
            event_loop: Some(event_loop),
            peers: HashMap::new(),
            pending_raft_groups: HashSet::new(),
            trans: trans,
        })
    }

    // Do something before store runs.
    fn prepare(&mut self) -> Result<()> {
        // Scan region meta to get saved regions.
        let start_key = keys::REGION_META_MIN_KEY;
        let end_key = keys::REGION_META_MAX_KEY;
        let engine = self.engine.clone();
        try!(engine.scan(start_key,
                         end_key,
                         &mut |key, value| -> Result<bool> {
                             let (_, suffix) = try!(keys::decode_region_meta_key(key));
                             if suffix != keys::REGION_INFO_SUFFIX {
                                 return Ok(true);
                             }

                             let region = try!(protobuf::parse_from_bytes::<metapb::Region>(value));
                             let region_id = region.get_region_id();
                             let peer = try!(Peer::create(self, region));
                             // TODO: check duplicated region id later?
                             self.peers.insert(region_id, peer);
                             Ok(true)
                         }));

        Ok(())
    }

    pub fn run(&mut self) -> Result<()> {
        try!(self.prepare());

        let mut event_loop = self.event_loop.take().unwrap();
        self.register_raft_base_tick(&mut event_loop);
        try!(event_loop.run(self));
        Ok(())
    }

    pub fn get_sender(&self) -> Sender {
        self.sender.clone()
    }

    pub fn get_engine(&self) -> Arc<DB> {
        self.engine.clone()
    }

    pub fn get_ident(&self) -> &StoreIdent {
        &self.ident
    }

    pub fn get_store_id(&self) -> u64 {
        self.ident.get_store_id()
    }

    pub fn get_config(&self) -> &Config {
        &self.cfg
    }

    fn register_raft_base_tick(&self, event_loop: &mut EventLoop<Self>) {
        // If we register raft base tick failed, the whole raft can't run correctly,
        // TODO: shutdown the store?
        if let Err(e) = register_timer(event_loop,
                                       Msg::RaftBaseTick,
                                       self.cfg.raft_base_tick_interval) {
            error!("register raft base tick err: {:?}", e);
        };
    }

    fn handle_raft_base_tick(&mut self, event_loop: &mut EventLoop<Self>) {
        for (region_id, peer) in &mut self.peers {
            peer.raft_group.tick();
            self.pending_raft_groups.insert(*region_id);
        }

        self.register_raft_base_tick(event_loop);
    }

    // Clippy doesn't allow hash_map contains_key followed by insert, and suggests
    // using entry().or_insert() instead, but we can't use this because creating peer
    // may fail, so we allow map_entry.
    #[allow(map_entry)]
    fn handle_raft_message(&mut self, msg: RaftMessage) -> Result<()> {
        let region_id = msg.get_region_id();
        let from_peer = msg.get_from_peer();
        let to_peer = msg.get_to_peer();
        debug!("handle raft message for region {}, from {} to {}",
               region_id,
               from_peer.get_peer_id(),
               to_peer.get_peer_id());

        {
            let mut trans = self.trans.write().unwrap();
            trans.cache_peer(from_peer.get_peer_id(), from_peer.clone());
            trans.cache_peer(to_peer.get_peer_id(), to_peer.clone());
        }

        if !self.peers.contains_key(&region_id) {
            let peer = try!(Peer::replicate(self, region_id, to_peer.get_peer_id()));
            self.peers.insert(region_id, peer);
        }

        let mut peer = self.peers.get_mut(&region_id).unwrap();

        try!(peer.raft_group.step(msg.get_message().clone()));

        // Add in pending raft group for later handing ready.
        self.pending_raft_groups.insert(region_id);

        Ok(())
    }

    fn handle_raft_ready(&mut self) -> Result<()> {
        let ids = self.pending_raft_groups.drain();

        for region_id in ids {
            if let Some(peer) = self.peers.get_mut(&region_id) {
                if let Err(e) = peer.handle_raft_ready(&self.trans) {
                    // TODO: should we panic here or shutdown the store?
                    error!("handle raft ready err: {:?}", e);
                    return Err(e);
                }
            }
        }

        Ok(())
    }

    fn propose_raft_command(&mut self, msg: RaftCommandRequest, cb: Callback) -> Result<()> {
        let mut resp: RaftCommandResponse;
        let uuid: Uuid = match Uuid::from_bytes(msg.get_header().get_uuid()) {
            None => {
                resp = cmd_resp::message_error("missing request uuid");
                return cb.call_box((resp,));
            }
            Some(uuid) => uuid,
        };

        if msg.has_status_request() {
            // For status commands, we handle it here directly.
            resp = match self.execute_status_command(msg) {
                Err(e) => cmd_resp::message_error(format!("{:?}", e)),
                Ok(resp) => resp,
            };
            bind_uuid(&mut resp, uuid);
            return cb.call_box((resp,));
        }

        let region_id = msg.get_header().get_region_id();
        let mut peer = match self.peers.get_mut(&region_id) {
            None => {
                resp = cmd_resp::region_not_found_error(region_id);
                bind_uuid(&mut resp, uuid);
                return cb.call_box((resp,));
            }
            Some(peer) => peer,
        };

        if !peer.is_leader() {
            let trans = self.trans.read().unwrap();
            resp = cmd_resp::not_leader_error(region_id, trans.get_peer(peer.get_leader()));
            bind_uuid(&mut resp, uuid);
            return cb.call_box((resp,));
        }

        let peer_id = msg.get_header().get_peer().get_peer_id();
        if peer.get_peer_id() != peer_id {
            resp = cmd_resp::message_error(format!("mismatch peer id {} != {}",
                                                   peer.get_peer_id(),
                                                   peer_id));
            bind_uuid(&mut resp, uuid);
            return cb.call_box((resp,));
        }

        if peer.pending_cmds.contains_key(&uuid) {
            resp = cmd_resp::message_error(format!("duplicated uuid {:?}", uuid));
            bind_uuid(&mut resp, uuid);
            return cb.call_box((resp,));
        }

        // Notice:
        // Here means the peer is leader, it can still step down to follower later,
        // but it doesn't matter, if the peer is not leader, the proposing command
        // log entry can't be committed.

        // TODO: support handing read-only commands later.
        // for read-only, if we don't care stale read, we can
        // execute these commands immediately in leader.

        let mut pending_cmd = PendingCmd {
            uuid: uuid,
            cb: None,
            cmd: Some(msg),
        };

        match peer.propose_pending_cmd(&mut pending_cmd) {
            Err(e) => {
                resp = cmd_resp::message_error(format!("{:?}", e));
                bind_uuid(&mut resp, uuid);
                return cb.call_box((resp,));
            }
            Ok(()) => (), // nothing to do.
        };

        // Keep the callback in pending_cmd so that we can call it later
        // after command applied.
        pending_cmd.cb = Some(cb);
        peer.pending_cmds.insert(uuid, pending_cmd);

        self.pending_raft_groups.insert(region_id);

        // TODO: add timeout, if the command is not applied after timeout,
        // we will call the callback with timeout error.

        Ok(())
    }
}

fn load_store_ident<T: Retriever>(r: &T) -> Result<Option<StoreIdent>> {
    let ident = try!(r.get_msg::<StoreIdent>(&keys::store_ident_key()));

    Ok(ident)
}

fn register_timer<T: Transport>(event_loop: &mut EventLoop<Store<T>>,
                                msg: Msg,
                                delay: u64)
                                -> Result<mio::Timeout> {
    // TODO: now mio TimerError doesn't implement Error trait,
    // so we can't use `try!` directly.
    event_loop.timeout_ms(msg, delay).map_err(|e| other(format!("register timer err: {:?}", e)))
}

impl<T: Transport> mio::Handler for Store<T> {
    type Timeout = Msg;
    type Message = Msg;

    fn notify(&mut self, event_loop: &mut EventLoop<Self>, msg: Msg) {
        match msg {
            Msg::RaftMessage(data) => {
                if let Err(e) = self.handle_raft_message(data) {
                    error!("handle raft message err: {:?}", e);
                }
            }
            Msg::RaftCommand{request, callback} => {
                if let Err(e) = self.propose_raft_command(request, callback) {
                    error!("propose raft command err: {:?}", e);
                }
            }
            Msg::Quit => {
                info!("receive quit message");
                event_loop.shutdown();
            }
            _ => panic!("invalid notify msg type {:?}", msg),
        }
    }

    fn timeout(&mut self, event_loop: &mut EventLoop<Self>, timeout: Msg) {
        match timeout {
            Msg::RaftBaseTick => self.handle_raft_base_tick(event_loop),
            _ => panic!("invalid timeout msg type {:?}", timeout),
        }
    }

    fn tick(&mut self, event_loop: &mut EventLoop<Self>) {
        if !event_loop.is_running() {
            // TODO: do some close here.
            return;
        }

        // We handle raft ready in event loop.
        if let Err(e) = self.handle_raft_ready() {
            // TODO: should we panic here or shutdown the store?
            error!("handle raft ready err: {:?}", e);
        }
    }
}

impl<T: Transport> Store<T> {
    // Handle status commands here, separate the logic, maybe we can move it
    // to another file later.
    // Unlike other commands (write or admin), status commands only show current
    // store status, so no need to handle it in raft group.
    fn execute_status_command(&mut self,
                              request: RaftCommandRequest)
                              -> Result<RaftCommandResponse> {
        let cmd_type = request.get_status_request().get_cmd_type();
        let mut response = try!(match cmd_type {
            cmd::StatusCommandType::RegionLeader => self.execute_region_leader(request),
            e => Err(other(format!("unsupported status command type {:?}", e))),
        });
        response.set_cmd_type(cmd_type);

        let mut resp = RaftCommandResponse::new();
        resp.set_status_response(response);
        Ok(resp)
    }

    fn execute_region_leader(&mut self,
                             request: RaftCommandRequest)
                             -> Result<cmd::StatusResponse> {
        let region_id = request.get_header().get_region_id();
        let peer = match self.peers.get_mut(&region_id) {
            None => return Err(other(format!("region {} not found", region_id))),
            Some(peer) => peer,
        };

        let mut resp = cmd::StatusResponse::new();
        let trans = self.trans.read().unwrap();
        if let Some(leader) = trans.get_peer(peer.get_leader()) {
            resp.mut_region_leader().set_leader(leader);
        }

        Ok(resp)
    }
}
<|endoftext|>


<commit_before>use core::convert::TryFrom;
use core::mem;

use crate::memory::Frame;
use crate::paging::{ActivePageTable, Page, PhysicalAddress, VirtualAddress};
use crate::paging::entry::EntryFlags;

/// RSDP
#[derive(Copy, Clone, Debug)]
#[repr(packed)]
pub struct RSDP {
    signature: [u8; 8],
    checksum: u8,
    oemid: [u8; 6],
    revision: u8,
    rsdt_address: u32,
    length: u32,
    xsdt_address: u64,
    extended_checksum: u8,
    reserved: [u8; 3]
}

impl RSDP {
    fn is_acpi_1_0(&self) -> bool {
        self.revision == 0
    }
    fn is_acpi_2_0(&self) -> bool {
        self.revision == 2
    }
    fn get_already_supplied_rsdps(area: &[u8]) -> Option<RSDP> {
        // the bootloader has already checked all the checksums for us, but we still need to
        // double-check.
        struct Iter<'a> {
            buf: &'a [u8],
        }
        impl<'a> Iterator for Iter<'a> {
            type Item = &'a [u8];

            fn next(&mut self) -> Option<Self::Item> {
                if self.buf.len() < 4 { return None }

                let length_bytes = <[u8; 4]>::try_from(&self.buf[..4]).ok()?;
                let length = u32::from_ne_bytes(length_bytes) as usize;

                if (4 + length as usize) > self.buf.len() { return None }

                let buf = &self.buf[4..4 + length];
                self.buf = &self.buf[4 + length..];

                Some(buf)
            }
        }
        fn slice_to_rsdp(slice: &[u8]) -> Option<&RSDP> {
            let ptr = slice.as_ptr() as usize;

            if slice.len() >= mem::size_of::<RSDP>() && ptr & (!0x7) == ptr {
                let rsdp = unsafe { &*(slice.as_ptr() as *const RSDP) };
                // TODO: Validate
                Some(rsdp)
            } else { None }
        }

        // first, find an RSDP for ACPI 2.0
        if let Some(rsdp_2_0) = (Iter { buf: area }.filter_map(slice_to_rsdp).find(|rsdp| rsdp.is_acpi_2_0())) {
            return Some(*rsdp_2_0);
        }

        // secondly, find an RSDP for ACPI 1.0
        if let Some(rsdp_1_0) = (Iter { buf: area }.filter_map(slice_to_rsdp).find(|rsdp| rsdp.is_acpi_1_0())) {
            return Some(*rsdp_1_0);
        }

        None
    }
    pub fn get_rsdp(active_table: &mut ActivePageTable, already_supplied_rsdps: Option<(u64, u64)>) -> Option<RSDP> {
        if let Some((base, size)) = already_supplied_rsdps {
            let area = unsafe { core::slice::from_raw_parts(base as usize as *const u8, size as usize) };
            Self::get_already_supplied_rsdps(area).or_else(|| Self::get_rsdp_by_searching(active_table))
        } else {
            Self::get_rsdp_by_searching(active_table)
        }
    }
    /// Search for the RSDP
    pub fn get_rsdp_by_searching(active_table: &mut ActivePageTable) -> Option<RSDP> {
        let start_addr = 0xE_0000;
        let end_addr = 0xF_FFFF;

        // Map all of the ACPI RSDP space
        {
            let start_frame = Frame::containing_address(PhysicalAddress::new(start_addr));
            let end_frame = Frame::containing_address(PhysicalAddress::new(end_addr));
            for frame in Frame::range_inclusive(start_frame, end_frame) {
                let page = Page::containing_address(VirtualAddress::new(frame.start_address().get()));
                let result = active_table.map_to(page, frame, EntryFlags::PRESENT | EntryFlags::NO_EXECUTE);
                result.flush(active_table);
            }
        }

        RSDP::search(start_addr, end_addr)
    }

    fn search(start_addr: usize, end_addr: usize) -> Option<RSDP> {
        for i in 0 .. (end_addr + 1 - start_addr)/16 {
            let rsdp = unsafe { &*((start_addr + i * 16) as *const RSDP) };
            if &rsdp.signature == b"RSD PTR " {
                return Some(*rsdp);
            }
        }
        None
    }

    /// Get the RSDT or XSDT address
    pub fn sdt_address(&self) -> usize {
        if self.revision >= 2 {
            self.xsdt_address as usize
        } else {
            self.rsdt_address as usize
        }
    }
}
<commit_msg>Only require alignment of 4 bytes for rsdp<commit_after>use core::convert::TryFrom;
use core::mem;

use crate::memory::Frame;
use crate::paging::{ActivePageTable, Page, PhysicalAddress, VirtualAddress};
use crate::paging::entry::EntryFlags;

/// RSDP
#[derive(Copy, Clone, Debug)]
#[repr(packed)]
pub struct RSDP {
    signature: [u8; 8],
    checksum: u8,
    oemid: [u8; 6],
    revision: u8,
    rsdt_address: u32,
    length: u32,
    xsdt_address: u64,
    extended_checksum: u8,
    reserved: [u8; 3]
}

impl RSDP {
    fn is_acpi_1_0(&self) -> bool {
        self.revision == 0
    }
    fn is_acpi_2_0(&self) -> bool {
        self.revision == 2
    }
    fn get_already_supplied_rsdps(area: &[u8]) -> Option<RSDP> {
        // the bootloader has already checked all the checksums for us, but we still need to
        // double-check.
        struct Iter<'a> {
            buf: &'a [u8],
        }
        impl<'a> Iterator for Iter<'a> {
            type Item = &'a [u8];

            fn next(&mut self) -> Option<Self::Item> {
                if self.buf.len() < 4 { return None }

                let length_bytes = <[u8; 4]>::try_from(&self.buf[..4]).ok()?;
                let length = u32::from_ne_bytes(length_bytes) as usize;

                if (4 + length as usize) > self.buf.len() { return None }

                let buf = &self.buf[4..4 + length];
                self.buf = &self.buf[4 + length..];

                Some(buf)
            }
        }
        fn slice_to_rsdp(slice: &[u8]) -> Option<&RSDP> {
            let ptr = slice.as_ptr() as usize;

            if slice.len() >= mem::size_of::<RSDP>() && ptr & (!0x3) == ptr {
                let rsdp = unsafe { &*(slice.as_ptr() as *const RSDP) };
                // TODO: Validate
                Some(rsdp)
            } else { None }
        }

        // first, find an RSDP for ACPI 2.0
        if let Some(rsdp_2_0) = (Iter { buf: area }.filter_map(slice_to_rsdp).find(|rsdp| rsdp.is_acpi_2_0())) {
            return Some(*rsdp_2_0);
        }

        // secondly, find an RSDP for ACPI 1.0
        if let Some(rsdp_1_0) = (Iter { buf: area }.filter_map(slice_to_rsdp).find(|rsdp| rsdp.is_acpi_1_0())) {
            return Some(*rsdp_1_0);
        }

        None
    }
    pub fn get_rsdp(active_table: &mut ActivePageTable, already_supplied_rsdps: Option<(u64, u64)>) -> Option<RSDP> {
        if let Some((base, size)) = already_supplied_rsdps {
            let area = unsafe { core::slice::from_raw_parts(base as usize as *const u8, size as usize) };
            Self::get_already_supplied_rsdps(area).or_else(|| Self::get_rsdp_by_searching(active_table))
        } else {
            Self::get_rsdp_by_searching(active_table)
        }
    }
    /// Search for the RSDP
    pub fn get_rsdp_by_searching(active_table: &mut ActivePageTable) -> Option<RSDP> {
        let start_addr = 0xE_0000;
        let end_addr = 0xF_FFFF;

        // Map all of the ACPI RSDP space
        {
            let start_frame = Frame::containing_address(PhysicalAddress::new(start_addr));
            let end_frame = Frame::containing_address(PhysicalAddress::new(end_addr));
            for frame in Frame::range_inclusive(start_frame, end_frame) {
                let page = Page::containing_address(VirtualAddress::new(frame.start_address().get()));
                let result = active_table.map_to(page, frame, EntryFlags::PRESENT | EntryFlags::NO_EXECUTE);
                result.flush(active_table);
            }
        }

        RSDP::search(start_addr, end_addr)
    }

    fn search(start_addr: usize, end_addr: usize) -> Option<RSDP> {
        for i in 0 .. (end_addr + 1 - start_addr)/16 {
            let rsdp = unsafe { &*((start_addr + i * 16) as *const RSDP) };
            if &rsdp.signature == b"RSD PTR " {
                return Some(*rsdp);
            }
        }
        None
    }

    /// Get the RSDT or XSDT address
    pub fn sdt_address(&self) -> usize {
        if self.revision >= 2 {
            self.xsdt_address as usize
        } else {
            self.rsdt_address as usize
        }
    }
}
<|endoftext|>