<commit_before>use std::io::prelude::*;

use crate::core::{resolver, Resolve, ResolveVersion, Workspace};
use crate::util::errors::{CargoResult, CargoResultExt};
use crate::util::toml as cargo_toml;
use crate::util::Filesystem;

pub fn load_pkg_lockfile(ws: &Workspace<'_>) -> CargoResult<Option<Resolve>> {
    if !ws.root().join("Cargo.lock").exists() {
        return Ok(None);
    }

    let root = Filesystem::new(ws.root().to_path_buf());
    let mut f = root.open_ro("Cargo.lock", ws.config(), "Cargo.lock file")?;

    let mut s = String::new();
    f.read_to_string(&mut s)
        .chain_err(|| format!("failed to read file: {}", f.path().display()))?;

    let resolve = (|| -> CargoResult<Option<Resolve>> {
        let resolve: toml::Value = cargo_toml::parse(&s, f.path(), ws.config())?;
        let v: resolver::EncodableResolve = resolve.try_into()?;
        Ok(Some(v.into_resolve(&s, ws)?))
    })()
    .chain_err(|| format!("failed to parse lock file at: {}", f.path().display()))?;
    Ok(resolve)
}

/// Generate a toml String of Cargo.lock from a Resolve.
pub fn resolve_to_string(ws: &Workspace<'_>, resolve: &mut Resolve) -> CargoResult<String> {
    let (_orig, out, _ws_root) = resolve_to_string_orig(ws, resolve)?;
    Ok(out)
}

pub fn write_pkg_lockfile(ws: &Workspace<'_>, resolve: &mut Resolve) -> CargoResult<()> {
    let (orig, mut out, ws_root) = resolve_to_string_orig(ws, resolve)?;

    // If the lock file contents haven't changed so don't rewrite it. This is
    // helpful on read-only filesystems.
    if let Some(orig) = &orig {
        if are_equal_lockfiles(orig, &out, ws) {
            return Ok(());
        }
    }

    if !ws.config().lock_update_allowed() {
        if ws.config().offline() {
            anyhow::bail!("can't update in the offline mode");
        }

        let flag = if ws.config().network_allowed() {
            "--locked"
        } else {
            "--frozen"
        };
        anyhow::bail!(
            "the lock file {} needs to be updated but {} was passed to prevent this\n\
             If you want to try to generate the lock file without accessing the network, \
             use the --offline flag.",
            ws.root().to_path_buf().join("Cargo.lock").display(),
            flag
        );
    }

    // While we're updating the lock file anyway go ahead and update its
    // encoding to whatever the latest default is. That way we can slowly roll
    // out lock file updates as they're otherwise already updated, and changes
    // which don't touch dependencies won't seemingly spuriously update the lock
    // file.
    if resolve.version() < ResolveVersion::default() {
        resolve.set_version(ResolveVersion::default());
        out = serialize_resolve(resolve, orig.as_deref());
    }

    // Ok, if that didn't work just write it out
    ws_root
        .open_rw("Cargo.lock", ws.config(), "Cargo.lock file")
        .and_then(|mut f| {
            f.file().set_len(0)?;
            f.write_all(out.as_bytes())?;
            Ok(())
        })
        .chain_err(|| format!("failed to write {}", ws.root().join("Cargo.lock").display()))?;
    Ok(())
}

fn resolve_to_string_orig(
    ws: &Workspace<'_>,
    resolve: &mut Resolve,
) -> CargoResult<(Option<String>, String, Filesystem)> {
    // Load the original lock file if it exists.
    let ws_root = Filesystem::new(ws.root().to_path_buf());
    let orig = ws_root.open_ro("Cargo.lock", ws.config(), "Cargo.lock file");
    let orig = orig.and_then(|mut f| {
        let mut s = String::new();
        f.read_to_string(&mut s)?;
        Ok(s)
    });
    let out = serialize_resolve(resolve, orig.as_ref().ok().map(|s| &**s));
    Ok((orig.ok(), out, ws_root))
}

fn serialize_resolve(resolve: &Resolve, orig: Option<&str>) -> String {
    let toml = toml::Value::try_from(resolve).unwrap();

    let mut out = String::new();

    // At the start of the file we notify the reader that the file is generated.
    // Specifically Phabricator ignores files containing "@generated", so we use that.
    let marker_line = "# This file is automatically @generated by Cargo.";
    let extra_line = "# It is not intended for manual editing.";
    out.push_str(marker_line);
    out.push('\n');
    out.push_str(extra_line);
    out.push('\n');
    // and preserve any other top comments
    if let Some(orig) = orig {
        let mut comments = orig.lines().take_while(|line| line.starts_with('#'));
        if let Some(first) = comments.next() {
            if first != marker_line {
                out.push_str(first);
                out.push('\n');
            }
            if let Some(second) = comments.next() {
                if second != extra_line {
                    out.push_str(second);
                    out.push('\n');
                }
                for line in comments {
                    out.push_str(line);
                    out.push('\n');
                }
            }
        }
    }

    if let Some(version) = toml.get("version") {
        out.push_str(&format!("version = {}\n\n", version));
    }

    let deps = toml["package"].as_array().unwrap();
    for dep in deps {
        let dep = dep.as_table().unwrap();

        out.push_str("[[package]]\n");
        emit_package(dep, &mut out);
    }

    if let Some(patch) = toml.get("patch") {
        let list = patch["unused"].as_array().unwrap();
        for entry in list {
            out.push_str("[[patch.unused]]\n");
            emit_package(entry.as_table().unwrap(), &mut out);
            out.push_str("\n");
        }
    }

    if let Some(meta) = toml.get("metadata") {
        out.push_str("[metadata]\n");
        out.push_str(&meta.to_string());
    }

    // Historical versions of Cargo in the old format accidentally left trailing
    // blank newlines at the end of files, so we just leave that as-is. For all
    // encodings going forward, though, we want to be sure that our encoded lock
    // file doesn't contain any trailing newlines so trim out the extra if
    // necessary.
    if resolve.version() >= ResolveVersion::V2 {
        while out.ends_with("\n\n") {
            out.pop();
        }
    }
    out
}

fn are_equal_lockfiles(orig: &str, current: &str, ws: &Workspace<'_>) -> bool {
    // If we want to try and avoid updating the lock file, parse both and
    // compare them; since this is somewhat expensive, don't do it in the
    // common case where we can update lock files.
    if !ws.config().lock_update_allowed() {
        let res: CargoResult<bool> = (|| {
            let old: resolver::EncodableResolve = toml::from_str(orig)?;
            let new: resolver::EncodableResolve = toml::from_str(current)?;
            Ok(old.into_resolve(orig, ws)? == new.into_resolve(current, ws)?)
        })();
        if let Ok(true) = res {
            return true;
        }
    }

    let mut orig_iter = orig.lines();
    let mut current_iter = current.lines();
    loop {
        match (orig_iter.next(), current_iter.next()) {
            (Some(o), Some(c)) => {
                if o != c {
                    return false;
                }
            }
            (Some(_), None) => return false,
            (None, Some(_)) => return false,
            (None, None) => return true,
        }
    }
}

fn emit_package(dep: &toml::value::Table, out: &mut String) {
    out.push_str(&format!("name = {}\n", &dep["name"]));
    out.push_str(&format!("version = {}\n", &dep["version"]));

    if dep.contains_key("source") {
        out.push_str(&format!("source = {}\n", &dep["source"]));
    }
    if dep.contains_key("checksum") {
        out.push_str(&format!("checksum = {}\n", &dep["checksum"]));
    }

    if let Some(s) = dep.get("dependencies") {
        let slice = s.as_array().unwrap();

        if !slice.is_empty() {
            out.push_str("dependencies = [\n");

            for child in slice.iter() {
                out.push_str(&format!(" {},\n", child));
            }

            out.push_str("]\n");
        }
        out.push_str("\n");
    } else if dep.contains_key("replace") {
        out.push_str(&format!("replace = {}\n\n", &dep["replace"]));
    }
}
<commit_msg>refactor: compare lockfile use Iterator.eq<commit_after>use std::io::prelude::*;

use crate::core::{resolver, Resolve, ResolveVersion, Workspace};
use crate::util::errors::{CargoResult, CargoResultExt};
use crate::util::toml as cargo_toml;
use crate::util::Filesystem;

pub fn load_pkg_lockfile(ws: &Workspace<'_>) -> CargoResult<Option<Resolve>> {
    if !ws.root().join("Cargo.lock").exists() {
        return Ok(None);
    }

    let root = Filesystem::new(ws.root().to_path_buf());
    let mut f = root.open_ro("Cargo.lock", ws.config(), "Cargo.lock file")?;

    let mut s = String::new();
    f.read_to_string(&mut s)
        .chain_err(|| format!("failed to read file: {}", f.path().display()))?;

    let resolve = (|| -> CargoResult<Option<Resolve>> {
        let resolve: toml::Value = cargo_toml::parse(&s, f.path(), ws.config())?;
        let v: resolver::EncodableResolve = resolve.try_into()?;
        Ok(Some(v.into_resolve(&s, ws)?))
    })()
    .chain_err(|| format!("failed to parse lock file at: {}", f.path().display()))?;
    Ok(resolve)
}

/// Generate a toml String of Cargo.lock from a Resolve.
pub fn resolve_to_string(ws: &Workspace<'_>, resolve: &mut Resolve) -> CargoResult<String> {
    let (_orig, out, _ws_root) = resolve_to_string_orig(ws, resolve)?;
    Ok(out)
}

pub fn write_pkg_lockfile(ws: &Workspace<'_>, resolve: &mut Resolve) -> CargoResult<()> {
    let (orig, mut out, ws_root) = resolve_to_string_orig(ws, resolve)?;

    // If the lock file contents haven't changed so don't rewrite it. This is
    // helpful on read-only filesystems.
    if let Some(orig) = &orig {
        if are_equal_lockfiles(orig, &out, ws) {
            return Ok(());
        }
    }

    if !ws.config().lock_update_allowed() {
        if ws.config().offline() {
            anyhow::bail!("can't update in the offline mode");
        }

        let flag = if ws.config().network_allowed() {
            "--locked"
        } else {
            "--frozen"
        };
        anyhow::bail!(
            "the lock file {} needs to be updated but {} was passed to prevent this\n\
             If you want to try to generate the lock file without accessing the network, \
             use the --offline flag.",
            ws.root().to_path_buf().join("Cargo.lock").display(),
            flag
        );
    }

    // While we're updating the lock file anyway go ahead and update its
    // encoding to whatever the latest default is. That way we can slowly roll
    // out lock file updates as they're otherwise already updated, and changes
    // which don't touch dependencies won't seemingly spuriously update the lock
    // file.
    if resolve.version() < ResolveVersion::default() {
        resolve.set_version(ResolveVersion::default());
        out = serialize_resolve(resolve, orig.as_deref());
    }

    // Ok, if that didn't work just write it out
    ws_root
        .open_rw("Cargo.lock", ws.config(), "Cargo.lock file")
        .and_then(|mut f| {
            f.file().set_len(0)?;
            f.write_all(out.as_bytes())?;
            Ok(())
        })
        .chain_err(|| format!("failed to write {}", ws.root().join("Cargo.lock").display()))?;
    Ok(())
}

fn resolve_to_string_orig(
    ws: &Workspace<'_>,
    resolve: &mut Resolve,
) -> CargoResult<(Option<String>, String, Filesystem)> {
    // Load the original lock file if it exists.
    let ws_root = Filesystem::new(ws.root().to_path_buf());
    let orig = ws_root.open_ro("Cargo.lock", ws.config(), "Cargo.lock file");
    let orig = orig.and_then(|mut f| {
        let mut s = String::new();
        f.read_to_string(&mut s)?;
        Ok(s)
    });
    let out = serialize_resolve(resolve, orig.as_ref().ok().map(|s| &**s));
    Ok((orig.ok(), out, ws_root))
}

