<commit_before>use std::io;
use byteorder::{ByteOrder, LittleEndian};
use key;

pub struct ChaCha20 {
    state: [u32; 16],
}

impl ChaCha20 {
    pub fn new(key: &[u8; 32]) -> io::Result<(Self, [u8; 12])> {
        let nonce = key::gen()?;
        let mut chacha20 = Self { state: [0; 16] };
        Self::setup_state(&mut chacha20.state, &key, &nonce);
        Ok((chacha20, nonce))
    }

    pub fn get_block(&self, counter: u32) -> [u8; 64] {
        Self::serialize_block(self.block(counter))
    }

    fn block(&self, counter: u32) -> [u32; 16] {
        let mut state = [0; 16];
        state.copy_from_slice(&self.state);
        state[12] = counter;
        for _ in 0..10 {
            Self::inner_block(&mut state);
        }
        for (x, &y) in state.iter_mut().zip(self.state.iter()) {
            *x = x.wrapping_add(y);
        }
        state[12] = state[12].wrapping_add(counter);
        state
    }

    fn setup_state(state: &mut [u32; 16], key: &[u8; 32], nonce: &[u8; 12]) {
        state[0] = 0x61707865;
        state[1] = 0x3320646e;
        state[2] = 0x79622d32;
        state[3] = 0x6b206574;
        state[4..12].copy_from_slice(&Self::convert_key(&key));
        state[13..].copy_from_slice(&Self::convert_nonce(&nonce));
    }

    fn inner_block(state: &mut [u32; 16]) {
        Self::quarter_round(state, 0, 4, 8, 12);
        Self::quarter_round(state, 1, 5, 9, 13);
        Self::quarter_round(state, 2, 6, 10, 14);
        Self::quarter_round(state, 3, 7, 11, 15);
        Self::quarter_round(state, 0, 5, 10, 15);
        Self::quarter_round(state, 1, 6, 11, 12);
        Self::quarter_round(state, 2, 7, 8, 13);
        Self::quarter_round(state, 3, 4, 9, 14);
    }

    fn quarter_round(state: &mut [u32; 16], i: usize, j: usize, k: usize, l: usize) {
        let mut a = state[i];
        let mut b = state[j];
        let mut c = state[k];
        let mut d = state[l];
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(16);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(12);
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(8);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(7);
        state[i] = a;
        state[j] = b;
        state[k] = c;
        state[l] = d;
    }

    fn convert_key(key: &[u8; 32]) -> [u32; 8] {
        let mut ret = [0; 8];
        for (byte, chunk) in ret.iter_mut().zip(key.chunks(4)) {
            *byte = LittleEndian::read_u32(chunk);
        }
        ret
    }

    fn convert_nonce(nonce: &[u8; 12]) -> [u32; 3] {
        let mut ret = [0; 3];
        for (byte, chunk) in ret.iter_mut().zip(nonce.chunks(4)) {
            *byte = LittleEndian::read_u32(chunk);
        }
        ret
    }

    fn serialize_block(block: [u32; 16]) -> [u8; 64] {
        let mut ret = [0; 64];
        for (chunk, &byte) in ret.chunks_mut(4).zip(block.iter()) {
            LittleEndian::write_u32(chunk, byte);
        }
        ret
    }
}

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

    const KEY: [u8; 32] = [
        0x00,
        0x01,
        0x02,
        0x03,
        0x04,
        0x05,
        0x06,
        0x07,
        0x08,
        0x09,
        0x0a,
        0x0b,
        0x0c,
        0x0d,
        0x0e,
        0x0f,
        0x10,
        0x11,
        0x12,
        0x13,
        0x14,
        0x15,
        0x16,
        0x17,
        0x18,
        0x19,
        0x1a,
        0x1b,
        0x1c,
        0x1d,
        0x1e,
        0x1f,
    ];
    const NONCE: [u8; 12] = [
        0x00,
        0x00,
        0x00,
        0x09,
        0x00,
        0x00,
        0x00,
        0x4a,
        0x00,
        0x00,
        0x00,
        0x00,
    ];
    const BLOCK_ONE: [u8; 64] = [
        0x10,
        0xf1,
        0xe7,
        0xe4,
        0xd1,
        0x3b,
        0x59,
        0x15,
        0x50,
        0x0f,
        0xdd,
        0x1f,
        0xa3,
        0x20,
        0x71,
        0xc4,
        0xc7,
        0xd1,
        0xf4,
        0xc7,
        0x33,
        0xc0,
        0x68,
        0x03,
        0x04,
        0x22,
        0xaa,
        0x9a,
        0xc3,
        0xd4,
        0x6c,
        0x4e,
        0xd2,
        0x82,
        0x64,
        0x46,
        0x07,
        0x9f,
        0xaa,
        0x09,
        0x14,
        0xc2,
        0xd7,
        0x05,
        0xd9,
        0x8b,
        0x02,
        0xa2,
        0xb5,
        0x12,
        0x9c,
        0xd1,
        0xde,
        0x16,
        0x4e,
        0xb9,
        0xcb,
        0xd0,
        0x83,
        0xe8,
        0xa2,
        0x50,
        0x3c,
        0x4e,
    ];
    const SETUP_STATE: [u32; 16] = [
        0x61707865,
        0x3320646e,
        0x79622d32,
        0x6b206574,
        0x03020100,
        0x07060504,
        0x0b0a0908,
        0x0f0e0d0c,
        0x13121110,
        0x17161514,
        0x1b1a1918,
        0x1f1e1d1c,
        0x00000000,
        0x09000000,
        0x4a000000,
        0x00000000,
    ];
    const FINAL_STATE: [u32; 16] = [
        0xe4e7f110,
        0x15593bd1,
        0x1fdd0f50,
        0xc47120a3,
        0xc7f4d1c7,
        0x0368c033,
        0x9aaa2204,
        0x4e6cd4c3,
        0x466482d2,
        0x09aa9f07,
        0x05d7c214,
        0xa2028bd9,
        0xd19c12b5,
        0xb94e16de,
        0xe883d0cb,
        0x4e3c50a2,
    ];

    #[test]
    fn test_new() {
        let (chacha20, nonce) = ChaCha20::new(&KEY).unwrap();
        let mut state = [0; 16];
        ChaCha20::setup_state(&mut state, &KEY, &nonce);
        assert_eq!(state, chacha20.state);
    }

    fn check_serialized_block(block: &[u8]) {
        assert_eq!(BLOCK_ONE.len(), block.len());
        for (lhs, rhs) in BLOCK_ONE.iter().zip(block) {
            assert_eq!(lhs, rhs);
        }
    }

    #[test]
    fn test_get_block() {
        let mut chacha20 = ChaCha20 { state: [0; 16] };
        ChaCha20::setup_state(&mut chacha20.state, &KEY, &NONCE);
        check_serialized_block(&chacha20.get_block(1));
    }

    #[test]
    fn test_block() {
        let mut chacha20 = ChaCha20 { state: [0; 16] };
        ChaCha20::setup_state(&mut chacha20.state, &KEY, &NONCE);
        assert_eq!(FINAL_STATE, chacha20.block(1));
    }

    #[test]
    fn test_setup_state() {
        let mut state = [0; 16];
        ChaCha20::setup_state(&mut state, &KEY, &NONCE);
        assert_eq!(SETUP_STATE, state);
    }

    #[test]
    fn test_inner_block() {
        let mut state = SETUP_STATE;
        state[12] = 1;
        for _ in 0..10 {
            ChaCha20::inner_block(&mut state);
        }
        assert_eq!(
            [
                0x837778ab,
                0xe238d763,
                0xa67ae21e,
                0x5950bb2f,
                0xc4f2d0c7,
                0xfc62bb2f,
                0x8fa018fc,
                0x3f5ec7b7,
                0x335271c2,
                0xf29489f3,
                0xeabda8fc,
                0x82e46ebd,
                0xd19c12b4,
                0xb04e16de,
                0x9e83d0cb,
                0x4e3c50a2,
            ],
            state
        );
    }

    #[test]
    fn test_quarter_round() {
        let mut state = [0; 16];
        state[0] = 0x11111111;
        state[1] = 0x01020304;
        state[2] = 0x9b8d6f43;
        state[3] = 0x01234567;
        ChaCha20::quarter_round(&mut state, 0, 1, 2, 3);
        assert_eq!([0xea2a92f4, 0xcb1cf8ce, 0x4581472e, 0x5881c4bb], state[..4]);

        state = [
            0x879531e0,
            0xc5ecf37d,
            0x516461b1,
            0xc9a62f8a,
            0x44c20ef3,
            0x3390af7f,
            0xd9fc690b,
            0x2a5f714c,
            0x53372767,
            0xb00a5631,
            0x974c541a,
            0x359e9963,
            0x5c971061,
            0x3d631689,
            0x2098d9d6,
            0x91dbd320,
        ];
        ChaCha20::quarter_round(&mut state, 2, 7, 8, 13);
        assert_eq!(
            [
                0x879531e0,
                0xc5ecf37d,
                0xbdb886dc,
                0xc9a62f8a,
                0x44c20ef3,
                0x3390af7f,
                0xd9fc690b,
                0xcfacafd2,
                0xe46bea80,
                0xb00a5631,
                0x974c541a,
                0x359e9963,
                0x5c971061,
                0xccc07c79,
                0x2098d9d6,
                0x91dbd320,
            ],
            state
        );
    }

    #[test]
    fn test_convert_key() {
        assert_eq!(
            [
                0x03020100,
                0x07060504,
                0x0b0a0908,
                0x0f0e0d0c,
                0x13121110,
                0x17161514,
                0x1b1a1918,
                0x1f1e1d1c,
            ],
            ChaCha20::convert_key(&KEY)
        );
    }

    #[test]
    fn test_convert_nonce() {
        assert_eq!(
            [0x09000000, 0x4a000000, 0x00000000],
            ChaCha20::convert_nonce(&NONCE)
        );
    }

    #[test]
    fn test_serialize_block() {
        check_serialized_block(&ChaCha20::serialize_block(FINAL_STATE));
    }
}
<commit_msg>chacha20: allow specifying a nonce + tiny changes<commit_after>use std::io;
use byteorder::{ByteOrder, LittleEndian};
use key;

pub struct ChaCha20 {
    state: [u32; 16],
}

pub fn gen_nonce() -> io::Result<[u8; 12]> {
    Ok(key::gen()?)
}

impl ChaCha20 {
    /// Be sure never to encrypt with a given (key, nonce) pair more than once
    pub fn new(key: &[u8; 32], nonce: &[u8; 12]) -> Self {
        let mut chacha20 = Self { state: [0; 16] };
        Self::setup_state(&mut chacha20.state, &key, &nonce);
        chacha20
    }

    pub fn get_block(&self, counter: u32) -> [u8; 64] {
        Self::serialize_block(self.block(counter))
    }

    fn block(&self, counter: u32) -> [u32; 16] {
        let mut state = [0; 16];
        state.copy_from_slice(&self.state);
        state[12] = counter;
        for _ in 0..10 {
            Self::inner_block(&mut state);
        }
        for (x, &y) in state.iter_mut().zip(&self.state) {
            *x = x.wrapping_add(y);
        }
        state[12] = state[12].wrapping_add(counter);
        state
    }

    fn setup_state(state: &mut [u32; 16], key: &[u8; 32], nonce: &[u8; 12]) {
        state[0] = 0x61707865;
        state[1] = 0x3320646e;
        state[2] = 0x79622d32;
        state[3] = 0x6b206574;
        state[4..12].copy_from_slice(&Self::convert_key(&key));
        state[13..].copy_from_slice(&Self::convert_nonce(&nonce));
    }

    fn inner_block(state: &mut [u32; 16]) {
        Self::quarter_round(state, 0, 4, 8, 12);
        Self::quarter_round(state, 1, 5, 9, 13);
        Self::quarter_round(state, 2, 6, 10, 14);
        Self::quarter_round(state, 3, 7, 11, 15);
        Self::quarter_round(state, 0, 5, 10, 15);
        Self::quarter_round(state, 1, 6, 11, 12);
        Self::quarter_round(state, 2, 7, 8, 13);
        Self::quarter_round(state, 3, 4, 9, 14);
    }

    fn quarter_round(state: &mut [u32; 16], i: usize, j: usize, k: usize, l: usize) {
        let mut a = state[i];
        let mut b = state[j];
        let mut c = state[k];
        let mut d = state[l];
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(16);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(12);
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(8);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(7);
        state[i] = a;
        state[j] = b;
        state[k] = c;
        state[l] = d;
    }

