<commit_before>extern crate sodiumoxide;

use self::sodiumoxide::crypto::aead::chacha20poly1305_ietf;
use self::sodiumoxide::utils;

use std::cmp;
use std::io;
use std::io::{Read, Write};

use utils::crypto::pwhash_argon2i13;
use errors::common::CommonError;
use domain::wallet::KeyDerivationMethod;

pub const KEYBYTES: usize = chacha20poly1305_ietf::KEYBYTES;
pub const NONCEBYTES: usize = chacha20poly1305_ietf::NONCEBYTES;
pub const TAGBYTES: usize = chacha20poly1305_ietf::TAGBYTES;

sodium_type!(Key, chacha20poly1305_ietf::Key, KEYBYTES);
sodium_type!(Nonce, chacha20poly1305_ietf::Nonce, NONCEBYTES);

impl Nonce {
    pub fn increment(&mut self) {
        utils::increment_le(&mut (self.0).0);
    }
}

pub fn gen_key() -> Key {
    Key(chacha20poly1305_ietf::gen_key())
}

pub fn derive_key(passphrase: &str, salt: &pwhash_argon2i13::Salt, key_derivation_method: &KeyDerivationMethod) -> Result<Key, CommonError> {
    let mut key_bytes = [0u8; chacha20poly1305_ietf::KEYBYTES];
    pwhash_argon2i13::pwhash(&mut key_bytes, passphrase.as_bytes(), salt, key_derivation_method)
        .map_err(|err| CommonError::InvalidStructure(format!("Can't derive key: {}", err)))?;
    Ok(Key::new(key_bytes))
}

pub fn gen_nonce() -> Nonce {
    Nonce(chacha20poly1305_ietf::gen_nonce())
}

pub fn gen_nonce_and_encrypt(data: &[u8], key: &Key) -> (Vec<u8>, Nonce) {
    let nonce = gen_nonce();

    let encrypted_data = chacha20poly1305_ietf::seal(
        data,
        None,
        &nonce.0,
        &key.0
    );

    (encrypted_data, nonce)
}

pub fn encrypt(data: &[u8], key: &Key, nonce: &Nonce) -> Vec<u8> {
    chacha20poly1305_ietf::seal(
        data,
        None,
        &nonce.0,
        &key.0,
    )
}

pub fn decrypt(data: &[u8], key: &Key, nonce: &Nonce) -> Result<Vec<u8>, CommonError> {
    chacha20poly1305_ietf::open(
        &data,
        None,
        &nonce.0,
        &key.0,
    )
        .map_err(|err| CommonError::InvalidStructure(format!("Unable to decrypt data: {:?}", err)))
}

pub struct Writer<W: Write> {
    buffer: Vec<u8>,
    chunk_size: usize,
    key: Key,
    nonce: Nonce,
    inner: W,
}

impl<W: Write> Writer<W> {
    pub fn new(inner: W, key: Key, nonce: Nonce, chunk_size: usize) -> Self {
        Writer {
            buffer: Vec::new(),
            chunk_size,
            key,
            nonce,
            inner,
        }
    }

    #[allow(unused)]
    pub fn into_inner(self) -> W {
        self.inner
    }
}

impl<W: Write> Write for Writer<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.buffer.write_all(buf)?; // TODO: Small optimizations are possible

        let mut chunk_start = 0;

        while self.buffer.len() >= chunk_start + self.chunk_size {
            let chunk = &self.buffer[chunk_start..chunk_start + self.chunk_size];
            self.inner.write_all(&encrypt(chunk, &self.key, &self.nonce))?;
            self.nonce.increment();
            chunk_start += self.chunk_size;
        }

        if chunk_start > 0 {
            self.buffer.drain(..chunk_start);
        }

        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        if self.buffer.len() > 0 {
            self.inner.write_all(&encrypt(&self.buffer, &self.key, &self.nonce))?;
            self.nonce.increment();
        }

        self.buffer.flush()
    }
}

pub struct Reader<R: Read> {
    rest_buffer: Vec<u8>,
    chunk_buffer: Vec<u8>,
    key: Key,
    nonce: Nonce,
    inner: R,
}

impl<R: Read> Reader<R> {
    pub fn new(inner: R, key: Key, nonce: Nonce, chunk_size: usize) -> Self {
        Reader {
            rest_buffer: Vec::new(),
            chunk_buffer: vec![0; chunk_size + TAGBYTES],
            key,
            nonce,
            inner,
        }
    }

    #[allow(unused)]
    pub fn into_inner(self) -> R {
        self.inner
    }

    fn _read_chunk(&mut self) -> io::Result<usize> {
        let mut read = 0;

        while read < self.chunk_buffer.len() {
            match self.inner.read(&mut self.chunk_buffer[read..]) {
                Ok(0) => break,
                Ok(n) => read += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
                Err(e) => Err(e)?
            }
        }

        if read == 0 {
            Err(io::Error::new(io::ErrorKind::UnexpectedEof, "No more crypto chucks to consume"))
        } else {
            Ok(read)
        }
    }
}

impl<R: Read> Read for Reader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut pos = 0;

        // Consume from rest buffer
        if self.rest_buffer.len() > 0 {
            let to_copy = cmp::min(self.rest_buffer.len(), buf.len() - pos);
            buf[pos..pos + to_copy].copy_from_slice(&self.rest_buffer[..to_copy]);
            pos = pos + to_copy;
            self.rest_buffer.drain(..to_copy);
        }

        // Consume from chunks
        while pos < buf.len() {
            let chunk_size = self._read_chunk()?;

            let chunk = decrypt(&self.chunk_buffer[..chunk_size], &self.key, &self.nonce)
                .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid data in crypto chunk"))?;

            self.nonce.increment();

            let to_copy = cmp::min(chunk.len(), buf.len() - pos);
            buf[pos..pos + to_copy].copy_from_slice(&chunk[..to_copy]);
            pos = pos + to_copy;

            // Save rest in rest buffer
            if pos == buf.len() && to_copy < chunk.len() {
                self.rest_buffer.extend(&chunk[to_copy..]);
            }
        }

        Ok(buf.len())
    }
}

#[cfg(test)]
mod tests {
    extern crate rmp_serde;

    use super::*;
    use utils::crypto::randombytes::randombytes;

    #[test]
    fn derivation_argon2i_mod_produces_expected_result() {
        let passphrase = "passphrase";
        let salt_bytes: [u8; 32] = [
            24, 62, 35, 31, 123, 241, 94, 24, 192, 110, 199, 143, 173, 20, 23, 102,
            184, 99, 221, 64, 247, 230, 11, 253, 10, 7, 80, 236, 185, 249, 110, 187
        ];
        let key_bytes: [u8; 32] = [
            148, 89, 76, 239, 127, 103, 13, 86, 84, 217, 216, 13, 223, 141, 225, 41,
            223, 126, 145, 138, 174, 31, 142, 199, 81, 12, 40, 201, 67, 8, 6, 251
        ];

        let res = derive_key(
            passphrase,
            &pwhash_argon2i13::Salt::from_slice(&salt_bytes).unwrap(),
            &KeyDerivationMethod::ARGON2I_MOD
        ).unwrap();

        assert_eq!(res, Key::new(key_bytes))
    }

    #[test]
    fn derivation_argon2i_int_produces_expected_result() {
        let passphrase = "passphrase";
        let salt_bytes: [u8; 32] = [
            24, 62, 35, 31, 123, 241, 94, 24, 192, 110, 199, 143, 173, 20, 23, 102,
            184, 99, 221, 64, 247, 230, 11, 253, 10, 7, 80, 236, 185, 249, 110, 187
        ];
        let key_bytes: [u8; 32] = [
            247, 55, 177, 252, 244, 130, 218, 129, 113, 206, 72, 44, 29, 68, 134, 215,
            249, 233, 131, 199, 38, 87, 69, 217, 156, 217, 10, 160, 30, 148, 80, 160
        ];

        let res = derive_key(
            passphrase,
            &pwhash_argon2i13::Salt::from_slice(&salt_bytes).unwrap(),
            &KeyDerivationMethod::ARGON2I_INT
        ).unwrap();

        assert_eq!(res, Key::new(key_bytes))
    }

    #[test]
    fn encrypt_decrypt_works() {
        let data = randombytes(100);
        let key = gen_key();

        let (c, nonce) = gen_nonce_and_encrypt(&data, &key);
        let u = decrypt(&c, &key, &nonce).unwrap();

        assert_eq!(data, u);
    }

    #[test]
    fn encrypt_decrypt_works_for_nonce() {
        let data = randombytes(16);

        let key = gen_key();
        let nonce = gen_nonce();
        let c = encrypt(&data, &key, &nonce);
        let u = decrypt(&c, &key, &nonce).unwrap();

        assert_eq!(data, u)
    }

    #[test]
    fn nonce_serialize_deserialize_works() {
        let nonce = gen_nonce();
        let serialized = rmp_serde::to_vec(&nonce).unwrap();
        let deserialized: Nonce = rmp_serde::from_slice(&serialized).unwrap();

        assert_eq!(serialized.len(), NONCEBYTES + 2);
        assert_eq!(nonce, deserialized)
    }

    #[test]
    fn key_serialize_deserialize_works() {
        let key = gen_key();
        let serialized = rmp_serde::to_vec(&key).unwrap();
        let deserialized: Key = rmp_serde::from_slice(&serialized).unwrap();

        assert_eq!(serialized.len(), KEYBYTES + 2);
        assert_eq!(key, deserialized)
    }

    #[test]
    fn writer_reader_works_for_less_than_one_chunk() {
        let plain = randombytes(7);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 7 + TAGBYTES);

        let mut decrypted = vec![0u8; 7];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }

    #[test]
    fn writer_reader_works_for_exact_one_chunk() {
        let plain = randombytes(10);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 10 + TAGBYTES);

        let mut decrypted = vec![0u8; 10];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }

    #[test]
    fn writer_reader_works_for_one_to_two_chunks() {
        let plain = randombytes(13);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 13 + 2 * TAGBYTES);

        let mut decrypted = vec![0u8; 13];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }

    #[test]
    fn writer_reader_works_for_exact_two_chunks() {
        let plain = randombytes(20);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 20 + 2 * TAGBYTES);

        let mut decrypted = vec![0u8; 20];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }
}
<commit_msg>added detached with aad functons<commit_after>extern crate sodiumoxide;

use self::sodiumoxide::crypto::aead::chacha20poly1305_ietf;
use self::sodiumoxide::utils;

use std::cmp;
use std::io;
use std::io::{Read, Write};

use utils::crypto::pwhash_argon2i13;
use errors::common::CommonError;
use domain::wallet::KeyDerivationMethod;

pub const KEYBYTES: usize = chacha20poly1305_ietf::KEYBYTES;
pub const NONCEBYTES: usize = chacha20poly1305_ietf::NONCEBYTES;
pub const TAGBYTES: usize = chacha20poly1305_ietf::TAGBYTES;

sodium_type!(Key, chacha20poly1305_ietf::Key, KEYBYTES);
sodium_type!(Nonce, chacha20poly1305_ietf::Nonce, NONCEBYTES);
sodium_type!(Tag, chacha20poly1305_ietf::Tag, TAGBYTES);

impl Nonce {
    pub fn increment(&mut self) {
        utils::increment_le(&mut (self.0).0);
    }
}

pub fn gen_key() -> Key {
    Key(chacha20poly1305_ietf::gen_key())
}