fn serialize_resolve(resolve: &Resolve, orig: Option<&str>) -> String {
    let toml = toml::Value::try_from(resolve).unwrap();

    let mut out = String::new();

    // At the start of the file we notify the reader that the file is generated.
    // Specifically Phabricator ignores files containing "@generated", so we use that.
    let marker_line = "# This file is automatically @generated by Cargo.";
    let extra_line = "# It is not intended for manual editing.";
    out.push_str(marker_line);
    out.push('\n');
    out.push_str(extra_line);
    out.push('\n');
    // and preserve any other top comments
    if let Some(orig) = orig {
        let mut comments = orig.lines().take_while(|line| line.starts_with('#'));
        if let Some(first) = comments.next() {
            if first != marker_line {
                out.push_str(first);
                out.push('\n');
            }
            if let Some(second) = comments.next() {
                if second != extra_line {
                    out.push_str(second);
                    out.push('\n');
                }
                for line in comments {
                    out.push_str(line);
                    out.push('\n');
                }
            }
        }
    }

    if let Some(version) = toml.get("version") {
        out.push_str(&format!("version = {}\n\n", version));
    }

    let deps = toml["package"].as_array().unwrap();
    for dep in deps {
        let dep = dep.as_table().unwrap();

        out.push_str("[[package]]\n");
        emit_package(dep, &mut out);
    }

    if let Some(patch) = toml.get("patch") {
        let list = patch["unused"].as_array().unwrap();
        for entry in list {
            out.push_str("[[patch.unused]]\n");
            emit_package(entry.as_table().unwrap(), &mut out);
            out.push_str("\n");
        }
    }

    if let Some(meta) = toml.get("metadata") {
        out.push_str("[metadata]\n");
        out.push_str(&meta.to_string());
    }

    // Historical versions of Cargo in the old format accidentally left trailing
    // blank newlines at the end of files, so we just leave that as-is. For all
    // encodings going forward, though, we want to be sure that our encoded lock
    // file doesn't contain any trailing newlines so trim out the extra if
    // necessary.
    if resolve.version() >= ResolveVersion::V2 {
        while out.ends_with("\n\n") {
            out.pop();
        }
    }
    out
}

fn are_equal_lockfiles(orig: &str, current: &str, ws: &Workspace<'_>) -> bool {
    // If we want to try and avoid updating the lock file, parse both and
    // compare them; since this is somewhat expensive, don't do it in the
    // common case where we can update lock files.
    if !ws.config().lock_update_allowed() {
        let res: CargoResult<bool> = (|| {
            let old: resolver::EncodableResolve = toml::from_str(orig)?;
            let new: resolver::EncodableResolve = toml::from_str(current)?;
            Ok(old.into_resolve(orig, ws)? == new.into_resolve(current, ws)?)
        })();
        if let Ok(true) = res {
            return true;
        }
    }

    orig.lines().eq(current.lines())
}

fn emit_package(dep: &toml::value::Table, out: &mut String) {
    out.push_str(&format!("name = {}\n", &dep["name"]));
    out.push_str(&format!("version = {}\n", &dep["version"]));

    if dep.contains_key("source") {
        out.push_str(&format!("source = {}\n", &dep["source"]));
    }
    if dep.contains_key("checksum") {
        out.push_str(&format!("checksum = {}\n", &dep["checksum"]));
    }

    if let Some(s) = dep.get("dependencies") {
        let slice = s.as_array().unwrap();

        if !slice.is_empty() {
            out.push_str("dependencies = [\n");

            for child in slice.iter() {
                out.push_str(&format!(" {},\n", child));
            }

            out.push_str("]\n");
        }
        out.push_str("\n");
    } else if dep.contains_key("replace") {
        out.push_str(&format!("replace = {}\n\n", &dep["replace"]));
    }
}
<|endoftext|>