    fn convert_key(key: &[u8; 32]) -> [u32; 8] {
        let mut ret = [0; 8];
        for (byte, chunk) in ret.iter_mut().zip(key.chunks(4)) {
            *byte = LittleEndian::read_u32(chunk);
        }
        ret
    }

    fn convert_nonce(nonce: &[u8; 12]) -> [u32; 3] {
        let mut ret = [0; 3];
        for (byte, chunk) in ret.iter_mut().zip(nonce.chunks(4)) {
            *byte = LittleEndian::read_u32(chunk);
        }
        ret
    }

    fn serialize_block(block: [u32; 16]) -> [u8; 64] {
        let mut ret = [0; 64];
        for (chunk, &byte) in ret.chunks_mut(4).zip(&block) {
            LittleEndian::write_u32(chunk, byte);
        }
        ret
    }
}

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

    const KEY: [u8; 32] = [
        0x00,
        0x01,
        0x02,
        0x03,
        0x04,
        0x05,
        0x06,
        0x07,
        0x08,
        0x09,
        0x0a,
        0x0b,
        0x0c,
        0x0d,
        0x0e,
        0x0f,
        0x10,
        0x11,
        0x12,
        0x13,
        0x14,
        0x15,
        0x16,
        0x17,
        0x18,
        0x19,
        0x1a,
        0x1b,
        0x1c,
        0x1d,
        0x1e,
        0x1f,
    ];
    const NONCE: [u8; 12] = [
        0x00,
        0x00,
        0x00,
        0x09,
        0x00,
        0x00,
        0x00,
        0x4a,
        0x00,
        0x00,
        0x00,
        0x00,
    ];
    const BLOCK_ONE: [u8; 64] = [
        0x10,
        0xf1,
        0xe7,
        0xe4,
        0xd1,
        0x3b,
        0x59,
        0x15,
        0x50,
        0x0f,
        0xdd,
        0x1f,
        0xa3,
        0x20,
        0x71,
        0xc4,
        0xc7,
        0xd1,
        0xf4,
        0xc7,
        0x33,
        0xc0,
        0x68,
        0x03,
        0x04,
        0x22,
        0xaa,
        0x9a,
        0xc3,
        0xd4,
        0x6c,
        0x4e,
        0xd2,
        0x82,
        0x64,
        0x46,
        0x07,
        0x9f,
        0xaa,
        0x09,
        0x14,
        0xc2,
        0xd7,
        0x05,
        0xd9,
        0x8b,
        0x02,
        0xa2,
        0xb5,
        0x12,
        0x9c,
        0xd1,
        0xde,
        0x16,
        0x4e,
        0xb9,
        0xcb,
        0xd0,
        0x83,
        0xe8,
        0xa2,
        0x50,
        0x3c,
        0x4e,
    ];
    const SETUP_STATE: [u32; 16] = [
        0x61707865,
        0x3320646e,
        0x79622d32,
        0x6b206574,
        0x03020100,
        0x07060504,
        0x0b0a0908,
        0x0f0e0d0c,
        0x13121110,
        0x17161514,
        0x1b1a1918,
        0x1f1e1d1c,
        0x00000000,
        0x09000000,
        0x4a000000,
        0x00000000,
    ];
    const FINAL_STATE: [u32; 16] = [
        0xe4e7f110,
        0x15593bd1,
        0x1fdd0f50,
        0xc47120a3,
        0xc7f4d1c7,
        0x0368c033,
        0x9aaa2204,
        0x4e6cd4c3,
        0x466482d2,
        0x09aa9f07,
        0x05d7c214,
        0xa2028bd9,
        0xd19c12b5,
        0xb94e16de,
        0xe883d0cb,
        0x4e3c50a2,
    ];

    #[test]
    fn test_new() {
        let chacha20 = ChaCha20::new(&KEY, &NONCE);
        let mut state = [0; 16];
        ChaCha20::setup_state(&mut state, &KEY, &NONCE);
        assert_eq!(state, chacha20.state);
    }

    fn check_serialized_block(block: &[u8]) {
        assert_eq!(BLOCK_ONE.len(), block.len());
        for (lhs, rhs) in BLOCK_ONE.iter().zip(block) {
            assert_eq!(lhs, rhs);
        }
    }

    #[test]
    fn test_get_block() {
        let mut chacha20 = ChaCha20 { state: [0; 16] };
        ChaCha20::setup_state(&mut chacha20.state, &KEY, &NONCE);
        check_serialized_block(&chacha20.get_block(1));
    }

    #[test]
    fn test_block() {
        let mut chacha20 = ChaCha20 { state: [0; 16] };
        ChaCha20::setup_state(&mut chacha20.state, &KEY, &NONCE);
        assert_eq!(FINAL_STATE, chacha20.block(1));
    }

    #[test]
    fn test_setup_state() {
        let mut state = [0; 16];
        ChaCha20::setup_state(&mut state, &KEY, &NONCE);
        assert_eq!(SETUP_STATE, state);
    }

    #[test]
    fn test_inner_block() {
        let mut state = SETUP_STATE;
        state[12] = 1;
        for _ in 0..10 {
            ChaCha20::inner_block(&mut state);
        }
        assert_eq!(
            [
                0x837778ab,
                0xe238d763,
                0xa67ae21e,
                0x5950bb2f,
                0xc4f2d0c7,
                0xfc62bb2f,
                0x8fa018fc,
                0x3f5ec7b7,
                0x335271c2,
                0xf29489f3,
                0xeabda8fc,
                0x82e46ebd,
                0xd19c12b4,
                0xb04e16de,
                0x9e83d0cb,
                0x4e3c50a2,
            ],
            state
        );
    }

    #[test]
    fn test_quarter_round() {
        let mut state = [0; 16];
        state[0] = 0x11111111;
        state[1] = 0x01020304;
        state[2] = 0x9b8d6f43;
        state[3] = 0x01234567;
        ChaCha20::quarter_round(&mut state, 0, 1, 2, 3);
        assert_eq!([0xea2a92f4, 0xcb1cf8ce, 0x4581472e, 0x5881c4bb], state[..4]);

        state = [
            0x879531e0,
            0xc5ecf37d,
            0x516461b1,
            0xc9a62f8a,
            0x44c20ef3,
            0x3390af7f,
            0xd9fc690b,
            0x2a5f714c,
            0x53372767,
            0xb00a5631,
            0x974c541a,
            0x359e9963,
            0x5c971061,
            0x3d631689,
            0x2098d9d6,
            0x91dbd320,
        ];
        ChaCha20::quarter_round(&mut state, 2, 7, 8, 13);
        assert_eq!(
            [
                0x879531e0,
                0xc5ecf37d,
                0xbdb886dc,
                0xc9a62f8a,
                0x44c20ef3,
                0x3390af7f,
                0xd9fc690b,
                0xcfacafd2,
                0xe46bea80,
                0xb00a5631,
                0x974c541a,
                0x359e9963,
                0x5c971061,
                0xccc07c79,
                0x2098d9d6,
                0x91dbd320,
            ],
            state
        );
    }

    #[test]
    fn test_convert_key() {
        assert_eq!(
            [
                0x03020100,
                0x07060504,
                0x0b0a0908,
                0x0f0e0d0c,
                0x13121110,
                0x17161514,
                0x1b1a1918,
                0x1f1e1d1c,
            ],
            ChaCha20::convert_key(&KEY)
        );
    }

    #[test]
    fn test_convert_nonce() {
        assert_eq!(
            [0x09000000, 0x4a000000, 0x00000000],
            ChaCha20::convert_nonce(&NONCE)
        );
    }

    #[test]
    fn test_serialize_block() {
        check_serialized_block(&ChaCha20::serialize_block(FINAL_STATE));
    }
}
<|endoftext|>


<commit_before>//! Implementation of the SPECK block cipher.
//!
//! SPECK is a really simple block cipher designed by the NSA. It is famous for its simple
//! structure and code size, which can fit in just a couple of lines, while still preserving
//! security.
#![feature(i128_type)]
#![no_std]
#![forbid(unsafe_code)]

/// The number of rounds.
const ROUNDS: u64 = 32;

/// A single round of SPECK.
///
/// This is a keyed ARX transformation.
macro_rules! round {
    ($x:ident, $y:ident, $k:ident) => {
        $x = $x.rotate_right(8);
        $x = $x.wrapping_add($y);
        $x ^= $k;
        $y = $y.rotate_left(3);
        $y ^= $x;
    }
}

/// Revert a SPECK round given some subkey.
macro_rules! inv_round {
    ($x:ident, $y:ident, $k:ident) => {
        $y ^= $x;
        $y = $y.rotate_right(3);
        $x ^= $k;
        $x = $x.wrapping_sub($y);
        $x = $x.rotate_left(8);
    }
}

/// Encrypt a block with key schedule generated on-the-go.
///
/// This works great for one-time use of a key (such as usages other than encryption), because it
/// should never read from memory (both the message and the keys are stored in the registers). As
/// such, this should be really fast for such usage.
///
/// If you want to reuse the key, however, it is recommended that you use the precomputed schedule
/// provided by the `Key` struct.
pub fn encrypt_block(m: u128, k: u128) -> u128 {
    let mut m1 = (m >> 64) as u64;
    let mut m2 = m as u64;
    let mut k1 = (k >> 64) as u64;
    let mut k2 = k as u64;

    // Run the initial round (similar to the loop below, but doesn't update the key schedule).
    round!(m1, m2, k2);

    for i in 0..ROUNDS - 1 {
        // Progress the key schedule.
        round!(k1, k2, i);
        // Run a round over the message.
        round!(m1, m2, k2);
    }

    m2 as u128 | (m1 as u128) << 64
}

/// A precomputed key.
///
/// This precomputes a key schedule, which can then be used for both encrypting and decrypting
/// messages.
#[derive(Debug, Copy, Clone)]
pub struct Key {
    /// The computed schedule.
    ///
    /// Each of these subkeys are used in a round of the cipher. The first subkey is used in the
    /// first round of the cipher and so on.
    schedule: [u64; ROUNDS as usize],
}

impl Key {
    /// Generate a new key from some seed.
    pub fn new(k: u128) -> Key {
        let mut k1 = (k >> 64) as u64;
        let mut k2 = k as u64;

        let mut ret = Key {
            schedule: [0; ROUNDS as usize],
        };

        // Run `ROUNDS - 1` rounds to generate the key's endpoint (the last key in the schedule).
        for i in 0..ROUNDS {
            // Insert the key into the schedule.
            ret.schedule[i as usize] = k2;

            // The beautiful thing about SPECK is that it reuses its round function to generate the
            // key schedule.
            round!(k1, k2, i);
        }

        ret
    }

    /// Encrypt a 128-bit block with this key.
    pub fn encrypt_block(&self, m: u128) -> u128 {
        let mut m1 = (m >> 64) as u64;
        let mut m2 = m as u64;

        // We run a round for every subkey in the generated key schedule.
        for &k in &self.schedule {
            // Run a round on the message.
            round!(m1, m2, k);
        }

        m2 as u128 | (m1 as u128) << 64
    }

    /// Decrypt a 128-bit block with this key.
    pub fn decrypt_block(&self, c: u128) -> u128 {
        let mut c1 = (c >> 64) as u64;
        let mut c2 = c as u64;

        // We run a round for every subkey in the generated key schedule.
        for &k in self.schedule.iter().rev() {
            // Run a round on the message.
            inv_round!(c1, c2, k);
        }

        c2 as u128 | (c1 as u128) << 64
    }
}

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

    #[test]
    fn encrypt_decrypt() {
        for mut x in 0u128..90000 {
            // <3
            x = x.wrapping_mul(0x6eed0e9da4d94a4f6eed0e9da4d94a4f);
            x ^= (x >> 6) >> (x >> 122);
            x = x.wrapping_mul(0x6eed0e9da4d94a4f6eed0e9da4d94a4f);

            let key = Key::new(!x);

            assert_eq!(key.decrypt_block(key.encrypt_block(x)), x);
            assert_eq!(key.encrypt_block(x), encrypt_block(x, !x));
        }
    }

    #[test]
    fn test_vectors() {
        // These test vectors are taken from the SPECK paper.
        assert_eq!(
            encrypt_block(
                0x6c617669757165207469206564616d20,
                0x0f0e0d0c0b0a09080706050403020100
            ),
            0xa65d9851797832657860fedf5c570d18
        );
    }
}
<commit_msg>Added `PartialEq`, `Eq`, and `Hash` to derived traits.<commit_after>//! Implementation of the SPECK block cipher.
//!
//! SPECK is a really simple block cipher designed by the NSA. It is famous for its simple
//! structure and code size, which can fit in just a couple of lines, while still preserving
//! security.
#![feature(i128_type)]
#![no_std]
#![forbid(unsafe_code)]