pub fn derive_key(passphrase: &str, salt: &pwhash_argon2i13::Salt, key_derivation_method: &KeyDerivationMethod) -> Result<Key, CommonError> {
    let mut key_bytes = [0u8; chacha20poly1305_ietf::KEYBYTES];
    pwhash_argon2i13::pwhash(&mut key_bytes, passphrase.as_bytes(), salt, key_derivation_method)
        .map_err(|err| CommonError::InvalidStructure(format!("Can't derive key: {}", err)))?;
    Ok(Key::new(key_bytes))
}

pub fn gen_nonce() -> Nonce {
    Nonce(chacha20poly1305_ietf::gen_nonce())
}

pub fn gen_nonce_and_encrypt(data: &[u8], key: &Key) -> (Vec<u8>, Nonce) {
    let nonce = gen_nonce();

    let encrypted_data = chacha20poly1305_ietf::seal(
        data,
        None,
        &nonce.0,
        &key.0
    );

    (encrypted_data, nonce)
}

pub fn gen_nonce_and_encrypt_detached(data: &[u8], aad: &[u8], key: &Key) -> (Vec<u8>, Nonce, Tag) {
    let nonce = gen_nonce();

    let mut plain = data.to_vec();
    let tag = chacha20poly1305_ietf::seal_detached(
        plain.as_mut_slice(),
        Some(aad),
        &nonce.0,
        &key.0
    );

    (plain.to_vec(), nonce, Tag(tag))
}

pub fn decrypt_detached(data: &[u8], key: &Key, nonce: &Nonce, tag: &Tag, ad: Option<&[u8]>) -> Result<Vec<u8>, CommonError> {
    let mut plain = data.to_vec();
    chacha20poly1305_ietf::open_detached(plain.as_mut_slice(),
        ad,
        &tag.0,
        &nonce.0,
        &key.0,
    ).map_err(|err| CommonError::InvalidStructure(format!("Unable to decrypt data: {:?}", err)))?;

    Ok(plain.to_vec())
}

pub fn encrypt(data: &[u8], key: &Key, nonce: &Nonce) -> Vec<u8> {
    chacha20poly1305_ietf::seal(
        data,
        None,
        &nonce.0,
        &key.0,
    )
}

pub fn decrypt(data: &[u8], key: &Key, nonce: &Nonce) -> Result<Vec<u8>, CommonError> {
    chacha20poly1305_ietf::open(
        &data,
        None,
        &nonce.0,
        &key.0,
    )
        .map_err(|err| CommonError::InvalidStructure(format!("Unable to decrypt data: {:?}", err)))
}

pub struct Writer<W: Write> {
    buffer: Vec<u8>,
    chunk_size: usize,
    key: Key,
    nonce: Nonce,
    inner: W,
}

impl<W: Write> Writer<W> {
    pub fn new(inner: W, key: Key, nonce: Nonce, chunk_size: usize) -> Self {
        Writer {
            buffer: Vec::new(),
            chunk_size,
            key,
            nonce,
            inner,
        }
    }

    #[allow(unused)]
    pub fn into_inner(self) -> W {
        self.inner
    }
}

impl<W: Write> Write for Writer<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.buffer.write_all(buf)?; // TODO: Small optimizations are possible

        let mut chunk_start = 0;

        while self.buffer.len() >= chunk_start + self.chunk_size {
            let chunk = &self.buffer[chunk_start..chunk_start + self.chunk_size];
            self.inner.write_all(&encrypt(chunk, &self.key, &self.nonce))?;
            self.nonce.increment();
            chunk_start += self.chunk_size;
        }

        if chunk_start > 0 {
            self.buffer.drain(..chunk_start);
        }

        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        if self.buffer.len() > 0 {
            self.inner.write_all(&encrypt(&self.buffer, &self.key, &self.nonce))?;
            self.nonce.increment();
        }

        self.buffer.flush()
    }
}

pub struct Reader<R: Read> {
    rest_buffer: Vec<u8>,
    chunk_buffer: Vec<u8>,
    key: Key,
    nonce: Nonce,
    inner: R,
}

impl<R: Read> Reader<R> {
    pub fn new(inner: R, key: Key, nonce: Nonce, chunk_size: usize) -> Self {
        Reader {
            rest_buffer: Vec::new(),
            chunk_buffer: vec![0; chunk_size + TAGBYTES],
            key,
            nonce,
            inner,
        }
    }

    #[allow(unused)]
    pub fn into_inner(self) -> R {
        self.inner
    }

    fn _read_chunk(&mut self) -> io::Result<usize> {
        let mut read = 0;

        while read < self.chunk_buffer.len() {
            match self.inner.read(&mut self.chunk_buffer[read..]) {
                Ok(0) => break,
                Ok(n) => read += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue,
                Err(e) => Err(e)?
            }
        }

        if read == 0 {
            Err(io::Error::new(io::ErrorKind::UnexpectedEof, "No more crypto chucks to consume"))
        } else {
            Ok(read)
        }
    }
}

impl<R: Read> Read for Reader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let mut pos = 0;

        // Consume from rest buffer
        if self.rest_buffer.len() > 0 {
            let to_copy = cmp::min(self.rest_buffer.len(), buf.len() - pos);
            buf[pos..pos + to_copy].copy_from_slice(&self.rest_buffer[..to_copy]);
            pos = pos + to_copy;
            self.rest_buffer.drain(..to_copy);
        }

        // Consume from chunks
        while pos < buf.len() {
            let chunk_size = self._read_chunk()?;

            let chunk = decrypt(&self.chunk_buffer[..chunk_size], &self.key, &self.nonce)
                .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "Invalid data in crypto chunk"))?;

            self.nonce.increment();

            let to_copy = cmp::min(chunk.len(), buf.len() - pos);
            buf[pos..pos + to_copy].copy_from_slice(&chunk[..to_copy]);
            pos = pos + to_copy;

            // Save rest in rest buffer
            if pos == buf.len() && to_copy < chunk.len() {
                self.rest_buffer.extend(&chunk[to_copy..]);
            }
        }

        Ok(buf.len())
    }
}

#[cfg(test)]
mod tests {
    extern crate rmp_serde;

    use super::*;
    use utils::crypto::randombytes::randombytes;

    #[test]
    fn derivation_argon2i_mod_produces_expected_result() {
        let passphrase = "passphrase";
        let salt_bytes: [u8; 32] = [
            24, 62, 35, 31, 123, 241, 94, 24, 192, 110, 199, 143, 173, 20, 23, 102,
            184, 99, 221, 64, 247, 230, 11, 253, 10, 7, 80, 236, 185, 249, 110, 187
        ];
        let key_bytes: [u8; 32] = [
            148, 89, 76, 239, 127, 103, 13, 86, 84, 217, 216, 13, 223, 141, 225, 41,
            223, 126, 145, 138, 174, 31, 142, 199, 81, 12, 40, 201, 67, 8, 6, 251
        ];

        let res = derive_key(
            passphrase,
            &pwhash_argon2i13::Salt::from_slice(&salt_bytes).unwrap(),
            &KeyDerivationMethod::ARGON2I_MOD
        ).unwrap();

        assert_eq!(res, Key::new(key_bytes))
    }

    #[test]
    fn derivation_argon2i_int_produces_expected_result() {
        let passphrase = "passphrase";
        let salt_bytes: [u8; 32] = [
            24, 62, 35, 31, 123, 241, 94, 24, 192, 110, 199, 143, 173, 20, 23, 102,
            184, 99, 221, 64, 247, 230, 11, 253, 10, 7, 80, 236, 185, 249, 110, 187
        ];
        let key_bytes: [u8; 32] = [
            247, 55, 177, 252, 244, 130, 218, 129, 113, 206, 72, 44, 29, 68, 134, 215,
            249, 233, 131, 199, 38, 87, 69, 217, 156, 217, 10, 160, 30, 148, 80, 160
        ];

        let res = derive_key(
            passphrase,
            &pwhash_argon2i13::Salt::from_slice(&salt_bytes).unwrap(),
            &KeyDerivationMethod::ARGON2I_INT
        ).unwrap();

        assert_eq!(res, Key::new(key_bytes))
    }

    #[test]
    fn gen_nonce_and_encrypt_decrypt_works() {
        let data = randombytes(100);
        let key = gen_key();

        let (c, nonce) = gen_nonce_and_encrypt(&data, &key);
        let u = decrypt(&c, &key, &nonce).unwrap();

        assert_eq!(data, u);
    }

    #[test]
    pub fn gen_nonce_and_encrypt_detached_decrypt_detached_works() {
        let data = randombytes(100);
        let key = gen_key();
        let aad= randombytes(100);

        let (c, nonce, tag) = gen_nonce_and_encrypt_detached(&data, aad.as_slice(), &key);
        let u = decrypt_detached(&c, &key, &nonce, &tag, Some(aad.as_slice())).unwrap();
        assert_eq!(data, u);
}

    #[test]
    fn encrypt_decrypt_works_for_nonce() {
        let data = randombytes(16);

        let key = gen_key();
        let nonce = gen_nonce();
        let c = encrypt(&data, &key, &nonce);
        let u = decrypt(&c, &key, &nonce).unwrap();

        assert_eq!(data, u)
    }

    #[test]
    fn nonce_serialize_deserialize_works() {
        let nonce = gen_nonce();
        let serialized = rmp_serde::to_vec(&nonce).unwrap();
        let deserialized: Nonce = rmp_serde::from_slice(&serialized).unwrap();

        assert_eq!(serialized.len(), NONCEBYTES + 2);
        assert_eq!(nonce, deserialized)
    }

    #[test]
    fn key_serialize_deserialize_works() {
        let key = gen_key();
        let serialized = rmp_serde::to_vec(&key).unwrap();
        let deserialized: Key = rmp_serde::from_slice(&serialized).unwrap();

        assert_eq!(serialized.len(), KEYBYTES + 2);
        assert_eq!(key, deserialized)
    }

    #[test]
    fn writer_reader_works_for_less_than_one_chunk() {
        let plain = randombytes(7);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 7 + TAGBYTES);

        let mut decrypted = vec![0u8; 7];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }

    #[test]
    fn writer_reader_works_for_exact_one_chunk() {
        let plain = randombytes(10);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 10 + TAGBYTES);

        let mut decrypted = vec![0u8; 10];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }

    #[test]
    fn writer_reader_works_for_one_to_two_chunks() {
        let plain = randombytes(13);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 13 + 2 * TAGBYTES);

        let mut decrypted = vec![0u8; 13];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }

    #[test]
    fn writer_reader_works_for_exact_two_chunks() {
        let plain = randombytes(20);
        let key = gen_key();
        let nonce = gen_nonce();

        let mut writer = Writer::new(Vec::<u8>::new(), key.clone(), nonce.clone(), 10);
        writer.write_all(&plain).unwrap();
        writer.flush().unwrap();

        let encrypted = writer.into_inner();
        assert_eq!(encrypted.len(), 20 + 2 * TAGBYTES);

        let mut decrypted = vec![0u8; 20];
        let mut reader = Reader::new(&encrypted[..], key, nonce, 10);
        reader.read_exact(&mut decrypted).unwrap();

        assert_eq!(plain, decrypted);
    }
}
<|endoftext|>


<commit_before>// A port of the simplistic benchmark from
//
//    http://github.com/PaulKeeble/ScalaVErlangAgents
//
// I *think* it's the same, more or less.

extern mod std;
extern mod extra;
extern mod zmq;

use std::comm;
use std::os;
use std::task;