<commit_before><commit_msg>Fix type-inference in sink::unfold() by specifying more of its types (#2311)<commit_after><|endoftext|>


<commit_before>// CITA
// Copyright 2016-2017 Cryptape Technologies LLC.

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

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

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

use byteorder::{BigEndian, ByteOrder};
use config;
use libproto::communication;
use parking_lot::RwLock;
use protobuf::Message;
use std::convert::AsRef;
use std::io::prelude::*;
use std::net::SocketAddr;
use std::net::TcpStream;
use std::sync::Arc;
use std::sync::mpsc::Receiver;
use std::thread;
use std::time::Duration;

const TIMEOUT: u64 = 15;

pub struct Connection {
    pub id_card: u32,
    pub peers_pair: Vec<(u32, SocketAddr, Arc<RwLock<Option<TcpStream>>>)>,
}

impl Connection {
    pub fn new(config: &config::NetConfig) -> Self {
        let peers = config.peers.as_ref().unwrap();
        let id_card = config.id_card.unwrap();
        let mut peers_pair = Vec::default();
        for peer in peers.iter() {
            let id_card: u32 = peer.id_card.unwrap();
            let addr = format!("{}:{}", peer.ip.clone().unwrap(), peer.port.unwrap());
            let addr = addr.parse::<SocketAddr>().unwrap();
            peers_pair.push((id_card, addr, Arc::new(RwLock::new(None))));
        }
        Connection { id_card, peers_pair }
    }
}

pub fn do_connect(con: &Connection) {
    for &(_, addr, ref stream) in &con.peers_pair {
        let stream_lock = stream.clone();
        thread::spawn(move || loop {
                          {
                              let stream_opt = &mut *stream_lock.as_ref().write();
                              if stream_opt.is_none() {
                                  trace!("connet {:?}", addr);
                                  let stream = TcpStream::connect(addr).ok();
                                  *stream_opt = stream;
                              }

                              let mut need_reconnect = false;
                              if let Some(ref mut stream) = stream_opt.as_mut() {
                                  trace!("handshake with {:?}!", addr);
                                  let mut header = [0; 8];
                                  BigEndian::write_u64(&mut header, 0xDEADBEEF00000000 as u64);
                                  let res = stream.write(&header);
                                  if res.is_err() {
                                      warn!("handshake with {:?} error!", addr);
                                      need_reconnect = true;
                                  }
                              }

                              if need_reconnect {
                                  *stream_opt = None;
                              }
                          }

                          let ten_sec = Duration::from_millis(TIMEOUT * 1000);
                          thread::sleep(ten_sec);
                          trace!("after sleep retry connect {:?}!", addr);
                      });
    }
}

pub fn broadcast(con: &Connection, mut msg: communication::Message) {
    let origin = msg.get_origin();
    let operate = msg.get_operate();
    msg.set_origin(con.id_card);

    trace!("broadcast msg {:?} ", msg);
    let msg = msg.write_to_bytes().unwrap();
    let request_id = 0xDEADBEEF00000000 + msg.len();
    let mut encoded_request_id = [0; 8];
    BigEndian::write_u64(&mut encoded_request_id, request_id as u64);
    let mut buf = Vec::new();
    buf.extend(&encoded_request_id);
    buf.extend(msg);
    let send_msg = move |stream: &Arc<RwLock<Option<TcpStream>>>| {
        let streams_lock = stream.clone();
        let stream_opt = &mut (*streams_lock.as_ref().write());
        if let Some(ref mut stream) = stream_opt.as_mut() {
            let _ = stream.write(&buf);
        }
    };
    let mut peers = vec![];
    for &(id_card, _, ref stream) in &con.peers_pair {
        if is_send(id_card, origin, operate) {
            peers.push(id_card);
            send_msg(stream);
        }
    }

    info!("{:?} broadcast msg to nodes {:?} {:?}", con.id_card, operate, peers);
}

pub fn is_send(id_card: u32, origin: u32, operate: communication::OperateType) -> bool {
    operate == communication::OperateType::BROADCAST || (operate == communication::OperateType::SINGLE && id_card == origin) || (operate == communication::OperateType::SUBTRACT && origin != id_card)
}

pub fn start_client(con: Arc<Connection>, rx: Receiver<communication::Message>) {
    thread::spawn(move || {
                      info!("start client!");
                      loop {
                          let msg = rx.recv().unwrap();
                          broadcast(&con, msg);
                      }
                  });
}

#[cfg(test)]
mod test {
    use super::is_send;
    use libproto::communication;
    #[test]
    fn is_seng_mag() {
        assert!(is_send(0, 0, communication::OperateType::BROADCAST));
        assert!(is_send(0, 1, communication::OperateType::BROADCAST));

        assert!(is_send(0, 0, communication::OperateType::SINGLE));
        assert!(!is_send(0, 1, communication::OperateType::SINGLE));

        assert!(!is_send(0, 0, communication::OperateType::SUBTRACT));
        assert!(is_send(0, 1, communication::OperateType::SUBTRACT));
    }
}
<commit_msg>修改单节点没有peers，可以正常启动 (#15)<commit_after>// CITA
// Copyright 2016-2017 Cryptape Technologies LLC.

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

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

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

use byteorder::{BigEndian, ByteOrder};
use config;
use libproto::communication;
use parking_lot::RwLock;
use protobuf::Message;
use std::convert::AsRef;
use std::io::prelude::*;
use std::net::SocketAddr;
use std::net::TcpStream;
use std::sync::Arc;
use std::sync::mpsc::Receiver;
use std::thread;
use std::time::Duration;

const TIMEOUT: u64 = 15;

pub struct Connection {
    pub id_card: u32,
    pub peers_pair: Vec<(u32, SocketAddr, Arc<RwLock<Option<TcpStream>>>)>,
}

impl Connection {
    pub fn new(config: &config::NetConfig) -> Self {
        let id_card = config.id_card.unwrap();
        let mut peers_pair = Vec::default();
        match config.peers.as_ref() {
            Some(peers) => {
                for peer in peers.iter() {
                    let id_card: u32 = peer.id_card.unwrap();
                    let addr = format!("{}:{}", peer.ip.clone().unwrap(), peer.port.unwrap());
                    let addr = addr.parse::<SocketAddr>().unwrap();
                    peers_pair.push((id_card, addr, Arc::new(RwLock::new(None))));
                }
            }
            None => (),

        }

        Connection { id_card, peers_pair }
    }
}

pub fn do_connect(con: &Connection) {
    for &(_, addr, ref stream) in &con.peers_pair {
        let stream_lock = stream.clone();
        thread::spawn(move || loop {
                          {
                              let stream_opt = &mut *stream_lock.as_ref().write();
                              if stream_opt.is_none() {
                                  trace!("connet {:?}", addr);
                                  let stream = TcpStream::connect(addr).ok();
                                  *stream_opt = stream;
                              }

                              let mut need_reconnect = false;
                              if let Some(ref mut stream) = stream_opt.as_mut() {
                                  trace!("handshake with {:?}!", addr);
                                  let mut header = [0; 8];
                                  BigEndian::write_u64(&mut header, 0xDEADBEEF00000000 as u64);
                                  let res = stream.write(&header);
                                  if res.is_err() {
                                      warn!("handshake with {:?} error!", addr);
                                      need_reconnect = true;
                                  }
                              }

                              if need_reconnect {
                                  *stream_opt = None;
                              }
                          }

                          let ten_sec = Duration::from_millis(TIMEOUT * 1000);
                          thread::sleep(ten_sec);
                          trace!("after sleep retry connect {:?}!", addr);
                      });
    }
}

pub fn broadcast(con: &Connection, mut msg: communication::Message) {
    let origin = msg.get_origin();
    let operate = msg.get_operate();
    msg.set_origin(con.id_card);

    trace!("broadcast msg {:?} ", msg);
    let msg = msg.write_to_bytes().unwrap();
    let request_id = 0xDEADBEEF00000000 + msg.len();
    let mut encoded_request_id = [0; 8];
    BigEndian::write_u64(&mut encoded_request_id, request_id as u64);
    let mut buf = Vec::new();
    buf.extend(&encoded_request_id);
    buf.extend(msg);
    let send_msg = move |stream: &Arc<RwLock<Option<TcpStream>>>| {
        let streams_lock = stream.clone();
        let stream_opt = &mut (*streams_lock.as_ref().write());
        if let Some(ref mut stream) = stream_opt.as_mut() {
            let _ = stream.write(&buf);
        }
    };
    let mut peers = vec![];
    for &(id_card, _, ref stream) in &con.peers_pair {
        if is_send(id_card, origin, operate) {
            peers.push(id_card);
            send_msg(stream);
        }
    }

    info!("{:?} broadcast msg to nodes {:?} {:?}", con.id_card, operate, peers);
}

pub fn is_send(id_card: u32, origin: u32, operate: communication::OperateType) -> bool {
    operate == communication::OperateType::BROADCAST || (operate == communication::OperateType::SINGLE && id_card == origin) || (operate == communication::OperateType::SUBTRACT && origin != id_card)
}

pub fn start_client(con: Arc<Connection>, rx: Receiver<communication::Message>) {
    thread::spawn(move || {
                      info!("start client!");
                      loop {
                          let msg = rx.recv().unwrap();
                          broadcast(&con, msg);
                      }
                  });
}

#[cfg(test)]
mod test {
    use super::is_send;
    use libproto::communication;
    #[test]
    fn is_seng_mag() {
        assert!(is_send(0, 0, communication::OperateType::BROADCAST));
        assert!(is_send(0, 1, communication::OperateType::BROADCAST));

        assert!(is_send(0, 0, communication::OperateType::SINGLE));
        assert!(!is_send(0, 1, communication::OperateType::SINGLE));

        assert!(!is_send(0, 0, communication::OperateType::SUBTRACT));
        assert!(is_send(0, 1, communication::OperateType::SUBTRACT));
    }
}
<|endoftext|>


<commit_before>extern crate llvm_sys as llvm;

use std::mem;

use llvm::core::*;
use llvm::execution_engine::*;
use llvm::target::*;

fn main() {
    unsafe {
        // Set up a context, module and builder in that context.
        let context = LLVMContextCreate();
        let module = LLVMModuleCreateWithName(b"sum\0".as_ptr() as *const _);
        let builder = LLVMCreateBuilderInContext(context);

        // get a type for sum function
        let i64t = LLVMInt64TypeInContext(context);
        let mut argts = [i64t, i64t, i64t];
        let function_type = LLVMFunctionType(
            i64t,
            argts.as_mut_ptr(),
            argts.len() as u32,
            0);

        // add it to our module
        let function = LLVMAddFunction(
            module,
            b"sum\0".as_ptr() as *const _,
            function_type);

        // Create a basic block in the function and set our builder to generate
        // code in it.
        let bb = LLVMAppendBasicBlockInContext(
            context,
            function,
            b"entry\0".as_ptr() as *const _);

        LLVMPositionBuilderAtEnd(builder, bb);

        // get the function's arguments
        let x = LLVMGetParam(function, 0);
        let y = LLVMGetParam(function, 1);
        let z = LLVMGetParam(function, 2);

        let sum = LLVMBuildAdd(builder, x, y, b"sum.1\0".as_ptr() as *const _);
        let sum = LLVMBuildAdd(builder, sum, z, b"sum.2\0".as_ptr() as *const _);

        // Emit a `ret void` into the function
        LLVMBuildRet(builder, sum);

        // done building
        LLVMDisposeBuilder(builder);

        // Dump the module as IR to stdout.
        LLVMDumpModule(module);

        // build an execution engine
        let mut ee = mem::uninitialized();
        let mut out = mem::zeroed();

        // robust code should check that these calls complete successfully
        // each of these calls is necessary to setup an execution engine which compiles to native
        // code
        LLVMLinkInMCJIT();
        LLVM_InitializeNativeTarget();
        LLVM_InitializeNativeAsmPrinter();

        // takes ownership of the module
        LLVMCreateExecutionEngineForModule(&mut ee, module, &mut out);

        let addr = LLVMGetFunctionAddress(ee, b"sum\0".as_ptr() as *const _);

        let f: extern "C" fn(u64, u64, u64) -> u64 = mem::transmute(addr);

        let x: u64 = 1;
        let y: u64 = 1;
        let z: u64 = 1;
        let res = f(x, y, z);

        println!("{} + {} + {} = {}", x, y, z, res);

        // Clean up the rest.
        LLVMDisposeExecutionEngine(ee);
        LLVMContextDispose(context);
    }
}
<commit_msg>Fix module context in JIT example.<commit_after>extern crate llvm_sys as llvm;

use std::mem;

use llvm::core::*;
use llvm::execution_engine::*;
use llvm::target::*;

fn main() {
    unsafe {
        // Set up a context, module and builder in that context.
        let context = LLVMContextCreate();
        let module = LLVMModuleCreateWithNameInContext(b"sum\0".as_ptr() as *const _,
                                                       context);
        let builder = LLVMCreateBuilderInContext(context);

        // get a type for sum function
        let i64t = LLVMInt64TypeInContext(context);
        let mut argts = [i64t, i64t, i64t];
        let function_type = LLVMFunctionType(
            i64t,
            argts.as_mut_ptr(),
            argts.len() as u32,
            0);

        // add it to our module
        let function = LLVMAddFunction(
            module,
            b"sum\0".as_ptr() as *const _,
            function_type);

        // Create a basic block in the function and set our builder to generate
        // code in it.
        let bb = LLVMAppendBasicBlockInContext(
            context,
            function,
            b"entry\0".as_ptr() as *const _);

        LLVMPositionBuilderAtEnd(builder, bb);

        // get the function's arguments
        let x = LLVMGetParam(function, 0);
        let y = LLVMGetParam(function, 1);
        let z = LLVMGetParam(function, 2);

        let sum = LLVMBuildAdd(builder, x, y, b"sum.1\0".as_ptr() as *const _);
        let sum = LLVMBuildAdd(builder, sum, z, b"sum.2\0".as_ptr() as *const _);

        // Emit a `ret void` into the function
        LLVMBuildRet(builder, sum);

        // done building
        LLVMDisposeBuilder(builder);

        // Dump the module as IR to stdout.
        LLVMDumpModule(module);

        // build an execution engine
        let mut ee = mem::uninitialized();
        let mut out = mem::zeroed();

        // robust code should check that these calls complete successfully
        // each of these calls is necessary to setup an execution engine which compiles to native
        // code
        LLVMLinkInMCJIT();
        LLVM_InitializeNativeTarget();
        LLVM_InitializeNativeAsmPrinter();

        // takes ownership of the module
        LLVMCreateExecutionEngineForModule(&mut ee, module, &mut out);

        let addr = LLVMGetFunctionAddress(ee, b"sum\0".as_ptr() as *const _);

        let f: extern "C" fn(u64, u64, u64) -> u64 = mem::transmute(addr);

        let x: u64 = 1;
        let y: u64 = 1;
        let z: u64 = 1;
        let res = f(x, y, z);

        println!("{} + {} + {} = {}", x, y, z, res);

        // Clean up the rest.
        LLVMDisposeExecutionEngine(ee);
        LLVMContextDispose(context);
    }
}
<|endoftext|>


<commit_before>//! Defines the render target for tray, where our image will be written too
//! during rendering
//! TODO: Reconstruction filters

use std::vec::Vec;
use std::{iter, cmp, f32};
use std::sync::Mutex;

use film::Colorf;
use film::filter::Filter;
use sampler::Region;

const FILTER_TABLE_SIZE: usize = 16;

/// A struct containing results of an image sample where a ray was fired through
/// continuous pixel coordinates [x, y] and color `color` was computed
pub struct ImageSample {
    pub x: f32,
    pub y: f32,
    pub color: Colorf,
}

impl ImageSample {
    pub fn new(x: f32, y: f32, color: Colorf) -> ImageSample {
        ImageSample { x: x, y: y, color: color }
    }
}

/// RenderTarget is a RGBF render target to write our image too while rendering
pub struct RenderTarget {
    width: usize,
    height: usize,
    pixels: Vec<Colorf>,
    pixels_locked: Vec<Mutex<Vec<Colorf>>>,
    lock_size: (i32, i32),
    filter: Box<Filter + Send + Sync>,
    filter_table: Vec<f32>,
    filter_pixel_width: (i32, i32),
}

impl RenderTarget {
    /// Create a render target with `width * height` pixels
    pub fn new(image_dim: (usize, usize), lock_size: (usize, usize), filter: Box<Filter + Send + Sync>) -> RenderTarget {
        if image_dim.0 % lock_size.0 != 0 || image_dim.1 % lock_size.1 != 0 {
            panic!("Image with dimension {:?} not evenly divided by blocks of {:?}", image_dim, lock_size);
        }
        let width = image_dim.0;
        let height = image_dim.1;
        let filter_pixel_width = (f32::floor(filter.width() / 0.5) as i32, f32::floor(filter.height() / 0.5) as i32);
        let mut filter_table: Vec<f32> = iter::repeat(0.0).take(FILTER_TABLE_SIZE * FILTER_TABLE_SIZE)
            .collect();
        for y in 0..FILTER_TABLE_SIZE {
            let fy = (y as f32 + 0.5) * filter.height() / FILTER_TABLE_SIZE as f32;
            for x in 0..FILTER_TABLE_SIZE {
                let fx = (x as f32 + 0.5) * filter.width() / FILTER_TABLE_SIZE as f32;
                filter_table[y * FILTER_TABLE_SIZE + x] = filter.weight(fx, fy);
            }
        }

        let x_blocks = width / lock_size.0;
        let y_blocks = height / lock_size.1;
        let mut pixels_locked = Vec::with_capacity(x_blocks * y_blocks);
        for _ in 0..x_blocks * y_blocks {
            pixels_locked.push(Mutex::new(iter::repeat(Colorf::broadcast(0.0)).take(lock_size.0 * lock_size.1).collect()));
        }

        RenderTarget { width: width, height: height,
            pixels: iter::repeat(Colorf::broadcast(0.0)).take(width * height).collect(),
            pixels_locked: pixels_locked,
            lock_size: (lock_size.0 as i32, lock_size.1 as i32),
            filter: filter,
            filter_table: filter_table,
            filter_pixel_width: filter_pixel_width,
        }
    }
    /// Write all the image samples to the render target
    pub fn write(&self, samples: &Vec<ImageSample>, region: &Region) {
        // Determine which blocks we touch with our set of samples
        let x_range = (cmp::max(region.start.0 as i32 - self.filter_pixel_width.0, 0),
                       cmp::min(region.end.0 as i32 + self.filter_pixel_width.0, self.width as i32 - 1));
        let y_range = (cmp::max(region.start.1 as i32 - self.filter_pixel_width.1, 0),
                       cmp::min(region.end.1 as i32 + self.filter_pixel_width.1, self.height as i32 - 1));

        if x_range.1 - x_range.0 < 0 || y_range.1 - y_range.0 < 0 {
                return;
        }
        let block_x_range = (x_range.0 / self.lock_size.0, x_range.1 / self.lock_size.0);
        let block_y_range = (y_range.0 / self.lock_size.1, y_range.1 / self.lock_size.1);
        // Temporary storage for filtered samples so we can compute the filtered results for
        // the block we're writing too without having to get the lock
        // TODO: Could each thread have this locally and pass it to the render target
        // instead of allocating it each time we write?
        let mut filtered_samples: Vec<_> = iter::repeat(Colorf::broadcast(0.0))
            .take((self.lock_size.0 * self.lock_size.1) as usize).collect();

        let blocks_per_row = self.width as i32 / self.lock_size.0;
        for y in block_y_range.0..block_y_range.1 + 1 {
            for x in block_x_range.0..block_x_range.1 + 1 {
                let block_x_start = x * self.lock_size.0;
                let block_y_start = y * self.lock_size.1;

                let x_write_range = (cmp::max(x_range.0, block_x_start),
                                     cmp::min(x_range.1 + 1, block_x_start + self.lock_size.0));
                let y_write_range = (cmp::max(y_range.0, block_y_start),
                                     cmp::min(y_range.1 + 1, block_y_start + self.lock_size.1));

                let block_samples = samples.iter().filter(|s| {
                    s.x >= (x_write_range.0 - self.filter_pixel_width.0) as f32
                        && s.x < (x_write_range.1 + self.filter_pixel_width.0) as f32
                        && s.y >= (y_write_range.0 - self.filter_pixel_width.1) as f32
                        && s.y < (y_write_range.1 + self.filter_pixel_width.1) as f32
                });

                for c in &mut filtered_samples {
                    *c = Colorf::broadcast(0.0);
                }

                // Compute the filtered samples for the block
                for c in block_samples {
                    let img_x = c.x - 0.5;
                    let img_y = c.y - 0.5;
                    for iy in y_write_range.0..y_write_range.1 {
                        let fy = f32::abs(iy as f32 - img_y) * self.filter.inv_height() * FILTER_TABLE_SIZE as f32;
                        let fy_idx = cmp::min(fy as usize, FILTER_TABLE_SIZE - 1);
                        for ix in x_write_range.0..x_write_range.1 {
                            let fx = f32::abs(ix as f32 - img_x) * self.filter.inv_width() * FILTER_TABLE_SIZE as f32;
                            let fx_idx = cmp::min(fx as usize, FILTER_TABLE_SIZE - 1);
                            let weight = self.filter_table[fy_idx * FILTER_TABLE_SIZE + fx_idx];
                            let px = ((iy - block_y_start) * self.lock_size.0 + ix - block_x_start) as usize;
                            // TODO: Can't currently overload the += operator
                            // Coming soon though, see RFC #953 https://github.com/rust-lang/rfcs/pull/953
                            filtered_samples[px].r += weight * c.color.r;
                            filtered_samples[px].g += weight * c.color.g;
                            filtered_samples[px].b += weight * c.color.b;
                            filtered_samples[px].a += weight;
                        }
                    }
                }

                // Acquire lock for the block and write the filtered samples
                let block_idx = (y * blocks_per_row + x) as usize;
                let mut pixels = self.pixels_locked[block_idx].lock().unwrap();
                for iy in y_write_range.0..y_write_range.1 {
                    for ix in x_write_range.0..x_write_range.1 {
                        let px = ((iy - block_y_start) * self.lock_size.0 + ix - block_x_start) as usize;
                        let c = &filtered_samples[px];
                        pixels[px].r += c.r;
                        pixels[px].g += c.g;
                        pixels[px].b += c.b;
                        pixels[px].a += c.a;
                    }
                }
            }
        }
    }
    /// Clear the render target to black
    pub fn clear(&mut self) {
        for p in self.pixels.iter_mut() {
            *p = Colorf::broadcast(0.0);
        }
    }
    /// Get the dimensions of the render target
    pub fn dimensions(&self) -> (usize, usize) {
        (self.width, self.height)
    }
    /// Convert the floating point color buffer to 24bpp sRGB for output to an image
    pub fn get_render(&self) -> Vec<u8> {
        let mut render: Vec<u8> = iter::repeat(0u8).take(self.width * self.height * 3).collect();
        let x_blocks = self.width / self.lock_size.0 as usize;
        let y_blocks = self.height / self.lock_size.1 as usize;
        for by in 0..y_blocks {
            for bx in 0..x_blocks {
                let block_x_start = bx * self.lock_size.0 as usize;
                let block_y_start = by * self.lock_size.1 as usize;
                let block_idx = (by * x_blocks + bx) as usize;
                let pixels = self.pixels_locked[block_idx].lock().unwrap();
                for y in 0..self.lock_size.1 as usize {
                    for x in 0..self.lock_size.0 as usize {
                        let c = &pixels[y * self.lock_size.0 as usize + x];
                        if c.a > 0.0 {
                            let cn = (*c / c.a).clamp().to_srgb();
                            let px = (y + block_y_start) * self.width * 3 + (x + block_x_start) * 3;
                            for i in 0..3 {
                                render[px + i] = (cn[i] * 255.0) as u8;
                            }
                        }
                    }
                }
            }
        }
        render
    }
}

<commit_msg>Tested passing a scratch buffer, seems slower which is surprising.<commit_after>//! Defines the render target for tray, where our image will be written too
//! during rendering
//! TODO: Reconstruction filters

use std::vec::Vec;
use std::{iter, cmp, f32};
use std::sync::Mutex;

use film::Colorf;
use film::filter::Filter;
use sampler::Region;

const FILTER_TABLE_SIZE: usize = 16;

/// A struct containing results of an image sample where a ray was fired through
/// continuous pixel coordinates [x, y] and color `color` was computed
pub struct ImageSample {
    pub x: f32,
    pub y: f32,
    pub color: Colorf,
}

impl ImageSample {
    pub fn new(x: f32, y: f32, color: Colorf) -> ImageSample {
        ImageSample { x: x, y: y, color: color }
    }
}

/// RenderTarget is a RGBF render target to write our image too while rendering
pub struct RenderTarget {
    width: usize,
    height: usize,
    pixels: Vec<Colorf>,
    pixels_locked: Vec<Mutex<Vec<Colorf>>>,
    lock_size: (i32, i32),
    filter: Box<Filter + Send + Sync>,
    filter_table: Vec<f32>,
    filter_pixel_width: (i32, i32),
}

impl RenderTarget {
    /// Create a render target with `width * height` pixels
    pub fn new(image_dim: (usize, usize), lock_size: (usize, usize), filter: Box<Filter + Send + Sync>) -> RenderTarget {
        if image_dim.0 % lock_size.0 != 0 || image_dim.1 % lock_size.1 != 0 {
            panic!("Image with dimension {:?} not evenly divided by blocks of {:?}", image_dim, lock_size);
        }
        let width = image_dim.0;
        let height = image_dim.1;
        let filter_pixel_width = (f32::floor(filter.width() / 0.5) as i32, f32::floor(filter.height() / 0.5) as i32);
        let mut filter_table: Vec<f32> = iter::repeat(0.0).take(FILTER_TABLE_SIZE * FILTER_TABLE_SIZE)
            .collect();
        for y in 0..FILTER_TABLE_SIZE {
            let fy = (y as f32 + 0.5) * filter.height() / FILTER_TABLE_SIZE as f32;
            for x in 0..FILTER_TABLE_SIZE {
                let fx = (x as f32 + 0.5) * filter.width() / FILTER_TABLE_SIZE as f32;
                filter_table[y * FILTER_TABLE_SIZE + x] = filter.weight(fx, fy);
            }
        }

        let x_blocks = width / lock_size.0;
        let y_blocks = height / lock_size.1;
        let mut pixels_locked = Vec::with_capacity(x_blocks * y_blocks);
        for _ in 0..x_blocks * y_blocks {
            pixels_locked.push(Mutex::new(iter::repeat(Colorf::broadcast(0.0)).take(lock_size.0 * lock_size.1).collect()));
        }

        RenderTarget { width: width, height: height,
            pixels: iter::repeat(Colorf::broadcast(0.0)).take(width * height).collect(),
            pixels_locked: pixels_locked,
            lock_size: (lock_size.0 as i32, lock_size.1 as i32),
            filter: filter,
            filter_table: filter_table,
            filter_pixel_width: filter_pixel_width,
        }
    }
    /// Write all the image samples to the render target
    pub fn write(&self, samples: &Vec<ImageSample>, region: &Region) {
        // Determine which blocks we touch with our set of samples
        let x_range = (cmp::max(region.start.0 as i32 - self.filter_pixel_width.0, 0),
                       cmp::min(region.end.0 as i32 + self.filter_pixel_width.0, self.width as i32 - 1));
        let y_range = (cmp::max(region.start.1 as i32 - self.filter_pixel_width.1, 0),
                       cmp::min(region.end.1 as i32 + self.filter_pixel_width.1, self.height as i32 - 1));

        if x_range.1 - x_range.0 < 0 || y_range.1 - y_range.0 < 0 {
                return;
        }
        let block_x_range = (x_range.0 / self.lock_size.0, x_range.1 / self.lock_size.0);
        let block_y_range = (y_range.0 / self.lock_size.1, y_range.1 / self.lock_size.1);
        // Temporary storage for filtered samples so we can compute the filtered results for
        // the block we're writing too without having to get the lock
        let mut filtered_samples: Vec<_> = iter::repeat(Colorf::broadcast(0.0))
            .take((self.lock_size.0 * self.lock_size.1) as usize).collect();

        let blocks_per_row = self.width as i32 / self.lock_size.0;
        for y in block_y_range.0..block_y_range.1 + 1 {
            for x in block_x_range.0..block_x_range.1 + 1 {
                let block_x_start = x * self.lock_size.0;
                let block_y_start = y * self.lock_size.1;

                let x_write_range = (cmp::max(x_range.0, block_x_start),
                                     cmp::min(x_range.1 + 1, block_x_start + self.lock_size.0));
                let y_write_range = (cmp::max(y_range.0, block_y_start),
                                     cmp::min(y_range.1 + 1, block_y_start + self.lock_size.1));

                let block_samples = samples.iter().filter(|s| {
                    s.x >= (x_write_range.0 - self.filter_pixel_width.0) as f32
                        && s.x < (x_write_range.1 + self.filter_pixel_width.0) as f32
                        && s.y >= (y_write_range.0 - self.filter_pixel_width.1) as f32
                        && s.y < (y_write_range.1 + self.filter_pixel_width.1) as f32
                });

                for c in &mut filtered_samples {
                    *c = Colorf::broadcast(0.0);
                }

                // Compute the filtered samples for the block
                for c in block_samples {
                    let img_x = c.x - 0.5;
                    let img_y = c.y - 0.5;
                    for iy in y_write_range.0..y_write_range.1 {
                        let fy = f32::abs(iy as f32 - img_y) * self.filter.inv_height() * FILTER_TABLE_SIZE as f32;
                        let fy_idx = cmp::min(fy as usize, FILTER_TABLE_SIZE - 1);
                        for ix in x_write_range.0..x_write_range.1 {
                            let fx = f32::abs(ix as f32 - img_x) * self.filter.inv_width() * FILTER_TABLE_SIZE as f32;
                            let fx_idx = cmp::min(fx as usize, FILTER_TABLE_SIZE - 1);
                            let weight = self.filter_table[fy_idx * FILTER_TABLE_SIZE + fx_idx];
                            let px = ((iy - block_y_start) * self.lock_size.0 + ix - block_x_start) as usize;
                            // TODO: Can't currently overload the += operator
                            // Coming soon though, see RFC #953 https://github.com/rust-lang/rfcs/pull/953
                            filtered_samples[px].r += weight * c.color.r;
                            filtered_samples[px].g += weight * c.color.g;
                            filtered_samples[px].b += weight * c.color.b;
                            filtered_samples[px].a += weight;
                        }
                    }
                }

                // Acquire lock for the block and write the filtered samples
                let block_idx = (y * blocks_per_row + x) as usize;
                let mut pixels = self.pixels_locked[block_idx].lock().unwrap();
                for iy in y_write_range.0..y_write_range.1 {
                    for ix in x_write_range.0..x_write_range.1 {
                        let px = ((iy - block_y_start) * self.lock_size.0 + ix - block_x_start) as usize;
                        let c = &filtered_samples[px];
                        pixels[px].r += c.r;
                        pixels[px].g += c.g;
                        pixels[px].b += c.b;
                        pixels[px].a += c.a;
                    }
                }
            }
        }
    }
    /// Clear the render target to black
    pub fn clear(&mut self) {
        for p in self.pixels.iter_mut() {
            *p = Colorf::broadcast(0.0);
        }
    }
    /// Get the dimensions of the render target
    pub fn dimensions(&self) -> (usize, usize) {
        (self.width, self.height)
    }
    /// Convert the floating point color buffer to 24bpp sRGB for output to an image
    pub fn get_render(&self) -> Vec<u8> {
        let mut render: Vec<u8> = iter::repeat(0u8).take(self.width * self.height * 3).collect();
        let x_blocks = self.width / self.lock_size.0 as usize;
        let y_blocks = self.height / self.lock_size.1 as usize;
        for by in 0..y_blocks {
            for bx in 0..x_blocks {
                let block_x_start = bx * self.lock_size.0 as usize;
                let block_y_start = by * self.lock_size.1 as usize;
                let block_idx = (by * x_blocks + bx) as usize;
                let pixels = self.pixels_locked[block_idx].lock().unwrap();
                for y in 0..self.lock_size.1 as usize {
                    for x in 0..self.lock_size.0 as usize {
                        let c = &pixels[y * self.lock_size.0 as usize + x];
                        if c.a > 0.0 {
                            let cn = (*c / c.a).clamp().to_srgb();
                            let px = (y + block_y_start) * self.width * 3 + (x + block_x_start) * 3;
                            for i in 0..3 {
                                render[px + i] = (cn[i] * 255.0) as u8;
                            }
                        }
                    }
                }
            }
        }
        render
    }
}

<|endoftext|>


<commit_before>//! Reports metrics to Datadog through their HTTPS API.

use chrono::UTC;
use hyper::client::{Client, RequestBuilder};
use hyper::header::ContentType;
use rustc_serialize::json::{self, Json, ToJson};
use std::collections::BTreeMap;

use super::Forwarder;
use super::super::metrics::AggregatedMetrics;

/// Forwards metrics to Datadog via its HTTPS API.
///
/// See [their documentation][] for more details.
///
/// [their documentation]: http://docs.datadoghq.com/api/
pub struct DatadogForwarder {
    pub api_key: String,
    pub base_url: String,
}

impl DatadogForwarder {
    pub fn new(api_key: &str) -> DatadogForwarder {
        DatadogForwarder {
            api_key: api_key.to_owned(),
            base_url: "https://app.datadoghq.com/api".to_owned(),
        }
    }

    fn serialize_metrics(metrics: AggregatedMetrics) -> Json {
        use super::super::metrics::AggregatedMetricType::*;

        let timestamp = UTC::now().timestamp();

        let series: Vec<Json> = metrics
            .iter()
            .map(|metric| {
                let mut object: BTreeMap<String, Json> = BTreeMap::new();

                let (ref metric_type, ref name, ref value) = *metric;

                let api_type = match *metric_type {
                    Count   => "count",
                    Measure => "gauge",
                    Sample  => "gauge",
                };

                object.insert("metric".to_owned(), name.to_json());
                object.insert("type".to_owned(), api_type.to_json());
                object.insert("points".to_owned(), Json::Array(vec![
                    Json::Array(vec![ timestamp.to_json(), value.to_json() ]),
                ]));

                object.to_json()
            })
            .collect();

        let mut data: BTreeMap<String, Json> = BTreeMap::new();
        data.insert("series".to_owned(), Json::Array(series));

        // Convert it to a `Json::Object`.
        data.to_json()
    }

    fn post<'a>(&'a self, client: &'a Client, path: &str) -> RequestBuilder {
        let path = format!("{}{}?api_key={}", self.base_url, path, self.api_key);

        client.post(&path)
            .header(ContentType::json())
    }
}

impl Forwarder for DatadogForwarder {
    fn forward_metrics(&self, metrics: AggregatedMetrics) {
        let body = json::encode(&DatadogForwarder::serialize_metrics(metrics)).unwrap();
        let client = Client::new();

        let res = self.post(&client, "/v1/series")
            .body(&body)
            .send()
            .unwrap();

        if !res.status.is_success() {
            println!("Datadog API Error: {:#?}", res);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::DatadogForwarder;
    use super::super::super::metrics::{
        AggregatedMetrics,
        AggregatedMetricType,
    };

    use rustc_serialize::json::ToJson;

    #[test]
    fn datadog_forwarder_serializes_metrics() {
        let metrics = AggregatedMetrics::with_metrics(vec![
            (AggregatedMetricType::Count, "test_count".to_owned(), 1.0),
        ]);
        let json = DatadogForwarder::serialize_metrics(metrics);

        let series = json.find("series").and_then(|s| s.as_array());
        assert_eq!(series.is_some(), true);

        let series = series.unwrap();
        assert_eq!(series.len(), 1);

        let item = series[0].as_object().unwrap();
        assert_eq!(item.get("metric"), Some(&"test_count".to_json()));
        assert_eq!(item.get("type"),   Some(&"count".to_json()));

        let points = item.get("points").unwrap().as_array().unwrap();
        assert_eq!(points.len(), 1);
        let point = points[0].as_array().unwrap();
        assert_eq!(point.len(), 2);
        let ref value = point[1];
        assert_eq!(value, &1.0.to_json());
    }
}
<commit_msg>Make Datadog forwarder gracefully handle HTTP errors<commit_after>//! Reports metrics to Datadog through their HTTPS API.

use chrono::UTC;
use hyper::client::{Client, RequestBuilder};
use hyper::header::ContentType;
use rustc_serialize::json::{self, Json, ToJson};
use std::collections::BTreeMap;

use super::Forwarder;
use super::super::metrics::AggregatedMetrics;

/// Forwards metrics to Datadog via its HTTPS API.
///
/// See [their documentation][] for more details.
///
/// [their documentation]: http://docs.datadoghq.com/api/
pub struct DatadogForwarder {
    pub api_key: String,
    pub base_url: String,
}

impl DatadogForwarder {
    pub fn new(api_key: &str) -> DatadogForwarder {
        DatadogForwarder {
            api_key: api_key.to_owned(),
            base_url: "https://app.datadoghq.com/api".to_owned(),
        }
    }

    fn serialize_metrics(metrics: AggregatedMetrics) -> Json {
        use super::super::metrics::AggregatedMetricType::*;

        let timestamp = UTC::now().timestamp();

        let series: Vec<Json> = metrics
            .iter()
            .map(|metric| {
                let mut object: BTreeMap<String, Json> = BTreeMap::new();

                let (ref metric_type, ref name, ref value) = *metric;

                let api_type = match *metric_type {
                    Count   => "count",
                    Measure => "gauge",
                    Sample  => "gauge",
                };

                object.insert("metric".to_owned(), name.to_json());
                object.insert("type".to_owned(), api_type.to_json());
                object.insert("points".to_owned(), Json::Array(vec![
                    Json::Array(vec![ timestamp.to_json(), value.to_json() ]),
                ]));

                object.to_json()
            })
            .collect();

        let mut data: BTreeMap<String, Json> = BTreeMap::new();
        data.insert("series".to_owned(), Json::Array(series));

        // Convert it to a `Json::Object`.
        data.to_json()
    }

    fn post<'a>(&'a self, client: &'a Client, path: &str) -> RequestBuilder {
        let path = format!("{}{}?api_key={}", self.base_url, path, self.api_key);

        client.post(&path)
            .header(ContentType::json())
    }
}

impl Forwarder for DatadogForwarder {
    fn forward_metrics(&self, metrics: AggregatedMetrics) {
        let body = json::encode(&DatadogForwarder::serialize_metrics(metrics)).unwrap();
        let client = Client::new();

        let res = self.post(&client, "/v1/series")
            .body(&body)
            .send();

        match res {
            Err(err) => {
                println!("Datadog HTTP Error: {:#?}", err)
            },
            Ok(res) => {
                if !res.status.is_success() {
                    println!("Datadog API Error: {:#?}", res);
                }
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::DatadogForwarder;
    use super::super::super::metrics::{
        AggregatedMetrics,
        AggregatedMetricType,
    };

    use rustc_serialize::json::ToJson;

    #[test]
    fn datadog_forwarder_serializes_metrics() {
        let metrics = AggregatedMetrics::with_metrics(vec![
            (AggregatedMetricType::Count, "test_count".to_owned(), 1.0),
        ]);
        let json = DatadogForwarder::serialize_metrics(metrics);

        let series = json.find("series").and_then(|s| s.as_array());
        assert_eq!(series.is_some(), true);

        let series = series.unwrap();
        assert_eq!(series.len(), 1);

        let item = series[0].as_object().unwrap();
        assert_eq!(item.get("metric"), Some(&"test_count".to_json()));
        assert_eq!(item.get("type"),   Some(&"count".to_json()));

        let points = item.get("points").unwrap().as_array().unwrap();
        assert_eq!(points.len(), 1);
        let point = points[0].as_array().unwrap();
        assert_eq!(point.len(), 2);
        let ref value = point[1];
        assert_eq!(value, &1.0.to_json());
    }
}
<|endoftext|>


<commit_before>use super::{ASTNode, SymbolTable};

use std::rc::Rc;
use std::borrow::Borrow;
use std::hash::Hash;

// pub type STCell<'a> = Rc<RefCell<SymbolTable<'a>>>;

pub struct Annotated<'a, T> where T: ASTNode {
    node: T, scope: Rc<SymbolTable<'a>>
}

impl<'a, T> Annotated<'a, T> where T: ASTNode {

    /// Get the type signature associated with the given name.
    ///
    /// This returns a borrowed reference to the type signature
    /// associated with that name in the current scope. The argument
    /// can be any type `Q` such that `String: Borrow<Q>` (i.e.
    /// you can pass an `&str` to this function).
    pub fn get_type<Q: ?Sized>(&self, name: &Q) -> Option<&Type>
    where String: Borrow<Q>, Q: Hash + Eq {
        self.scope.get(name)
    }

    /// Check if the given name is defined in this scope.
    ///
    /// The argument can be any type `Q` such that `String: Borrow<Q>`
    /// (i.e. you can pass an `&str` to this function).
    pub fn is_defined_here<Q: ?Sized>(&self, name: &Q) -> bool
    where String: Borrow<Q>, Q: Hash + Eq {
        self.scope.contains_key(name)
    }

}

#[derive(Debug, Clone)]
pub enum Type { Ref(Reference)
              , Prim(Primitive)
              , Algebraic //TODO: make me a thing
              , Function { params: Vec<Type>
                         , rt: Rc<Type> }
              }

/// Reference types
///
/// TODO: how will lifetime analysis actually work?
#[derive(Debug, Clone)]
pub enum Reference { Borrowed(Rc<Type>) // TODO: borrowed from where?
                   , Moved(Rc<Type>) // TODO: moved from where?
                   , Owned(Rc<Type>)
                   , Raw(Rc<Type>)
                   }

#[derive(Debug, Clone)]
pub enum Primitive { Int
                   , Uint
                   , Byte
                   , Char
                   , Str
                   , Bool
                   , Symbol(String)
                   // TODO: finish
                   }
<commit_msg>Symbol really isn't a primitive, is it?<commit_after>use super::{ASTNode, SymbolTable};

use std::rc::Rc;
use std::borrow::Borrow;
use std::hash::Hash;

// pub type STCell<'a> = Rc<RefCell<SymbolTable<'a>>>;

pub struct Annotated<'a, T> where T: ASTNode {
    node: T, scope: Rc<SymbolTable<'a>>
}

impl<'a, T> Annotated<'a, T> where T: ASTNode {

    /// Get the type signature associated with the given name.
    ///
    /// This returns a borrowed reference to the type signature
    /// associated with that name in the current scope. The argument
    /// can be any type `Q` such that `String: Borrow<Q>` (i.e.
    /// you can pass an `&str` to this function).
    pub fn get_type<Q: ?Sized>(&self, name: &Q) -> Option<&Type>
    where String: Borrow<Q>, Q: Hash + Eq {
        self.scope.get(name)
    }

    /// Check if the given name is defined in this scope.
    ///
    /// The argument can be any type `Q` such that `String: Borrow<Q>`
    /// (i.e. you can pass an `&str` to this function).
    pub fn is_defined_here<Q: ?Sized>(&self, name: &Q) -> bool
    where String: Borrow<Q>, Q: Hash + Eq {
        self.scope.contains_key(name)
    }

}

#[derive(Debug, Clone)]
pub enum Type { Ref(Reference)
              , Prim(Primitive)
              , Algebraic //TODO: make me a thing
              , Function { params: Vec<Type>
                         , rt: Rc<Type> }
              , Symbol(String)
              }

/// Reference types
///
/// TODO: how will lifetime analysis actually work?
#[derive(Debug, Clone)]
pub enum Reference { Borrowed(Rc<Type>) // TODO: borrowed from where?
                   , Moved(Rc<Type>) // TODO: moved from where?
                   , Owned(Rc<Type>)
                   , Raw(Rc<Type>)
                   }

#[derive(Debug, Clone, Copy)]
pub enum Primitive { Int
                   , Uint
                   , Byte
                   , Char
                   , Str
                   , Bool
                   // TODO: finish
                   }
<|endoftext|>


<commit_before>use types::NumPlies;

use super::super::color::{White, Black};
use super::super::piece::Piece::{self, WK, WR, BK, BR};
use super::super::piece::Pawn;
use super::super::square::{Square, File, Rank};
use super::super::moves::Move;
use super::super::castle::{Kingside, Queenside};

use super::{Position, ExtraData};

pub fn make_move(p: &mut Position, m: &Move) {

    debug!("make_move(): {:?}", *m);

    debug!("before: {:?}", *p);

    let from = m.from();
    let to = m.to();
    let curr_piece = p.at(from).unwrap();
    let curr_color = p.side_to_move();

    if m.capture_normal().is_some() || m.is_en_passant() || curr_piece.piece_type() == Pawn {
        p.set_ply_count(NumPlies(0));
    } else {
        let NumPlies(temp) = p.ply_count();
        p.set_ply_count(NumPlies(temp+1));
    }

    if let Some(castle_side) = m.castle() {

        //no en passant
        p.set_en_passant(None);

        let (rook_from, rook_to) = match (curr_color, castle_side) {
            (White, Kingside ) => (Square::new(File(7), Rank(0)), Square::new(File(5), Rank(0))),
            (White, Queenside) => (Square::new(File(0), Rank(0)), Square::new(File(3), Rank(0))),
            (Black, Kingside ) => (Square::new(File(7), Rank(7)), Square::new(File(5), Rank(7))),
            (Black, Queenside) => (Square::new(File(0), Rank(7)), Square::new(File(3), Rank(7))),
        };

        p.remove_at(from, curr_piece);
        p.set_at(to, curr_piece);
        let curr_rook = p.at(rook_from).unwrap();
        p.remove_at(rook_from, curr_rook);
        p.set_at(rook_to, curr_rook);

    } else if m.is_en_passant() {

        //no en passant after this
        p.set_en_passant(None);

        let captured = Square::new(to.file(), from.rank());
        let captured_piece = Piece::new(p.side_to_move().invert(), Pawn);

        p.remove_at(captured, captured_piece);
        p.remove_at(from, curr_piece);
        p.set_at(to, curr_piece);

    } else if let Some(promote_piece) = m.promote() {
        //no en passant
        p.set_en_passant(None);

        //change the board and promote
        if let Some(captured_piece) = m.capture_normal() {
            p.remove_at(to, captured_piece);
        }
        p.remove_at(from, curr_piece);
        p.set_at(to, Piece::new(curr_color, promote_piece));
    } else {

        if m.is_pawn_double_move() {
            p.set_en_passant(Some(to.file()));
        } else {
            p.set_en_passant(None);
        }

        //change the board
        if let Some(captured_piece) = m.capture_normal() {
            p.remove_at(to, captured_piece);
        }
        p.remove_at(from, curr_piece);
        p.set_at(to, curr_piece);
    }

    //castling
    if p.at(Square::new(File(4), Rank(0))) != Some(WK) {
        p.set_castle(Queenside, White, false);
        p.set_castle(Kingside , White, false);
    }
    if p.at(Square::new(File(0), Rank(0))) != Some(WR) {
        p.set_castle(Queenside, White, false);
    }
    if p.at(Square::new(File(7), Rank(0))) != Some(WR) {
        p.set_castle(Kingside , White, false);
    }
    if p.at(Square::new(File(4), Rank(7))) != Some(BK) {
        p.set_castle(Queenside, Black, false);
        p.set_castle(Kingside , Black, false);
    }
    if p.at(Square::new(File(0), Rank(7))) != Some(BR) {
        p.set_castle(Queenside, Black, false);
    }
    if p.at(Square::new(File(7), Rank(7))) != Some(BR) {
        p.set_castle(Kingside , Black, false);
    }

    p.swap_side_to_move();

    debug!("after : {:?}", *p);
}

pub fn unmake_move(p: &mut Position, m: &Move, extra_data: ExtraData) {

    debug!("unmake_move(): {:?}", *m);
    debug!("{:?}", extra_data);

    debug!("before: {:?}", *p);

    let from = m.from();
    let to = m.to();
    let curr_piece = p.at(to).unwrap();
    let curr_color = p.side_to_move().invert();

    if let Some(castle_side) = m.castle() {

        let (rook_from, rook_to) = match (curr_color, castle_side) {
            (White, Kingside ) => (Square::new(File(7), Rank(0)), Square::new(File(5), Rank(0))),
            (White, Queenside) => (Square::new(File(0), Rank(0)), Square::new(File(3), Rank(0))),
            (Black, Kingside ) => (Square::new(File(7), Rank(7)), Square::new(File(5), Rank(7))),
            (Black, Queenside) => (Square::new(File(0), Rank(7)), Square::new(File(3), Rank(7))),
        };

        p.remove_at(to, curr_piece);
        p.set_at(from, curr_piece);
        let curr_rook = p.at(rook_to).unwrap();
        p.remove_at(rook_to, curr_rook);
        p.set_at(rook_from, curr_rook);

    } else if m.is_en_passant() {

        let captured = Square::new(to.file(), from.rank());
        let captured_piece = Piece::new(curr_color.invert(), Pawn);

        p.set_at(captured, captured_piece);
        p.remove_at(to, curr_piece);
        p.set_at(from, curr_piece);

    } else if let Some(promote_piece) = m.promote() {

        p.remove_at(to, Piece::new(curr_color, promote_piece));
        p.set_at(from, curr_piece);
        if let Some(captured_piece) = m.capture_normal() {
            p.set_at(to, captured_piece);
        }

    } else {

        p.remove_at(to, curr_piece);
        p.set_at(from, curr_piece);
        if let Some(captured_piece) = m.capture_normal() {
            p.set_at(to, captured_piece);
        }

    }

    p.set_extra_data(extra_data);
    p.swap_side_to_move();

    debug!("after : {:?}", *p);

}
<commit_msg>Fix more promotion bugs<commit_after>use types::NumPlies;

use super::super::color::{White, Black};
use super::super::piece::Piece::{self, WK, WR, BK, BR};
use super::super::piece::Pawn;
use super::super::square::{Square, File, Rank};
use super::super::moves::Move;
use super::super::castle::{Kingside, Queenside};

use super::{Position, ExtraData};

pub fn make_move(p: &mut Position, m: &Move) {

    debug!("make_move(): {:?}", *m);

    debug!("before: {:?}", *p);

    let from = m.from();
    let to = m.to();
    let curr_piece = p.at(from).unwrap();
    let curr_color = p.side_to_move();

    if m.capture_normal().is_some() || m.is_en_passant() || curr_piece.piece_type() == Pawn {
        p.set_ply_count(NumPlies(0));
    } else {
        let NumPlies(temp) = p.ply_count();
        p.set_ply_count(NumPlies(temp+1));
    }

    if let Some(castle_side) = m.castle() {

        //no en passant
        p.set_en_passant(None);

        let (rook_from, rook_to) = match (curr_color, castle_side) {
            (White, Kingside ) => (Square::new(File(7), Rank(0)), Square::new(File(5), Rank(0))),
            (White, Queenside) => (Square::new(File(0), Rank(0)), Square::new(File(3), Rank(0))),
            (Black, Kingside ) => (Square::new(File(7), Rank(7)), Square::new(File(5), Rank(7))),
            (Black, Queenside) => (Square::new(File(0), Rank(7)), Square::new(File(3), Rank(7))),
        };

        p.remove_at(from, curr_piece);
        p.set_at(to, curr_piece);
        let curr_rook = p.at(rook_from).unwrap();
        p.remove_at(rook_from, curr_rook);
        p.set_at(rook_to, curr_rook);

    } else if m.is_en_passant() {

        //no en passant after this
        p.set_en_passant(None);

        let captured = Square::new(to.file(), from.rank());
        let captured_piece = Piece::new(p.side_to_move().invert(), Pawn);

        p.remove_at(captured, captured_piece);
        p.remove_at(from, curr_piece);
        p.set_at(to, curr_piece);

    } else if let Some(promote_piece) = m.promote() {
        //no en passant
        p.set_en_passant(None);

        //change the board and promote
        if let Some(captured_piece) = m.capture_normal() {
            p.remove_at(to, captured_piece);
        }
        p.remove_at(from, curr_piece);
        p.set_at(to, Piece::new(curr_color, promote_piece));
    } else {

        if m.is_pawn_double_move() {
            p.set_en_passant(Some(to.file()));
        } else {
            p.set_en_passant(None);
        }

        //change the board
        if let Some(captured_piece) = m.capture_normal() {
            p.remove_at(to, captured_piece);
        }
        p.remove_at(from, curr_piece);
        p.set_at(to, curr_piece);
    }

    //castling
    if p.at(Square::new(File(4), Rank(0))) != Some(WK) {
        p.set_castle(Queenside, White, false);
        p.set_castle(Kingside , White, false);
    }
    if p.at(Square::new(File(0), Rank(0))) != Some(WR) {
        p.set_castle(Queenside, White, false);
    }
    if p.at(Square::new(File(7), Rank(0))) != Some(WR) {
        p.set_castle(Kingside , White, false);
    }
    if p.at(Square::new(File(4), Rank(7))) != Some(BK) {
        p.set_castle(Queenside, Black, false);
        p.set_castle(Kingside , Black, false);
    }
    if p.at(Square::new(File(0), Rank(7))) != Some(BR) {
        p.set_castle(Queenside, Black, false);
    }
    if p.at(Square::new(File(7), Rank(7))) != Some(BR) {
        p.set_castle(Kingside , Black, false);
    }

    p.swap_side_to_move();

    debug!("after : {:?}", *p);
}

pub fn unmake_move(p: &mut Position, m: &Move, extra_data: ExtraData) {

    debug!("unmake_move(): {:?}", *m);
    debug!("{:?}", extra_data);

    debug!("before: {:?}", *p);

    let from = m.from();
    let to = m.to();
    let curr_piece = p.at(to).unwrap();
    let curr_color = p.side_to_move().invert();

    if let Some(castle_side) = m.castle() {

        let (rook_from, rook_to) = match (curr_color, castle_side) {
            (White, Kingside ) => (Square::new(File(7), Rank(0)), Square::new(File(5), Rank(0))),
            (White, Queenside) => (Square::new(File(0), Rank(0)), Square::new(File(3), Rank(0))),
            (Black, Kingside ) => (Square::new(File(7), Rank(7)), Square::new(File(5), Rank(7))),
            (Black, Queenside) => (Square::new(File(0), Rank(7)), Square::new(File(3), Rank(7))),
        };

        p.remove_at(to, curr_piece);
        p.set_at(from, curr_piece);
        let curr_rook = p.at(rook_to).unwrap();
        p.remove_at(rook_to, curr_rook);
        p.set_at(rook_from, curr_rook);

    } else if m.is_en_passant() {

        let captured = Square::new(to.file(), from.rank());
        let captured_piece = Piece::new(curr_color.invert(), Pawn);

        p.set_at(captured, captured_piece);
        p.remove_at(to, curr_piece);
        p.set_at(from, curr_piece);

    } else if let Some(promote_piece) = m.promote() {

        p.remove_at(to, Piece::new(curr_color, promote_piece));
        p.set_at(from, Piece::new(curr_color, Pawn));
        if let Some(captured_piece) = m.capture_normal() {
            p.set_at(to, captured_piece);
        }

    } else {

        p.remove_at(to, curr_piece);
        p.set_at(from, curr_piece);
        if let Some(captured_piece) = m.capture_normal() {
            p.set_at(to, captured_piece);
        }

    }

    p.set_extra_data(extra_data);
    p.swap_side_to_move();

    debug!("after : {:?}", *p);

}
<|endoftext|>


<commit_before>extern crate bignum;
extern crate num_traits;
extern crate rsa;

use num_traits::NumOps;

pub struct DsaPublic<'a, T: 'a> {
    pub params: &'a DsaParams<T>,
    pub y: T,
}

impl<'a, T> DsaPublic<'a, T>
where
    T: bignum::BigNumTrait,
    for<'a1, 'a2> &'a1 T: NumOps<&'a2 T, T>,
{
    pub fn generate(private: &'a DsaPrivate<T>) -> Self {
        let params = private.params;
        let y = params.g.mod_exp(&private.x, &params.p);
        DsaPublic {
            params: params,
            y: y,
        }
    }

    pub fn secret_key_from_k(&self, m: &T, &Signature { ref r, ref s }: &Signature<T>, k: &T) -> T {
        let q = &self.params.q;
        let x = &(&(s * k) - m) * &r.invmod(q).unwrap();
        x.mod_math(q)
    }

    pub fn secret_key_from_two_signatures_with_same_k(
        &self,
        m1: &T,
        s1: &Signature<T>,
        m2: &T,
        s2: &Signature<T>,
    ) -> T {
        assert_eq!(s1.r, s2.r);
        let q = &self.params.q;
        let k = &(m1 - m2).mod_math(q) * &T::invmod(&T::mod_math(&(&s1.s - &s2.s), q), q).unwrap();
        self.secret_key_from_k(m1, s1, &k)
    }

    pub fn verify_signature(&self, m: &T, &Signature { ref r, ref s }: &Signature<T>) -> bool {
        let zero = T::zero();
        let p = &self.params.p;
        let q = &self.params.q;
        let g = &self.params.g;
        if r <= &zero || r >= q || s <= &zero || s >= q {
            return false;
        }
        let w = T::invmod(s, q).unwrap();
        let u1 = T::mod_math(&(m * &w), q);
        let u2 = T::mod_math(&(r * &w), q);
        let v1 = g.mod_exp(&u1, p);
        let v2 = self.y.mod_exp(&u2, p);
        let v = T::mod_math(&T::mod_math(&(&v1 * &v2), p), q);
        &v == r
    }
}

pub struct DsaPrivate<'a, T: 'a> {
    pub params: &'a DsaParams<T>,
    pub x: T,
}

impl<'a, T> DsaPrivate<'a, T>
where
    T: bignum::BigNumTrait,
    for<'a1, 'a2> &'a1 T: NumOps<&'a2 T, T>,
{
    pub fn generate(params: &'a DsaParams<T>) -> Self {
        let x = rand_range_safe(&params.q);
        DsaPrivate {
            params: params,
            x: x,
        }
    }

    pub fn sign(&self, m: &T) -> (Signature<T>, T) {
        let zero = T::zero();
        let one = T::one();
        let p = &self.params.p;
        let q = &self.params.q;
        let g = &self.params.g;

        let mut k: T;
        let mut r: T;
        let mut s: T;
        loop {
            k = T::gen_below(q);
            if k == zero || k == one {
                continue;
            }
            r = g.mod_exp(&k, p);
            r = r.mod_math(q);
            if r == zero {
                continue;
            }
            s = &k.invmod(q).unwrap() * &(m + &(&r * &self.x));
            s = s.mod_math(q);
            if s == zero {
                continue;
            }
            break;
        }
        (Signature { r: r, s: s }, k)
    }
}

pub fn rand_range_safe<T: bignum::BigNumTrait>(q: &T) -> T {
    let zero = T::zero();
    let one = T::one();
    let mut x = T::clone(&zero);
    while x == zero || x == one {
        x = T::gen_below(q);
    }
    x
}

pub fn fake_signature<T>(public: &DsaPublic<T>) -> Signature<T>
where
    T: bignum::BigNumTrait,
    for<'b> &'b T: NumOps<&'b T, T>,
{
    let p = &public.params.p;
    let q = &public.params.q;
    let z = rand_range_safe(q);
    let mut r = T::mod_exp(&public.y, &z, p);
    r = T::mod_math(&r, q);
    let mut s = &r * &T::invmod(&z, q).unwrap();
    s = T::mod_math(&s, q);
    Signature { r: r, s: s }
}

pub struct DsaParams<T> {
    pub p: T,
    pub q: T,
    pub g: T,
}

impl<T: bignum::BigNumTrait> DsaParams<T> {
    pub fn generate() -> Self {
        let p = T::from_hex_str(
            "800000000000000089e1855218a0e7dac38136ffafa72eda7859f2171e25e65eac698c1702578b07dc2a1076da241c76c62d374d8389ea5aeffd3226a0530cc565f3bf6b50929139ebeac04f48c3c84afb796d61e5a4f9a8fda812ab59494232c7d2b4deb50aa18ee9e132bfa85ac4374d7f9091abc3d015efc871a584471bb1"
            ).unwrap();
        let q = T::from_hex_str("f4f47f05794b256174bba6e9b396a7707e563c5b").unwrap();
        let g = T::from_hex_str("5958c9d3898b224b12672c0b98e06c60df923cb8bc999d119458fef538b8fa4046c8db53039db620c094c9fa077ef389b5322a559946a71903f990f1f7e0e025e2d7f7cf494aff1a0470f5b64c36b625a097f1651fe775323556fe00b3608c887892878480e99041be601a62166ca6894bdd41a7054ec89f756ba9fc95302291").unwrap();
        DsaParams { p: p, q: q, g: g }
    }
}

pub struct Signature<T> {
    pub r: T,
    pub s: T,
}
<commit_msg>Fix build<commit_after>extern crate bignum;
extern crate num_traits;
extern crate rsa;

use num_traits::NumOps;

pub struct DsaPublic<'a, T: 'a> {
    pub params: &'a DsaParams<T>,
    pub y: T,
}

impl<'a, T> DsaPublic<'a, T>
where
    T: bignum::BigNumTrait+bignum::BigNumExt,
    for<'a1, 'a2> &'a1 T: NumOps<&'a2 T, T>,
{
    pub fn generate(private: &'a DsaPrivate<T>) -> Self {
        let params = private.params;
        let y = params.g.mod_exp(&private.x, &params.p);
        DsaPublic {
            params: params,
            y: y,
        }
    }

    pub fn secret_key_from_k(&self, m: &T, &Signature { ref r, ref s }: &Signature<T>, k: &T) -> T {
        let q = &self.params.q;
        let x = &(&(s * k) - m) * &r.invmod(q).unwrap();
        x.mod_math(q)
    }

    pub fn secret_key_from_two_signatures_with_same_k(
        &self,
        m1: &T,
        s1: &Signature<T>,
        m2: &T,
        s2: &Signature<T>,
    ) -> T {
        assert_eq!(s1.r, s2.r);
        let q = &self.params.q;
        let k = &(m1 - m2).mod_math(q) * &T::invmod(&T::mod_math(&(&s1.s - &s2.s), q), q).unwrap();
        self.secret_key_from_k(m1, s1, &k)
    }

    pub fn verify_signature(&self, m: &T, &Signature { ref r, ref s }: &Signature<T>) -> bool {
        let zero = T::zero();
        let p = &self.params.p;
        let q = &self.params.q;
        let g = &self.params.g;
        if r <= &zero || r >= q || s <= &zero || s >= q {
            return false;
        }
        let w = T::invmod(s, q).unwrap();
        let u1 = T::mod_math(&(m * &w), q);
        let u2 = T::mod_math(&(r * &w), q);
        let v1 = g.mod_exp(&u1, p);
        let v2 = self.y.mod_exp(&u2, p);
        let v = T::mod_math(&T::mod_math(&(&v1 * &v2), p), q);
        &v == r
    }
}

pub struct DsaPrivate<'a, T: 'a> {
    pub params: &'a DsaParams<T>,
    pub x: T,
}

impl<'a, T> DsaPrivate<'a, T>
where
    T: bignum::BigNumTrait+bignum::BigNumExt,
    for<'a1, 'a2> &'a1 T: NumOps<&'a2 T, T>,
{
    pub fn generate(params: &'a DsaParams<T>) -> Self {
        let x = rand_range_safe(&params.q);
        DsaPrivate {
            params: params,
            x: x,
        }
    }

    pub fn sign(&self, m: &T) -> (Signature<T>, T) {
        let zero = T::zero();
        let one = T::one();
        let p = &self.params.p;
        let q = &self.params.q;
        let g = &self.params.g;

        let mut k: T;
        let mut r: T;
        let mut s: T;
        loop {
            k = T::gen_below(q);
            if k == zero || k == one {
                continue;
            }
            r = g.mod_exp(&k, p);
            r = r.mod_math(q);
            if r == zero {
                continue;
            }
            s = &k.invmod(q).unwrap() * &(m + &(&r * &self.x));
            s = s.mod_math(q);
            if s == zero {
                continue;
            }
            break;
        }
        (Signature { r: r, s: s }, k)
    }
}

pub fn rand_range_safe<T: bignum::BigNumTrait>(q: &T) -> T {
    let zero = T::zero();
    let one = T::one();
    let mut x = T::clone(&zero);
    while x == zero || x == one {
        x = T::gen_below(q);
    }
    x
}

pub fn fake_signature<T>(public: &DsaPublic<T>) -> Signature<T>
where
    T: bignum::BigNumTrait+bignum::BigNumExt,
    for<'b> &'b T: NumOps<&'b T, T>,
{
    let p = &public.params.p;
    let q = &public.params.q;
    let z = rand_range_safe(q);
    let mut r = T::mod_exp(&public.y, &z, p);
    r = T::mod_math(&r, q);
    let mut s = &r * &T::invmod(&z, q).unwrap();
    s = T::mod_math(&s, q);
    Signature { r: r, s: s }
}

pub struct DsaParams<T> {
    pub p: T,
    pub q: T,
    pub g: T,
}

impl<T: bignum::BigNumTrait> DsaParams<T> {
    pub fn generate() -> Self {
        let p = T::from_hex_str(
            "800000000000000089e1855218a0e7dac38136ffafa72eda7859f2171e25e65eac698c1702578b07dc2a1076da241c76c62d374d8389ea5aeffd3226a0530cc565f3bf6b50929139ebeac04f48c3c84afb796d61e5a4f9a8fda812ab59494232c7d2b4deb50aa18ee9e132bfa85ac4374d7f9091abc3d015efc871a584471bb1"
            ).unwrap();
        let q = T::from_hex_str("f4f47f05794b256174bba6e9b396a7707e563c5b").unwrap();
        let g = T::from_hex_str("5958c9d3898b224b12672c0b98e06c60df923cb8bc999d119458fef538b8fa4046c8db53039db620c094c9fa077ef389b5322a559946a71903f990f1f7e0e025e2d7f7cf494aff1a0470f5b64c36b625a097f1651fe775323556fe00b3608c887892878480e99041be601a62166ca6894bdd41a7054ec89f756ba9fc95302291").unwrap();
        DsaParams { p: p, q: q, g: g }
    }
}

pub struct Signature<T> {
    pub r: T,
    pub s: T,
}
<|endoftext|>


<commit_before>use sodiumoxide::crypto::hash::sha256;
use sodiumoxide::crypto::sign;
use sodiumoxide::randombytes;
use bs58;

mod error;
pub use self::error::*;

const PUB_ADDRESS_PREFIX: &str = "GOD";
const PRIV_BUF_PREFIX: u8 = 0x01;
const PUB_BUF_PREFIX: u8 = 0x02;

pub trait Wif<T> {
    fn from_wif(s: &str) -> Result<T, WifError>;
    fn to_wif(&self) -> Box<str>;
}

#[derive(Debug, Clone, PartialEq)]
pub struct PublicKey {
    key: sign::PublicKey
}

impl PublicKey {
    #[inline]
    pub fn as_bytes(&self) -> &[u8] {
        self.key.as_ref()
    }

    #[inline]
    pub fn verify(&self, msg: &[u8], sig: &sign::Signature) -> bool {
        sign::verify_detached(sig, msg, &self.key)
    }

    #[inline]
    pub fn from_bytes(bytes: &[u8]) -> Option<PublicKey> {
        let key = sign::PublicKey::from_slice(bytes)?;
        Some(PublicKey { key })
    }
}

impl Wif<PublicKey> for PublicKey {
    fn from_wif(s: &str) -> Result<PublicKey, WifError> {
        if s.len() < 3 || &s[0..3] != PUB_ADDRESS_PREFIX {
            return Err(WifError::new(WifErrorKind::InvalidPrefix))
        }
        let raw = match bs58::decode(&s[3..]).into_vec() {
            Ok(bytes) => bytes,
            Err(_) => return Err(WifError::new(WifErrorKind::InvalidBs58Encoding))
        };
        if raw.len() != 37 {
            return Err(WifError::new(WifErrorKind::InvalidLen))
        } else if raw[0] != PUB_BUF_PREFIX {
            return Err(WifError::new(WifErrorKind::InvalidPrefix))
        }

        let prefixed_key = &raw[0..raw.len() - 4];
        {
            let checksum_a = &raw[raw.len() - 4 .. raw.len()];
            let checksum_b = &double_sha256(prefixed_key)[0..4];
            if checksum_a != checksum_b {
                return Err(WifError::new(WifErrorKind::InvalidChecksum))
            }
        }

        let key = &prefixed_key[1 .. prefixed_key.len()];
        Ok(PublicKey {
            key: sign::PublicKey::from_slice(key).unwrap()
        })
    }

    fn to_wif(&self) -> Box<str> {
        let mut buf: Vec<u8> = Vec::<u8>::with_capacity(37);
        buf.push(PUB_BUF_PREFIX);
        buf.extend_from_slice(self.key.as_ref());

        let checksum = &double_sha256(&buf)[0..4];
        buf.extend_from_slice(checksum);

        let mut s = bs58::encode(buf).into_string();
        s.insert_str(0, PUB_ADDRESS_PREFIX);
        s.into_boxed_str()
    }
}

#[derive(Debug, Clone)]
pub struct PrivateKey {
    seed: sign::Seed,
    key: sign::SecretKey
}

impl PrivateKey {
    pub fn as_bytes(&self) -> ([u8; sign::SEEDBYTES], [u8; sign::SECRETKEYBYTES]) {
        (self.seed.0, self.key.0)
    }

    #[inline]
    pub fn sign(&self, msg: &[u8]) -> sign::Signature {
        sign::sign_detached(msg, &self.key)
    }
}

impl Wif<KeyPair> for PrivateKey {
    fn from_wif(s: &str) -> Result<KeyPair, WifError> {
        let raw = match bs58::decode(s).into_vec() {
            Ok(bytes) => bytes,
            Err(_) => return Err(WifError::new(WifErrorKind::InvalidBs58Encoding))
        };
        if raw.len() != 37 {
            return Err(WifError::new(WifErrorKind::InvalidLen))
        } else if raw[0] != PRIV_BUF_PREFIX {
            return Err(WifError::new(WifErrorKind::InvalidPrefix))
        }

        let key = &raw[0..raw.len() - 4];
        {
            let checksum_a = &raw[raw.len() - 4 .. raw.len()];
            let checksum_b = &double_sha256(key)[0..4];
            if checksum_a != checksum_b {
                return Err(WifError::new(WifErrorKind::InvalidChecksum))
            }
        }

        let seed = sign::Seed::from_slice(&key[1..]).unwrap();
        let (pk, sk) = sign::keypair_from_seed(&seed);
        Ok(KeyPair(PublicKey {
            key: pk
        }, PrivateKey {
            seed,
            key: sk
        }))
    }

    fn to_wif(&self) -> Box<str> {
        let mut buf = Vec::<u8>::with_capacity(37);
        buf.push(PRIV_BUF_PREFIX);
        buf.extend_from_slice(&self.seed.0);

        let checksum = &double_sha256(&buf)[0..4];
        buf.extend_from_slice(checksum);

        bs58::encode(buf).into_string().into_boxed_str()
    }
}

#[derive(Debug, Clone)]
pub struct KeyPair(pub PublicKey, pub PrivateKey);

impl KeyPair {
    #[inline]
    pub fn sign(&self, msg: &[u8]) -> SigPair {
        SigPair {
            pub_key: self.0.clone(),
            signature: self.1.sign(msg)
        }
    }

    #[inline]
    pub fn verify(&self, msg: &[u8], sig: &sign::Signature) -> bool {
        PublicKey::verify(&self.0, msg, sig)
    }

    pub fn gen_keypair() -> KeyPair {
        let mut raw_seed: [u8; sign::SEEDBYTES] = [0; sign::SEEDBYTES];
        randombytes::randombytes_into(&mut raw_seed);
        let seed = sign::Seed::from_slice(&raw_seed).unwrap();
        let (pk, sk) = sign::keypair_from_seed(&seed);
        KeyPair(PublicKey {
            key: pk
        }, PrivateKey {
            seed,
            key: sk
        })
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct SigPair {
    pub pub_key: PublicKey,
    pub signature: sign::Signature
}

impl SigPair {
    #[inline]
    pub fn verify(&self, msg: &[u8]) -> bool {
        sign::verify_detached(&self.signature, msg, &self.pub_key.key)
    }
}

#[inline]
pub fn double_sha256(buf: &[u8]) -> sha256::Digest {
    sha256::hash(sha256::hash(buf).as_ref())
}

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

    #[test]
    fn test_create_and_recover_keys() {
        let kp = KeyPair::gen_keypair();

        let pk = &*kp.0.to_wif();
        assert_eq!(&*PublicKey::from_wif(pk).unwrap().to_wif(), pk);

        let sk = &*kp.1.to_wif();
        assert_eq!(&*PrivateKey::from_wif(sk).unwrap().1.to_wif(), sk);
    }

    #[test]
    fn test_import_keys_from_wif() {
        let kp = PrivateKey::from_wif("3GAD3otqozDorfu1iDpMQJ1gzWp8PRFEjVHZivZdedKW3i3KtM").unwrap();
        assert_eq!(&*kp.1.to_wif(), "3GAD3otqozDorfu1iDpMQJ1gzWp8PRFEjVHZivZdedKW3i3KtM");
        assert_eq!(&*kp.0.to_wif(), "GOD52QZDBUStV5CudxvKf6bPsQeN7oeKTkEm2nAU1vAUqNVexGTb8");
    }

    #[test]
    fn test_invalid_prefix() {
        let mut bytes = bs58::decode("3GAD3otqozDorfu1iDpMQJ1gzWp8PRFEjVHZivZdedKW3i3KtM").into_vec().unwrap();
        bytes[0] = 255;
        let wif = bs58::encode(bytes).into_string();
        assert_eq!(PrivateKey::from_wif(&wif).unwrap_err().kind, WifErrorKind::InvalidPrefix);

        let mut bytes = bs58::decode("52QZDBUStV5CudxvKf6bPsQeN7oeKTkEm2nAU1vAUqNVexGTb8").into_vec().unwrap();
        bytes[0] = 255;
        let mut wif = bs58::encode(bytes).into_string();
        wif.insert_str(0, PUB_ADDRESS_PREFIX);
        assert_eq!(PublicKey::from_wif(&wif).unwrap_err().kind, WifErrorKind::InvalidPrefix);
    }

    #[test]
    fn test_invalid_checksum() {
        let mut bytes = bs58::decode("3GAD3otqozDorfu1iDpMQJ1gzWp8PRFEjVHZivZdedKW3i3KtM").into_vec().unwrap();
        let len = bytes.len();
        for i in 1..5 {
            bytes[len - i] = 0;
        }
        let wif = bs58::encode(bytes).into_string();
        assert_eq!(PrivateKey::from_wif(&wif).unwrap_err().kind, WifErrorKind::InvalidChecksum);

        let mut bytes = bs58::decode("52QZDBUStV5CudxvKf6bPsQeN7oeKTkEm2nAU1vAUqNVexGTb8").into_vec().unwrap();
        let len = bytes.len();
        for i in 1..5 {
            bytes[len - i] = 0;
        }
        let mut wif = bs58::encode(bytes).into_string();
        wif.insert_str(0, PUB_ADDRESS_PREFIX);
        assert_eq!(PublicKey::from_wif(&wif).unwrap_err().kind, WifErrorKind::InvalidChecksum);
    }

    #[test]
    fn test_sign_message() {
        let msg = "Hello world!".as_bytes();
        let kp = KeyPair::gen_keypair();

        let sig = &kp.1.sign(msg);
        assert!(kp.0.verify(msg, sig));

        let pair = SigPair {
            pub_key: kp.0,
            signature: *sig
        };
        assert!(pair.verify(msg));

        // Test bad keys
        let kp = KeyPair::gen_keypair();
        assert!(!kp.verify(msg, sig));
    }
}
<commit_msg>Add from_bytes to allow directly creating a private key<commit_after>use sodiumoxide::crypto::hash::sha256;
use sodiumoxide::crypto::sign;
use sodiumoxide::randombytes;
use bs58;

mod error;
pub use self::error::*;

const PUB_ADDRESS_PREFIX: &str = "GOD";
const PRIV_BUF_PREFIX: u8 = 0x01;
const PUB_BUF_PREFIX: u8 = 0x02;

pub trait Wif<T> {
    fn from_wif(s: &str) -> Result<T, WifError>;
    fn to_wif(&self) -> Box<str>;
}

#[derive(Debug, Clone, PartialEq)]
pub struct PublicKey {
    key: sign::PublicKey
}

impl PublicKey {
    #[inline]
    pub fn as_bytes(&self) -> &[u8] {
        self.key.as_ref()
    }

    #[inline]
    pub fn verify(&self, msg: &[u8], sig: &sign::Signature) -> bool {
        sign::verify_detached(sig, msg, &self.key)
    }

    #[inline]
    pub fn from_bytes(bytes: &[u8]) -> Option<PublicKey> {
        let key = sign::PublicKey::from_slice(bytes)?;
        Some(PublicKey { key })
    }
}

impl Wif<PublicKey> for PublicKey {
    fn from_wif(s: &str) -> Result<PublicKey, WifError> {
        if s.len() < 3 || &s[0..3] != PUB_ADDRESS_PREFIX {
            return Err(WifError::new(WifErrorKind::InvalidPrefix))
        }
        let raw = match bs58::decode(&s[3..]).into_vec() {
            Ok(bytes) => bytes,
            Err(_) => return Err(WifError::new(WifErrorKind::InvalidBs58Encoding))
        };
        if raw.len() != 37 {
            return Err(WifError::new(WifErrorKind::InvalidLen))
        } else if raw[0] != PUB_BUF_PREFIX {
            return Err(WifError::new(WifErrorKind::InvalidPrefix))
        }

        let prefixed_key = &raw[0..raw.len() - 4];
        {
            let checksum_a = &raw[raw.len() - 4 .. raw.len()];
            let checksum_b = &double_sha256(prefixed_key)[0..4];
            if checksum_a != checksum_b {
                return Err(WifError::new(WifErrorKind::InvalidChecksum))
            }
        }

        let key = &prefixed_key[1 .. prefixed_key.len()];
        Ok(PublicKey {
            key: sign::PublicKey::from_slice(key).unwrap()
        })
    }

    fn to_wif(&self) -> Box<str> {
        let mut buf: Vec<u8> = Vec::<u8>::with_capacity(37);
        buf.push(PUB_BUF_PREFIX);
        buf.extend_from_slice(self.key.as_ref());

        let checksum = &double_sha256(&buf)[0..4];
        buf.extend_from_slice(checksum);

        let mut s = bs58::encode(buf).into_string();
        s.insert_str(0, PUB_ADDRESS_PREFIX);
        s.into_boxed_str()
    }
}

#[derive(Debug, Clone)]
pub struct PrivateKey {
    seed: sign::Seed,
    key: sign::SecretKey
}

impl PrivateKey {
    #[inline]
    pub fn sign(&self, msg: &[u8]) -> sign::Signature {
        sign::sign_detached(msg, &self.key)
    }

    #[inline]
    pub fn as_bytes(&self) -> ([u8; sign::SEEDBYTES], [u8; sign::SECRETKEYBYTES]) {
        (self.seed.0, self.key.0)
    }

    pub fn from_bytes(&self,
                        seed: &[u8; sign::SEEDBYTES],
                        key: &[u8; sign::SECRETKEYBYTES]) -> Option<PrivateKey> {
        Some(PrivateKey {
            seed: sign::Seed::from_slice(seed)?,
            key: sign::SecretKey::from_slice(key)?
        })
    }
}

impl Wif<KeyPair> for PrivateKey {
    fn from_wif(s: &str) -> Result<KeyPair, WifError> {
        let raw = match bs58::decode(s).into_vec() {
            Ok(bytes) => bytes,
            Err(_) => return Err(WifError::new(WifErrorKind::InvalidBs58Encoding))
        };
        if raw.len() != 37 {
            return Err(WifError::new(WifErrorKind::InvalidLen))
        } else if raw[0] != PRIV_BUF_PREFIX {
            return Err(WifError::new(WifErrorKind::InvalidPrefix))
        }

        let key = &raw[0..raw.len() - 4];
        {
            let checksum_a = &raw[raw.len() - 4 .. raw.len()];
            let checksum_b = &double_sha256(key)[0..4];
            if checksum_a != checksum_b {
                return Err(WifError::new(WifErrorKind::InvalidChecksum))
            }
        }

        let seed = sign::Seed::from_slice(&key[1..]).unwrap();
        let (pk, sk) = sign::keypair_from_seed(&seed);
        Ok(KeyPair(PublicKey {
            key: pk
        }, PrivateKey {
            seed,
            key: sk
        }))
    }

    fn to_wif(&self) -> Box<str> {
        let mut buf = Vec::<u8>::with_capacity(37);
        buf.push(PRIV_BUF_PREFIX);
        buf.extend_from_slice(&self.seed.0);

        let checksum = &double_sha256(&buf)[0..4];
        buf.extend_from_slice(checksum);

        bs58::encode(buf).into_string().into_boxed_str()
    }
}

#[derive(Debug, Clone)]
pub struct KeyPair(pub PublicKey, pub PrivateKey);

impl KeyPair {
    #[inline]
    pub fn sign(&self, msg: &[u8]) -> SigPair {
        SigPair {
            pub_key: self.0.clone(),
            signature: self.1.sign(msg)
        }
    }

    #[inline]
    pub fn verify(&self, msg: &[u8], sig: &sign::Signature) -> bool {
        PublicKey::verify(&self.0, msg, sig)
    }

    pub fn gen_keypair() -> KeyPair {
        let mut raw_seed: [u8; sign::SEEDBYTES] = [0; sign::SEEDBYTES];
        randombytes::randombytes_into(&mut raw_seed);
        let seed = sign::Seed::from_slice(&raw_seed).unwrap();
        let (pk, sk) = sign::keypair_from_seed(&seed);
        KeyPair(PublicKey {
            key: pk
        }, PrivateKey {
            seed,
            key: sk
        })
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct SigPair {
    pub pub_key: PublicKey,
    pub signature: sign::Signature
}

impl SigPair {
    #[inline]
    pub fn verify(&self, msg: &[u8]) -> bool {
        sign::verify_detached(&self.signature, msg, &self.pub_key.key)
    }
}

#[inline]
pub fn double_sha256(buf: &[u8]) -> sha256::Digest {
    sha256::hash(sha256::hash(buf).as_ref())
}

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

    #[test]
    fn test_create_and_recover_keys() {
        let kp = KeyPair::gen_keypair();

        let pk = &*kp.0.to_wif();
        assert_eq!(&*PublicKey::from_wif(pk).unwrap().to_wif(), pk);

        let sk = &*kp.1.to_wif();
        assert_eq!(&*PrivateKey::from_wif(sk).unwrap().1.to_wif(), sk);
    }

    #[test]
    fn test_import_keys_from_wif() {
        let kp = PrivateKey::from_wif("3GAD3otqozDorfu1iDpMQJ1gzWp8PRFEjVHZivZdedKW3i3KtM").unwrap();
        assert_eq!(&*kp.1.to_wif(), "3GAD3otqozDorfu1iDpMQJ1gzWp8PRFEjVHZivZdedKW3i3KtM");
        assert_eq!(&*kp.0.to_wif(), "GOD52QZDBUStV5CudxvKf6bPsQeN7oeKTkEm2nAU1vAUqNVexGTb8");
    }

    #[test]
    fn test_invalid_prefix() {
        let mut bytes = bs58::decode("3GAD3otqozDorfu1iDpMQJ1gzWp8PRFEjVHZivZdedKW3i3KtM").into_vec().unwrap();
        bytes[0] = 255;
        let wif = bs58::encode(bytes).into_string();
        assert_eq!(PrivateKey::from_wif(&wif).unwrap_err().kind, WifErrorKind::InvalidPrefix);

        let mut bytes = bs58::decode("52QZDBUStV5CudxvKf6bPsQeN7oeKTkEm2nAU1vAUqNVexGTb8").into_vec().unwrap();
        bytes[0] = 255;
        let mut wif = bs58::encode(bytes).into_string();
        wif.insert_str(0, PUB_ADDRESS_PREFIX);
        assert_eq!(PublicKey::from_wif(&wif).unwrap_err().kind, WifErrorKind::InvalidPrefix);
    }

    #[test]
    fn test_invalid_checksum() {
        let mut bytes = bs58::decode("3GAD3otqozDorfu1iDpMQJ1gzWp8PRFEjVHZivZdedKW3i3KtM").into_vec().unwrap();
        let len = bytes.len();
        for i in 1..5 {
            bytes[len - i] = 0;
        }
        let wif = bs58::encode(bytes).into_string();
        assert_eq!(PrivateKey::from_wif(&wif).unwrap_err().kind, WifErrorKind::InvalidChecksum);

        let mut bytes = bs58::decode("52QZDBUStV5CudxvKf6bPsQeN7oeKTkEm2nAU1vAUqNVexGTb8").into_vec().unwrap();
        let len = bytes.len();
        for i in 1..5 {
            bytes[len - i] = 0;
        }
        let mut wif = bs58::encode(bytes).into_string();
        wif.insert_str(0, PUB_ADDRESS_PREFIX);
        assert_eq!(PublicKey::from_wif(&wif).unwrap_err().kind, WifErrorKind::InvalidChecksum);
    }

    #[test]
    fn test_sign_message() {
        let msg = "Hello world!".as_bytes();
        let kp = KeyPair::gen_keypair();

        let sig = &kp.1.sign(msg);
        assert!(kp.0.verify(msg, sig));

        let pair = SigPair {
            pub_key: kp.0,
            signature: *sig
        };
        assert!(pair.verify(msg));

        // Test bad keys
        let kp = KeyPair::gen_keypair();
        assert!(!kp.verify(msg, sig));
    }
}
<|endoftext|>