/// The number of rounds.
const ROUNDS: u64 = 32;

/// A single round of SPECK.
///
/// This is a keyed ARX transformation.
macro_rules! round {
    ($x:ident, $y:ident, $k:ident) => {
        $x = $x.rotate_right(8);
        $x = $x.wrapping_add($y);
        $x ^= $k;
        $y = $y.rotate_left(3);
        $y ^= $x;
    }
}

/// Revert a SPECK round given some subkey.
macro_rules! inv_round {
    ($x:ident, $y:ident, $k:ident) => {
        $y ^= $x;
        $y = $y.rotate_right(3);
        $x ^= $k;
        $x = $x.wrapping_sub($y);
        $x = $x.rotate_left(8);
    }
}

/// Encrypt a block with key schedule generated on-the-go.
///
/// This works great for one-time use of a key (such as usages other than encryption), because it
/// should never read from memory (both the message and the keys are stored in the registers). As
/// such, this should be really fast for such usage.
///
/// If you want to reuse the key, however, it is recommended that you use the precomputed schedule
/// provided by the `Key` struct.
pub fn encrypt_block(m: u128, k: u128) -> u128 {
    let mut m1 = (m >> 64) as u64;
    let mut m2 = m as u64;
    let mut k1 = (k >> 64) as u64;
    let mut k2 = k as u64;

    // Run the initial round (similar to the loop below, but doesn't update the key schedule).
    round!(m1, m2, k2);

    for i in 0..ROUNDS - 1 {
        // Progress the key schedule.
        round!(k1, k2, i);
        // Run a round over the message.
        round!(m1, m2, k2);
    }

    m2 as u128 | (m1 as u128) << 64
}

/// A precomputed key.
///
/// This precomputes a key schedule, which can then be used for both encrypting and decrypting
/// messages.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct Key {
    /// The computed schedule.
    ///
    /// Each of these subkeys are used in a round of the cipher. The first subkey is used in the
    /// first round of the cipher and so on.
    schedule: [u64; ROUNDS as usize],
}

impl Key {
    /// Generate a new key from some seed.
    pub fn new(k: u128) -> Key {
        let mut k1 = (k >> 64) as u64;
        let mut k2 = k as u64;

        let mut ret = Key {
            schedule: [0; ROUNDS as usize],
        };

        // Run `ROUNDS - 1` rounds to generate the key's endpoint (the last key in the schedule).
        for i in 0..ROUNDS {
            // Insert the key into the schedule.
            ret.schedule[i as usize] = k2;

            // The beautiful thing about SPECK is that it reuses its round function to generate the
            // key schedule.
            round!(k1, k2, i);
        }

        ret
    }

    /// Encrypt a 128-bit block with this key.
    pub fn encrypt_block(&self, m: u128) -> u128 {
        let mut m1 = (m >> 64) as u64;
        let mut m2 = m as u64;

        // We run a round for every subkey in the generated key schedule.
        for &k in &self.schedule {
            // Run a round on the message.
            round!(m1, m2, k);
        }

        m2 as u128 | (m1 as u128) << 64
    }

    /// Decrypt a 128-bit block with this key.
    pub fn decrypt_block(&self, c: u128) -> u128 {
        let mut c1 = (c >> 64) as u64;
        let mut c2 = c as u64;

        // We run a round for every subkey in the generated key schedule.
        for &k in self.schedule.iter().rev() {
            // Run a round on the message.
            inv_round!(c1, c2, k);
        }

        c2 as u128 | (c1 as u128) << 64
    }
}

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

    #[test]
    fn encrypt_decrypt() {
        for mut x in 0u128..90000 {
            // <3
            x = x.wrapping_mul(0x6eed0e9da4d94a4f6eed0e9da4d94a4f);
            x ^= (x >> 6) >> (x >> 122);
            x = x.wrapping_mul(0x6eed0e9da4d94a4f6eed0e9da4d94a4f);

            let key = Key::new(!x);

            assert_eq!(key.decrypt_block(key.encrypt_block(x)), x);
            assert_eq!(key.encrypt_block(x), encrypt_block(x, !x));
        }
    }

    #[test]
    fn test_vectors() {
        // These test vectors are taken from the SPECK paper.
        assert_eq!(
            encrypt_block(
                0x6c617669757165207469206564616d20,
                0x0f0e0d0c0b0a09080706050403020100
            ),
            0xa65d9851797832657860fedf5c570d18
        );
    }
}
<|endoftext|>


<commit_before>#![deny(intra_doc_link_resolution_failure)]

extern crate proc_macro;
extern crate protobuf;

use std::collections::hash_map::HashMap;
use std::fmt::Write as FmtWrite;
use std::fs::File;
use std::io;
use std::io::Write;
use std::path::{Path, PathBuf};

use protobuf::descriptor::*;
use protobuf::Message;

use protobuf::prelude::*;

mod amend_io_error_util;
pub mod case_convert;
mod compiler_plugin;
mod customize;
mod enums;
mod extensions;
mod field;
mod file;
pub(crate) mod file_and_mod;
mod file_descriptor;
mod inside;
mod map;
mod message;
mod oneof;
mod protobuf_abs_path;
mod protobuf_ident;
mod protobuf_rel_path;
mod rust_name;
mod rust_types_values;
mod serde;
mod well_known_types;

pub(crate) mod rust;
pub(crate) mod scope;
pub(crate) mod strx;
pub(crate) mod syntax;

use customize::customize_from_rustproto_for_file;
pub use customize::Customize;

pub mod code_writer;

use self::code_writer::CodeWriter;
use self::enums::*;
use self::extensions::*;
use self::message::*;
#[doc(hidden)]
pub use amend_io_error_util::amend_io_error;
use file::proto_path_to_rust_mod;
use map::map_entry;
use scope::FileScope;
use scope::RootScope;

use inside::protobuf_crate_path;
pub use protobuf_abs_path::ProtobufAbsolutePath;
pub use protobuf_ident::ProtobufIdent;
pub use protobuf_rel_path::ProtobufRelativePath;

fn escape_byte(s: &mut String, b: u8) {
    if b == b'\n' {
        write!(s, "\\n").unwrap();
    } else if b == b'\r' {
        write!(s, "\\r").unwrap();
    } else if b == b'\t' {
        write!(s, "\\t").unwrap();
    } else if b == b'\\' || b == b'"' {
        write!(s, "\\{}", b as char).unwrap();
    } else if b == b'\0' {
        write!(s, "\\0").unwrap();
    // ASCII printable except space
    } else if b > 0x20 && b < 0x7f {
        write!(s, "{}", b as char).unwrap();
    } else {
        write!(s, "\\x{:02x}", b).unwrap();
    }
}

fn write_file_descriptor_data(
    file: &FileDescriptorProto,
    customize: &Customize,
    w: &mut CodeWriter,
) {
    let fdp_bytes = file.write_to_bytes().unwrap();
    w.write_line("static file_descriptor_proto_data: &'static [u8] = b\"\\");
    w.indented(|w| {
        const MAX_LINE_LEN: usize = 72;

        let mut s = String::new();
        for &b in &fdp_bytes {
            let prev_len = s.len();
            escape_byte(&mut s, b);
            let truncate = s.len() > MAX_LINE_LEN;
            if truncate {
                s.truncate(prev_len);
            }
            if truncate || s.len() == MAX_LINE_LEN {
                write!(s, "\\").unwrap();
                w.write_line(&s);
                s.clear();
            }
            if truncate {
                escape_byte(&mut s, b);
            }
        }
        if !s.is_empty() {
            write!(s, "\\").unwrap();
            w.write_line(&s);
            s.clear();
        }
    });
    w.write_line("\";");
    w.write_line("");
    w.lazy_static(
        "file_descriptor_proto_lazy",
        &format!(
            "{}::descriptor::FileDescriptorProto",
            protobuf_crate_path(customize)
        ),
        protobuf_crate_path(customize),
    );
    w.write_line("");
    w.def_fn(
        &format!(
            "parse_descriptor_proto() -> {}::descriptor::FileDescriptorProto",
            protobuf_crate_path(customize)
        ),
        |w| {
            w.write_line(&format!(
                "{}::parse_from_bytes(file_descriptor_proto_data).unwrap()",
                protobuf_crate_path(customize)
            ));
        },
    );
    w.write_line("");
    w.write_line("/// `FileDescriptorProto` object which was a source for this generated file");
    w.pub_fn(
        &format!(
            "file_descriptor_proto() -> &'static {}::descriptor::FileDescriptorProto",
            protobuf_crate_path(customize)
        ),
        |w| {
            w.block("file_descriptor_proto_lazy.get(|| {", "})", |w| {
                w.write_line("parse_descriptor_proto()");
            });
        },
    );
}

fn gen_file(
    file: &FileDescriptorProto,
    _files_map: &HashMap<&Path, &FileDescriptorProto>,
    root_scope: &RootScope,
    customize: &Customize,
    parser: &str,
) -> Option<compiler_plugin::GenResult> {
    // TODO: use it
    let mut customize = customize.clone();
    // options specified in invocation have precedence over options specified in file
    customize.update_with(&customize_from_rustproto_for_file(
        file.options.get_message(),
    ));

    let scope = FileScope {
        file_descriptor: file,
    }
    .to_scope();
    let lite_runtime = customize.lite_runtime.unwrap_or_else(|| {
        file.options.get_message().get_optimize_for() == file_options::OptimizeMode::LITE_RUNTIME
    });

    let mut v = Vec::new();

    {
        let mut w = CodeWriter::new(&mut v);

        w.write_generated_by("rust-protobuf", env!("CARGO_PKG_VERSION"), parser);

        w.write_line("");
        w.write_line(&format!("//! Generated file from `{}`", file.get_name()));
        if customize.inside_protobuf != Some(true) {
            w.write_line("");
            w.write_line("/// Generated files are compatible only with the same version");
            w.write_line("/// of protobuf runtime.");
            w.write_line(&format!(
                "const _PROTOBUF_VERSION_CHECK: () = {}::{};",
                protobuf_crate_path(&customize),
                protobuf::VERSION_IDENT
            ));
        }

        static NESTED_TYPE_NUMBER: protobuf::rt::LazyV2<i32> = protobuf::rt::LazyV2::INIT;
        let message_type_number = *NESTED_TYPE_NUMBER.get(|| {
            protobuf::reflect::MessageDescriptor::for_type::<FileDescriptorProto>()
                .get_field_by_name("message_type")
                .expect("`message_type` must exist")
                .proto()
                .get_number()
        });

        let mut path = vec![message_type_number, 0];
        for (id, message) in scope.get_messages().iter().enumerate() {
            // ignore map entries, because they are not used in map fields
            if map_entry(message).is_none() {
                path[1] = id as i32;

                w.write_line("");
                MessageGen::new(
                    message,
                    &root_scope,
                    &customize,
                    &path,
                    file.source_code_info.as_ref(),
                )
                .write(&mut w);
            }
        }

        static ENUM_TYPE_NUMBER: protobuf::rt::LazyV2<i32> = protobuf::rt::LazyV2::INIT;
        let enum_type_number = *ENUM_TYPE_NUMBER.get(|| {
            protobuf::reflect::MessageDescriptor::for_type::<FileDescriptorProto>()
                .get_field_by_name("enum_type")
                .expect("`enum_type` must exist")
                .proto()
                .get_number()
        });

        let mut path = vec![enum_type_number, 0];
        for (id, enum_type) in scope.get_enums().iter().enumerate() {
            path[1] = id as i32;

            w.write_line("");
            EnumGen::new(
                enum_type,
                &customize,
                root_scope,
                &path,
                file.source_code_info.as_ref(),
            )
            .write(&mut w);
        }

        write_extensions(file, &root_scope, &mut w, &customize);

        if !lite_runtime {
            w.write_line("");
            write_file_descriptor_data(file, &customize, &mut w);
        }
    }

    Some(compiler_plugin::GenResult {
        name: format!("{}.rs", proto_path_to_rust_mod(file.get_name())),
        content: v,
    })
}