fn server(pull_socket: zmq::Socket, push_socket: zmq::Socket, mut workers: uint) {
    let mut count = 0u;
    let mut msg = zmq::Message::new();

    while workers != 0 {
        match pull_socket.recv(&mut msg, 0) {
            Err(e) => fail!(e.to_str()),
            Ok(()) => {
                msg.with_str(|s| {
                    if s == "" {
                        workers -= 1;
                    } else {
                        count += from_str::<uint>(s).unwrap();
                    }
                })
            }
        }
    }

    match push_socket.send_str(count.to_str(), 0) {
        Ok(()) => { }
        Err(e) => fail!(e.to_str()),
    }
}

fn spawn_server(ctx: &zmq::Context, workers: uint) -> comm::Chan<()> {
    let pull_socket = ctx.socket(zmq::PULL).unwrap();
    let push_socket = ctx.socket(zmq::PUSH).unwrap();

    pull_socket.bind("inproc://server-pull").unwrap();
    push_socket.bind("inproc://server-push").unwrap();
    //pull_socket.bind("tcp://127.0.0.1:3456").unwrap();
    //push_socket.bind("tcp://127.0.0.1:3457").unwrap();

    // Spawn the server.
    let (ready_po, ready_ch) = comm::stream();
    let (start_po, start_ch) = comm::stream();

    let mut task = task::task();
    task.sched_mode(task::SingleThreaded);
    do task.spawn {
        // Let the main thread know we're ready.
        ready_ch.send(());

        // Wait until we need to start.
        start_po.recv();

        server(pull_socket, push_socket, workers);
    }

    // Wait for the server to start.
    ready_po.recv();

    start_ch
}

fn worker(push_socket: zmq::Socket, count: uint) {
    for _ in range(0, count) {
        push_socket.send_str(100u.to_str(), 0).unwrap();
    }

    // Let the server know we're done.
    push_socket.send_str("", 0).unwrap();
}

fn spawn_worker(ctx: &zmq::Context, count: uint) -> comm::Port<()> {
    let push_socket = ctx.socket(zmq::PUSH).unwrap();

    push_socket.connect("inproc://server-pull").unwrap();
    //push_socket.connect("tcp://127.0.0.1:3456").unwrap();

    // Spawn the worker.
    let (po, ch) = comm::stream();
    let mut task = task::task();
    task.sched_mode(task::SingleThreaded);
    do task.spawn {
        // Let the main thread we're ready.
        ch.send(());

        worker(push_socket, count);

        ch.send(());
    }

    // Wait for the worker to start.
    po.recv();

    po
}

fn run(ctx: zmq::Context, size: uint, workers: uint) {
    let start_ch = spawn_server(&ctx, workers);

    // Create some command/control sockets.
    let push_socket = ctx.socket(zmq::PUSH).unwrap();
    let pull_socket = ctx.socket(zmq::PULL).unwrap();

    push_socket.connect("inproc://server-pull").unwrap();
    pull_socket.connect("inproc://server-push").unwrap();
    //push_socket.connect("tcp://127.0.0.1:3456").unwrap();
    //pull_socket.connect("tcp://127.0.0.1:3457").unwrap();

    // Spawn all the workers.
    let mut worker_results = ~[];
    for _ in range(0, workers) {
        worker_results.push(spawn_worker(&ctx, size / workers));
    }

    let start = extra::time::precise_time_s();

    start_ch.send(());

    // Block until all the workers finish.
    for po in worker_results.iter() {
        po.recv();
    }

    // Receive the final count.
    let result = match pull_socket.recv_msg(0) {
        Ok(msg) => msg.with_str(|s| from_str::<uint>(s).unwrap()),
        Err(e) => fail!(e.to_str()),
    };

    let end = extra::time::precise_time_s();
    let elapsed = end - start;

    println!("Count is {}", result);
    println!("Test took {} seconds", elapsed);
    let thruput = ((size / workers * workers) as f64) / (elapsed as f64);
    println!("Throughput={:f} per sec", thruput);
}

fn main() {
    let args = os::args();

    let args = if os::getenv("RUST_BENCH").is_some() {
        ~[~"", ~"1000000", ~"10000"]
    } else if args.len() <= 1u {
        ~[~"", ~"10000", ~"4"]
    } else {
        args
    };

    let size = from_str::<uint>(args[1]).unwrap();
    let workers = from_str::<uint>(args[2]).unwrap();

    let ctx = zmq::Context::new();

    run(ctx, size, workers);
}
<commit_msg>Update to rust 0.9-pre (09b840)<commit_after>// A port of the simplistic benchmark from
//
//    http://github.com/PaulKeeble/ScalaVErlangAgents
//
// I *think* it's the same, more or less.

extern mod std;
extern mod extra;
extern mod zmq;

use std::comm;
use std::os;
use std::task;

fn server(pull_socket: zmq::Socket, push_socket: zmq::Socket, mut workers: uint) {
    let mut count = 0u;
    let mut msg = zmq::Message::new();

    while workers != 0 {
        match pull_socket.recv(&mut msg, 0) {
            Err(e) => fail!(e.to_str()),
            Ok(()) => {
                msg.with_str(|s| {
                    if s == "" {
                        workers -= 1;
                    } else {
                        count += from_str::<uint>(s).unwrap();
                    }
                })
            }
        }
    }

    match push_socket.send_str(count.to_str(), 0) {
        Ok(()) => { }
        Err(e) => fail!(e.to_str()),
    }
}

fn spawn_server(ctx: &zmq::Context, workers: uint) -> comm::Chan<()> {
    let pull_socket = ctx.socket(zmq::PULL).unwrap();
    let push_socket = ctx.socket(zmq::PUSH).unwrap();

    pull_socket.bind("inproc://server-pull").unwrap();
    push_socket.bind("inproc://server-push").unwrap();
    //pull_socket.bind("tcp://127.0.0.1:3456").unwrap();
    //push_socket.bind("tcp://127.0.0.1:3457").unwrap();

    // Spawn the server.
    let (ready_po, ready_ch) = comm::Chan::new();
    let (start_po, start_ch) = comm::Chan::new();

    let mut task = task::task();
    task.sched_mode(task::SingleThreaded);
    do task.spawn {
        // Let the main thread know we're ready.
        ready_ch.send(());

        // Wait until we need to start.
        start_po.recv();

        server(pull_socket, push_socket, workers);
    }

    // Wait for the server to start.
    ready_po.recv();

    start_ch
}

fn worker(push_socket: zmq::Socket, count: uint) {
    for _ in range(0, count) {
        push_socket.send_str(100u.to_str(), 0).unwrap();
    }

    // Let the server know we're done.
    push_socket.send_str("", 0).unwrap();
}

fn spawn_worker(ctx: &zmq::Context, count: uint) -> comm::Port<()> {
    let push_socket = ctx.socket(zmq::PUSH).unwrap();

    push_socket.connect("inproc://server-pull").unwrap();
    //push_socket.connect("tcp://127.0.0.1:3456").unwrap();

    // Spawn the worker.
    let (po, ch) = comm::Chan::new();
    let mut task = task::task();
    task.sched_mode(task::SingleThreaded);
    do task.spawn {
        // Let the main thread we're ready.
        ch.send(());

        worker(push_socket, count);

        ch.send(());
    }

    // Wait for the worker to start.
    po.recv();

    po
}

fn run(ctx: zmq::Context, size: uint, workers: uint) {
    let start_ch = spawn_server(&ctx, workers);

    // Create some command/control sockets.
    let push_socket = ctx.socket(zmq::PUSH).unwrap();
    let pull_socket = ctx.socket(zmq::PULL).unwrap();

    push_socket.connect("inproc://server-pull").unwrap();
    pull_socket.connect("inproc://server-push").unwrap();
    //push_socket.connect("tcp://127.0.0.1:3456").unwrap();
    //pull_socket.connect("tcp://127.0.0.1:3457").unwrap();

    // Spawn all the workers.
    let mut worker_results = ~[];
    for _ in range(0, workers) {
        worker_results.push(spawn_worker(&ctx, size / workers));
    }

    let start = extra::time::precise_time_s();

    start_ch.send(());

    // Block until all the workers finish.
    for po in worker_results.iter() {
        po.recv();
    }

    // Receive the final count.
    let result = match pull_socket.recv_msg(0) {
        Ok(msg) => msg.with_str(|s| from_str::<uint>(s).unwrap()),
        Err(e) => fail!(e.to_str()),
    };

    let end = extra::time::precise_time_s();
    let elapsed = end - start;

    println!("Count is {}", result);
    println!("Test took {} seconds", elapsed);
    let thruput = ((size / workers * workers) as f64) / (elapsed as f64);
    println!("Throughput={:f} per sec", thruput);
}

fn main() {
    let args = os::args();

    let args = if os::getenv("RUST_BENCH").is_some() {
        ~[~"", ~"1000000", ~"10000"]
    } else if args.len() <= 1u {
        ~[~"", ~"10000", ~"4"]
    } else {
        args
    };

    let size = from_str::<uint>(args[1]).unwrap();
    let workers = from_str::<uint>(args[2]).unwrap();

    let ctx = zmq::Context::new();

    run(ctx, size, workers);
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate askama;

use askama::Template;

#[derive(Template)] // this will generate the code...
#[template(path = "hello.html", print = "all")] // using the template in this path, relative
                                 // to the templates dir in the crate root
struct HelloTemplate<'a> { // the name of the struct can be anything
    name: &'a str, // the field name should match the variable name
                   // in your template
}

#[test]
fn main() {
    let hello = HelloTemplate { name: "world" }; // instantiate your struct
    assert_eq!("Hello, world!", hello.render().unwrap()); // then render it.
}
<commit_msg>Disable debug output for example by default<commit_after>#[macro_use]
extern crate askama;

use askama::Template;

#[derive(Template)] // this will generate the code...
#[template(path = "hello.html")] // using the template in this path, relative
                                 // to the templates dir in the crate root
struct HelloTemplate<'a> { // the name of the struct can be anything
    name: &'a str, // the field name should match the variable name
                   // in your template
}

#[test]
fn main() {
    let hello = HelloTemplate { name: "world" }; // instantiate your struct
    assert_eq!("Hello, world!", hello.render().unwrap()); // then render it.
}
<|endoftext|>


<commit_before>// Copyright 2021 Vector 35 Inc.
//
// 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.

use std::env;
use std::ffi::{CStr, CString, OsStr};
use std::mem;
use std::os::raw;
use std::path::PathBuf;

#[repr(C)]
struct DlInfo {
    dli_fname: *const raw::c_char,
    dli_fbase: *mut raw::c_void,
    dli_sname: *const raw::c_char,
    dli_saddr: *mut raw::c_void,
}

#[cfg(not(target_os = "windows"))]
fn binja_path() -> PathBuf {
    use std::os::unix::ffi::OsStrExt;

    if let Ok(p) = env::var("BINJA_DIR") {
        return PathBuf::from(p);
    }

    extern "C" {
        fn dladdr(addr: *mut raw::c_void, info: *mut DlInfo) -> raw::c_int;
    }

    unsafe {
        let mut info: DlInfo = mem::zeroed();

        if dladdr(BNSetBundledPluginDirectory as *mut _, &mut info) == 0 {
            panic!("Failed to find libbinaryninjacore path!");
        }

        if info.dli_fname.is_null() {
            panic!("Failed to find libbinaryninjacore path!");
        }

        let path = CStr::from_ptr(info.dli_fname);
        let path = OsStr::from_bytes(path.to_bytes());
        let mut path = PathBuf::from(path);

        path.pop();
        path
    }
}

#[cfg(target_os = "windows")]
fn binja_path() -> PathBuf {
    PathBuf::from(env::var("PROGRAMFILES").unwrap()).join("Vector35\\BinaryNinja\\")
}

use binaryninjacore_sys::{BNInitPlugins, BNInitRepoPlugins, BNSetBundledPluginDirectory};

pub fn init() {
    unsafe {
        let path = binja_path().join("plugins").into_os_string();
        let path = CString::new(path.into_string().unwrap()).unwrap();

        BNSetBundledPluginDirectory(path.as_ptr());
        BNInitPlugins(true);
        BNInitRepoPlugins();
    }
}

pub fn shutdown() {
    //! Unloads plugins, stops all worker threads, and closes open logs
    //! Important! : Must be called at the end of scripts

    unsafe { binaryninjacore_sys::BNShutdown() };
}

pub fn script_helper(func: fn()) {
    //! Prelued-postlued helper function:
    //! ```
    //! fn main() {
    //!     binaryninja::headless::script_helper(|| {
    //!         binaryninja::open_view("/bin/cat")
    //!             .expect("Couldn't open `/bin/cat`")
    //!             .iter()
    //!             .for_each(|func| println!("  `{}`", func.symbol().full_name()));
    //!     });
    //! }
    //! ```

    init();
    func();
    shutdown();
}
<commit_msg>rust: fix warnings on windows<commit_after>// Copyright 2021 Vector 35 Inc.
//
// 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.

use std::env;
use std::ffi::CString;
use std::path::PathBuf;

#[cfg(not(target_os = "windows"))]
fn binja_path() -> PathBuf {
    use std::ffi::{CStr, OsStr};
    use std::mem;
    use std::os::raw;
    use std::os::unix::ffi::OsStrExt;

    #[repr(C)]
    struct DlInfo {
        dli_fname: *const raw::c_char,
        dli_fbase: *mut raw::c_void,
        dli_sname: *const raw::c_char,
        dli_saddr: *mut raw::c_void,
    }

    if let Ok(p) = env::var("BINJA_DIR") {
        return PathBuf::from(p);
    }

    extern "C" {
        fn dladdr(addr: *mut raw::c_void, info: *mut DlInfo) -> raw::c_int;
    }

    unsafe {
        let mut info: DlInfo = mem::zeroed();

        if dladdr(BNSetBundledPluginDirectory as *mut _, &mut info) == 0 {
            panic!("Failed to find libbinaryninjacore path!");
        }

        if info.dli_fname.is_null() {
            panic!("Failed to find libbinaryninjacore path!");
        }

        let path = CStr::from_ptr(info.dli_fname);
        let path = OsStr::from_bytes(path.to_bytes());
        let mut path = PathBuf::from(path);

        path.pop();
        path
    }
}

#[cfg(target_os = "windows")]
fn binja_path() -> PathBuf {
    PathBuf::from(env::var("PROGRAMFILES").unwrap()).join("Vector35\\BinaryNinja\\")
}

use binaryninjacore_sys::{BNInitPlugins, BNInitRepoPlugins, BNSetBundledPluginDirectory};

pub fn init() {
    unsafe {
        let path = binja_path().join("plugins").into_os_string();
        let path = CString::new(path.into_string().unwrap()).unwrap();

        BNSetBundledPluginDirectory(path.as_ptr());
        BNInitPlugins(true);
        BNInitRepoPlugins();
    }
}

pub fn shutdown() {
    //! Unloads plugins, stops all worker threads, and closes open logs
    //! Important! : Must be called at the end of scripts

    unsafe { binaryninjacore_sys::BNShutdown() };
}

pub fn script_helper(func: fn()) {
    //! Prelued-postlued helper function:
    //! ```
    //! fn main() {
    //!     binaryninja::headless::script_helper(|| {
    //!         binaryninja::open_view("/bin/cat")
    //!             .expect("Couldn't open `/bin/cat`")
    //!             .iter()
    //!             .for_each(|func| println!("  `{}`", func.symbol().full_name()));
    //!     });
    //! }
    //! ```

    init();
    func();
    shutdown();
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use common::*;
use arena::*;
use layout;

pub struct ReaderOptions {
    traversalLimitInWords : u64,
    nestingLimit : uint
}

pub static DEFAULT_READER_OPTIONS : ReaderOptions =
    ReaderOptions { traversalLimitInWords : 8 * 1024 * 1024, nestingLimit : 64 };

pub struct MessageReader<'a> {
    segments : &'a [ &'a [Word]],
    options : ReaderOptions,
    arena : ReaderArena<'a>
}

type SegmentId = u32;

impl <'a> MessageReader<'a> {

    #[inline]
    pub fn get_options<'b>(&'b self) -> &'b ReaderOptions {
        return &self.options;
    }
}

impl <'a> MessageReader<'a> {

    pub fn new<'b>(segments : &'b [&'b [Word]], options : ReaderOptions) -> ~MessageReader<'b> {

        assert!(segments.len() > 0);

        let mut result = ~MessageReader {
            segments : segments,
            arena : ReaderArena {
                segment0 : SegmentReader {
                    arena : Null,
                    ptr : unsafe { segments[0].unsafe_ref(0) },
                    size : segments[0].len()
                },
                more_segments : None
            },
            options : options
        };

        let arena_ptr = ReaderArenaPtr (std::ptr::to_unsafe_ptr(&result.arena));

        result.arena.segment0.arena = arena_ptr;

        if (segments.len() > 1 ) {
            let mut moreSegmentReaders = ~[];
            for segment in segments.slice_from(1).iter() {
                let segmentReader = SegmentReader {
                    arena : arena_ptr,
                    ptr : unsafe { segment.unsafe_ref(0) },
                    size : segment.len()
                };
                moreSegmentReaders.push(segmentReader);
            }
            result.arena.more_segments = Some(moreSegmentReaders);
        }

        result

    }

    pub fn get_root<T : layout::FromStructReader<'a>>(&self) -> T {
        unsafe {
            let segment : *SegmentReader<'a> = std::ptr::to_unsafe_ptr(&self.arena.segment0);

            let pointer_reader = layout::PointerReader::get_root::<'a>(
                segment, (*segment).get_start_ptr(), self.options.nestingLimit as int);

            let result : T = layout::FromStructReader::from_struct_reader(
                pointer_reader.get_struct::<'a>(std::ptr::null()));

            result
        }
    }
}

pub enum AllocationStrategy {
    FIXED_SIZE,
    GROW_HEURISTICALLY
}

pub static SUGGESTED_FIRST_SEGMENT_WORDS : uint = 1024;
pub static SUGGESTED_ALLOCATION_STRATEGY : AllocationStrategy = GROW_HEURISTICALLY;

pub struct MessageBuilder<'a> {
    nextSize : uint,
    allocation_strategy : AllocationStrategy,
    arena : ~BuilderArena<'a>,
    segments : ~[~[Word]]
}

impl <'a>MessageBuilder<'a> {

    // TODO: maybe when Rust issue #5121 is fixed we can safely get away with not passing
    //  a closure here.
    pub fn new<T>(firstSegmentWords : uint,
                  allocationStrategy : AllocationStrategy,
                  cont : |&mut MessageBuilder<'a>| -> T) -> T {

        let mut segments = ~[];
        segments.push(allocate_zeroed_words(firstSegmentWords));
        let mut arena = ~BuilderArena::<'a> {
            message : std::ptr::mut_null(),
            segment0 : SegmentBuilder {
                reader : SegmentReader {
                    ptr : unsafe { segments[0].unsafe_ref(0) },
                    size : segments[0].len(),
                    arena : Null },
                id : 0,
                pos : unsafe { segments[0].unsafe_mut_ref(0) }
            },
            more_segments : None };

        let arena_ptr = std::ptr::to_mut_unsafe_ptr(arena);
        arena.segment0.reader.arena = BuilderArenaPtr(arena_ptr);

        let mut result = ~MessageBuilder {
            nextSize : firstSegmentWords,
            allocation_strategy : allocationStrategy,
            segments : segments,
            arena : arena
        };

        (*result.arena).message = std::ptr::to_mut_unsafe_ptr(result);

        cont(result)
    }

    pub fn new_default<T>(cont : |&mut MessageBuilder| -> T) -> T {
        MessageBuilder::new(SUGGESTED_FIRST_SEGMENT_WORDS, SUGGESTED_ALLOCATION_STRATEGY, cont)
    }

    pub fn allocate_segment(&mut self, minimumSize : WordCount) -> (*mut Word, WordCount) {
        let size = std::cmp::max(minimumSize, self.nextSize);
        let mut new_words = allocate_zeroed_words(size);
        let ptr = unsafe { new_words.unsafe_mut_ref(0) };
        self.segments.push(new_words);

        match self.allocation_strategy {
            GROW_HEURISTICALLY => { self.nextSize += size; }
            _ => { }
        }
        (ptr, size)
    }

    pub fn get_segments_for_output<T>(&self, cont : |&[&[Word]]| -> T) -> T {
        self.arena.get_segments_for_output(cont)
    }
}

impl <'a>MessageBuilder<'a> {
    // Note: This type signature ought to prevent a MessageBuilder
    // from being initted twice simultaneously. It currently does not
    // fulfill that goal, perhaps due to Rust issue #5121.
    pub fn init_root<T : layout::HasStructSize + layout::FromStructBuilder<'a>>(&mut self) -> T {
        // Rolled in this stuff form getRootSegment.
        let rootSegment = std::ptr::to_mut_unsafe_ptr(&mut self.arena.segment0);