// This function is also used externally by cargo plugin
// https://github.com/plietar/rust-protobuf-build
// So be careful changing its signature.
pub fn gen(
    file_descriptors: &[FileDescriptorProto],
    parser: &str,
    files_to_generate: &[PathBuf],
    customize: &Customize,
) -> Vec<compiler_plugin::GenResult> {
    let root_scope = RootScope {
        file_descriptors: file_descriptors,
    };

    let mut results: Vec<compiler_plugin::GenResult> = Vec::new();
    let files_map: HashMap<&Path, &FileDescriptorProto> = file_descriptors
        .iter()
        .map(|f| (Path::new(f.get_name()), f))
        .collect();

    for file_name in files_to_generate {
        let file = files_map.get(file_name.as_path()).expect(&format!(
            "file not found in file descriptors: {:?}, files: {:?}",
            file_name,
            files_map.keys()
        ));
        results.extend(gen_file(file, &files_map, &root_scope, customize, parser));
    }
    results
}

pub fn gen_and_write(
    file_descriptors: &[FileDescriptorProto],
    parser: &str,
    files_to_generate: &[PathBuf],
    out_dir: &Path,
    customize: &Customize,
) -> io::Result<()> {
    match out_dir.metadata() {
        Ok(m) => {
            if !m.is_dir() {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    format!("{} is not a directory", out_dir.display()),
                ));
            }
        }
        Err(e) => {
            return Err(amend_io_error(
                e,
                format!("{} does not exist or not accessible", out_dir.display()),
            ));
        }
    }

    let results = gen(file_descriptors, parser, files_to_generate, customize);

    for r in &results {
        let mut file_path = out_dir.to_owned();
        file_path.push(&r.name);
        let mut file_writer = File::create(&file_path)
            .map_err(|e| amend_io_error(e, format!("failed to create {:?}", file_path)))?;
        file_writer
            .write_all(&r.content)
            .map_err(|e| amend_io_error(e, format!("failed to write to {:?}", file_path)))?;
        file_writer
            .flush()
            .map_err(|e| amend_io_error(e, format!("failed to flush {:?}", file_path)))?;
    }

    Ok(())
}

pub fn protoc_gen_rust_main() {
    compiler_plugin::plugin_main(|r| {
        let customize = Customize::parse_from_parameter(r.parameter).expect("parse options");
        gen(
            r.file_descriptors,
            "protoc --rust-out=...",
            r.files_to_generate,
            &customize,
        )
    });
}

/// Used in protobuf-codegen-identical-test
#[doc(hidden)]
pub fn proto_name_to_rs(name: &str) -> String {
    format!("{}.rs", proto_path_to_rust_mod(name))
}
<commit_msg>Move file_descriptor_proto_lazy out of global scope in generated files<commit_after>#![deny(intra_doc_link_resolution_failure)]

extern crate proc_macro;
extern crate protobuf;

use std::collections::hash_map::HashMap;
use std::fmt::Write as FmtWrite;
use std::fs::File;
use std::io;
use std::io::Write;
use std::path::{Path, PathBuf};

use protobuf::descriptor::*;
use protobuf::Message;

use protobuf::prelude::*;

mod amend_io_error_util;
pub mod case_convert;
mod compiler_plugin;
mod customize;
mod enums;
mod extensions;
mod field;
mod file;
pub(crate) mod file_and_mod;
mod file_descriptor;
mod inside;
mod map;
mod message;
mod oneof;
mod protobuf_abs_path;
mod protobuf_ident;
mod protobuf_rel_path;
mod rust_name;
mod rust_types_values;
mod serde;
mod well_known_types;

pub(crate) mod rust;
pub(crate) mod scope;
pub(crate) mod strx;
pub(crate) mod syntax;

use customize::customize_from_rustproto_for_file;
pub use customize::Customize;

pub mod code_writer;

use self::code_writer::CodeWriter;
use self::enums::*;
use self::extensions::*;
use self::message::*;
#[doc(hidden)]
pub use amend_io_error_util::amend_io_error;
use file::proto_path_to_rust_mod;
use map::map_entry;
use scope::FileScope;
use scope::RootScope;

use inside::protobuf_crate_path;
pub use protobuf_abs_path::ProtobufAbsolutePath;
pub use protobuf_ident::ProtobufIdent;
pub use protobuf_rel_path::ProtobufRelativePath;

fn escape_byte(s: &mut String, b: u8) {
    if b == b'\n' {
        write!(s, "\\n").unwrap();
    } else if b == b'\r' {
        write!(s, "\\r").unwrap();
    } else if b == b'\t' {
        write!(s, "\\t").unwrap();
    } else if b == b'\\' || b == b'"' {
        write!(s, "\\{}", b as char).unwrap();
    } else if b == b'\0' {
        write!(s, "\\0").unwrap();
    // ASCII printable except space
    } else if b > 0x20 && b < 0x7f {
        write!(s, "{}", b as char).unwrap();
    } else {
        write!(s, "\\x{:02x}", b).unwrap();
    }
}

fn write_file_descriptor_data(
    file: &FileDescriptorProto,
    customize: &Customize,
    w: &mut CodeWriter,
) {
    let fdp_bytes = file.write_to_bytes().unwrap();
    w.write_line("static file_descriptor_proto_data: &'static [u8] = b\"\\");
    w.indented(|w| {
        const MAX_LINE_LEN: usize = 72;

        let mut s = String::new();
        for &b in &fdp_bytes {
            let prev_len = s.len();
            escape_byte(&mut s, b);
            let truncate = s.len() > MAX_LINE_LEN;
            if truncate {
                s.truncate(prev_len);
            }
            if truncate || s.len() == MAX_LINE_LEN {
                write!(s, "\\").unwrap();
                w.write_line(&s);
                s.clear();
            }
            if truncate {
                escape_byte(&mut s, b);
            }
        }
        if !s.is_empty() {
            write!(s, "\\").unwrap();
            w.write_line(&s);
            s.clear();
        }
    });
    w.write_line("\";");
    w.write_line("");
    w.def_fn(
        &format!(
            "parse_descriptor_proto() -> {}::descriptor::FileDescriptorProto",
            protobuf_crate_path(customize)
        ),
        |w| {
            w.write_line(&format!(
                "{}::parse_from_bytes(file_descriptor_proto_data).unwrap()",
                protobuf_crate_path(customize)
            ));
        },
    );
    w.write_line("");
    w.write_line("/// `FileDescriptorProto` object which was a source for this generated file");
    w.pub_fn(
        &format!(
            "file_descriptor_proto() -> &'static {}::descriptor::FileDescriptorProto",
            protobuf_crate_path(customize)
        ),
        |w| {
            w.lazy_static(
                "file_descriptor_proto_lazy",
                &format!(
                    "{}::descriptor::FileDescriptorProto",
                    protobuf_crate_path(customize)
                ),
                protobuf_crate_path(customize),
            );
            w.block("file_descriptor_proto_lazy.get(|| {", "})", |w| {
                w.write_line("parse_descriptor_proto()");
            });
        },
    );
}

fn gen_file(
    file: &FileDescriptorProto,
    _files_map: &HashMap<&Path, &FileDescriptorProto>,
    root_scope: &RootScope,
    customize: &Customize,
    parser: &str,
) -> Option<compiler_plugin::GenResult> {
    // TODO: use it
    let mut customize = customize.clone();
    // options specified in invocation have precedence over options specified in file
    customize.update_with(&customize_from_rustproto_for_file(
        file.options.get_message(),
    ));

    let scope = FileScope {
        file_descriptor: file,
    }
    .to_scope();
    let lite_runtime = customize.lite_runtime.unwrap_or_else(|| {
        file.options.get_message().get_optimize_for() == file_options::OptimizeMode::LITE_RUNTIME
    });

    let mut v = Vec::new();

    {
        let mut w = CodeWriter::new(&mut v);

        w.write_generated_by("rust-protobuf", env!("CARGO_PKG_VERSION"), parser);

        w.write_line("");
        w.write_line(&format!("//! Generated file from `{}`", file.get_name()));
        if customize.inside_protobuf != Some(true) {
            w.write_line("");
            w.write_line("/// Generated files are compatible only with the same version");
            w.write_line("/// of protobuf runtime.");
            w.write_line(&format!(
                "const _PROTOBUF_VERSION_CHECK: () = {}::{};",
                protobuf_crate_path(&customize),
                protobuf::VERSION_IDENT
            ));
        }

        static NESTED_TYPE_NUMBER: protobuf::rt::LazyV2<i32> = protobuf::rt::LazyV2::INIT;
        let message_type_number = *NESTED_TYPE_NUMBER.get(|| {
            protobuf::reflect::MessageDescriptor::for_type::<FileDescriptorProto>()
                .get_field_by_name("message_type")
                .expect("`message_type` must exist")
                .proto()
                .get_number()
        });

        let mut path = vec![message_type_number, 0];
        for (id, message) in scope.get_messages().iter().enumerate() {
            // ignore map entries, because they are not used in map fields
            if map_entry(message).is_none() {
                path[1] = id as i32;

                w.write_line("");
                MessageGen::new(
                    message,
                    &root_scope,
                    &customize,
                    &path,
                    file.source_code_info.as_ref(),
                )
                .write(&mut w);
            }
        }

        static ENUM_TYPE_NUMBER: protobuf::rt::LazyV2<i32> = protobuf::rt::LazyV2::INIT;
        let enum_type_number = *ENUM_TYPE_NUMBER.get(|| {
            protobuf::reflect::MessageDescriptor::for_type::<FileDescriptorProto>()
                .get_field_by_name("enum_type")
                .expect("`enum_type` must exist")
                .proto()
                .get_number()
        });

        let mut path = vec![enum_type_number, 0];
        for (id, enum_type) in scope.get_enums().iter().enumerate() {
            path[1] = id as i32;

            w.write_line("");
            EnumGen::new(
                enum_type,
                &customize,
                root_scope,
                &path,
                file.source_code_info.as_ref(),
            )
            .write(&mut w);
        }

        write_extensions(file, &root_scope, &mut w, &customize);

        if !lite_runtime {
            w.write_line("");
            write_file_descriptor_data(file, &customize, &mut w);
        }
    }

    Some(compiler_plugin::GenResult {
        name: format!("{}.rs", proto_path_to_rust_mod(file.get_name())),
        content: v,
    })
}

// This function is also used externally by cargo plugin
// https://github.com/plietar/rust-protobuf-build
// So be careful changing its signature.
pub fn gen(
    file_descriptors: &[FileDescriptorProto],
    parser: &str,
    files_to_generate: &[PathBuf],
    customize: &Customize,
) -> Vec<compiler_plugin::GenResult> {
    let root_scope = RootScope {
        file_descriptors: file_descriptors,
    };

    let mut results: Vec<compiler_plugin::GenResult> = Vec::new();
    let files_map: HashMap<&Path, &FileDescriptorProto> = file_descriptors
        .iter()
        .map(|f| (Path::new(f.get_name()), f))
        .collect();

    for file_name in files_to_generate {
        let file = files_map.get(file_name.as_path()).expect(&format!(
            "file not found in file descriptors: {:?}, files: {:?}",
            file_name,
            files_map.keys()
        ));
        results.extend(gen_file(file, &files_map, &root_scope, customize, parser));
    }
    results
}

pub fn gen_and_write(
    file_descriptors: &[FileDescriptorProto],
    parser: &str,
    files_to_generate: &[PathBuf],
    out_dir: &Path,
    customize: &Customize,
) -> io::Result<()> {
    match out_dir.metadata() {
        Ok(m) => {
            if !m.is_dir() {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    format!("{} is not a directory", out_dir.display()),
                ));
            }
        }
        Err(e) => {
            return Err(amend_io_error(
                e,
                format!("{} does not exist or not accessible", out_dir.display()),
            ));
        }
    }

    let results = gen(file_descriptors, parser, files_to_generate, customize);

    for r in &results {
        let mut file_path = out_dir.to_owned();
        file_path.push(&r.name);
        let mut file_writer = File::create(&file_path)
            .map_err(|e| amend_io_error(e, format!("failed to create {:?}", file_path)))?;
        file_writer
            .write_all(&r.content)
            .map_err(|e| amend_io_error(e, format!("failed to write to {:?}", file_path)))?;
        file_writer
            .flush()
            .map_err(|e| amend_io_error(e, format!("failed to flush {:?}", file_path)))?;
    }

    Ok(())
}