        match self.arena.segment0.allocate(WORDS_PER_POINTER) {
            None => {fail!("could not allocate root pointer") }
            Some(location) => {
                //assert!(location == 0,
                //        "First allocated word of new segment was not at offset 0");

                let pb = layout::PointerBuilder::get_root(rootSegment, location);

                return layout::FromStructBuilder::from_struct_builder(
                    pb.init_struct(layout::HasStructSize::struct_size(None::<T>)));
            }
        }
    }

}
<commit_msg>use calloc in MessageBuilder<commit_after>/*
 * Copyright (c) 2013, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use std;
use common::*;
use arena::*;
use layout;

pub struct ReaderOptions {
    traversalLimitInWords : u64,
    nestingLimit : uint
}

pub static DEFAULT_READER_OPTIONS : ReaderOptions =
    ReaderOptions { traversalLimitInWords : 8 * 1024 * 1024, nestingLimit : 64 };

pub struct MessageReader<'a> {
    segments : &'a [ &'a [Word]],
    options : ReaderOptions,
    arena : ReaderArena<'a>
}

type SegmentId = u32;

impl <'a> MessageReader<'a> {

    #[inline]
    pub fn get_options<'b>(&'b self) -> &'b ReaderOptions {
        return &self.options;
    }
}

impl <'a> MessageReader<'a> {

    pub fn new<'b>(segments : &'b [&'b [Word]], options : ReaderOptions) -> ~MessageReader<'b> {

        assert!(segments.len() > 0);

        let mut result = ~MessageReader {
            segments : segments,
            arena : ReaderArena {
                segment0 : SegmentReader {
                    arena : Null,
                    ptr : unsafe { segments[0].unsafe_ref(0) },
                    size : segments[0].len()
                },
                more_segments : None
            },
            options : options
        };

        let arena_ptr = ReaderArenaPtr (std::ptr::to_unsafe_ptr(&result.arena));

        result.arena.segment0.arena = arena_ptr;

        if (segments.len() > 1 ) {
            let mut moreSegmentReaders = ~[];
            for segment in segments.slice_from(1).iter() {
                let segmentReader = SegmentReader {
                    arena : arena_ptr,
                    ptr : unsafe { segment.unsafe_ref(0) },
                    size : segment.len()
                };
                moreSegmentReaders.push(segmentReader);
            }
            result.arena.more_segments = Some(moreSegmentReaders);
        }

        result

    }

    pub fn get_root<T : layout::FromStructReader<'a>>(&self) -> T {
        unsafe {
            let segment : *SegmentReader<'a> = std::ptr::to_unsafe_ptr(&self.arena.segment0);

            let pointer_reader = layout::PointerReader::get_root::<'a>(
                segment, (*segment).get_start_ptr(), self.options.nestingLimit as int);

            let result : T = layout::FromStructReader::from_struct_reader(
                pointer_reader.get_struct::<'a>(std::ptr::null()));

            result
        }
    }
}

pub enum AllocationStrategy {
    FIXED_SIZE,
    GROW_HEURISTICALLY
}

pub static SUGGESTED_FIRST_SEGMENT_WORDS : uint = 1024;
pub static SUGGESTED_ALLOCATION_STRATEGY : AllocationStrategy = GROW_HEURISTICALLY;

pub struct MessageBuilder<'a> {
    nextSize : uint,
    allocation_strategy : AllocationStrategy,
    arena : ~BuilderArena<'a>,
    segments : ~[*mut Word]
}

impl <'a> Drop for MessageBuilder<'a> {
    fn drop(&mut self) {
        for &segment_ptr in self.segments.iter() {
            unsafe { std::libc::free(std::cast::transmute(segment_ptr)); }
        }
    }
}

impl <'a>MessageBuilder<'a> {

    // TODO: maybe when Rust issue #5121 is fixed we can safely get away with not passing
    //  a closure here.
    pub fn new<T>(firstSegmentWords : uint,
                  allocationStrategy : AllocationStrategy,
                  cont : |&mut MessageBuilder<'a>| -> T) -> T {

        let mut segments : ~[*mut Word] = ~[];
        segments.push(unsafe { std::cast::transmute(std::libc::calloc(firstSegmentWords as std::libc::size_t,
                                                                      BYTES_PER_WORD as std::libc::size_t))});

        let mut arena = ~BuilderArena::<'a> {
            message : std::ptr::mut_null(),
            segment0 : SegmentBuilder {
                reader : SegmentReader {
                    ptr : segments[0] as * Word,
                    size : firstSegmentWords,
                    arena : Null },
                id : 0,
                pos : segments[0]
            },
            more_segments : None };

        let arena_ptr = std::ptr::to_mut_unsafe_ptr(arena);
        arena.segment0.reader.arena = BuilderArenaPtr(arena_ptr);

        let mut result = ~MessageBuilder {
            nextSize : firstSegmentWords,
            allocation_strategy : allocationStrategy,
            segments : segments,
            arena : arena
        };

        (*result.arena).message = std::ptr::to_mut_unsafe_ptr(result);

        cont(result)
    }

    pub fn new_default<T>(cont : |&mut MessageBuilder| -> T) -> T {
        MessageBuilder::new(SUGGESTED_FIRST_SEGMENT_WORDS, SUGGESTED_ALLOCATION_STRATEGY, cont)
    }

    pub fn allocate_segment(&mut self, minimumSize : WordCount) -> (*mut Word, WordCount) {
        let size = std::cmp::max(minimumSize, self.nextSize);
        let new_words : *mut Word = unsafe {
            std::cast::transmute(std::libc::calloc(size as std::libc::size_t,
                                                   BYTES_PER_WORD as std::libc::size_t)) };
        self.segments.push(new_words);

        match self.allocation_strategy {
            GROW_HEURISTICALLY => { self.nextSize += size; }
            _ => { }
        }
        (new_words, size)
    }

    pub fn get_segments_for_output<T>(&self, cont : |&[&[Word]]| -> T) -> T {
        self.arena.get_segments_for_output(cont)
    }
}

impl <'a>MessageBuilder<'a> {
    // Note: This type signature ought to prevent a MessageBuilder
    // from being initted twice simultaneously. It currently does not
    // fulfill that goal, perhaps due to Rust issue #5121.
    pub fn init_root<T : layout::HasStructSize + layout::FromStructBuilder<'a>>(&mut self) -> T {
        // Rolled in this stuff form getRootSegment.
        let rootSegment = std::ptr::to_mut_unsafe_ptr(&mut self.arena.segment0);

        match self.arena.segment0.allocate(WORDS_PER_POINTER) {
            None => {fail!("could not allocate root pointer") }
            Some(location) => {
                //assert!(location == 0,
                //        "First allocated word of new segment was not at offset 0");

                let pb = layout::PointerBuilder::get_root(rootSegment, location);

                return layout::FromStructBuilder::from_struct_builder(
                    pb.init_struct(layout::HasStructSize::struct_size(None::<T>)));
            }
        }
    }

}
<|endoftext|>


<commit_before>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

//! # Cap'n Proto Runtime Library
//!
//! This crate contains basic facilities for reading and writing
//! [Cap'n Proto](https://capnproto.org) messages in Rust. It is intended to
//! be used in conjunction with code generated by the
//! [capnpc-rust](https://github.com/dwrensha/capnpc-rust) crate.

extern crate byteorder;

#[cfg(any(feature="quickcheck", test))]
extern crate quickcheck;

#[cfg(feature = "rpc")]
extern crate futures;

/// Constructs a [`Word`](struct.Word.html) from its constituent bytes, accounting
/// for endianness. This macro can be used to construct constants. In the future, once
/// Rust supports [constant functions](https://github.com/rust-lang/rust/issues/24111),
/// this macro will be replaced by such a function.
#[cfg(target_endian = "little")]
#[macro_export]
macro_rules! capnp_word {
  ($b0:expr, $b1:expr, $b2:expr, $b3:expr,
   $b4:expr, $b5:expr, $b6:expr, $b7:expr) => (
    $crate::Word {
        raw_content: (($b0 as u64) << 0) + (($b1 as u64) << 8) +
                     (($b2 as u64) << 16) + (($b3 as u64) << 24) +
                     (($b4 as u64) << 32) + (($b5 as u64) << 40) +
                     (($b6 as u64) << 48) + (($b7 as u64) << 56)
    }
  )
}

#[cfg(target_endian = "big")]
#[macro_export]
macro_rules! capnp_word {
  ($b0:expr, $b1:expr, $b2:expr, $b3:expr,
   $b4:expr, $b5:expr, $b6:expr, $b7:expr) => (
     $crate::Word {
         raw_content: (($b7 as u64) << 0) + (($b6 as u64) << 8) +
                      (($b5 as u64) << 16) + (($b4 as u64) << 24) +
                      (($b3 as u64) << 32) + (($b2 as u64) << 40) +
                      (($b1 as u64) << 48) + (($b0 as u64) << 56)
     }
  )
}

pub mod any_pointer;
pub mod capability;
pub mod capability_list;
pub mod constant;
pub mod data;
pub mod data_list;
pub mod enum_list;
pub mod list_list;
pub mod message;
pub mod primitive_list;
pub mod private;
pub mod serialize;
pub mod serialize_packed;
pub mod struct_list;
pub mod text;
pub mod text_list;
pub mod traits;

mod util;

/// Eight bytes of memory with opaque interior. Use [`capnp_word!()`](macro.capnp_word!.html)
/// to construct one of these.
///
/// This type is used to ensure that the data of a message is properly aligned.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(C)]
pub struct Word {
    #[doc(hidden)]
    pub raw_content: u64,
}

impl Word {
    /// Does this, but faster:
    /// `::std::iter::repeat(Word(0)).take(length).collect()`
    pub fn allocate_zeroed_vec(length: usize) -> Vec<Word> {
        let mut result : Vec<Word> = Vec::with_capacity(length);
        unsafe {
            result.set_len(length);
            let p : *mut u8 = result.as_mut_ptr() as *mut u8;
            ::std::ptr::write_bytes(p, 0u8, length * ::std::mem::size_of::<Word>());
        }
        result
    }

    /// Converts a byte slice into a `Word` slice. Unsafe due to possible alignment issues.
    /// Only call this if you know that either
    ///    1. `bytes.to_ptr()` falls on an eight-byte boundary, or
    ///    2. your processor is okay with unaligned reads.
    pub unsafe fn bytes_to_words<'a>(bytes: &'a [u8]) -> &'a [Word] {
        ::std::slice::from_raw_parts(bytes.as_ptr() as *const Word, bytes.len() / 8)
    }

    /// Converts a mutable byte slice into a mutable `Word` slice. Unsafe due to possible
    /// alignment issues. Only call this if you know that either
    ///    1. `bytes.to_ptr()` falls on an eight-byte boundary, or
    ///    2. your processor is okay with unaligned reads and writes
    pub unsafe fn bytes_to_words_mut<'a>(bytes: &'a mut [u8]) -> &'a mut [Word] {
        ::std::slice::from_raw_parts_mut(bytes.as_ptr() as *mut Word, bytes.len() / 8)
    }

    pub fn words_to_bytes<'a>(words: &'a [Word]) -> &'a [u8] {
        unsafe {
            ::std::slice::from_raw_parts(words.as_ptr() as *const u8, words.len() * 8)
        }
    }

    pub fn words_to_bytes_mut<'a>(words: &'a mut [Word]) -> &'a mut [u8] {
        unsafe {
            ::std::slice::from_raw_parts_mut(words.as_mut_ptr() as *mut u8, words.len() * 8)
        }
    }
}

#[cfg(any(feature="quickcheck", test))]
impl quickcheck::Arbitrary for Word {
    fn arbitrary<G: quickcheck::Gen>(g: &mut G) -> Word {
        Word { raw_content: quickcheck::Arbitrary::arbitrary(g) }
    }
}

/// Size of a message. Every generated struct has a method `.total_size()` that returns this.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct MessageSize {
    pub word_count: u64,

    /// Size of the capability table.
    pub cap_count: u32
}

impl MessageSize {
    pub fn plus_eq(&mut self, other : MessageSize) {
        self.word_count += other.word_count;
        self.cap_count += other.cap_count;
    }
}

/// An enum value or union discriminant that was not found among those defined in a schema.
#[derive(PartialEq, Clone, Copy, Debug)]
pub struct NotInSchema(pub u16);

impl ::std::fmt::Display for NotInSchema {
    fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
        write!(fmt, "Enum value or union discriminant {} was not present in the schema.", self.0)
    }
}

impl ::std::error::Error for NotInSchema {
    fn description<'a>(&'a self) -> &'a str {
        "Enum value or union disriminant was not present in schema."
    }
}

/// Because messages are lazily validated, the return type of any method that reads a pointer field
/// must be wrapped in a Result.
pub type Result<T> = ::std::result::Result<T, Error>;

/// Describes an arbitrary error that prevented an operation from completing.
#[derive(Debug, Clone)]
pub struct Error {
    /// The general kind of the error. Code that decides how to respond to an error
    /// should read only this field in making its decision.
    pub kind: ErrorKind,

    /// Human-readable failure description.
    pub description: String,
}

/// The general nature of an error. The purpose of this enum is not to describe the error itself,
/// but rather to describe how the client might want to respond to the error.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ErrorKind {
    /// Something went wrong. This is the usual error kind. It includes decoding errors.
    Failed,

    /// The call failed because of a temporary lack of resources. This could be space resources
    /// (out of memory, out of disk space) or time resources (request queue overflow, operation
    /// timed out).
    //
    /// The operation might work if tried again, but it should NOT be repeated immediately as this
    /// may simply exacerbate the problem.
    Overloaded,

    /// The call required communication over a connection that has been lost. The callee will need
    /// to re-establish connections and try again.
    Disconnected,

    /// The requested method is not implemented. The caller may wish to revert to a fallback
    /// approach based on other methods.
    Unimplemented,
}

impl Error {
    pub fn failed(description: String) -> Error {
        Error { description: description, kind: ErrorKind::Failed }
    }
    pub fn overloaded(description: String) -> Error {
        Error { description: description, kind: ErrorKind::Overloaded }
    }
    pub fn disconnected(description: String) -> Error {
        Error { description: description, kind: ErrorKind::Disconnected }
    }
    pub fn unimplemented(description: String) -> Error {
        Error { description: description, kind: ErrorKind::Unimplemented }
    }
}

impl ::std::convert::From<::std::io::Error> for Error {
    fn from(err: ::std::io::Error) -> Error {
        use std::io;
        let kind = match err.kind() {
            io::ErrorKind::TimedOut => ErrorKind::Overloaded,
            io::ErrorKind::BrokenPipe |
            io::ErrorKind::ConnectionRefused |
            io::ErrorKind::ConnectionReset |
            io::ErrorKind::ConnectionAborted |
            io::ErrorKind::NotConnected  => ErrorKind::Disconnected,
            _ => ErrorKind::Failed,
        };
        Error { description: format!("{}", err), kind: kind }
    }
}

impl ::std::convert::From<::std::string::FromUtf8Error> for Error {
    fn from(err: ::std::string::FromUtf8Error) -> Error {
        Error::failed(format!("{}", err))
    }
}

impl ::std::convert::From<::std::str::Utf8Error> for Error {
    fn from(err: ::std::str::Utf8Error) -> Error {
        Error::failed(format!("{}", err))
    }
}

#[cfg(feature = "rpc")]
impl ::std::convert::From<futures::sync::oneshot::Canceled> for Error {
    fn from(_e: futures::sync::oneshot::Canceled) -> Error {
        Error::failed(format!("oneshot was canceled"))
    }
}

impl ::std::convert::From<NotInSchema> for Error {
    fn from(e: NotInSchema) -> Error {
        Error::failed(format!("Enum value or union discriminant {} was not present in schema.", e.0))
    }
}

impl ::std::fmt::Display for Error {
    fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
        write!(fmt, "{:?}: {}", self.kind, self.description)
    }
}

impl ::std::error::Error for Error {
    fn description(&self) -> &str {
        &self.description
    }
    fn cause(&self) -> Option<&::std::error::Error> {
        None
    }
}

/// Helper struct that allows `MessageBuilder::get_segments_for_output()` to avoid heap allocations
/// in the single-segment case.
pub enum OutputSegments<'a> {
    SingleSegment([&'a [Word]; 1]),
    MultiSegment(Vec<&'a [Word]>),
}

impl <'a> ::std::ops::Deref for OutputSegments<'a> {
    type Target = [&'a [Word]];
    fn deref<'b>(&'b self) -> &'b [&'a [Word]] {
        match *self {
            OutputSegments::SingleSegment(ref s) => {
                s
            }
            OutputSegments::MultiSegment(ref v) => {
                &*v
            }
        }
    }
}

impl<'s> message::ReaderSegments for OutputSegments<'s> {
    fn get_segment<'a>(&'a self, id: u32) -> Option<&'a [Word]> {
        match *self {
            OutputSegments::SingleSegment(ref s) => {
                s.get(id as usize).map(|slice| *slice)
            }
            OutputSegments::MultiSegment(ref v) => {
                v.get(id as usize).map(|slice| *slice)
            }
        }
    }
}
<commit_msg>remove unnecessary ref & deref operations<commit_after>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

//! # Cap'n Proto Runtime Library
//!
//! This crate contains basic facilities for reading and writing
//! [Cap'n Proto](https://capnproto.org) messages in Rust. It is intended to
//! be used in conjunction with code generated by the
//! [capnpc-rust](https://github.com/dwrensha/capnpc-rust) crate.

extern crate byteorder;

#[cfg(any(feature="quickcheck", test))]
extern crate quickcheck;

#[cfg(feature = "rpc")]
extern crate futures;

/// Constructs a [`Word`](struct.Word.html) from its constituent bytes, accounting
/// for endianness. This macro can be used to construct constants. In the future, once
/// Rust supports [constant functions](https://github.com/rust-lang/rust/issues/24111),
/// this macro will be replaced by such a function.
#[cfg(target_endian = "little")]
#[macro_export]
macro_rules! capnp_word {
  ($b0:expr, $b1:expr, $b2:expr, $b3:expr,
   $b4:expr, $b5:expr, $b6:expr, $b7:expr) => (
    $crate::Word {
        raw_content: (($b0 as u64) << 0) + (($b1 as u64) << 8) +
                     (($b2 as u64) << 16) + (($b3 as u64) << 24) +
                     (($b4 as u64) << 32) + (($b5 as u64) << 40) +
                     (($b6 as u64) << 48) + (($b7 as u64) << 56)
    }
  )
}

#[cfg(target_endian = "big")]
#[macro_export]
macro_rules! capnp_word {
  ($b0:expr, $b1:expr, $b2:expr, $b3:expr,
   $b4:expr, $b5:expr, $b6:expr, $b7:expr) => (
     $crate::Word {
         raw_content: (($b7 as u64) << 0) + (($b6 as u64) << 8) +
                      (($b5 as u64) << 16) + (($b4 as u64) << 24) +
                      (($b3 as u64) << 32) + (($b2 as u64) << 40) +
                      (($b1 as u64) << 48) + (($b0 as u64) << 56)
     }
  )
}

pub mod any_pointer;
pub mod capability;
pub mod capability_list;
pub mod constant;
pub mod data;
pub mod data_list;
pub mod enum_list;
pub mod list_list;
pub mod message;
pub mod primitive_list;
pub mod private;
pub mod serialize;
pub mod serialize_packed;
pub mod struct_list;
pub mod text;
pub mod text_list;
pub mod traits;

mod util;

/// Eight bytes of memory with opaque interior. Use [`capnp_word!()`](macro.capnp_word!.html)
/// to construct one of these.
///
/// This type is used to ensure that the data of a message is properly aligned.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(C)]
pub struct Word {
    #[doc(hidden)]
    pub raw_content: u64,
}

impl Word {
    /// Does this, but faster:
    /// `::std::iter::repeat(Word(0)).take(length).collect()`
    pub fn allocate_zeroed_vec(length: usize) -> Vec<Word> {
        let mut result : Vec<Word> = Vec::with_capacity(length);
        unsafe {
            result.set_len(length);
            let p : *mut u8 = result.as_mut_ptr() as *mut u8;
            ::std::ptr::write_bytes(p, 0u8, length * ::std::mem::size_of::<Word>());
        }
        result
    }

    /// Converts a byte slice into a `Word` slice. Unsafe due to possible alignment issues.
    /// Only call this if you know that either
    ///    1. `bytes.to_ptr()` falls on an eight-byte boundary, or
    ///    2. your processor is okay with unaligned reads.
    pub unsafe fn bytes_to_words<'a>(bytes: &'a [u8]) -> &'a [Word] {
        ::std::slice::from_raw_parts(bytes.as_ptr() as *const Word, bytes.len() / 8)
    }

    /// Converts a mutable byte slice into a mutable `Word` slice. Unsafe due to possible
    /// alignment issues. Only call this if you know that either
    ///    1. `bytes.to_ptr()` falls on an eight-byte boundary, or
    ///    2. your processor is okay with unaligned reads and writes
    pub unsafe fn bytes_to_words_mut<'a>(bytes: &'a mut [u8]) -> &'a mut [Word] {
        ::std::slice::from_raw_parts_mut(bytes.as_ptr() as *mut Word, bytes.len() / 8)
    }

    pub fn words_to_bytes<'a>(words: &'a [Word]) -> &'a [u8] {
        unsafe {
            ::std::slice::from_raw_parts(words.as_ptr() as *const u8, words.len() * 8)
        }
    }

    pub fn words_to_bytes_mut<'a>(words: &'a mut [Word]) -> &'a mut [u8] {
        unsafe {
            ::std::slice::from_raw_parts_mut(words.as_mut_ptr() as *mut u8, words.len() * 8)
        }
    }
}

#[cfg(any(feature="quickcheck", test))]
impl quickcheck::Arbitrary for Word {
    fn arbitrary<G: quickcheck::Gen>(g: &mut G) -> Word {
        Word { raw_content: quickcheck::Arbitrary::arbitrary(g) }
    }
}

/// Size of a message. Every generated struct has a method `.total_size()` that returns this.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct MessageSize {
    pub word_count: u64,

    /// Size of the capability table.
    pub cap_count: u32
}

impl MessageSize {
    pub fn plus_eq(&mut self, other : MessageSize) {
        self.word_count += other.word_count;
        self.cap_count += other.cap_count;
    }
}

/// An enum value or union discriminant that was not found among those defined in a schema.
#[derive(PartialEq, Clone, Copy, Debug)]
pub struct NotInSchema(pub u16);

impl ::std::fmt::Display for NotInSchema {
    fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
        write!(fmt, "Enum value or union discriminant {} was not present in the schema.", self.0)
    }
}

impl ::std::error::Error for NotInSchema {
    fn description<'a>(&'a self) -> &'a str {
        "Enum value or union disriminant was not present in schema."
    }
}

/// Because messages are lazily validated, the return type of any method that reads a pointer field
/// must be wrapped in a Result.
pub type Result<T> = ::std::result::Result<T, Error>;

/// Describes an arbitrary error that prevented an operation from completing.
#[derive(Debug, Clone)]
pub struct Error {
    /// The general kind of the error. Code that decides how to respond to an error
    /// should read only this field in making its decision.
    pub kind: ErrorKind,

    /// Human-readable failure description.
    pub description: String,
}

/// The general nature of an error. The purpose of this enum is not to describe the error itself,
/// but rather to describe how the client might want to respond to the error.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ErrorKind {
    /// Something went wrong. This is the usual error kind. It includes decoding errors.
    Failed,

    /// The call failed because of a temporary lack of resources. This could be space resources
    /// (out of memory, out of disk space) or time resources (request queue overflow, operation
    /// timed out).
    //
    /// The operation might work if tried again, but it should NOT be repeated immediately as this
    /// may simply exacerbate the problem.
    Overloaded,

    /// The call required communication over a connection that has been lost. The callee will need
    /// to re-establish connections and try again.
    Disconnected,

    /// The requested method is not implemented. The caller may wish to revert to a fallback
    /// approach based on other methods.
    Unimplemented,
}

impl Error {
    pub fn failed(description: String) -> Error {
        Error { description: description, kind: ErrorKind::Failed }
    }
    pub fn overloaded(description: String) -> Error {
        Error { description: description, kind: ErrorKind::Overloaded }
    }
    pub fn disconnected(description: String) -> Error {
        Error { description: description, kind: ErrorKind::Disconnected }
    }
    pub fn unimplemented(description: String) -> Error {
        Error { description: description, kind: ErrorKind::Unimplemented }
    }
}

impl ::std::convert::From<::std::io::Error> for Error {
    fn from(err: ::std::io::Error) -> Error {
        use std::io;
        let kind = match err.kind() {
            io::ErrorKind::TimedOut => ErrorKind::Overloaded,
            io::ErrorKind::BrokenPipe |
            io::ErrorKind::ConnectionRefused |
            io::ErrorKind::ConnectionReset |
            io::ErrorKind::ConnectionAborted |
            io::ErrorKind::NotConnected  => ErrorKind::Disconnected,
            _ => ErrorKind::Failed,
        };
        Error { description: format!("{}", err), kind: kind }
    }
}

impl ::std::convert::From<::std::string::FromUtf8Error> for Error {
    fn from(err: ::std::string::FromUtf8Error) -> Error {
        Error::failed(format!("{}", err))
    }
}

impl ::std::convert::From<::std::str::Utf8Error> for Error {
    fn from(err: ::std::str::Utf8Error) -> Error {
        Error::failed(format!("{}", err))
    }
}

#[cfg(feature = "rpc")]
impl ::std::convert::From<futures::sync::oneshot::Canceled> for Error {
    fn from(_e: futures::sync::oneshot::Canceled) -> Error {
        Error::failed(format!("oneshot was canceled"))
    }
}

impl ::std::convert::From<NotInSchema> for Error {
    fn from(e: NotInSchema) -> Error {
        Error::failed(format!("Enum value or union discriminant {} was not present in schema.", e.0))
    }
}

impl ::std::fmt::Display for Error {
    fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
        write!(fmt, "{:?}: {}", self.kind, self.description)
    }
}

impl ::std::error::Error for Error {
    fn description(&self) -> &str {
        &self.description
    }
    fn cause(&self) -> Option<&::std::error::Error> {
        None
    }
}

/// Helper struct that allows `MessageBuilder::get_segments_for_output()` to avoid heap allocations
/// in the single-segment case.
pub enum OutputSegments<'a> {
    SingleSegment([&'a [Word]; 1]),
    MultiSegment(Vec<&'a [Word]>),
}

impl <'a> ::std::ops::Deref for OutputSegments<'a> {
    type Target = [&'a [Word]];
    fn deref<'b>(&'b self) -> &'b [&'a [Word]] {
        match *self {
            OutputSegments::SingleSegment(ref s) => {
                s
            }
            OutputSegments::MultiSegment(ref v) => {
                v
            }
        }
    }
}

impl<'s> message::ReaderSegments for OutputSegments<'s> {
    fn get_segment<'a>(&'a self, id: u32) -> Option<&'a [Word]> {
        match *self {
            OutputSegments::SingleSegment(ref s) => {
                s.get(id as usize).map(|slice| *slice)
            }
            OutputSegments::MultiSegment(ref v) => {
                v.get(id as usize).map(|slice| *slice)
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Handles helper utilities and global methods.

pub use self::alignment::*;
pub use self::angle::*;
pub use self::border::*;
pub use self::brush::*;
pub use self::color::*;
pub use self::constraint::*;
pub use self::dirty_size::*;
pub use self::expression::*;
pub use self::f32_cmp::*;
pub use self::f64_cmp::*;
pub use self::filter::*;
pub use self::gradients::*;
pub use self::number::*;
pub use self::orientation::*;
pub use self::point::*;
pub use self::rectangle::*;
pub use self::relative_direction::*;
pub use self::selection_mode::*;
pub use self::size::*;
pub use self::string16::*;
pub use self::text_alignment::*;
pub use self::text_baseline::*;
pub use self::thickness::*;
pub use self::value::*;
pub use self::visibility::*;

mod alignment;
mod angle;
mod border;
mod brush;
mod color;
mod constraint;
mod dirty_size;
mod expression;
mod f32_cmp;
mod f64_cmp;
mod filter;
mod gradients;
mod number;
mod orientation;
mod point;
pub mod prelude;
mod rectangle;
mod relative_direction;
mod selection_mode;
mod size;
mod spacer;
mod string16;
mod text_alignment;
mod text_baseline;
mod thickness;
mod value;
mod visibility;
<commit_msg>utils: update documentation strings<commit_after>//! Handles helper utilities and global methods.

pub use self::alignment::*;
pub use self::angle::*;
pub use self::border::*;
pub use self::brush::*;
pub use self::color::*;
pub use self::constraint::*;
pub use self::dirty_size::*;
pub use self::expression::*;
pub use self::f32_cmp::*;
pub use self::f64_cmp::*;
pub use self::filter::*;
pub use self::gradients::*;
pub use self::number::*;
pub use self::orientation::*;
pub use self::point::*;
pub use self::rectangle::*;
pub use self::relative_direction::*;
pub use self::selection_mode::*;
pub use self::size::*;
pub use self::string16::*;
pub use self::text_alignment::*;
pub use self::text_baseline::*;
pub use self::thickness::*;
pub use self::value::*;
pub use self::visibility::*;

mod alignment;
mod angle;
mod border;
mod brush;
mod color;
mod constraint;
mod dirty_size;
mod expression;
mod f32_cmp;
mod f64_cmp;
mod filter;
mod gradients;
mod number;
mod orientation;
mod point;

/// This module pre-selects commonly used OrbTk crates and put them into scope.
pub mod prelude;

mod rectangle;
mod relative_direction;
mod selection_mode;
mod size;
mod spacer;
mod string16;
mod text_alignment;
mod text_baseline;
mod thickness;
mod value;
mod visibility;
<|endoftext|>


<commit_before>use crate::term2;
use rustup::dist::Notification as In;
use rustup::utils::tty;
use rustup::utils::Notification as Un;
use rustup::Notification;
use std::collections::VecDeque;
use std::fmt;
use std::io::Write;
use term::Terminal;
use time::precise_time_s;

/// Keep track of this many past download amounts
const DOWNLOAD_TRACK_COUNT: usize = 5;

/// Tracks download progress and displays information about it to a terminal.
pub struct DownloadTracker {
    /// Content-Length of the to-be downloaded object.
    content_len: Option<usize>,
    /// Total data downloaded in bytes.
    total_downloaded: usize,
    /// Data downloaded this second.
    downloaded_this_sec: usize,
    /// Keeps track of amount of data downloaded every last few secs.
    /// Used for averaging the download speed. NB: This does not necessarily
    /// represent adjacent seconds; thus it may not show the average at all.
    downloaded_last_few_secs: VecDeque<usize>,
    /// Time stamp of the last second
    last_sec: Option<f64>,
    /// Time stamp of the start of the download
    start_sec: f64,
    /// The terminal we write the information to.
    /// XXX: Could be a term trait, but with #1818 on the horizon that
    ///      is a pointless change to make - better to let that transition
    ///      happen and take stock after that.
    term: term2::StdoutTerminal,
    /// Whether we displayed progress for the download or not.
    ///
    /// If the download is quick enough, we don't have time to
    /// display the progress info.
    /// In that case, we do not want to do some cleanup stuff we normally do.
    ///
    /// If we have displayed progress, this is the number of characters we
    /// rendered, so we can erase it cleanly.
    displayed_charcount: Option<usize>,
    /// What units to show progress in
    units: Vec<String>,
}

impl DownloadTracker {
    /// Creates a new DownloadTracker.
    pub fn new() -> Self {
        DownloadTracker {
            content_len: None,
            total_downloaded: 0,
            downloaded_this_sec: 0,
            downloaded_last_few_secs: VecDeque::with_capacity(DOWNLOAD_TRACK_COUNT),
            start_sec: precise_time_s(),
            last_sec: None,
            term: term2::stdout(),
            displayed_charcount: None,
            units: vec!["B".into(); 1],
        }
    }

    pub fn handle_notification(&mut self, n: &Notification<'_>) -> bool {
        match *n {
            Notification::Install(In::Utils(Un::DownloadContentLengthReceived(content_len))) => {
                self.content_length_received(content_len);

                true
            }
            Notification::Install(In::Utils(Un::DownloadDataReceived(data))) => {
                if tty::stdout_isatty() {
                    self.data_received(data.len());
                }
                true
            }
            Notification::Install(In::Utils(Un::DownloadFinished)) => {
                self.download_finished();
                true
            }
            Notification::Install(In::Utils(Un::DownloadPushUnits(units))) => {
                self.push_units(units.into());
                true
            }
            Notification::Install(In::Utils(Un::DownloadPopUnits)) => {
                self.pop_units();
                true
            }

            _ => false,
        }
    }

    /// Notifies self that Content-Length information has been received.
    pub fn content_length_received(&mut self, content_len: u64) {
        self.content_len = Some(content_len as usize);
    }

    /// Notifies self that data of size `len` has been received.
    pub fn data_received(&mut self, len: usize) {
        self.total_downloaded += len;
        self.downloaded_this_sec += len;

        let current_time = precise_time_s();

        match self.last_sec {
            None => self.last_sec = Some(current_time),
            Some(prev) => {
                let elapsed = current_time - prev;
                if elapsed >= 1.0 {
                    self.display();
                    self.last_sec = Some(current_time);
                    if self.downloaded_last_few_secs.len() == DOWNLOAD_TRACK_COUNT {
                        self.downloaded_last_few_secs.pop_back();
                    }
                    self.downloaded_last_few_secs
                        .push_front(self.downloaded_this_sec);
                    self.downloaded_this_sec = 0;
                }
            }
        }
    }
    /// Notifies self that the download has finished.
    pub fn download_finished(&mut self) {
        if self.displayed_charcount.is_some() {
            // Display the finished state
            self.display();
            let _ = writeln!(self.term);
        }
        self.prepare_for_new_download();
    }
    // we're doing modular arithmetic, treat as integer
    pub fn from_seconds(sec: u32) -> (u32, u32, u32, u32) {
        let d = sec / (24 * 3600);
        let h = sec % (24 * 3600) / 3600;
        let min = sec % 3600 / 60;
        let sec = sec % 60;

        (d, h, min, sec)
    }
    /// Resets the state to be ready for a new download.
    fn prepare_for_new_download(&mut self) {
        self.content_len = None;
        self.total_downloaded = 0;
        self.downloaded_this_sec = 0;
        self.downloaded_last_few_secs.clear();
        self.start_sec = precise_time_s();
        self.last_sec = None;
        self.displayed_charcount = None;
    }
    /// Display the tracked download information to the terminal.
    fn display(&mut self) {
        // Panic if someone pops the default bytes unit...
        let units = &self.units.last().unwrap();
        let total_h = Size(self.total_downloaded, units);
        let sum = self.downloaded_last_few_secs.iter().fold(0, |a, &v| a + v);
        let len = self.downloaded_last_few_secs.len();
        let speed = if len > 0 { sum / len } else { 0 };
        let speed_h = Size(speed, units);
        let elapsed_h = Duration(precise_time_s() - self.start_sec);

        // First, move to the start of the current line and clear it.
        let _ = self.term.carriage_return();
        // We'd prefer to use delete_line() but on Windows it seems to
        // sometimes do unusual things
        // let _ = self.term.as_mut().unwrap().delete_line();
        // So instead we do:
        if let Some(n) = self.displayed_charcount {
            // This is not ideal as very narrow terminals might mess up,
            // but it is more likely to succeed until term's windows console
            // fixes whatever's up with delete_line().
            let _ = write!(self.term, "{}", " ".repeat(n));
            let _ = self.term.flush();
            let _ = self.term.carriage_return();
        }

        let output = match self.content_len {
            Some(content_len) => {
                let content_len_h = Size(content_len, units);
                let content_len = content_len as f64;
                let percent = (self.total_downloaded as f64 / content_len) * 100.;
                let remaining = content_len - self.total_downloaded as f64;
                let eta_h = Duration(remaining / speed as f64);
                format!(
                    "{} / {} ({:3.0} %) {}/s in {} ETA: {}",
                    total_h, content_len_h, percent, speed_h, elapsed_h, eta_h
                )
            }
            None => format!(
                "Total: {} Speed: {}/s Elapsed: {}",
                total_h, speed_h, elapsed_h
            ),
        };

        let _ = write!(self.term, "{}", output);
        // Since stdout is typically line-buffered and we don't print a newline, we manually flush.
        let _ = self.term.flush();
        self.displayed_charcount = Some(output.chars().count());
    }

    pub fn push_units(&mut self, new_units: String) {
        self.units.push(new_units);
    }

    pub fn pop_units(&mut self) {
        self.units.pop();
    }
}

/// Human readable representation of duration(seconds).
struct Duration(f64);

impl fmt::Display for Duration {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // repurposing the alternate mode for ETA
        let sec = self.0;

        if sec.is_infinite() {
            write!(f, "Unknown")
        } else {
            match DownloadTracker::from_seconds(sec as u32) {
                (d, h, m, s) if d > 0 => write!(f, "{:3.0}d {:2.0}h {:2.0}m {:2.0}s", d, h, m, s),
                (0, h, m, s) if h > 0 => write!(f, "{:2.0}h {:2.0}m {:2.0}s", h, m, s),
                (0, 0, m, s) if m > 0 => write!(f, "{:2.0}m {:2.0}s", m, s),
                (_, _, _, s) => write!(f, "{:2.0}s", s),
            }
        }
    }
}

/// Human readable size (some units)
struct Size<'a>(usize, &'a str);

impl<'a> fmt::Display for Size<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        const KI: f64 = 1024.0;
        const MI: f64 = KI * KI;
        let size = self.0 as f64;

        if size >= MI {
            write!(f, "{:5.1} Mi{}", size / MI, self.1) // XYZ.P Mi
        } else if size >= KI {
            write!(f, "{:5.1} Ki{}", size / KI, self.1)
        } else {
            write!(f, "{:3.0} {}", size, self.1)
        }
    }
}

#[cfg(test)]
mod tests {

    #[test]
    fn download_tracker_from_seconds_test() {
        use crate::download_tracker::DownloadTracker;
        assert_eq!(DownloadTracker::from_seconds(2), (0, 0, 0, 2));

        assert_eq!(DownloadTracker::from_seconds(60), (0, 0, 1, 0));

        assert_eq!(DownloadTracker::from_seconds(3600), (0, 1, 0, 0));

        assert_eq!(DownloadTracker::from_seconds(3600 * 24), (1, 0, 0, 0));

        assert_eq!(DownloadTracker::from_seconds(52292), (0, 14, 31, 32));

        assert_eq!(DownloadTracker::from_seconds(222292), (2, 13, 44, 52));
    }

}
<commit_msg>Fix clippy: unnecessary_fold<commit_after>use crate::term2;
use rustup::dist::Notification as In;
use rustup::utils::tty;
use rustup::utils::Notification as Un;
use rustup::Notification;
use std::collections::VecDeque;
use std::fmt;
use std::io::Write;
use term::Terminal;
use time::precise_time_s;

/// Keep track of this many past download amounts
const DOWNLOAD_TRACK_COUNT: usize = 5;

/// Tracks download progress and displays information about it to a terminal.
pub struct DownloadTracker {
    /// Content-Length of the to-be downloaded object.
    content_len: Option<usize>,
    /// Total data downloaded in bytes.
    total_downloaded: usize,
    /// Data downloaded this second.
    downloaded_this_sec: usize,
    /// Keeps track of amount of data downloaded every last few secs.
    /// Used for averaging the download speed. NB: This does not necessarily
    /// represent adjacent seconds; thus it may not show the average at all.
    downloaded_last_few_secs: VecDeque<usize>,
    /// Time stamp of the last second
    last_sec: Option<f64>,
    /// Time stamp of the start of the download
    start_sec: f64,
    /// The terminal we write the information to.
    /// XXX: Could be a term trait, but with #1818 on the horizon that
    ///      is a pointless change to make - better to let that transition
    ///      happen and take stock after that.
    term: term2::StdoutTerminal,
    /// Whether we displayed progress for the download or not.
    ///
    /// If the download is quick enough, we don't have time to
    /// display the progress info.
    /// In that case, we do not want to do some cleanup stuff we normally do.
    ///
    /// If we have displayed progress, this is the number of characters we
    /// rendered, so we can erase it cleanly.
    displayed_charcount: Option<usize>,
    /// What units to show progress in
    units: Vec<String>,
}

impl DownloadTracker {
    /// Creates a new DownloadTracker.
    pub fn new() -> Self {
        DownloadTracker {
            content_len: None,
            total_downloaded: 0,
            downloaded_this_sec: 0,
            downloaded_last_few_secs: VecDeque::with_capacity(DOWNLOAD_TRACK_COUNT),
            start_sec: precise_time_s(),
            last_sec: None,
            term: term2::stdout(),
            displayed_charcount: None,
            units: vec!["B".into(); 1],
        }
    }

    pub fn handle_notification(&mut self, n: &Notification<'_>) -> bool {
        match *n {
            Notification::Install(In::Utils(Un::DownloadContentLengthReceived(content_len))) => {
                self.content_length_received(content_len);

                true
            }
            Notification::Install(In::Utils(Un::DownloadDataReceived(data))) => {
                if tty::stdout_isatty() {
                    self.data_received(data.len());
                }
                true
            }
            Notification::Install(In::Utils(Un::DownloadFinished)) => {
                self.download_finished();
                true
            }
            Notification::Install(In::Utils(Un::DownloadPushUnits(units))) => {
                self.push_units(units.into());
                true
            }
            Notification::Install(In::Utils(Un::DownloadPopUnits)) => {
                self.pop_units();
                true
            }

            _ => false,
        }
    }

    /// Notifies self that Content-Length information has been received.
    pub fn content_length_received(&mut self, content_len: u64) {
        self.content_len = Some(content_len as usize);
    }

    /// Notifies self that data of size `len` has been received.
    pub fn data_received(&mut self, len: usize) {
        self.total_downloaded += len;
        self.downloaded_this_sec += len;

        let current_time = precise_time_s();

        match self.last_sec {
            None => self.last_sec = Some(current_time),
            Some(prev) => {
                let elapsed = current_time - prev;
                if elapsed >= 1.0 {
                    self.display();
                    self.last_sec = Some(current_time);
                    if self.downloaded_last_few_secs.len() == DOWNLOAD_TRACK_COUNT {
                        self.downloaded_last_few_secs.pop_back();
                    }
                    self.downloaded_last_few_secs
                        .push_front(self.downloaded_this_sec);
                    self.downloaded_this_sec = 0;
                }
            }
        }
    }
    /// Notifies self that the download has finished.
    pub fn download_finished(&mut self) {
        if self.displayed_charcount.is_some() {
            // Display the finished state
            self.display();
            let _ = writeln!(self.term);
        }
        self.prepare_for_new_download();
    }
    // we're doing modular arithmetic, treat as integer
    pub fn from_seconds(sec: u32) -> (u32, u32, u32, u32) {
        let d = sec / (24 * 3600);
        let h = sec % (24 * 3600) / 3600;
        let min = sec % 3600 / 60;
        let sec = sec % 60;

        (d, h, min, sec)
    }
    /// Resets the state to be ready for a new download.
    fn prepare_for_new_download(&mut self) {
        self.content_len = None;
        self.total_downloaded = 0;
        self.downloaded_this_sec = 0;
        self.downloaded_last_few_secs.clear();
        self.start_sec = precise_time_s();
        self.last_sec = None;
        self.displayed_charcount = None;
    }
    /// Display the tracked download information to the terminal.
    fn display(&mut self) {
        // Panic if someone pops the default bytes unit...
        let units = &self.units.last().unwrap();
        let total_h = Size(self.total_downloaded, units);
        let sum: usize = self.downloaded_last_few_secs.iter().sum();
        let len = self.downloaded_last_few_secs.len();
        let speed = if len > 0 { sum / len } else { 0 };
        let speed_h = Size(speed, units);
        let elapsed_h = Duration(precise_time_s() - self.start_sec);

        // First, move to the start of the current line and clear it.
        let _ = self.term.carriage_return();
        // We'd prefer to use delete_line() but on Windows it seems to
        // sometimes do unusual things
        // let _ = self.term.as_mut().unwrap().delete_line();
        // So instead we do:
        if let Some(n) = self.displayed_charcount {
            // This is not ideal as very narrow terminals might mess up,
            // but it is more likely to succeed until term's windows console
            // fixes whatever's up with delete_line().
            let _ = write!(self.term, "{}", " ".repeat(n));
            let _ = self.term.flush();
            let _ = self.term.carriage_return();
        }

        let output = match self.content_len {
            Some(content_len) => {
                let content_len_h = Size(content_len, units);
                let content_len = content_len as f64;
                let percent = (self.total_downloaded as f64 / content_len) * 100.;
                let remaining = content_len - self.total_downloaded as f64;
                let eta_h = Duration(remaining / speed as f64);
                format!(
                    "{} / {} ({:3.0} %) {}/s in {} ETA: {}",
                    total_h, content_len_h, percent, speed_h, elapsed_h, eta_h
                )
            }
            None => format!(
                "Total: {} Speed: {}/s Elapsed: {}",
                total_h, speed_h, elapsed_h
            ),
        };

        let _ = write!(self.term, "{}", output);
        // Since stdout is typically line-buffered and we don't print a newline, we manually flush.
        let _ = self.term.flush();
        self.displayed_charcount = Some(output.chars().count());
    }

    pub fn push_units(&mut self, new_units: String) {
        self.units.push(new_units);
    }

    pub fn pop_units(&mut self) {
        self.units.pop();
    }
}

/// Human readable representation of duration(seconds).
struct Duration(f64);

impl fmt::Display for Duration {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // repurposing the alternate mode for ETA
        let sec = self.0;

        if sec.is_infinite() {
            write!(f, "Unknown")
        } else {
            match DownloadTracker::from_seconds(sec as u32) {
                (d, h, m, s) if d > 0 => write!(f, "{:3.0}d {:2.0}h {:2.0}m {:2.0}s", d, h, m, s),
                (0, h, m, s) if h > 0 => write!(f, "{:2.0}h {:2.0}m {:2.0}s", h, m, s),
                (0, 0, m, s) if m > 0 => write!(f, "{:2.0}m {:2.0}s", m, s),
                (_, _, _, s) => write!(f, "{:2.0}s", s),
            }
        }
    }
}

/// Human readable size (some units)
struct Size<'a>(usize, &'a str);

impl<'a> fmt::Display for Size<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        const KI: f64 = 1024.0;
        const MI: f64 = KI * KI;
        let size = self.0 as f64;

        if size >= MI {
            write!(f, "{:5.1} Mi{}", size / MI, self.1) // XYZ.P Mi
        } else if size >= KI {
            write!(f, "{:5.1} Ki{}", size / KI, self.1)
        } else {
            write!(f, "{:3.0} {}", size, self.1)
        }
    }
}

#[cfg(test)]
mod tests {

    #[test]
    fn download_tracker_from_seconds_test() {
        use crate::download_tracker::DownloadTracker;
        assert_eq!(DownloadTracker::from_seconds(2), (0, 0, 0, 2));

        assert_eq!(DownloadTracker::from_seconds(60), (0, 0, 1, 0));

        assert_eq!(DownloadTracker::from_seconds(3600), (0, 1, 0, 0));

        assert_eq!(DownloadTracker::from_seconds(3600 * 24), (1, 0, 0, 0));

        assert_eq!(DownloadTracker::from_seconds(52292), (0, 14, 31, 32));

        assert_eq!(DownloadTracker::from_seconds(222292), (2, 13, 44, 52));
    }

}
<|endoftext|>