pub fn protoc_gen_rust_main() {
    compiler_plugin::plugin_main(|r| {
        let customize = Customize::parse_from_parameter(r.parameter).expect("parse options");
        gen(
            r.file_descriptors,
            "protoc --rust-out=...",
            r.files_to_generate,
            &customize,
        )
    });
}

/// Used in protobuf-codegen-identical-test
#[doc(hidden)]
pub fn proto_name_to_rs(name: &str) -> String {
    format!("{}.rs", proto_path_to_rust_mod(name))
}
<|endoftext|>


<commit_before>extern crate redis = "redis#0.1";
extern crate time;
extern crate native;

use redis::Client;
use native::task;

fn bench_set(tid: uint, n: uint) {
  let mut redis = Client::new("127.0.0.1:6379");

  for _ in range(0, n) {
    redis.set("key", "12");
  }

  println!("Thread {} finished", tid);
}

fn main() {
  let before = time::precise_time_ns();

  let concurrency: uint = from_str(std::os::args()[1]).unwrap();
  let repeats: uint = from_str(std::os::args()[2]).unwrap();
  let per_thread: uint = repeats / concurrency;
  let total_reqs = per_thread * concurrency;

  let mut threads = ~[];

  for tid in range(0, concurrency) {
    println!("Thread {} started", tid);

    let (port, chan) = Chan::new();
    task::spawn(proc() {
      bench_set(tid, per_thread);
      chan.send(());
    });
    threads.push(port);
  }

  println!("Waiting for all clients to terminate");
  for port in threads.iter() {
      port.recv();
  }

  let after = time::precise_time_ns();

  let time = ((after - before) / 1_000_000) as f64 / 1000f64;

  println!("Concurrency: {}", concurrency);
  println!("Total requests: {}", total_reqs);
  println!("Total time: {}", time);
  let reqs_per_s = total_reqs as f64 / time;
  println!("Requests per second: {}", reqs_per_s);
}
<commit_msg>Chan::new renamed to channel()<commit_after>extern crate redis = "redis#0.1";
extern crate time;
extern crate native;

use redis::Client;
use native::task;
use std::comm::channel;

fn bench_set(tid: uint, n: uint) {
  let mut redis = Client::new("127.0.0.1:6379");

  for _ in range(0, n) {
    redis.set("key", "12");
  }

  println!("Thread {} finished", tid);
}

fn main() {
  let before = time::precise_time_ns();

  let concurrency: uint = from_str(std::os::args()[1]).unwrap();
  let repeats: uint = from_str(std::os::args()[2]).unwrap();
  let per_thread: uint = repeats / concurrency;
  let total_reqs = per_thread * concurrency;

  let mut threads = ~[];

  for tid in range(0, concurrency) {
    println!("Thread {} started", tid);

    let (sender, receiver) = channel();
    task::spawn(proc() {
      bench_set(tid, per_thread);
      sender.send(());
    });
    threads.push(receiver);
  }

  println!("Waiting for all clients to terminate");
  for receiver in threads.iter() {
      receiver.recv();
  }

  let after = time::precise_time_ns();

  let time = ((after - before) / 1_000_000) as f64 / 1000f64;

  println!("Concurrency: {}", concurrency);
  println!("Total requests: {}", total_reqs);
  println!("Total time: {}", time);
  let reqs_per_s = total_reqs as f64 / time;
  println!("Requests per second: {}", reqs_per_s);
}
<|endoftext|>


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

//! Basic TUI to better output the overall system status and status
//! of various subsystems

use chrono::prelude::Utc;
use cursive::direction::Orientation;
use cursive::theme::BaseColor::{Black, Blue, Cyan, White};
use cursive::theme::Color::Dark;
use cursive::theme::PaletteColor::{
	Background, Highlight, HighlightInactive, Primary, Shadow, View,
};
use cursive::theme::{BaseColor, BorderStyle, Color, Theme};
use cursive::traits::Boxable;
use cursive::traits::Identifiable;
use cursive::utils::markup::StyledString;
use cursive::views::{BoxedView, CircularFocus, Dialog, LinearLayout, Panel, StackView, TextView};
use cursive::Cursive;
use std::sync::mpsc;

use crate::built_info;
use crate::servers::Server;
use crate::tui::constants::ROOT_STACK;
use crate::tui::types::{TUIStatusListener, UIMessage};
use crate::tui::{logs, menu, mining, peers, status, version};
use grin_util::logger::LogEntry;

pub struct UI {
	cursive: Cursive,
	ui_rx: mpsc::Receiver<UIMessage>,
	ui_tx: mpsc::Sender<UIMessage>,
	controller_tx: mpsc::Sender<ControllerMessage>,
	logs_rx: mpsc::Receiver<LogEntry>,
}

fn modify_theme(theme: &mut Theme) {
	theme.shadow = false;
	theme.borders = BorderStyle::Simple;
	theme.palette[Background] = Dark(Black);
	theme.palette[Shadow] = Dark(Black);
	theme.palette[View] = Dark(Black);
	theme.palette[Primary] = Dark(White);
	theme.palette[Highlight] = Dark(Cyan);
	theme.palette[HighlightInactive] = Dark(Blue);
	// also secondary, tertiary, TitlePrimary, TitleSecondary
}

impl UI {
	/// Create a new UI
	pub fn new(
		controller_tx: mpsc::Sender<ControllerMessage>,
		logs_rx: mpsc::Receiver<LogEntry>,
	) -> UI {
		let (ui_tx, ui_rx) = mpsc::channel::<UIMessage>();

		let mut grin_ui = UI {
			cursive: Cursive::default(),
			ui_tx,
			ui_rx,
			controller_tx,
			logs_rx,
		};

		// Create UI objects, etc
		let status_view = status::TUIStatusView::create();
		let mining_view = mining::TUIMiningView::create();
		let peer_view = peers::TUIPeerView::create();
		let logs_view = logs::TUILogsView::create();
		let version_view = version::TUIVersionView::create();

		let main_menu = menu::create();

		let root_stack = StackView::new()
			.layer(version_view)
			.layer(mining_view)
			.layer(peer_view)
			.layer(logs_view)
			.layer(status_view)
			.with_name(ROOT_STACK)
			.full_height();

		let mut title_string = StyledString::new();
		title_string.append(StyledString::styled(
			format!(
				"Grin Version {} (proto: {})",
				built_info::PKG_VERSION,
				Server::protocol_version()
			),
			Color::Dark(BaseColor::Green),
		));

		let main_layer = LinearLayout::new(Orientation::Vertical)
			.child(Panel::new(TextView::new(title_string).full_width()))
			.child(
				LinearLayout::new(Orientation::Horizontal)
					.child(Panel::new(BoxedView::new(main_menu)))
					.child(Panel::new(root_stack)),
			);

		//set theme
		let mut theme = grin_ui.cursive.current_theme().clone();
		modify_theme(&mut theme);
		grin_ui.cursive.set_theme(theme);
		grin_ui.cursive.add_fullscreen_layer(main_layer);

		// Configure a callback (shutdown, for the first test)
		let controller_tx_clone = grin_ui.controller_tx.clone();
		grin_ui.cursive.add_global_callback('q', move |c| {
			let content = StyledString::styled("Shutting down...", Color::Light(BaseColor::Yellow));
			c.add_layer(CircularFocus::wrap_tab(Dialog::around(TextView::new(
				content,
			))));
			controller_tx_clone
				.send(ControllerMessage::Shutdown)
				.unwrap();
		});
		grin_ui.cursive.set_fps(3);
		grin_ui
	}

	/// Step the UI by calling into Cursive's step function, then
	/// processing any UI messages
	pub fn step(&mut self) -> bool {
		if !self.cursive.is_running() {
			return false;
		}

		while let Some(message) = self.logs_rx.try_iter().next() {
			logs::TUILogsView::update(&mut self.cursive, message);
		}

		// Process any pending UI messages
		while let Some(message) = self.ui_rx.try_iter().next() {
			match message {
				UIMessage::UpdateStatus(update) => {
					status::TUIStatusView::update(&mut self.cursive, &update);
					mining::TUIMiningView::update(&mut self.cursive, &update);
					peers::TUIPeerView::update(&mut self.cursive, &update);
					version::TUIVersionView::update(&mut self.cursive, &update);
				}
			}
		}

		// Step the UI
		self.cursive.step();
		true
	}

	/// Stop the UI
	pub fn stop(&mut self) {
		self.cursive.quit();
	}
}

pub struct Controller {
	rx: mpsc::Receiver<ControllerMessage>,
	ui: UI,
}

pub enum ControllerMessage {
	Shutdown,
}

impl Controller {
	/// Create a new controller
	pub fn new(logs_rx: mpsc::Receiver<LogEntry>) -> Result<Controller, String> {
		let (tx, rx) = mpsc::channel::<ControllerMessage>();
		Ok(Controller {
			rx,
			ui: UI::new(tx, logs_rx),
		})
	}

	/// Run the controller
	pub fn run(&mut self, server: Server) {
		let stat_update_interval = 1;
		let mut next_stat_update = Utc::now().timestamp() + stat_update_interval;
		while self.ui.step() {
			if let Some(message) = self.rx.try_iter().next() {
				match message {
					ControllerMessage::Shutdown => {
						warn!("Shutdown in progress, please wait");
						self.ui.stop();
						server.stop();
						return;
					}
				}
			}

			if Utc::now().timestamp() > next_stat_update {
				next_stat_update = Utc::now().timestamp() + stat_update_interval;
				if let Ok(stats) = server.get_server_stats() {
					self.ui.ui_tx.send(UIMessage::UpdateStatus(stats)).unwrap();
				}
			}
		}
		server.stop();
	}
}
<commit_msg>Introduce a pause after calling Cursive.step (#3326)<commit_after>// Copyright 2020 The Grin Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Basic TUI to better output the overall system status and status
//! of various subsystems

use chrono::prelude::Utc;
use cursive::direction::Orientation;
use cursive::theme::BaseColor::{Black, Blue, Cyan, White};
use cursive::theme::Color::Dark;
use cursive::theme::PaletteColor::{
	Background, Highlight, HighlightInactive, Primary, Shadow, View,
};
use cursive::theme::{BaseColor, BorderStyle, Color, Theme};
use cursive::traits::Boxable;
use cursive::traits::Identifiable;
use cursive::utils::markup::StyledString;
use cursive::views::{BoxedView, CircularFocus, Dialog, LinearLayout, Panel, StackView, TextView};
use cursive::Cursive;
use std::sync::mpsc;
use std::{thread, time};

use crate::built_info;
use crate::servers::Server;
use crate::tui::constants::ROOT_STACK;
use crate::tui::types::{TUIStatusListener, UIMessage};
use crate::tui::{logs, menu, mining, peers, status, version};
use grin_util::logger::LogEntry;

pub struct UI {
	cursive: Cursive,
	ui_rx: mpsc::Receiver<UIMessage>,
	ui_tx: mpsc::Sender<UIMessage>,
	controller_tx: mpsc::Sender<ControllerMessage>,
	logs_rx: mpsc::Receiver<LogEntry>,
}

fn modify_theme(theme: &mut Theme) {
	theme.shadow = false;
	theme.borders = BorderStyle::Simple;
	theme.palette[Background] = Dark(Black);
	theme.palette[Shadow] = Dark(Black);
	theme.palette[View] = Dark(Black);
	theme.palette[Primary] = Dark(White);
	theme.palette[Highlight] = Dark(Cyan);
	theme.palette[HighlightInactive] = Dark(Blue);
	// also secondary, tertiary, TitlePrimary, TitleSecondary
}

impl UI {
	/// Create a new UI
	pub fn new(
		controller_tx: mpsc::Sender<ControllerMessage>,
		logs_rx: mpsc::Receiver<LogEntry>,
	) -> UI {
		let (ui_tx, ui_rx) = mpsc::channel::<UIMessage>();

		let mut grin_ui = UI {
			cursive: Cursive::default(),
			ui_tx,
			ui_rx,
			controller_tx,
			logs_rx,
		};

		// Create UI objects, etc
		let status_view = status::TUIStatusView::create();
		let mining_view = mining::TUIMiningView::create();
		let peer_view = peers::TUIPeerView::create();
		let logs_view = logs::TUILogsView::create();
		let version_view = version::TUIVersionView::create();

		let main_menu = menu::create();

		let root_stack = StackView::new()
			.layer(version_view)
			.layer(mining_view)
			.layer(peer_view)
			.layer(logs_view)
			.layer(status_view)
			.with_name(ROOT_STACK)
			.full_height();

		let mut title_string = StyledString::new();
		title_string.append(StyledString::styled(
			format!(
				"Grin Version {} (proto: {})",
				built_info::PKG_VERSION,
				Server::protocol_version()
			),
			Color::Dark(BaseColor::Green),
		));

		let main_layer = LinearLayout::new(Orientation::Vertical)
			.child(Panel::new(TextView::new(title_string).full_width()))
			.child(
				LinearLayout::new(Orientation::Horizontal)
					.child(Panel::new(BoxedView::new(main_menu)))
					.child(Panel::new(root_stack)),
			);

		//set theme
		let mut theme = grin_ui.cursive.current_theme().clone();
		modify_theme(&mut theme);
		grin_ui.cursive.set_theme(theme);
		grin_ui.cursive.add_fullscreen_layer(main_layer);

		// Configure a callback (shutdown, for the first test)
		let controller_tx_clone = grin_ui.controller_tx.clone();
		grin_ui.cursive.add_global_callback('q', move |c| {
			let content = StyledString::styled("Shutting down...", Color::Light(BaseColor::Yellow));
			c.add_layer(CircularFocus::wrap_tab(Dialog::around(TextView::new(
				content,
			))));
			controller_tx_clone
				.send(ControllerMessage::Shutdown)
				.unwrap();
		});
		grin_ui.cursive.set_fps(3);
		grin_ui
	}

	/// Step the UI by calling into Cursive's step function, then
	/// processing any UI messages
	pub fn step(&mut self) -> bool {
		if !self.cursive.is_running() {
			return false;
		}

		while let Some(message) = self.logs_rx.try_iter().next() {
			logs::TUILogsView::update(&mut self.cursive, message);
		}

		// Process any pending UI messages
		while let Some(message) = self.ui_rx.try_iter().next() {
			match message {
				UIMessage::UpdateStatus(update) => {
					status::TUIStatusView::update(&mut self.cursive, &update);
					mining::TUIMiningView::update(&mut self.cursive, &update);
					peers::TUIPeerView::update(&mut self.cursive, &update);
					version::TUIVersionView::update(&mut self.cursive, &update);
				}
			}
		}

		// Step the UI
		self.cursive.step();
		true
	}

	/// Stop the UI
	pub fn stop(&mut self) {
		self.cursive.quit();
	}
}

pub struct Controller {
	rx: mpsc::Receiver<ControllerMessage>,
	ui: UI,
}

pub enum ControllerMessage {
	Shutdown,
}

impl Controller {
	/// Create a new controller
	pub fn new(logs_rx: mpsc::Receiver<LogEntry>) -> Result<Controller, String> {
		let (tx, rx) = mpsc::channel::<ControllerMessage>();
		Ok(Controller {
			rx,
			ui: UI::new(tx, logs_rx),
		})
	}

	/// Run the controller
	pub fn run(&mut self, server: Server) {
		let stat_update_interval = 1;
		let mut next_stat_update = Utc::now().timestamp() + stat_update_interval;
		let delay = time::Duration::from_millis(50);
		while self.ui.step() {
			if let Some(message) = self.rx.try_iter().next() {
				match message {
					ControllerMessage::Shutdown => {
						warn!("Shutdown in progress, please wait");
						self.ui.stop();
						server.stop();
						return;
					}
				}
			}

			if Utc::now().timestamp() > next_stat_update {
				next_stat_update = Utc::now().timestamp() + stat_update_interval;
				if let Ok(stats) = server.get_server_stats() {
					self.ui.ui_tx.send(UIMessage::UpdateStatus(stats)).unwrap();
				}
			}
			thread::sleep(delay);
		}
		server.stop();
	}
}
<|endoftext|>


<commit_before>use super::bit_vec::BitVec;
use std::usize;
use rand::Rng;
use super::{Probability, ProbabilityValue};

pub trait BitRepr {
    fn nbits(&self) -> usize;
    fn get_bit(&self, pos: usize) -> bool;
}

impl BitRepr for usize {
    fn nbits(&self) -> usize { usize::BITS }
    fn get_bit(&self, pos: usize) -> bool {
        assert!(pos < self.nbits());
        (*self >> pos) & 1 == 1
    }
}

impl BitRepr for bool {
    fn nbits(&self) -> usize { 1 }
    fn get_bit(&self, pos: usize) -> bool {
        assert!(pos == 0);
        *self
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BitString {
    bits: BitVec,
}

impl BitString {
    /// Append the lowest `nbits` bits of `value` to the BitString, msb first.
    pub fn append_bits_msb_from<T:BitRepr>(&mut self, value: &T, nbits: usize) {
        assert!(nbits <= value.nbits());
        self.bits.reserve(nbits);
        for i in (0 .. nbits).rev() {
            self.bits.push(value.get_bit(i));
        }
    }

    /// Flip bits with probability `prob`.
    /// The default rng for floats returns numbers in [0, 1)
    pub fn flip_bits_randomly<R:Rng>(&mut self, rng: &mut R, prob: Probability) {
        for i in 0..self.len() {
            if rng.gen::<ProbabilityValue>().is_probable_with(prob) {
                self.flip(i);
            }
        }
    }

    pub fn flip(&mut self, pos: usize) {
        let old = self.get(pos);
        self.set(pos, !old);
    }

    pub fn set(&mut self, pos: usize, val: bool) {
        self.bits.set(pos, val);
    }

    pub fn get(&self, pos: usize) -> bool {
        self.bits[pos]
    }

    pub fn push(&mut self, bit: bool) {
        self.bits.push(bit);
    }

    pub fn count(&self, val: bool) -> usize {
        let mut cnt = 0;
        for v in self.bits.iter() {
            if v == val {
                cnt += 1;
            }
        }
        return cnt;
    }

    pub fn len(&self) -> usize { self.bits.len() }

    pub fn with_capacity(cap: usize) -> BitString {
        BitString {bits: BitVec::with_capacity(cap)}
    }

    pub fn new() -> BitString {
        BitString {bits: BitVec::new()}
    }

    pub fn from_elem(len: usize, val: bool) -> BitString {
        BitString {bits: BitVec::from_elem(len, val)}
    }

    pub fn from_iter<I>(iter: I) -> BitString where I:Iterator<Item=bool> {
        let mut bs;
        match iter.size_hint() {
            (lower, Some(upper)) if upper > lower => {
                bs = BitString::with_capacity(upper);
            }
            (lower, _) => {
                bs = BitString::with_capacity(lower);
            }
        }
        for val in iter {
            bs.push(val);
        }
        return bs;
    }

    pub fn to_vec(&self) -> Vec<bool> {
        self.bits.iter().collect()
    }
}

/// One-point crossover
pub fn crossover_one_point(point: usize, parents: (&BitString, &BitString)) -> (BitString, BitString) {
    let (pa, pb) = parents;
    assert!(point <= pa.len());
    assert!(point <= pb.len());

    // `ca` will contain pa[..point], pb[point..]
    // `cb` will contain pb[..point], pa[point..]
    let mut ca = BitString::with_capacity(pb.len());
    let mut cb = BitString::with_capacity(pa.len());

    for i in 0..point {
        ca.push(pa.get(i));
        cb.push(pb.get(i));
    }
    for i in point..pb.len() {
        ca.push(pb.get(i));
    }
    for i in point..pa.len() {
        cb.push(pa.get(i));
    }

    assert!(ca.len() == pb.len());
    assert!(cb.len() == pa.len());

    return (ca, cb);
}

#[test]
fn test_bitstring() {
    let bs = BitString::new();
    assert_eq!(0, bs.len());

    let bs = BitString::from_elem(4, true);
    assert_eq!(4, bs.len());
    assert_eq!(vec!(true, true, true, true), bs.to_vec());

    let mut bs = BitString::new();
    bs.append_bits_msb_from(&0b11001, 4);
    assert_eq!(4, bs.len());
    assert_eq!(vec!(true, false, false, true), bs.to_vec());

}

#[test]
fn test_crossover_one_point() {
    {
        let mut pa = BitString::new();
        pa.append_bits_msb_from(&0b1011usize, 4);
        let mut pb = BitString::new();
        pb.append_bits_msb_from(&0b0010usize, 4);

        let (ca, cb) = crossover_one_point(2, (&pa, &pb));
        assert_eq!(vec!(true, false, true, false), ca.to_vec());
        assert_eq!(vec!(false, false, true, true), cb.to_vec());
    }
    {
        let mut pa = BitString::new();
        pa.append_bits_msb_from(&0b1011usize, 4);
        let pb = BitString::new();

        let (ca, cb) = crossover_one_point(0, (&pa, &pb));
        assert!(ca.to_vec().is_empty());
        assert_eq!(vec!(true, false, true, true), cb.to_vec());

    }
}
<commit_msg>Add some inlines<commit_after>use super::bit_vec::BitVec;
use std::usize;
use rand::Rng;
use super::{Probability, ProbabilityValue};

pub trait BitRepr {
    fn nbits(&self) -> usize;
    fn get_bit(&self, pos: usize) -> bool;
}

impl BitRepr for usize {
    #[inline(always)]
    fn nbits(&self) -> usize { usize::BITS }
    #[inline(always)]
    fn get_bit(&self, pos: usize) -> bool {
        assert!(pos < self.nbits());
        (*self >> pos) & 1 == 1
    }
}

impl BitRepr for bool {
    #[inline(always)]
    fn nbits(&self) -> usize { 1 }
    #[inline(always)]
    fn get_bit(&self, pos: usize) -> bool {
        assert!(pos == 0);
        *self
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BitString {
    bits: BitVec,
}

impl BitString {
    /// Append the lowest `nbits` bits of `value` to the BitString, msb first.
    pub fn append_bits_msb_from<T:BitRepr>(&mut self, value: &T, nbits: usize) {
        assert!(nbits <= value.nbits());
        self.bits.reserve(nbits);
        for i in (0 .. nbits).rev() {
            self.bits.push(value.get_bit(i));
        }
    }

    /// Flip bits with probability `prob`.
    /// The default rng for floats returns numbers in [0, 1)
    pub fn flip_bits_randomly<R:Rng>(&mut self, rng: &mut R, prob: Probability) {
        for i in 0..self.len() {
            if rng.gen::<ProbabilityValue>().is_probable_with(prob) {
                self.flip(i);
            }
        }
    }

    #[inline]
    pub fn flip(&mut self, pos: usize) {
        let old = self.get(pos);
        self.set(pos, !old);
    }

    #[inline(always)]
    pub fn set(&mut self, pos: usize, val: bool) {
        self.bits.set(pos, val);
    }

    #[inline(always)]
    pub fn get(&self, pos: usize) -> bool {
        self.bits[pos]
    }

    #[inline(always)]
    pub fn push(&mut self, bit: bool) {
        self.bits.push(bit);
    }

    pub fn count(&self, val: bool) -> usize {
        let mut cnt = 0;
        for v in self.bits.iter() {
            if v == val {
                cnt += 1;
            }
        }
        return cnt;
    }

    #[inline(always)]
    pub fn len(&self) -> usize { self.bits.len() }

    #[inline]
    pub fn with_capacity(cap: usize) -> BitString {
        BitString {bits: BitVec::with_capacity(cap)}
    }

    #[inline]
    pub fn new() -> BitString {
        BitString {bits: BitVec::new()}
    }

    #[inline]
    pub fn from_elem(len: usize, val: bool) -> BitString {
        BitString {bits: BitVec::from_elem(len, val)}
    }

    pub fn from_iter<I>(iter: I) -> BitString where I:Iterator<Item=bool> {
        let mut bs;
        match iter.size_hint() {
            (lower, Some(upper)) if upper > lower => {
                bs = BitString::with_capacity(upper);
            }
            (lower, _) => {
                bs = BitString::with_capacity(lower);
            }
        }
        for val in iter {
            bs.push(val);
        }
        return bs;
    }

    #[inline]
    pub fn to_vec(&self) -> Vec<bool> {
        self.bits.iter().collect()
    }
}

/// One-point crossover
pub fn crossover_one_point(point: usize, parents: (&BitString, &BitString)) -> (BitString, BitString) {
    let (pa, pb) = parents;
    assert!(point <= pa.len());
    assert!(point <= pb.len());

    // `ca` will contain pa[..point], pb[point..]
    // `cb` will contain pb[..point], pa[point..]
    let mut ca = BitString::with_capacity(pb.len());
    let mut cb = BitString::with_capacity(pa.len());

    for i in 0..point {
        ca.push(pa.get(i));
        cb.push(pb.get(i));
    }
    for i in point..pb.len() {
        ca.push(pb.get(i));
    }
    for i in point..pa.len() {
        cb.push(pa.get(i));
    }

    assert!(ca.len() == pb.len());
    assert!(cb.len() == pa.len());

    return (ca, cb);
}

#[test]
fn test_bitstring() {
    let bs = BitString::new();
    assert_eq!(0, bs.len());

    let bs = BitString::from_elem(4, true);
    assert_eq!(4, bs.len());
    assert_eq!(vec!(true, true, true, true), bs.to_vec());

    let mut bs = BitString::new();
    bs.append_bits_msb_from(&0b11001, 4);
    assert_eq!(4, bs.len());
    assert_eq!(vec!(true, false, false, true), bs.to_vec());

}

#[test]
fn test_crossover_one_point() {
    {
        let mut pa = BitString::new();
        pa.append_bits_msb_from(&0b1011usize, 4);
        let mut pb = BitString::new();
        pb.append_bits_msb_from(&0b0010usize, 4);

        let (ca, cb) = crossover_one_point(2, (&pa, &pb));
        assert_eq!(vec!(true, false, true, false), ca.to_vec());
        assert_eq!(vec!(false, false, true, true), cb.to_vec());
    }
    {
        let mut pa = BitString::new();
        pa.append_bits_msb_from(&0b1011usize, 4);
        let pb = BitString::new();

        let (ca, cb) = crossover_one_point(0, (&pa, &pb));
        assert!(ca.to_vec().is_empty());
        assert_eq!(vec!(true, false, true, true), cb.to_vec());

    }
}
<|endoftext|>


<commit_before>use std::env;
use std::fs;
use std::io::prelude::*;
use std::io;
use std::process;
use toml;
use toml::de::Error;
use config::Config;

const DEFAULT_ARGS: &'static str = "\"--with-zlib --with-ssl --enable-mruby --without-libstemmer \
                                    --disable-zeromq\"";

#[derive(Debug, PartialEq, Eq,Serialize, Deserialize)]
pub struct Configuration {
    settings: BuildConfig,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize,Deserialize)]
pub struct BuildConfig {
    args: String,
}

pub fn write_conf(config: &Config) {
    if !config.build_conf.exists() {
        let mut f = fs::File::create(&config.build_conf)
            .expect("Cound not create build configuration file.");
        let contents = format!("[settings]\nargs = {}", DEFAULT_ARGS);
        match f.write_all(&contents.as_bytes()) {
            Ok(_) => (),
            Err(e) => {
                println!("{}", e);
                process::exit(2)
            }
        }
    }
}

pub fn parse_toml(config_content: String) -> BuildConfig {
    println!("config:\n{}", config_content);
    let configuration: Result<Configuration, Error> = toml::from_str(&*config_content);
    let config = match configuration {
        Ok(config) => config,
        Err(_) => panic!("Couldn't decode toml with Configuration struct"),
    };
    config.settings
}

pub fn read_build_args(config: &Config) -> Result<String, io::Error> {
    let mut f = fs::File::open(&config.build_conf)?;
    let mut buffer = String::new();
    f.read_to_string(&mut buffer)?;
    let build_conf = parse_toml(buffer);
    Ok(build_conf.args)
}

// mainly used for macOS.
fn openssl_pkg_config_path() -> String {
    let key = "OPENSSL_PKG_CONFIG_PATH";
    env::var(key).unwrap_or("/usr/local/opt/openssl/lib/pkgconfig".to_string())
}

pub fn build_args(config: &Config, groonga_dir: String) -> Result<Vec<String>, io::Error> {
    let conf_args = read_build_args(config)?;
    let build_args: Vec<String> =
        conf_args.split_whitespace().into_iter().map(|e| e.to_owned()).collect();
    println!("{} with args: {:?}",
             config.versions_dir.join(groonga_dir.clone()).display(),
             build_args.clone());
    let mut args = vec![format!("--prefix={}",
                                config.versions_dir.join(groonga_dir.clone()).display()),
                        format!("PKG_CONFIG_PATH={}", openssl_pkg_config_path())];
    args.extend(build_args.iter().cloned());
    Ok(args)
}

#[cfg(not(target_os = "windows"))]
pub fn is_program_in_path(program: &str) -> bool {
    if let Ok(path) = env::var("PATH") {
        for p in path.split(":") {
            let p_str = format!("{}/{}", p, program);
            if fs::metadata(p_str).is_ok() {
                return true;
            }
        }
    }
    false
}

#[cfg(target_os = "windows")]
pub fn is_program_in_path(program: &str) -> bool {
    use std::process::{Command, Stdio};

    let mut args = vec!["-Command".to_string(), "Get-Command".to_string()];
    args.push(program.to_string());
    Command::new("powershell")
        .args(&*args)
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()
        .is_ok()
}

pub fn make() -> Option<&'static str> {
    if is_program_in_path("gmake") {
        Some("gmake")
    } else if is_program_in_path("make") {
        Some("make")
    } else {
        None
    }
}

pub fn cmake() -> Option<&'static str> {
    if is_program_in_path("cmake") {
        Some("cmake")
    } else {
        None
    }
}

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

    #[test]
    #[cfg(target_os = "windows")]
    fn test_is_program_in_path() {
        let cmd = is_program_in_path("cmd");
        assert_eq!(true, cmd);
    }

    #[test]
    #[cfg(target_os = "linux")]
    fn test_make() {
        let cmake = make();
        assert_eq!(Some("make"), cmake);
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_make() {
        let cmake = make();
        assert_eq!(Some("make"), cmake);
    }

    #[test]
    fn test_parse_toml() {
        let toml = r#"[settings]
args = "an example arguments"
"#;
        let settings = parse_toml(toml.to_string());
        let expected = BuildConfig { args: "an example arguments".to_string() };
        assert_eq!(expected, settings);
    }
}
<commit_msg>rustfmt<commit_after>use std::env;
use std::fs;
use std::io::prelude::*;
use std::io;
use std::process;
use toml;
use toml::de::Error;
use config::Config;

const DEFAULT_ARGS: &'static str = "\"--with-zlib --with-ssl --enable-mruby --without-libstemmer \
                                    --disable-zeromq\"";

#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Configuration {
    settings: BuildConfig,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BuildConfig {
    args: String,
}

pub fn write_conf(config: &Config) {
    if !config.build_conf.exists() {
        let mut f = fs::File::create(&config.build_conf)
            .expect("Cound not create build configuration file.");
        let contents = format!("[settings]\nargs = {}", DEFAULT_ARGS);
        match f.write_all(&contents.as_bytes()) {
            Ok(_) => (),
            Err(e) => {
                println!("{}", e);
                process::exit(2)
            }
        }
    }
}

pub fn parse_toml(config_content: String) -> BuildConfig {
    println!("config:\n{}", config_content);
    let configuration: Result<Configuration, Error> = toml::from_str(&*config_content);
    let config = match configuration {
        Ok(config) => config,
        Err(_) => panic!("Couldn't decode toml with Configuration struct"),
    };
    config.settings
}

pub fn read_build_args(config: &Config) -> Result<String, io::Error> {
    let mut f = fs::File::open(&config.build_conf)?;
    let mut buffer = String::new();
    f.read_to_string(&mut buffer)?;
    let build_conf = parse_toml(buffer);
    Ok(build_conf.args)
}

// mainly used for macOS.
fn openssl_pkg_config_path() -> String {
    let key = "OPENSSL_PKG_CONFIG_PATH";
    env::var(key).unwrap_or("/usr/local/opt/openssl/lib/pkgconfig".to_string())
}

pub fn build_args(config: &Config, groonga_dir: String) -> Result<Vec<String>, io::Error> {
    let conf_args = read_build_args(config)?;
    let build_args: Vec<String> =
        conf_args.split_whitespace().into_iter().map(|e| e.to_owned()).collect();
    println!("{} with args: {:?}",
             config.versions_dir.join(groonga_dir.clone()).display(),
             build_args.clone());
    let mut args = vec![format!("--prefix={}",
                                config.versions_dir.join(groonga_dir.clone()).display()),
                        format!("PKG_CONFIG_PATH={}", openssl_pkg_config_path())];
    args.extend(build_args.iter().cloned());
    Ok(args)
}

#[cfg(not(target_os = "windows"))]
pub fn is_program_in_path(program: &str) -> bool {
    if let Ok(path) = env::var("PATH") {
        for p in path.split(":") {
            let p_str = format!("{}/{}", p, program);
            if fs::metadata(p_str).is_ok() {
                return true;
            }
        }
    }
    false
}

#[cfg(target_os = "windows")]
pub fn is_program_in_path(program: &str) -> bool {
    use std::process::{Command, Stdio};

    let mut args = vec!["-Command".to_string(), "Get-Command".to_string()];
    args.push(program.to_string());
    Command::new("powershell")
        .args(&*args)
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()
        .is_ok()
}

pub fn make() -> Option<&'static str> {
    if is_program_in_path("gmake") {
        Some("gmake")
    } else if is_program_in_path("make") {
        Some("make")
    } else {
        None
    }
}

pub fn cmake() -> Option<&'static str> {
    if is_program_in_path("cmake") {
        Some("cmake")
    } else {
        None
    }
}

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

    #[test]
    #[cfg(target_os = "windows")]
    fn test_is_program_in_path() {
        let cmd = is_program_in_path("cmd");
        assert_eq!(true, cmd);
    }

    #[test]
    #[cfg(target_os = "linux")]
    fn test_make() {
        let cmake = make();
        assert_eq!(Some("make"), cmake);
    }

    #[test]
    #[cfg(target_os = "macos")]
    fn test_make() {
        let cmake = make();
        assert_eq!(Some("make"), cmake);
    }

    #[test]
    fn test_parse_toml() {
        let toml = r#"[settings]
args = "an example arguments"
"#;
        let settings = parse_toml(toml.to_string());
        let expected = BuildConfig { args: "an example arguments".to_string() };
        assert_eq!(expected, settings);
    }
}
<|endoftext|>


<commit_before>use std::io::Read;

use url::Url;
use hyper::{self, Client};
use rustc_serialize::json::{self, Json, Object};

use errors::WeChatError;

pub struct WeChatClient {
    pub appid: String,
    pub secret: String,
    pub access_token: String,
}

impl WeChatClient {
    pub fn new(appid: &str, secret: &str) -> WeChatClient {
        WeChatClient {
            appid: appid.to_owned(),
            secret: secret.to_owned(),
            access_token: "".to_owned(),
        }
    }

    pub fn with_access_token(appid: &str, secret: &str, access_token: &str) -> WeChatClient {
        WeChatClient {
            appid: appid.to_owned(),
            secret: secret.to_owned(),
            access_token: access_token.to_owned(),
        }
    }

    pub fn post(&self, url: &str, params: Vec<(&str, &str)>, data: &Object) -> Result<Json, WeChatError> {
        let mut http_url = Url::parse(url).unwrap();
        http_url.set_query_from_pairs(params.into_iter());
        let body = match json::encode(data) {
            Ok(text) => text,
            Err(_) => "".to_owned(),
        };
        let mut client = Client::new();
        let mut res = match client.post(http_url).body(&body).send() {
            Ok(_res) => _res,
            Err(_) => { return Err(WeChatError::ClientError { errcode: -1, errmsg: "".to_owned() }); }
        };
        if res.status != hyper::Ok {
            return Err(WeChatError::ClientError { errcode: -2, errmsg: "Request status error".to_owned() })
        }
        let mut buf = String::new();
        match res.read_to_string(&mut buf) {
            Ok(_) => {},
            Err(_) => { return Err(WeChatError::ClientError { errcode: -3, errmsg: "Read response error".to_owned() }); }
        };
        let obj = match Json::from_str(&buf) {
            Ok(decoded) => { decoded },
            Err(_) => { return Err(WeChatError::ClientError { errcode: -4, errmsg: "Json decode error".to_owned() }); }
        };
        Ok(obj)
    }

    pub fn get(&self, url: &str, params: Vec<(&str, &str)>) -> Result<Json, WeChatError> {
        let mut http_url = Url::parse(url).unwrap();
        http_url.set_query_from_pairs(params.into_iter());
        let mut client = Client::new();
        let mut res = match client.post(http_url).send() {
            Ok(_res) => _res,
            Err(_) => { return Err(WeChatError::ClientError { errcode: -1, errmsg: "".to_owned() }); }
        };
        if res.status != hyper::Ok {
            return Err(WeChatError::ClientError { errcode: -2, errmsg: "Request status error".to_owned() })
        }
        let mut buf = String::new();
        match res.read_to_string(&mut buf) {
            Ok(_) => {},
            Err(_) => { return Err(WeChatError::ClientError { errcode: -3, errmsg: "Read response error".to_owned() }); }
        };
        let obj = match Json::from_str(&buf) {
            Ok(decoded) => { decoded },
            Err(_) => { return Err(WeChatError::ClientError { errcode: -4, errmsg: "Json decode error".to_owned() }); }
        };
        Ok(obj)
    }
}<commit_msg>Use Json::from_read instead of Json::from_str to simplify code<commit_after>use url::Url;
use hyper::{self, Client};
use rustc_serialize::json::{self, Json, Object};

use errors::WeChatError;

pub struct WeChatClient {
    pub appid: String,
    pub secret: String,
    pub access_token: String,
}

impl WeChatClient {
    pub fn new(appid: &str, secret: &str) -> WeChatClient {
        WeChatClient {
            appid: appid.to_owned(),
            secret: secret.to_owned(),
            access_token: "".to_owned(),
        }
    }

    pub fn with_access_token(appid: &str, secret: &str, access_token: &str) -> WeChatClient {
        WeChatClient {
            appid: appid.to_owned(),
            secret: secret.to_owned(),
            access_token: access_token.to_owned(),
        }
    }

    pub fn post(&self, url: &str, params: Vec<(&str, &str)>, data: &Object) -> Result<Json, WeChatError> {
        let mut http_url = Url::parse(url).unwrap();
        http_url.set_query_from_pairs(params.into_iter());
        let body = match json::encode(data) {
            Ok(text) => text,
            Err(_) => "".to_owned(),
        };
        let mut client = Client::new();
        let mut res = match client.post(http_url).body(&body).send() {
            Ok(_res) => _res,
            Err(_) => { return Err(WeChatError::ClientError { errcode: -1, errmsg: "".to_owned() }); }
        };
        if res.status != hyper::Ok {
            return Err(WeChatError::ClientError { errcode: -2, errmsg: "Request status error".to_owned() })
        }
        let obj = match Json::from_reader(&mut res) {
            Ok(decoded) => { decoded },
            Err(_) => { return Err(WeChatError::ClientError { errcode: -3, errmsg: "Json decode error".to_owned() }); }
        };
        Ok(obj)
    }

    pub fn get(&self, url: &str, params: Vec<(&str, &str)>) -> Result<Json, WeChatError> {
        let mut http_url = Url::parse(url).unwrap();
        http_url.set_query_from_pairs(params.into_iter());
        let mut client = Client::new();
        let mut res = match client.post(http_url).send() {
            Ok(_res) => _res,
            Err(_) => { return Err(WeChatError::ClientError { errcode: -1, errmsg: "".to_owned() }); }
        };
        if res.status != hyper::Ok {
            return Err(WeChatError::ClientError { errcode: -2, errmsg: "Request status error".to_owned() })
        }
        let obj = match Json::from_reader(&mut res) {
            Ok(decoded) => { decoded },
            Err(_) => { return Err(WeChatError::ClientError { errcode: -3, errmsg: "Json decode error".to_owned() }); }
        };
        Ok(obj)
    }
}<|endoftext|>


<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str, ptr};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

extern "C" {
    fn register_writer_tx(tx: *mut Sender<Vec<u8>>);
    fn register_stop_tx(tx: *mut Sender<()>);
}

/// Connects to the provided address, (eg "123.123.123.123:3000") and registers
/// the on data received handler
#[no_mangle]
pub extern "C" fn connect(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    1 as c_int;

    // let mut r_address;
    // unsafe {
    //     r_address = CStr::from_ptr(address);
    // }
    // let s_address = r_address.to_bytes();
    // let host_address = match str::from_utf8(s_address) {
    //     Ok(safe_str) => safe_str,
    //     Err(_) => {
    //         println!("Invalid host address");
    //         return -1 as c_int;
    //     }
    // };
    //
    // // Create and register a way to kill this client
    // let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    // let kill_tx = k_tx.clone();
    // let mut k_tx_ptr = Box::new(k_tx);
    // unsafe {
    //     register_stop_tx(&mut *k_tx_ptr);
    // }
    //
    // // Writer thread's channel
    // let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    // let mut w_tx_ptr = Box::new(w_tx);
    // unsafe {
    //     register_writer_tx(&mut *w_tx_ptr);
    // }
    //
    // let result = TcpStream::connect(host_address);
    // if result.is_err() {
    //     println!("Error connecting to {} - {}", host_address, result.unwrap_err());
    //     return -1 as c_int;
    // }
    // println!("Connected");
    // on_connect_handler();
    //
    // let stream = result.unwrap();
    // let client = Bstream::new(stream);
    //
    // let r_client = client.clone();
    // let r_kill_tx = kill_tx.clone();
    //
    // let w_client = client.clone();
    // let w_kill_tx = kill_tx.clone();
    //
    // // Start the reader thread
    // thread::Builder::new()
    //     .name("ReaderThread".to_string())
    //     .spawn(move||{
    //         reader_thread(r_client, handler, r_kill_tx)
    //     }).unwrap();
    //
    // // Start the writer thread
    // thread::Builder::new()
    //     .name("WriterThread".to_string())
    //     .spawn(move||{
    //         writer_thread(w_rx, w_client, w_kill_tx)
    //     }).unwrap();
    //
    // // Wait for the kill signal
    // match kill_rx.recv() {
    //     Ok(_) => { }
    //     Err(e) => {
    //         println!("Error on kill channel: {}", e);
    //         return -1 as c_int;
    //     }
    // };
    // on_disconnect_handler();
    //
    // // Exit out in standard C fashion
    // 0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn send_to_writer(w_tx: *mut Sender<Vec<u8>>,
                                 buffer: *const c_char,
                                 count: c_int,
                                 k_tx: *mut Sender<()>) -> c_int {
    if count < 1 {
        println!("Error - count must be greater than zero");
        return -1 as c_int;
    }

    let num_elts = count as usize;
    let mut n_buffer = Vec::<u8>::with_capacity(num_elts);
    for x in 0..num_elts as isize {
        unsafe {
            n_buffer.push(ptr::read(buffer.offset(x)) as u8);
        }
    }

    unsafe {
        match (*w_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                println!("Error sending buffer: {}", e);
                let _ = (*k_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream,
                 event_handler: extern fn(*const c_char, c_int),
                 kill_tx: Sender<()>) {

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        event_handler(c_buffer.as_ptr(), buffer.len() as c_int);
                    }).unwrap();
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Reader thread finished");
    let _ = kill_tx.send(());
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream, kill_tx: Sender<()>) {
    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        println!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Writer thread finished");
    let _ = kill_tx.send(());
}
<commit_msg>Debugging stuff<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the
// terms of the Mozilla Public License, v.
// 2.0. If a copy of the MPL was not
// distributed with this file, You can
// obtain one at
// http://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible
// With Secondary Licenses", as defined by
// the Mozilla Public License, v. 2.0.

use std::{str, ptr};
use std::ffi::{CStr, CString};
use std::thread;
use std::net::TcpStream;
use std::sync::mpsc::{channel, Sender, Receiver};

use super::libc::{c_int, c_char};
use super::simple_stream::bstream::Bstream;

extern "C" {
    fn register_writer_tx(tx: *mut Sender<Vec<u8>>);
    fn register_stop_tx(tx: *mut Sender<()>);
}

/// Connects to the provided address, (eg "123.123.123.123:3000") and registers
/// the on data received handler
#[no_mangle]
pub extern "C" fn connect(address: *const c_char,
    handler: extern fn(*const c_char, c_int),
    on_connect_handler: extern fn(),
    on_disconnect_handler: extern fn()) -> c_int {

    1 as c_int

    // let mut r_address;
    // unsafe {
    //     r_address = CStr::from_ptr(address);
    // }
    // let s_address = r_address.to_bytes();
    // let host_address = match str::from_utf8(s_address) {
    //     Ok(safe_str) => safe_str,
    //     Err(_) => {
    //         println!("Invalid host address");
    //         return -1 as c_int;
    //     }
    // };
    //
    // // Create and register a way to kill this client
    // let (k_tx, kill_rx): (Sender<()>, Receiver<()>) = channel();
    // let kill_tx = k_tx.clone();
    // let mut k_tx_ptr = Box::new(k_tx);
    // unsafe {
    //     register_stop_tx(&mut *k_tx_ptr);
    // }
    //
    // // Writer thread's channel
    // let (w_tx, w_rx): (Sender<Vec<u8>>, Receiver<Vec<u8>>) = channel();
    // let mut w_tx_ptr = Box::new(w_tx);
    // unsafe {
    //     register_writer_tx(&mut *w_tx_ptr);
    // }
    //
    // let result = TcpStream::connect(host_address);
    // if result.is_err() {
    //     println!("Error connecting to {} - {}", host_address, result.unwrap_err());
    //     return -1 as c_int;
    // }
    // println!("Connected");
    // on_connect_handler();
    //
    // let stream = result.unwrap();
    // let client = Bstream::new(stream);
    //
    // let r_client = client.clone();
    // let r_kill_tx = kill_tx.clone();
    //
    // let w_client = client.clone();
    // let w_kill_tx = kill_tx.clone();
    //
    // // Start the reader thread
    // thread::Builder::new()
    //     .name("ReaderThread".to_string())
    //     .spawn(move||{
    //         reader_thread(r_client, handler, r_kill_tx)
    //     }).unwrap();
    //
    // // Start the writer thread
    // thread::Builder::new()
    //     .name("WriterThread".to_string())
    //     .spawn(move||{
    //         writer_thread(w_rx, w_client, w_kill_tx)
    //     }).unwrap();
    //
    // // Wait for the kill signal
    // match kill_rx.recv() {
    //     Ok(_) => { }
    //     Err(e) => {
    //         println!("Error on kill channel: {}", e);
    //         return -1 as c_int;
    //     }
    // };
    // on_disconnect_handler();
    //
    // // Exit out in standard C fashion
    // 0 as c_int
}

/// Writes the complete contents of buffer to the server
/// Returns -1 on error
#[no_mangle]
pub extern "C" fn send_to_writer(w_tx: *mut Sender<Vec<u8>>,
                                 buffer: *const c_char,
                                 count: c_int,
                                 k_tx: *mut Sender<()>) -> c_int {
    if count < 1 {
        println!("Error - count must be greater than zero");
        return -1 as c_int;
    }

    let num_elts = count as usize;
    let mut n_buffer = Vec::<u8>::with_capacity(num_elts);
    for x in 0..num_elts as isize {
        unsafe {
            n_buffer.push(ptr::read(buffer.offset(x)) as u8);
        }
    }

    unsafe {
        match (*w_tx).send(n_buffer) {
            Ok(_) => { }
            Err(e) => {
                println!("Error sending buffer: {}", e);
                let _ = (*k_tx).send(());
                return -1 as c_int;
            }
        };
    }

    0 as c_int
}

/// Forever listens to incoming data and when a complete message is received,
/// the passed callback is hit
fn reader_thread(client: Bstream,
                 event_handler: extern fn(*const c_char, c_int),
                 kill_tx: Sender<()>) {

    let mut reader = client.clone();
    loop {
        match reader.read() {
            Ok(buffer) => {
                // Launch the handler in a new thread
                thread::Builder::new()
                    .name("Reader-Worker".to_string())
                    .spawn(move||{
                        let slice = &buffer[..];
                        let c_buffer = CString::new(slice).unwrap();
                        event_handler(c_buffer.as_ptr(), buffer.len() as c_int);
                    }).unwrap();
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Reader thread finished");
    let _ = kill_tx.send(());
}

/// Forever listens to Receiver<Vec<u8>> waiting on messages to come in
/// Once available, blocks until the entire message has been written
fn writer_thread(rx: Receiver<Vec<u8>>, client: Bstream, kill_tx: Sender<()>) {
    let mut writer = client.clone();
    loop {
        match rx.recv() {
            Ok(ref mut buffer) => {
                match writer.write(buffer) {
                    Ok(_) => { }
                    Err(e) => {
                        println!("Error: {}", e);
                        break;
                    }
                };
            }
            Err(e) => {
                println!("Error: {}", e);
                break;
            }
        };
    }
    println!("Writer thread finished");
    let _ = kill_tx.send(());
}
<|endoftext|>