<commit_before>//! A minimal implementation of SHA1 for rust.
//!
//! Example:
//!
//! ```rust
//! extern crate sha1;
//! # fn main() {
//!
//! let mut m = sha1::Sha1::new();
//! m.update(b"Hello World!");
//! assert_eq!(m.digest().to_string(),
//!            "2ef7bde608ce5404e97d5f042f95f89f1c232871");
//! # }
//! ```

#![no_std]
#![deny(missing_docs)]

use core::cmp;
use core::fmt;

/// Represents a Sha1 hash object in memory.
#[derive(Clone)]
pub struct Sha1 {
    state: Sha1State,
    blocks: Blocks,
    len: u64,
}

struct Blocks {
    len: u32,
    block: [u8; 64],
}

#[derive(Copy, Clone)]
struct Sha1State {
    state: [u32; 5],
}

/// Digest generated from a `Sha1` instance.
///
/// A digest can be formatted to view the digest as a hex string, or the bytes
/// can be extracted for later processing.
pub struct Digest {
    data: Sha1State,
}

const DEFAULT_STATE: Sha1State = Sha1State {
    state: [0x67452301, 0xefcdab89, 0x98badcfe, 0x10325476, 0xc3d2e1f0],
};

impl Sha1 {
    /// Creates an fresh sha1 hash object.
    pub fn new() -> Sha1 {
        Sha1 {
            state: DEFAULT_STATE,
            len: 0,
            blocks: Blocks {
                len: 0,
                block: [0; 64],
            },
        }
    }

    /// Resets the hash object to it's initial state.
    pub fn reset(&mut self) {
        self.state = DEFAULT_STATE;
        self.len = 0;
        self.blocks.len = 0;
    }

    /// Update hash with input data.
    pub fn update(&mut self, data: &[u8]) {
        let len = &mut self.len;
        let state = &mut self.state;
        self.blocks.input(data, |chunk| {
            *len += 64;
            state.process(chunk);
        })
    }

    /// Retrieve digest result.
    pub fn digest(&self) -> Digest {
        let mut state = self.state;
        let bits = (self.len + (self.blocks.len as u64)) * 8;
        let extra = [
            (bits >> 56) as u8,
            (bits >> 48) as u8,
            (bits >> 40) as u8,
            (bits >> 32) as u8,
            (bits >> 24) as u8,
            (bits >> 16) as u8,
            (bits >>  8) as u8,
            (bits >>  0) as u8,
        ];
        let mut last = [0; 128];
        let blocklen = self.blocks.len as usize;
        last[..blocklen].clone_from_slice(&self.blocks.block[..blocklen]);
        last[blocklen] = 0x80;

        if blocklen < 56 {
            last[56..64].clone_from_slice(&extra);
            state.process(&last[0..64]);
        } else {
            last[120..128].clone_from_slice(&extra);
            state.process(&last[0..64]);
            state.process(&last[64..128]);
        }

        Digest { data: state }
    }
}

impl Digest {
    /// Returns the 160 bit (20 byte) digest as a byte array.
    pub fn bytes(&self) -> [u8; 20] {
        [
            (self.data.state[0] >> 24) as u8,
            (self.data.state[0] >> 16) as u8,
            (self.data.state[0] >>  8) as u8,
            (self.data.state[0] >>  0) as u8,
            (self.data.state[1] >> 24) as u8,
            (self.data.state[1] >> 16) as u8,
            (self.data.state[1] >>  8) as u8,
            (self.data.state[1] >>  0) as u8,
            (self.data.state[2] >> 24) as u8,
            (self.data.state[2] >> 16) as u8,
            (self.data.state[2] >>  8) as u8,
            (self.data.state[2] >>  0) as u8,
            (self.data.state[3] >> 24) as u8,
            (self.data.state[3] >> 16) as u8,
            (self.data.state[3] >>  8) as u8,
            (self.data.state[3] >>  0) as u8,
            (self.data.state[4] >> 24) as u8,
            (self.data.state[4] >> 16) as u8,
            (self.data.state[4] >>  8) as u8,
            (self.data.state[4] >>  0) as u8,
        ]
    }
}

impl Blocks {
    fn input<F>(&mut self, mut input: &[u8], mut f: F) where F: FnMut(&[u8]) {
        if self.len > 0 {
            let len = self.len as usize;
            let amt = cmp::min(input.len(), self.block.len() - len);
            self.block[len..len + amt].clone_from_slice(&input[..amt]);
            if len + amt == self.block.len() {
                f(&self.block);
                self.len = 0;
                input = &input[amt..];
            } else {
                self.len += amt as u32;
                return
            }
        }
        assert_eq!(self.len, 0);
        for chunk in input.chunks(64) {
            if chunk.len() == 64 {
                f(chunk)
            } else {
                self.block[..chunk.len()].clone_from_slice(chunk);
                self.len = chunk.len() as u32;
            }
        }
    }
}

impl Sha1State {
    fn process(&mut self, block: &[u8]) {
        let mut words = [0u32; 80];
        for (i, chunk) in block.chunks(4).enumerate() {
            words[i] = (chunk[3] as u32) |
                       ((chunk[2] as u32) << 8) |
                       ((chunk[1] as u32) << 16) |
                       ((chunk[0] as u32) << 24);
        }

        fn ff(b: u32, c: u32, d: u32) -> u32 { d ^ (b & (c ^ d)) }
        fn gg(b: u32, c: u32, d: u32) -> u32 { b ^ c ^ d }
        fn hh(b: u32, c: u32, d: u32) -> u32 { (b & c) | (d & (b | c)) }
        fn ii(b: u32, c: u32, d: u32) -> u32 { b ^ c ^ d }

        for i in 16..80 {
            let n = words[i - 3] ^ words[i - 8] ^ words[i - 14] ^ words[i - 16];
            words[i] = n.rotate_left(1);
        }

        let mut a = self.state[0];
        let mut b = self.state[1];
        let mut c = self.state[2];
        let mut d = self.state[3];
        let mut e = self.state[4];

        for i in 0..80 {
            let (f, k) = match i {
                0 ... 19 => (ff(b, c, d), 0x5a827999),
                20 ... 39 => (gg(b, c, d), 0x6ed9eba1),
                40 ... 59 => (hh(b, c, d), 0x8f1bbcdc),
                60 ... 79 => (ii(b, c, d), 0xca62c1d6),
                _ => (0, 0),
            };

            let tmp = a.rotate_left(5)
                       .wrapping_add(f)
                       .wrapping_add(e)
                       .wrapping_add(k)
                       .wrapping_add(words[i]);
            e = d;
            d = c;
            c = b.rotate_left(30);
            b = a;
            a = tmp;
        }

        self.state[0] = self.state[0].wrapping_add(a);
        self.state[1] = self.state[1].wrapping_add(b);
        self.state[2] = self.state[2].wrapping_add(c);
        self.state[3] = self.state[3].wrapping_add(d);
        self.state[4] = self.state[4].wrapping_add(e);
    }
}

impl Clone for Blocks {
    fn clone(&self) -> Blocks {
        Blocks { ..*self }
    }
}

impl fmt::Display for Digest {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for byte in self.data.state.iter() {
            try!(write!(f, "{:08x}", byte));
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    extern crate std;
    extern crate rand;
    extern crate openssl;

    use self::std::prelude::v1::*;
    use self::std::io::Write;

    use Sha1;

    #[test]
    fn test_simple() {
        let mut m = Sha1::new();

        let tests = [
            ("The quick brown fox jumps over the lazy dog",
             "2fd4e1c67a2d28fced849ee1bb76e7391b93eb12"),
            ("The quick brown fox jumps over the lazy cog",
             "de9f2c7fd25e1b3afad3e85a0bd17d9b100db4b3"),
            ("", "da39a3ee5e6b4b0d3255bfef95601890afd80709"),
            ("testing\n", "9801739daae44ec5293d4e1f53d3f4d2d426d91c"),
            ("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
             "025ecbd5d70f8fb3c5457cd96bab13fda305dc59"),
            ("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
             "4300320394f7ee239bcdce7d3b8bcee173a0cd5c"),
            ("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
             "cef734ba81a024479e09eb5a75b6ddae62e6abf1"),
        ];

        for &(s, ref h) in tests.iter() {
            let data = s.as_bytes();

            m.reset();
            m.update(data);
            let hh = m.digest().to_string();

            assert_eq!(hh.len(), h.len());
            assert_eq!(hh, *h);
        }
    }

    #[test]
    fn test_multiple_updates() {
        let mut m = Sha1::new();

        m.reset();
        m.update("The quick brown ".as_bytes());
        m.update("fox jumps over ".as_bytes());
        m.update("the lazy dog".as_bytes());
        let hh = m.digest().to_string();

        let h = "2fd4e1c67a2d28fced849ee1bb76e7391b93eb12";
        assert_eq!(hh.len(), h.len());
        assert_eq!(hh, &*h);
    }

    #[test]
    fn test_sha1_loop() {
        let mut m = Sha1::new();
        let s = "The quick brown fox jumps over the lazy dog.";
        let n = 1000u64;

        for _ in 0..3 {
            m.reset();
            for _ in 0..n {
                m.update(s.as_bytes());
            }
            assert_eq!(m.digest().to_string(),
                       "7ca27655f67fceaa78ed2e645a81c7f1d6e249d2");
        }
    }

    #[test]
    fn spray_and_pray() {
        use self::rand::Rng;

        let mut rng = rand::thread_rng();
        let mut m = Sha1::new();
        let mut bytes = [0; 512];

        for _ in 0..20 {
            let ty = openssl::crypto::hash::Type::SHA1;
            let mut r = openssl::crypto::hash::Hasher::new(ty);
            m.reset();
            for _ in 0..50 {
                let len = rng.gen::<usize>() % bytes.len();
                rng.fill_bytes(&mut bytes[..len]);
                m.update(&bytes[..len]);
                r.write(&bytes[..len]).unwrap();
            }
            assert_eq!(r.finish(), m.digest().bytes());
        }
    }
}
<commit_msg>Use more accurate variable name for the loop<commit_after>//! A minimal implementation of SHA1 for rust.
//!
//! Example:
//!
//! ```rust
//! extern crate sha1;
//! # fn main() {
//!
//! let mut m = sha1::Sha1::new();
//! m.update(b"Hello World!");
//! assert_eq!(m.digest().to_string(),
//!            "2ef7bde608ce5404e97d5f042f95f89f1c232871");
//! # }
//! ```

#![no_std]
#![deny(missing_docs)]

use core::cmp;
use core::fmt;

/// Represents a Sha1 hash object in memory.
#[derive(Clone)]
pub struct Sha1 {
    state: Sha1State,
    blocks: Blocks,
    len: u64,
}

struct Blocks {
    len: u32,
    block: [u8; 64],
}

#[derive(Copy, Clone)]
struct Sha1State {
    state: [u32; 5],
}

/// Digest generated from a `Sha1` instance.
///
/// A digest can be formatted to view the digest as a hex string, or the bytes
/// can be extracted for later processing.
pub struct Digest {
    data: Sha1State,
}

const DEFAULT_STATE: Sha1State = Sha1State {
    state: [0x67452301, 0xefcdab89, 0x98badcfe, 0x10325476, 0xc3d2e1f0],
};

impl Sha1 {
    /// Creates an fresh sha1 hash object.
    pub fn new() -> Sha1 {
        Sha1 {
            state: DEFAULT_STATE,
            len: 0,
            blocks: Blocks {
                len: 0,
                block: [0; 64],
            },
        }
    }

    /// Resets the hash object to it's initial state.
    pub fn reset(&mut self) {
        self.state = DEFAULT_STATE;
        self.len = 0;
        self.blocks.len = 0;
    }

    /// Update hash with input data.
    pub fn update(&mut self, data: &[u8]) {
        let len = &mut self.len;
        let state = &mut self.state;
        self.blocks.input(data, |chunk| {
            *len += 64;
            state.process(chunk);
        })
    }

    /// Retrieve digest result.
    pub fn digest(&self) -> Digest {
        let mut state = self.state;
        let bits = (self.len + (self.blocks.len as u64)) * 8;
        let extra = [
            (bits >> 56) as u8,
            (bits >> 48) as u8,
            (bits >> 40) as u8,
            (bits >> 32) as u8,
            (bits >> 24) as u8,
            (bits >> 16) as u8,
            (bits >>  8) as u8,
            (bits >>  0) as u8,
        ];
        let mut last = [0; 128];
        let blocklen = self.blocks.len as usize;
        last[..blocklen].clone_from_slice(&self.blocks.block[..blocklen]);
        last[blocklen] = 0x80;

        if blocklen < 56 {
            last[56..64].clone_from_slice(&extra);
            state.process(&last[0..64]);
        } else {
            last[120..128].clone_from_slice(&extra);
            state.process(&last[0..64]);
            state.process(&last[64..128]);
        }

        Digest { data: state }
    }
}

impl Digest {
    /// Returns the 160 bit (20 byte) digest as a byte array.
    pub fn bytes(&self) -> [u8; 20] {
        [
            (self.data.state[0] >> 24) as u8,
            (self.data.state[0] >> 16) as u8,
            (self.data.state[0] >>  8) as u8,
            (self.data.state[0] >>  0) as u8,
            (self.data.state[1] >> 24) as u8,
            (self.data.state[1] >> 16) as u8,
            (self.data.state[1] >>  8) as u8,
            (self.data.state[1] >>  0) as u8,
            (self.data.state[2] >> 24) as u8,
            (self.data.state[2] >> 16) as u8,
            (self.data.state[2] >>  8) as u8,
            (self.data.state[2] >>  0) as u8,
            (self.data.state[3] >> 24) as u8,
            (self.data.state[3] >> 16) as u8,
            (self.data.state[3] >>  8) as u8,
            (self.data.state[3] >>  0) as u8,
            (self.data.state[4] >> 24) as u8,
            (self.data.state[4] >> 16) as u8,
            (self.data.state[4] >>  8) as u8,
            (self.data.state[4] >>  0) as u8,
        ]
    }
}

impl Blocks {
    fn input<F>(&mut self, mut input: &[u8], mut f: F) where F: FnMut(&[u8]) {
        if self.len > 0 {
            let len = self.len as usize;
            let amt = cmp::min(input.len(), self.block.len() - len);
            self.block[len..len + amt].clone_from_slice(&input[..amt]);
            if len + amt == self.block.len() {
                f(&self.block);
                self.len = 0;
                input = &input[amt..];
            } else {
                self.len += amt as u32;
                return
            }
        }
        assert_eq!(self.len, 0);
        for chunk in input.chunks(64) {
            if chunk.len() == 64 {
                f(chunk)
            } else {
                self.block[..chunk.len()].clone_from_slice(chunk);
                self.len = chunk.len() as u32;
            }
        }
    }
}

impl Sha1State {
    fn process(&mut self, block: &[u8]) {
        let mut words = [0u32; 80];
        for (i, chunk) in block.chunks(4).enumerate() {
            words[i] = (chunk[3] as u32) |
                       ((chunk[2] as u32) << 8) |
                       ((chunk[1] as u32) << 16) |
                       ((chunk[0] as u32) << 24);
        }

        fn ff(b: u32, c: u32, d: u32) -> u32 { d ^ (b & (c ^ d)) }
        fn gg(b: u32, c: u32, d: u32) -> u32 { b ^ c ^ d }
        fn hh(b: u32, c: u32, d: u32) -> u32 { (b & c) | (d & (b | c)) }
        fn ii(b: u32, c: u32, d: u32) -> u32 { b ^ c ^ d }

        for i in 16..80 {
            let n = words[i - 3] ^ words[i - 8] ^ words[i - 14] ^ words[i - 16];
            words[i] = n.rotate_left(1);
        }

        let mut a = self.state[0];
        let mut b = self.state[1];
        let mut c = self.state[2];
        let mut d = self.state[3];
        let mut e = self.state[4];

        for i in 0..80 {
            let (f, k) = match i {
                0 ... 19 => (ff(b, c, d), 0x5a827999),
                20 ... 39 => (gg(b, c, d), 0x6ed9eba1),
                40 ... 59 => (hh(b, c, d), 0x8f1bbcdc),
                60 ... 79 => (ii(b, c, d), 0xca62c1d6),
                _ => (0, 0),
            };

            let tmp = a.rotate_left(5)
                       .wrapping_add(f)
                       .wrapping_add(e)
                       .wrapping_add(k)
                       .wrapping_add(words[i]);
            e = d;
            d = c;
            c = b.rotate_left(30);
            b = a;
            a = tmp;
        }

        self.state[0] = self.state[0].wrapping_add(a);
        self.state[1] = self.state[1].wrapping_add(b);
        self.state[2] = self.state[2].wrapping_add(c);
        self.state[3] = self.state[3].wrapping_add(d);
        self.state[4] = self.state[4].wrapping_add(e);
    }
}

impl Clone for Blocks {
    fn clone(&self) -> Blocks {
        Blocks { ..*self }
    }
}

impl fmt::Display for Digest {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for i in self.data.state.iter() {
            try!(write!(f, "{:08x}", i));
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    extern crate std;
    extern crate rand;
    extern crate openssl;

    use self::std::prelude::v1::*;
    use self::std::io::Write;

    use Sha1;

    #[test]
    fn test_simple() {
        let mut m = Sha1::new();

        let tests = [
            ("The quick brown fox jumps over the lazy dog",
             "2fd4e1c67a2d28fced849ee1bb76e7391b93eb12"),
            ("The quick brown fox jumps over the lazy cog",
             "de9f2c7fd25e1b3afad3e85a0bd17d9b100db4b3"),
            ("", "da39a3ee5e6b4b0d3255bfef95601890afd80709"),
            ("testing\n", "9801739daae44ec5293d4e1f53d3f4d2d426d91c"),
            ("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
             "025ecbd5d70f8fb3c5457cd96bab13fda305dc59"),
            ("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
             "4300320394f7ee239bcdce7d3b8bcee173a0cd5c"),
            ("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
             "cef734ba81a024479e09eb5a75b6ddae62e6abf1"),
        ];

        for &(s, ref h) in tests.iter() {
            let data = s.as_bytes();

            m.reset();
            m.update(data);
            let hh = m.digest().to_string();

            assert_eq!(hh.len(), h.len());
            assert_eq!(hh, *h);
        }
    }

    #[test]
    fn test_multiple_updates() {
        let mut m = Sha1::new();

        m.reset();
        m.update("The quick brown ".as_bytes());
        m.update("fox jumps over ".as_bytes());
        m.update("the lazy dog".as_bytes());
        let hh = m.digest().to_string();

        let h = "2fd4e1c67a2d28fced849ee1bb76e7391b93eb12";
        assert_eq!(hh.len(), h.len());
        assert_eq!(hh, &*h);
    }

    #[test]
    fn test_sha1_loop() {
        let mut m = Sha1::new();
        let s = "The quick brown fox jumps over the lazy dog.";
        let n = 1000u64;

        for _ in 0..3 {
            m.reset();
            for _ in 0..n {
                m.update(s.as_bytes());
            }
            assert_eq!(m.digest().to_string(),
                       "7ca27655f67fceaa78ed2e645a81c7f1d6e249d2");
        }
    }

    #[test]
    fn spray_and_pray() {
        use self::rand::Rng;

        let mut rng = rand::thread_rng();
        let mut m = Sha1::new();
        let mut bytes = [0; 512];

        for _ in 0..20 {
            let ty = openssl::crypto::hash::Type::SHA1;
            let mut r = openssl::crypto::hash::Hasher::new(ty);
            m.reset();
            for _ in 0..50 {
                let len = rng.gen::<usize>() % bytes.len();
                rng.fill_bytes(&mut bytes[..len]);
                m.update(&bytes[..len]);
                r.write(&bytes[..len]).unwrap();
            }
            assert_eq!(r.finish(), m.digest().bytes());
        }
    }
}
<|endoftext|>


<commit_before>extern crate log;
extern crate simplelog;

use goblin::pe::PE;
use goblin::Object;
use log::{debug, error, info, trace};
use rayon::prelude::*;
use std::env;
use std::fs;
use std::io::prelude::*;
use zydis;

// TODO: this really shouldn't be public. only for testing.
// TODO: what is the `cfg(doctest)` setting?
//#[cfg(test)]
pub mod rsrc;

pub struct Config {
    pub filename: String,
}

impl Config {
    pub fn from_args(args: env::Args) -> Result<Config, &'static str> {
        let args: Vec<String> = args.collect();

        if args.len() < 2 {
            return Err("not enough arguments");
        }

        let filename = args[1].clone();
        trace!("config: parsed filename: {:?}", filename);

        Ok(Config { filename })
    }
}

pub fn setup_logging(_args: &Config) {
    simplelog::TermLogger::init(simplelog::LevelFilter::Info, simplelog::Config::default())
        .expect("failed to setup logging");
}

#[derive(Debug)]
pub enum Error {
    FileAccess,
    FileFormat,
    NotImplemented,
    ValueError,
}

fn align(i: usize, b: usize) -> usize {
    let rem = i % b;
    if rem == 0 {
        i
    } else {
        i + (b - rem)
    }
}

pub fn hexdump_ascii(b: u8) -> char {
    if b.is_ascii_graphic() || b == b' ' {
        b as char
    } else {
        '.'
    }
}

pub fn hexdump(buf: &[u8], offset: usize) -> String {
    // 01234567:  00 01 02 03 04 05 06 07  ...............
    // <prefix>   <hex col>                <ascii col>

    let padding = "  ";

    let padding_size = 2;
    let hex_col_size = 3;
    let ascii_col_size = 1;
    let prefix_size = 8 + 1;
    let newline_size = 1;
    let line_size = prefix_size
        + padding_size
        + 16 * hex_col_size
        + padding_size
        + 16 * ascii_col_size
        + newline_size;
    let line_count = align(buf.len(), 0x10) / 0x10;

    let mut ret = String::with_capacity(line_count * line_size);

    let mut line = String::with_capacity(line_size);
    let mut remaining_count = buf.len();
    for line_index in 0..line_count {
        let line_elem_count = 0x10.min(remaining_count);
        let padding_elem_count = 0x10 - line_elem_count;

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        // ^^^^^^^^^
        line.push_str(format!("{:08x}:", offset + 0x10 * line_index).as_str());

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //          ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //            ^^^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push_str(format!("{:02x} ", elem).as_str());
        }
        for _ in 0..padding_elem_count {
            line.push_str("   ");
        }

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                   ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                     ^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push(hexdump_ascii(*elem))
        }
        for _ in 0..padding_elem_count {
            line.push(' ');
        }
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                                    ^
        line.push('\n');

        ret.push_str(line.as_str());
        line.truncate(0x0);
        remaining_count -= line_elem_count;
    }

    ret
}

fn foo(ws: &Workspace) -> Result<(), Error> {
    if let Object::PE(pe) = ws.get_obj().unwrap() {
        info!("module: {}", pe.name.unwrap_or("(unknown)"));

        info!("bitness: {}", if pe.is_64 { "64" } else { "32" });
        info!("image base: 0x{:x}", pe.image_base);
        info!("entry rva: 0x{:x}", pe.entry);

        // like:
        //
        //     sections:
        //       - .text
        //         raw size:     0x18aa00
        //         virtual size: 0x18a9a8
        info!("sections:");
        for section in ws.sections.iter() {
            info!("  - {}", section.name);
            info!("    virtual address: 0x{:x}", section.addr);
            info!("    virtual size: 0x{:x}", section.buf.len());

            info!(
                "\n{}",
                hexdump(&section.buf[..0x1C], pe.image_base + section.addr as usize)
            );

            let decoder =
                zydis::Decoder::new(zydis::MachineMode::Long64, zydis::AddressWidth::_64).unwrap();
            let insns: Vec<_> = section
                .buf
                .par_windows(0x10)
                .map(|ibuf| decoder.decode(ibuf))
                .collect();

            info!("total instructions: {}", insns.len());

            info!(
                "successful disassembles: {}",
                insns.par_iter().filter(|insn| insn.is_ok()).count()
            );

            info!(
                "valid instructions: {}",
                insns
                    .par_iter()
                    .filter(|insn| match insn {
                        Ok(Some(_)) => true,
                        _ => false,
                    })
                    .count()
            );
        }
    }

    Ok(())
}

pub fn read_file(filename: &str) -> Result<Vec<u8>, Error> {
    debug!("read_file: {:?}", filename);

    let mut buf = Vec::new();
    {
        debug!("reading file: {}", filename);
        let mut f = match fs::File::open(filename) {
            Ok(f) => f,
            Err(_) => {
                error!("failed to open file: {}", filename);
                return Err(Error::FileAccess);
            }
        };
        let bytes_read = match f.read_to_end(&mut buf) {
            Ok(c) => c,
            Err(_) => {
                error!("failed to read entire file: {}", filename);
                return Err(Error::FileAccess);
            }
        };
        debug!("read {} bytes", bytes_read);
        if bytes_read < 0x10 {
            error!("file too small: {}", filename);
            return Err(Error::FileFormat);
        }
    }

    Ok(buf)
}

pub struct Section {
    pub name: String,
    pub addr: u32,
    pub buf: Vec<u8>,
}

pub struct Workspace {
    pub filename: String,
    pub buf: Vec<u8>,
    pub sections: Vec<Section>,
}

impl Workspace {
    /// Parse the given file into its object.
    ///
    /// # Errors
    ///   - `Error::FileFormat`: when not able to be parsed by Goblin.
    ///
    /// # Examples
    ///
    /// ```
    /// use goblin::Object;
    /// use matches::matches;
    /// use lancelot::rsrc::*;
    /// let ws = get_workspace(Rsrc::K32);
    /// assert!(matches!(ws.get_obj().unwrap(), Object::PE(_)));
    /// ```
    ///
    /// you might be tempted to maintain a method `get_pe`,
    /// however, i don't think this is a good idea:
    /// its fragile because the file type may not be PE.
    /// therefore, force clients to be explicit:
    ///
    /// ```
    /// use goblin::Object;
    /// use lancelot::rsrc::*;
    /// let ws = get_workspace(Rsrc::K32);
    /// if let Object::PE(_) = ws.get_obj().unwrap() {
    ///     // everyone is happy!
    /// }
    /// ```
    ///
    /// TODO: demonstrate `Error::FileFormat`.
    pub fn get_obj(&self) -> Result<Object, Error> {
        let obj = match Object::parse(&self.buf) {
            Ok(o) => o,
            Err(e) => {
                error!("failed to parse file: {} error: {:?}", self.filename, e);
                return Err(Error::FileFormat);
            }
        };

        match obj {
            Object::Unknown(_) => {
                error!(
                    "unknown file format, magic: | {:02X} {:02X} | '{}{}' ",
                    self.buf[0],
                    self.buf[1],
                    hexdump_ascii(self.buf[0]),
                    hexdump_ascii(self.buf[1])
                );

                Err(Error::FileFormat)
            }
            _ => Ok(obj),
        }
    }

    pub fn get_pe(&self) -> Result<PE, Error> {
        match self.get_obj()? {
            Object::PE(pe) => {
                return Ok(pe);
            }
            _ => {
                error!("not a PE file");
                return Err(Error::ValueError);
            }
        }
    }

    /// Construct a workspace from the module with the given contents.
    ///
    /// # Errors
    ///   - `Error::FileFormat`: when not able to be parsed by Goblin.
    ///   - `Error::NotImplemented`: when not a PE file.
    ///
    /// # Examples
    ///
    /// ```
    /// use lancelot::*;
    /// use lancelot::rsrc::*;
    /// let buf = get_buf(Rsrc::K32);
    /// let ws = Workspace::from_buf("kernel32.dll", buf).unwrap();
    /// ```
    ///
    /// TODO: demonstrate ELF file behavior.
    /// TODO: demonstrate MachO file behavior.
    /// TODO: demonstrate unknown file behavior.
    pub fn from_buf(filename: &str, buf: Vec<u8>) -> Result<Workspace, Error> {
        let mut ws = Workspace {
            filename: filename.to_string(),
            buf: buf.clone(),
            sections: vec![],
        };

        match ws.get_obj()? {
            Object::PE(pe) => {
                info!("found PE file");

                ws.sections
                    .extend(pe.sections.iter().map(|section| -> Section {
                        // TODO: i'm sure this can be abused.
                        // TODO: add tests for weird section names.
                        let name = String::from_utf8_lossy(&section.name[..]).into_owned();

                        // TODO: figure out if we will work with usize, or u64, or what,
                        // then assert usize is ok.
                        // ref: `usize::max_value()`
                        let mut secbuf = vec![0; align(section.virtual_size as usize, 0x200)];

                        {
                            let secsize = section.size_of_raw_data as usize;
                            let rawbuf = &mut secbuf[..secsize];
                            let pstart = section.pointer_to_raw_data as usize;
                            rawbuf.copy_from_slice(&buf[pstart..pstart + secsize]);
                        }

                        Section {
                            name: name,
                            addr: section.virtual_address,
                            buf: secbuf,
                        }
                    }));
            }
            Object::Elf(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Mach(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Archive(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Unknown(_) => {
                return Err(Error::NotImplemented);
            }
        }

        Ok(ws)
    }

    /// Construct a workspace from the module at given a file path.
    ///
    /// ```
    /// use lancelot::*;
    /// use lancelot::rsrc::*;
    /// let path = get_path(Rsrc::K32);
    /// // This test resource file is mangled. Needs to be fixed before parsing.
    /// // Otherwise, the following would work:
    /// // let ws = Workspace::from_file(&path).unwrap();
    /// ```
    pub fn from_file(filename: &str) -> Result<Workspace, Error> {
        let buf = read_file(filename)?;
        Workspace::from_buf(filename, buf)
    }
}

pub fn run(args: &Config) -> Result<(), Error> {
    debug!("filename: {:?}", args.filename);
    let ws = Workspace::from_file(&args.filename)?;

    if let Object::PE(_) = ws.get_obj()? {
        foo(&ws).expect("failed to foo");
    }

    Ok(())
}
<commit_msg>lib: add tests for align<commit_after>extern crate log;
extern crate simplelog;

use goblin::pe::PE;
use goblin::Object;
use log::{debug, error, info, trace};
use rayon::prelude::*;
use std::env;
use std::fs;
use std::io::prelude::*;
use zydis;

// TODO: this really shouldn't be public. only for testing.
// TODO: what is the `cfg(doctest)` setting?
//#[cfg(test)]
pub mod rsrc;

pub struct Config {
    pub filename: String,
}

impl Config {
    pub fn from_args(args: env::Args) -> Result<Config, &'static str> {
        let args: Vec<String> = args.collect();

        if args.len() < 2 {
            return Err("not enough arguments");
        }

        let filename = args[1].clone();
        trace!("config: parsed filename: {:?}", filename);

        Ok(Config { filename })
    }
}

pub fn setup_logging(_args: &Config) {
    simplelog::TermLogger::init(simplelog::LevelFilter::Info, simplelog::Config::default())
        .expect("failed to setup logging");
}

#[derive(Debug)]
pub enum Error {
    FileAccess,
    FileFormat,
    NotImplemented,
    ValueError,
}

/// Round the given value up to the next multiple of the given base.
///
/// # Panics
///
///   - Base `b` must be at least `2`.
///
/// # Examples
///
/// ```
/// use lancelot::*;
/// assert_eq!(align(0, 2), 0);
/// assert_eq!(align(1, 2), 2);
/// assert_eq!(align(2, 2), 2);
/// assert_eq!(align(3, 2), 4);
/// assert_eq!(align(4, 2), 4);
/// ```
pub fn align(i: usize, b: usize) -> usize {
    if b < 2 {
        panic!("base `b` must be at least: 2");
    }
    let rem = i % b;
    if rem == 0 {
        i
    } else {
        i + (b - rem)
    }
}

pub fn hexdump_ascii(b: u8) -> char {
    if b.is_ascii_graphic() || b == b' ' {
        b as char
    } else {
        '.'
    }
}

pub fn hexdump(buf: &[u8], offset: usize) -> String {
    // 01234567:  00 01 02 03 04 05 06 07  ...............
    // <prefix>   <hex col>                <ascii col>

    let padding = "  ";

    let padding_size = 2;
    let hex_col_size = 3;
    let ascii_col_size = 1;
    let prefix_size = 8 + 1;
    let newline_size = 1;
    let line_size = prefix_size
        + padding_size
        + 16 * hex_col_size
        + padding_size
        + 16 * ascii_col_size
        + newline_size;
    let line_count = align(buf.len(), 0x10) / 0x10;

    let mut ret = String::with_capacity(line_count * line_size);

    let mut line = String::with_capacity(line_size);
    let mut remaining_count = buf.len();
    for line_index in 0..line_count {
        let line_elem_count = 0x10.min(remaining_count);
        let padding_elem_count = 0x10 - line_elem_count;

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        // ^^^^^^^^^
        line.push_str(format!("{:08x}:", offset + 0x10 * line_index).as_str());

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //          ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //            ^^^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push_str(format!("{:02x} ", elem).as_str());
        }
        for _ in 0..padding_elem_count {
            line.push_str("   ");
        }

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                   ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                     ^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push(hexdump_ascii(*elem))
        }
        for _ in 0..padding_elem_count {
            line.push(' ');
        }
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                                    ^
        line.push('\n');

        ret.push_str(line.as_str());
        line.truncate(0x0);
        remaining_count -= line_elem_count;
    }

    ret
}

fn foo(ws: &Workspace) -> Result<(), Error> {
    if let Object::PE(pe) = ws.get_obj().unwrap() {
        info!("module: {}", pe.name.unwrap_or("(unknown)"));

        info!("bitness: {}", if pe.is_64 { "64" } else { "32" });
        info!("image base: 0x{:x}", pe.image_base);
        info!("entry rva: 0x{:x}", pe.entry);

        // like:
        //
        //     sections:
        //       - .text
        //         raw size:     0x18aa00
        //         virtual size: 0x18a9a8
        info!("sections:");
        for section in ws.sections.iter() {
            info!("  - {}", section.name);
            info!("    virtual address: 0x{:x}", section.addr);
            info!("    virtual size: 0x{:x}", section.buf.len());

            info!(
                "\n{}",
                hexdump(&section.buf[..0x1C], pe.image_base + section.addr as usize)
            );

            let decoder =
                zydis::Decoder::new(zydis::MachineMode::Long64, zydis::AddressWidth::_64).unwrap();
            let insns: Vec<_> = section
                .buf
                .par_windows(0x10)
                .map(|ibuf| decoder.decode(ibuf))
                .collect();

            info!("total instructions: {}", insns.len());

            info!(
                "successful disassembles: {}",
                insns.par_iter().filter(|insn| insn.is_ok()).count()
            );

            info!(
                "valid instructions: {}",
                insns
                    .par_iter()
                    .filter(|insn| match insn {
                        Ok(Some(_)) => true,
                        _ => false,
                    })
                    .count()
            );
        }
    }

    Ok(())
}

pub fn read_file(filename: &str) -> Result<Vec<u8>, Error> {
    debug!("read_file: {:?}", filename);

    let mut buf = Vec::new();
    {
        debug!("reading file: {}", filename);
        let mut f = match fs::File::open(filename) {
            Ok(f) => f,
            Err(_) => {
                error!("failed to open file: {}", filename);
                return Err(Error::FileAccess);
            }
        };
        let bytes_read = match f.read_to_end(&mut buf) {
            Ok(c) => c,
            Err(_) => {
                error!("failed to read entire file: {}", filename);
                return Err(Error::FileAccess);
            }
        };
        debug!("read {} bytes", bytes_read);
        if bytes_read < 0x10 {
            error!("file too small: {}", filename);
            return Err(Error::FileFormat);
        }
    }

    Ok(buf)
}

pub struct Section {
    pub name: String,
    pub addr: u32,
    pub buf: Vec<u8>,
}

pub struct Workspace {
    pub filename: String,
    pub buf: Vec<u8>,
    pub sections: Vec<Section>,
}

impl Workspace {
    /// Parse the given file into its object.
    ///
    /// # Errors
    ///   - `Error::FileFormat`: when not able to be parsed by Goblin.
    ///
    /// # Examples
    ///
    /// ```
    /// use goblin::Object;
    /// use matches::matches;
    /// use lancelot::rsrc::*;
    /// let ws = get_workspace(Rsrc::K32);
    /// assert!(matches!(ws.get_obj().unwrap(), Object::PE(_)));
    /// ```
    ///
    /// you might be tempted to maintain a method `get_pe`,
    /// however, i don't think this is a good idea:
    /// its fragile because the file type may not be PE.
    /// therefore, force clients to be explicit:
    ///
    /// ```
    /// use goblin::Object;
    /// use lancelot::rsrc::*;
    /// let ws = get_workspace(Rsrc::K32);
    /// if let Object::PE(_) = ws.get_obj().unwrap() {
    ///     // everyone is happy!
    /// }
    /// ```
    ///
    /// TODO: demonstrate `Error::FileFormat`.
    pub fn get_obj(&self) -> Result<Object, Error> {
        let obj = match Object::parse(&self.buf) {
            Ok(o) => o,
            Err(e) => {
                error!("failed to parse file: {} error: {:?}", self.filename, e);
                return Err(Error::FileFormat);
            }
        };

        match obj {
            Object::Unknown(_) => {
                error!(
                    "unknown file format, magic: | {:02X} {:02X} | '{}{}' ",
                    self.buf[0],
                    self.buf[1],
                    hexdump_ascii(self.buf[0]),
                    hexdump_ascii(self.buf[1])
                );

                Err(Error::FileFormat)
            }
            _ => Ok(obj),
        }
    }

    pub fn get_pe(&self) -> Result<PE, Error> {
        match self.get_obj()? {
            Object::PE(pe) => {
                return Ok(pe);
            }
            _ => {
                error!("not a PE file");
                return Err(Error::ValueError);
            }
        }
    }

    /// Construct a workspace from the module with the given contents.
    ///
    /// # Errors
    ///   - `Error::FileFormat`: when not able to be parsed by Goblin.
    ///   - `Error::NotImplemented`: when not a PE file.
    ///
    /// # Examples
    ///
    /// ```
    /// use lancelot::*;
    /// use lancelot::rsrc::*;
    /// let buf = get_buf(Rsrc::K32);
    /// let ws = Workspace::from_buf("kernel32.dll", buf).unwrap();
    /// ```
    ///
    /// TODO: demonstrate ELF file behavior.
    /// TODO: demonstrate MachO file behavior.
    /// TODO: demonstrate unknown file behavior.
    pub fn from_buf(filename: &str, buf: Vec<u8>) -> Result<Workspace, Error> {
        let mut ws = Workspace {
            filename: filename.to_string(),
            buf: buf.clone(),
            sections: vec![],
        };

        match ws.get_obj()? {
            Object::PE(pe) => {
                info!("found PE file");

                // TODO: load PE header, too

                ws.sections
                    .extend(pe.sections.iter().map(|section| -> Section {
                        // TODO: i'm sure this can be abused.
                        // TODO: add tests for weird section names.
                        let name = String::from_utf8_lossy(&section.name[..]).into_owned();

                        // TODO: figure out if we will work with usize, or u64, or what,
                        // then assert usize is ok.
                        // ref: `usize::max_value()`
                        let mut secbuf = vec![0; align(section.virtual_size as usize, 0x200)];

                        {
                            let secsize = section.size_of_raw_data as usize;
                            let rawbuf = &mut secbuf[..secsize];
                            let pstart = section.pointer_to_raw_data as usize;
                            rawbuf.copy_from_slice(&buf[pstart..pstart + secsize]);
                        }

                        Section {
                            name: name,
                            addr: section.virtual_address,
                            buf: secbuf,
                        }
                    }));
            }
            Object::Elf(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Mach(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Archive(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Unknown(_) => {
                return Err(Error::NotImplemented);
            }
        }

        Ok(ws)
    }

    /// Construct a workspace from the module at given a file path.
    ///
    /// ```
    /// use lancelot::*;
    /// use lancelot::rsrc::*;
    /// let path = get_path(Rsrc::K32);
    /// // This test resource file is mangled. Needs to be fixed before parsing.
    /// // Otherwise, the following would work:
    /// // let ws = Workspace::from_file(&path).unwrap();
    /// ```
    pub fn from_file(filename: &str) -> Result<Workspace, Error> {
        let buf = read_file(filename)?;
        Workspace::from_buf(filename, buf)
    }
}

pub fn run(args: &Config) -> Result<(), Error> {
    debug!("filename: {:?}", args.filename);
    let ws = Workspace::from_file(&args.filename)?;

    if let Object::PE(_) = ws.get_obj()? {
        foo(&ws).expect("failed to foo");
    }

    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![cfg_attr(feature = "unstable", feature(fn_must_use))]
#![cfg_attr(not(test), no_std)]

//! A collection of strongly typed math tools for computer graphics with an inclination
//! towards 2d graphics and layout.
//!
//! All types are generic over the scalar type of their component (`f32`, `i32`, etc.),
//! and tagged with a generic Unit parameter which is useful to prevent mixing
//! values from different spaces. For example it should not be legal to translate
//! a screen-space position by a world-space vector and this can be expressed using
//! the generic Unit parameter.
//!
//! This unit system is not mandatory and all Typed* structures have an alias
//! with the default unit: `UnknownUnit`.
//! for example ```Point2D<T>``` is equivalent to ```TypedPoint2D<T, UnknownUnit>```.
//! Client code typically creates a set of aliases for each type and doesn't need
//! to deal with the specifics of typed units further. For example:
//!
//! ```rust
//! use euclid::*;
//! pub struct ScreenSpace;
//! pub type ScreenPoint = TypedPoint2D<f32, ScreenSpace>;
//! pub type ScreenSize = TypedSize2D<f32, ScreenSpace>;
//! pub struct WorldSpace;
//! pub type WorldPoint = TypedPoint3D<f32, WorldSpace>;
//! pub type ProjectionMatrix = TypedTransform3D<f32, WorldSpace, ScreenSpace>;
//! // etc...
//! ```
//!
//! All euclid types are marked `#[repr(C)]` in order to facilitate exposing them to
//! foreign function interfaces (provided the underlying scalar type is also `repr(C)`).
//!
//! Components are accessed in their scalar form by default for convenience, and most
//! types additionally implement strongly typed accessors which return typed ```Length``` wrappers.
//! For example:
//!
//! ```rust
//! # use euclid::*;
//! # pub struct WorldSpace;
//! # pub type WorldPoint = TypedPoint3D<f32, WorldSpace>;
//! let p = WorldPoint::new(0.0, 1.0, 1.0);
//! // p.x is an f32.
//! println!("p.x = {:?} ", p.x);
//! // p.x is a Length<f32, WorldSpace>.
//! println!("p.x_typed() = {:?} ", p.x_typed());
//! // Length::get returns the scalar value (f32).
//! assert_eq!(p.x, p.x_typed().get());
//! ```

#[cfg(feature = "serde")]
#[macro_use]
extern crate serde;

#[cfg(feature = "mint")]
pub extern crate mint;
extern crate num_traits;
#[cfg(test)]
extern crate rand;
#[cfg(test)]
use std as core;

pub use box2d::{TypedBox2D, Box2D};
pub use length::Length;
pub use scale::TypedScale;
pub use transform2d::{Transform2D, TypedTransform2D};
pub use transform3d::{Transform3D, TypedTransform3D};
pub use point::{Point2D, Point3D, TypedPoint2D, TypedPoint3D, point2, point3};
pub use vector::{TypedVector2D, TypedVector3D, Vector2D, Vector3D, vec2, vec3};
pub use vector::{BoolVector2D, BoolVector3D, bvec2, bvec3};
pub use homogen::HomogeneousVector;

pub use rect::{rect, Rect, TypedRect};
pub use rigid::{RigidTransform3D, TypedRigidTransform3D};
pub use box3d::{box3d, Box3D, TypedBox3D};
pub use translation::{TypedTranslation2D, TypedTranslation3D};
pub use rotation::{Angle, Rotation2D, Rotation3D, TypedRotation2D, TypedRotation3D};
pub use side_offsets::{SideOffsets2D, TypedSideOffsets2D};
pub use size::{Size2D, TypedSize2D, size2};
pub use trig::Trig;

#[macro_use]
mod macros;

pub mod approxeq;
pub mod approxord;
mod box2d;
mod homogen;
pub mod num;
mod length;
mod point;
mod rect;
mod rigid;
mod rotation;
mod scale;
mod side_offsets;
mod size;
mod transform2d;
mod transform3d;
mod translation;
mod trig;
mod vector;
mod box3d;

/// The default unit.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct UnknownUnit;

/// Temporary alias to facilitate the transition to the new naming scheme
#[deprecated]
pub type Matrix2D<T> = Transform2D<T>;

/// Temporary alias to facilitate the transition to the new naming scheme
#[deprecated]
pub type TypedMatrix2D<T, Src, Dst> = TypedTransform2D<T, Src, Dst>;

/// Temporary alias to facilitate the transition to the new naming scheme
#[deprecated]
pub type Matrix4D<T> = Transform3D<T>;

/// Temporary alias to facilitate the transition to the new naming scheme
#[deprecated]
pub type TypedMatrix4D<T, Src, Dst> = TypedTransform3D<T, Src, Dst>;

/// Temporary alias to facilitate the transition to the new naming scheme
#[deprecated]
pub type ScaleFactor<T, Src, Dst> = TypedScale<T, Src, Dst>;

/// Temporary alias to facilitate the transition to the new naming scheme
#[deprecated]
pub use Angle as Radians;
<commit_msg>Add a default module with type aliases for the default unknown units.<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![cfg_attr(feature = "unstable", feature(fn_must_use))]
#![cfg_attr(not(test), no_std)]

//! A collection of strongly typed math tools for computer graphics with an inclination
//! towards 2d graphics and layout.
//!
//! All types are generic over the scalar type of their component (`f32`, `i32`, etc.),
//! and tagged with a generic Unit parameter which is useful to prevent mixing
//! values from different spaces. For example it should not be legal to translate
//! a screen-space position by a world-space vector and this can be expressed using
//! the generic Unit parameter.
//!
//! This unit system is not mandatory and all Typed* structures have an alias
//! with the default unit: `UnknownUnit`.
//! for example ```Point2D<T>``` is equivalent to ```TypedPoint2D<T, UnknownUnit>```.
//! Client code typically creates a set of aliases for each type and doesn't need
//! to deal with the specifics of typed units further. For example:
//!
//! ```rust
//! use euclid::*;
//! pub struct ScreenSpace;
//! pub type ScreenPoint = TypedPoint2D<f32, ScreenSpace>;
//! pub type ScreenSize = TypedSize2D<f32, ScreenSpace>;
//! pub struct WorldSpace;
//! pub type WorldPoint = TypedPoint3D<f32, WorldSpace>;
//! pub type ProjectionMatrix = TypedTransform3D<f32, WorldSpace, ScreenSpace>;
//! // etc...
//! ```
//!
//! All euclid types are marked `#[repr(C)]` in order to facilitate exposing them to
//! foreign function interfaces (provided the underlying scalar type is also `repr(C)`).
//!
//! Components are accessed in their scalar form by default for convenience, and most
//! types additionally implement strongly typed accessors which return typed ```Length``` wrappers.
//! For example:
//!
//! ```rust
//! # use euclid::*;
//! # pub struct WorldSpace;
//! # pub type WorldPoint = TypedPoint3D<f32, WorldSpace>;
//! let p = WorldPoint::new(0.0, 1.0, 1.0);
//! // p.x is an f32.
//! println!("p.x = {:?} ", p.x);
//! // p.x is a Length<f32, WorldSpace>.
//! println!("p.x_typed() = {:?} ", p.x_typed());
//! // Length::get returns the scalar value (f32).
//! assert_eq!(p.x, p.x_typed().get());
//! ```

#[cfg(feature = "serde")]
#[macro_use]
extern crate serde;

#[cfg(feature = "mint")]
pub extern crate mint;
extern crate num_traits;
#[cfg(test)]
extern crate rand;
#[cfg(test)]
use std as core;

pub use box2d::{TypedBox2D, Box2D};
pub use length::Length;
pub use scale::TypedScale;
pub use transform2d::{Transform2D, TypedTransform2D};
pub use transform3d::{Transform3D, TypedTransform3D};
pub use point::{Point2D, Point3D, TypedPoint2D, TypedPoint3D, point2, point3};
pub use vector::{TypedVector2D, TypedVector3D, Vector2D, Vector3D, vec2, vec3};
pub use vector::{BoolVector2D, BoolVector3D, bvec2, bvec3};
pub use homogen::HomogeneousVector;

pub use rect::{rect, Rect, TypedRect};
pub use rigid::{RigidTransform3D, TypedRigidTransform3D};
pub use box3d::{box3d, Box3D, TypedBox3D};
pub use translation::{TypedTranslation2D, TypedTranslation3D};
pub use rotation::{Angle, Rotation2D, Rotation3D, TypedRotation2D, TypedRotation3D};
pub use side_offsets::{SideOffsets2D, TypedSideOffsets2D};
pub use size::{Size2D, TypedSize2D, TypedSize3D, size2, size3};
pub use trig::Trig;

#[macro_use]
mod macros;

pub mod approxeq;
pub mod approxord;
mod box2d;
mod homogen;
pub mod num;
mod length;
mod point;
mod rect;
mod rigid;
mod rotation;
mod scale;
mod side_offsets;
mod size;
mod transform2d;
mod transform3d;
mod translation;
mod trig;
mod vector;
mod box3d;

/// The default unit.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct UnknownUnit;

/// Temporary alias to facilitate the transition to the new naming scheme
#[deprecated]
pub type Matrix2D<T> = Transform2D<T>;

/// Temporary alias to facilitate the transition to the new naming scheme
#[deprecated]
pub type TypedMatrix2D<T, Src, Dst> = TypedTransform2D<T, Src, Dst>;

/// Temporary alias to facilitate the transition to the new naming scheme
#[deprecated]
pub type Matrix4D<T> = Transform3D<T>;

/// Temporary alias to facilitate the transition to the new naming scheme
#[deprecated]
pub type TypedMatrix4D<T, Src, Dst> = TypedTransform3D<T, Src, Dst>;

/// Temporary alias to facilitate the transition to the new naming scheme
#[deprecated]
pub type ScaleFactor<T, Src, Dst> = TypedScale<T, Src, Dst>;

/// Temporary alias to facilitate the transition to the new naming scheme
#[deprecated]
pub use Angle as Radians;

pub mod default {
    use super::*;

    pub type Point2D<T> = TypedPoint2D<T, UnknownUnit>;
    pub type Point3D<T> = TypedPoint3D<T, UnknownUnit>;
    pub type Vector2D<T> = TypedVector2D<T, UnknownUnit>;
    pub type Vector3D<T> = TypedVector3D<T, UnknownUnit>;
    pub type Vector4D<T> = HomogeneousVector<T, UnknownUnit>;
    pub type Size2D<T> = TypedSize2D<T, UnknownUnit>;
    pub type Size3D<T> = TypedSize3D<T, UnknownUnit>;
    pub type Rect<T> = TypedRect<T, UnknownUnit>;
    pub type Box2D<T> = TypedBox2D<T, UnknownUnit>;
    pub type Box3D<T> = TypedBox3D<T, UnknownUnit>;
    pub type SideOffsets2D<T> = TypedSideOffsets2D<T, UnknownUnit>;
    pub type Tranform2D<T> = TypedTransform2D<T, UnknownUnit, UnknownUnit>;
    pub type Tranform3D<T> = TypedTransform3D<T, UnknownUnit, UnknownUnit>;
    pub type Rotation2D<T> = TypedRotation2D<T, UnknownUnit, UnknownUnit>;
    pub type Rotation3D<T> = TypedRotation3D<T, UnknownUnit, UnknownUnit>;
    pub type Translation3D<T> = TypedTranslation3D<T, UnknownUnit, UnknownUnit>;
    pub type Scale<T> = TypedScale<T, UnknownUnit, UnknownUnit>;
    pub type RigidTransform3D<T> = TypedRigidTransform3D<T, UnknownUnit, UnknownUnit>;
}
<|endoftext|>


<commit_before>use crate::config::{Config, EnvConfig};
use crate::error::{Error, ErrorKind, Result};
use crate::page::parse_page;
use crate::service;
use std::process::{Command, Stdio};

pub fn build_url(cfg: &Config) -> Result<String> {
    let page = parse_page(&cfg)?;
    service::build_page_url(&page, &cfg)
}

fn browse_with_cmd(url: &str, cmd: &str) -> Result<()> {
    let out = Command::new(cmd)
        .arg(url)
        .stdout(Stdio::inherit())
        .output()?;
    if out.status.success() {
        Ok(())
    } else {
        Error::err(ErrorKind::UserBrowseCommandFailed {
            cmd: cmd.to_string(),
            url: url.to_string(),
            msg: String::from_utf8_lossy(&out.stderr)
                .trim()
                .replace('\n', " "),
        })
    }
}

pub fn browse(url: &str, env: &EnvConfig) -> Result<()> {
    if let Some(ref cmd) = env.browse_command {
        return browse_with_cmd(url, cmd);
    }

    match open::that(url) {
        Ok(status) if status.success() => Ok(()),
        Ok(status) => {
            let url = url.to_string();
            let msg = if let Some(code) = status.code() {
                format!("Command exited with non-zero status {}", code)
            } else if cfg!(unix) {
                use std::os::unix::process::ExitStatusExt;
                if let Some(sig) = status.signal() {
                    format!("Command terminated by signal {}", sig)
                } else {
                    "Command terminated by signal".to_string()
                }
            } else {
                "Command terminated by signal".to_string()
            };
            Error::err(ErrorKind::OpenUrlFailure { url, msg })
        }
        Err(e) => Error::err(ErrorKind::OpenUrlFailure {
            url: url.to_string(),
            msg: format!("{}", e),
        }),
    }
}
<commit_msg>fix build on Windows<commit_after>use crate::config::{Config, EnvConfig};
use crate::error::{Error, ErrorKind, Result};
use crate::page::parse_page;
use crate::service;
use std::process::{Command, ExitStatus, Stdio};

pub fn build_url(cfg: &Config) -> Result<String> {
    let page = parse_page(&cfg)?;
    service::build_page_url(&page, &cfg)
}

fn browse_with_cmd(url: &str, cmd: &str) -> Result<()> {
    let out = Command::new(cmd)
        .arg(url)
        .stdout(Stdio::inherit())
        .output()?;
    if out.status.success() {
        Ok(())
    } else {
        Error::err(ErrorKind::UserBrowseCommandFailed {
            cmd: cmd.to_string(),
            url: url.to_string(),
            msg: String::from_utf8_lossy(&out.stderr)
                .trim()
                .replace('\n', " "),
        })
    }
}

#[cfg(unix)]
fn error_without_status(status: ExitStatus) -> String {
    use std::os::unix::process::ExitStatusExt;
    if let Some(sig) = status.signal() {
        format!("Command terminated by signal {}", sig)
    } else {
        "Command terminated by signal".to_string()
    }
}

#[cfg(not(unix))]
fn error_without_status(status: ExitStatus) -> String {
    "Command terminated by signal".to_string()
}

pub fn browse(url: &str, env: &EnvConfig) -> Result<()> {
    if let Some(ref cmd) = env.browse_command {
        return browse_with_cmd(url, cmd);
    }

    match open::that(url) {
        Ok(status) if status.success() => Ok(()),
        Ok(status) => {
            let url = url.to_string();
            let msg = if let Some(code) = status.code() {
                format!("Command exited with non-zero status {}", code)
            } else {
                error_without_status(status)
            };
            Error::err(ErrorKind::OpenUrlFailure { url, msg })
        }
        Err(e) => Error::err(ErrorKind::OpenUrlFailure {
            url: url.to_string(),
            msg: format!("{}", e),
        }),
    }
}
<|endoftext|>


<commit_before>//! The master module provides the Master struct which instructs Workers which
//! portions of the image they should render and collects their results to combine
//! into the final image.

use std::path::PathBuf;
use std::io::prelude::*;
use std::collections::HashMap;
use std::net::ToSocketAddrs;
use std::iter;

use bincode::rustc_serialize::decode;
use image;
use mio::tcp::{TcpStream, Shutdown};
use mio::*;

use film::Image;
use exec::Config;
use exec::distrib::{worker, Instructions};
use sampler::BlockQueue;

pub struct DistributedFrame {
    pub frame: usize,
    // The number of workers who have reported results for this
    // frame so far
    pub num_reporting: usize,
    pub render: Image,
    pub completed: bool,
}

#[derive(Clone, Debug)]
struct WorkerBuffer {
    pub buf: Vec<u8>,
    pub expected_size: usize,
    pub currently_read: usize,
}

impl WorkerBuffer {
    pub fn new() -> WorkerBuffer {
        WorkerBuffer { buf: Vec::new(), expected_size: 8, currently_read: 0 }
    }
}

pub struct Master {
    // Hostnames of the workers to send work too
    workers: Vec<String>,
    connections: Vec<TcpStream>,
    // Temporary buffers to store worker results in as they're
    // read in over TCP
    worker_buffers: Vec<WorkerBuffer>,
    config: Config,
    frames: HashMap<usize, DistributedFrame>,
    selected_frames: Option<(usize, usize)>,
    img_dim: (usize, usize),
    blocks_per_worker: usize,
    blocks_remainder: usize,
}

impl Master {
    /// Create a new master that will contact the worker nodes passed and
    /// send instructions on what parts of the scene to start rendering
    pub fn start_workers(workers: Vec<String>, config: Config, img_dim: (usize, usize),
                         frames: Option<(usize, usize)>) -> (Master, EventLoop<Master>) {
        // Figure out how many blocks we have for this image and assign them to our workers
        let queue = BlockQueue::new((img_dim.0 as u32, img_dim.1 as u32), (8, 8), (0, 0));
        let blocks_per_worker = queue.len() / workers.len();
        let blocks_remainder = queue.len() % workers.len();
        println!("Have {} workers, each worker does {} blocks w/ {} remainder",
                 workers.len(), blocks_per_worker, blocks_remainder);

        let mut event_loop = EventLoop::<Master>::new().unwrap();
        let mut connections = Vec::new();

        // Connect to each worker and send instructions on what to render
        // I guess this should also be queue'd up and the writes actually performed asynchronously?
        // Otherwise we need to loop on write_all while waiting for the connection to open
        for (i, host) in workers.iter().enumerate() {
            let addr = (&host[..], worker::PORT).to_socket_addrs().unwrap().next().unwrap();
            println!("Connecting too {:?}", addr);
            match TcpStream::connect(&addr) {
                Ok(stream) => {
                    println!("Connected and writing");
                    println!("Adding stream to event loop");
                    // Each worker is identified in the event loop by their index in the vec
                    match event_loop.register(&stream, Token(i)) {
                        Err(e) => println!("Error registering stream from {}: {}", host, e),
                        Ok(_) => {},
                    }
                    connections.push(stream);
                },
                Err(e) => println!("Failed to contact worker {}: {:?}", host, e),
            }
        }
        let worker_buffers: Vec<_> = iter::repeat(WorkerBuffer::new()).take(workers.len()).collect();
        let master = Master { workers: workers, connections: connections,
                              worker_buffers: worker_buffers, config: config,
                              frames: HashMap::new(), selected_frames: frames,
                              img_dim: img_dim,
                              blocks_per_worker: blocks_per_worker,
                              blocks_remainder: blocks_remainder };
        (master, event_loop)
    }
    /// Read a result frame from a worker and save it into the list of frames we're collecting from
    /// all workers. Will save out the final render if all workers have reported results for this
    /// frame.
    fn save_results(&mut self, frame: worker::Frame) {
        // Find the frame this worker is referring to and add its results
        let frame_num = frame.frame as usize;
        println!("Worker reporting frame {}", frame_num);
        if let Some(df) = self.frames.get_mut(&frame_num) {
            df.render.add_blocks(frame.block_size, &frame.blocks, &frame.pixels);
            df.num_reporting += 1;
            println!("We have parts of this frame already");
        }
        // TODO: Better way here? We get stuck because self.frames is counted as borrowed
        // from the if above, making this really awkward.
        if self.frames.get_mut(&frame_num).is_none() {
            // TODO: I need to seperate the render target from the actual image as the
            // filter here isn't used in this path and it doesn't make sense to have one
            let render = Image::new(self.img_dim);
            let mut df = DistributedFrame { frame: frame_num, num_reporting: 1, render: render, completed: false };
            df.render.add_blocks(frame.block_size, &frame.blocks, &frame.pixels);
            self.frames.insert(df.frame, df);
            println!("This is a new frame");
        }
        if let Some(df) = self.frames.get_mut(&frame_num) {
            if df.num_reporting == self.workers.len() {
                df.completed = true;
                println!("We have completed frame {}", frame_num);
                let out_file = match self.config.out_path.extension() {
                    Some(_) => self.config.out_path.clone(),
                    None => self.config.out_path.join(PathBuf::from(format!("frame{:05}.png", df.frame))),
                };
                println!("Frame {}: rendered to '{}'\n--------------------", frame_num, out_file.display());
                let img = df.render.get_srgb8();
                let dim = df.render.dimensions();
                match image::save_buffer(&out_file.as_path(), &img[..], dim.0 as u32, dim.1 as u32, image::RGB(8)) {
                    Ok(_) => {},
                    Err(e) => println!("Error saving image, {}", e),
                };
            }
        }
    }
    fn read_worker_buffer(&mut self, worker: usize) -> bool {
        let mut buf = &mut self.worker_buffers[worker];
        // If we haven't read the size of data being sent, read that now
        if buf.currently_read < 8 {
            // First 8 bytes are a u64 specifying the number of bytes being sent
            buf.buf.extend(iter::repeat(0u8).take(8));
            println!("Reading size header...");
            match self.connections[worker].read(&mut buf.buf[buf.currently_read..]) {
                Ok(n) => buf.currently_read += n,
                Err(e) => println!("Error reading results from worker {}: {}", self.workers[worker], e),
            }
            println!("Read {} bytes so far, must read {} in total", buf.currently_read, buf.expected_size);
            if buf.currently_read == buf.expected_size {
                println!("Got size header");
                // How many bytes we expect to get from the worker for a frame
                buf.expected_size = decode(&buf.buf[..]).unwrap();
                println!("Expecting {} bytes", buf.expected_size);
                // Extend the Vec so we've got enough room for the remaning bytes, minus the 8 for the
                // encoded size header
                buf.buf.extend(iter::repeat(0u8).take(buf.expected_size - 8));
            }
        }
        // If we've finished reading the size header we can now start reading the frame data
        if buf.currently_read >= 8 {
            println!("Reading from worker");
            match self.connections[worker].read(&mut buf.buf[buf.currently_read..]) {
                Ok(n) => buf.currently_read += n,
                Err(e) => println!("Error reading results from worker {}: {}", self.workers[worker], e),
            }
            println!("Read {} bytes so far, must read {} in total", buf.currently_read, buf.expected_size);
        }
        buf.currently_read == buf.expected_size
    }
}

impl Handler for Master {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Master>, token: Token, event: EventSet) {
        let worker = token.as_usize();
        println!("Event from {}: {:?}", worker, event);
        if event.is_error() {
            println!("Error connecting too {}", self.workers[worker]);
            match self.connections[worker].shutdown(Shutdown::Both) {
                Err(e) => println!("Error shutting down worker {}: {}", worker, e),
                Ok(_) => {},
            }
            // Remove the connection from the event loop
            match event_loop.deregister(&self.connections[worker]) {
                Err(e) => println!("Error deregistering worker {}: {}", worker, e),
                Ok(_) => {},
            }
        }
        // A worker is ready to receive instructions from us
        if event.is_writable() {
            let b_start = worker * self.blocks_per_worker;
            let b_count =
                if worker == self.workers.len() - 1 {
                    self.blocks_per_worker + self.blocks_remainder
                } else {
                    self.blocks_per_worker
                };
            let instr = Instructions::new(&self.config.scene_file, self.selected_frames, b_start, b_count);
            println!("Sending instructions {:?} to {}", instr, self.workers[worker]);
            let bytes = instr.to_json().into_bytes();
            // Loop until we successfully write the byes
            match self.connections[worker].write_all(&bytes[..]) {
                Err(e) => {
                    println!("Failed to send instructions to {}: {:?}", self.workers[worker], e);
                    return;
                },
                Ok(_) => println!("Instructions sent"),
            }
            // Register that we no longer care about writable events on this connection
            event_loop.reregister(&self.connections[worker], token,
                                  EventSet::readable() | EventSet::error() | EventSet::hup(),
                                  PollOpt::level()).expect("Re-registering failed");
            // We no longer need to write anything, so close the write end
            match self.connections[worker].shutdown(Shutdown::Write) {
                Err(e) => panic!(format!("Failed to shutdown write end to worker {}: {}",
                                         self.workers[worker], e)),
                Ok(_) => {},
            }
        }
        // Some results are available from a worker
        if event.is_readable() {
            println!("Readable event from worker {}", worker);
            if self.read_worker_buffer(worker) {
                println!("Got all frame data, now parsing");
                let frame: worker::Frame = decode(&self.worker_buffers[worker].buf[..]).unwrap();
                println!("Saving frame results");
                self.save_results(frame);
                // Clean up our read buffer
                self.worker_buffers[worker].buf.clear();
                self.worker_buffers[worker].expected_size = 8;
                self.worker_buffers[worker].currently_read = 0;
            }
        }
        // If the worker has terminated, shutdown the read end of the connection
        if event.is_hup() {
            println!("Worker {} has hung up", worker);
            match self.connections[worker].shutdown(Shutdown::Both) {
                Err(e) => println!("Error shutting down worker {}: {}", worker, e),
                Ok(_) => {},
            }
            // Remove the connection from the event loop
            match event_loop.deregister(&self.connections[worker]) {
                Err(e) => println!("Error deregistering worker {}: {}", worker, e),
                Ok(_) => {},
            }
        }
        // After getting results from the worker we check if we've completed all our frames
        // and exit if so
        if self.frames.len() == self.config.frame_info.frames
            && self.frames.values().fold(true, |all, v| all && v.completed) {
                println!("All frames have been rendered, master exiting");
                event_loop.shutdown();
        }
    }
}

<commit_msg>Just need to shutdown read end since we already closed write<commit_after>//! The master module provides the Master struct which instructs Workers which
//! portions of the image they should render and collects their results to combine
//! into the final image.

use std::path::PathBuf;
use std::io::prelude::*;
use std::collections::HashMap;
use std::net::ToSocketAddrs;
use std::iter;

use bincode::rustc_serialize::decode;
use image;
use mio::tcp::{TcpStream, Shutdown};
use mio::*;

use film::Image;
use exec::Config;
use exec::distrib::{worker, Instructions};
use sampler::BlockQueue;

pub struct DistributedFrame {
    pub frame: usize,
    // The number of workers who have reported results for this
    // frame so far
    pub num_reporting: usize,
    pub render: Image,
    pub completed: bool,
}

#[derive(Clone, Debug)]
struct WorkerBuffer {
    pub buf: Vec<u8>,
    pub expected_size: usize,
    pub currently_read: usize,
}

impl WorkerBuffer {
    pub fn new() -> WorkerBuffer {
        WorkerBuffer { buf: Vec::new(), expected_size: 8, currently_read: 0 }
    }
}

pub struct Master {
    // Hostnames of the workers to send work too
    workers: Vec<String>,
    connections: Vec<TcpStream>,
    // Temporary buffers to store worker results in as they're
    // read in over TCP
    worker_buffers: Vec<WorkerBuffer>,
    config: Config,
    frames: HashMap<usize, DistributedFrame>,
    selected_frames: Option<(usize, usize)>,
    img_dim: (usize, usize),
    blocks_per_worker: usize,
    blocks_remainder: usize,
}

impl Master {
    /// Create a new master that will contact the worker nodes passed and
    /// send instructions on what parts of the scene to start rendering
    pub fn start_workers(workers: Vec<String>, config: Config, img_dim: (usize, usize),
                         frames: Option<(usize, usize)>) -> (Master, EventLoop<Master>) {
        // Figure out how many blocks we have for this image and assign them to our workers
        let queue = BlockQueue::new((img_dim.0 as u32, img_dim.1 as u32), (8, 8), (0, 0));
        let blocks_per_worker = queue.len() / workers.len();
        let blocks_remainder = queue.len() % workers.len();
        println!("Have {} workers, each worker does {} blocks w/ {} remainder",
                 workers.len(), blocks_per_worker, blocks_remainder);

        let mut event_loop = EventLoop::<Master>::new().unwrap();
        let mut connections = Vec::new();

        // Connect to each worker and send instructions on what to render
        // I guess this should also be queue'd up and the writes actually performed asynchronously?
        // Otherwise we need to loop on write_all while waiting for the connection to open
        for (i, host) in workers.iter().enumerate() {
            let addr = (&host[..], worker::PORT).to_socket_addrs().unwrap().next().unwrap();
            println!("Connecting too {:?}", addr);
            match TcpStream::connect(&addr) {
                Ok(stream) => {
                    println!("Connected and writing");
                    println!("Adding stream to event loop");
                    // Each worker is identified in the event loop by their index in the vec
                    match event_loop.register(&stream, Token(i)) {
                        Err(e) => println!("Error registering stream from {}: {}", host, e),
                        Ok(_) => {},
                    }
                    connections.push(stream);
                },
                Err(e) => println!("Failed to contact worker {}: {:?}", host, e),
            }
        }
        let worker_buffers: Vec<_> = iter::repeat(WorkerBuffer::new()).take(workers.len()).collect();
        let master = Master { workers: workers, connections: connections,
                              worker_buffers: worker_buffers, config: config,
                              frames: HashMap::new(), selected_frames: frames,
                              img_dim: img_dim,
                              blocks_per_worker: blocks_per_worker,
                              blocks_remainder: blocks_remainder };
        (master, event_loop)
    }
    /// Read a result frame from a worker and save it into the list of frames we're collecting from
    /// all workers. Will save out the final render if all workers have reported results for this
    /// frame.
    fn save_results(&mut self, frame: worker::Frame) {
        // Find the frame this worker is referring to and add its results
        let frame_num = frame.frame as usize;
        println!("Worker reporting frame {}", frame_num);
        if let Some(df) = self.frames.get_mut(&frame_num) {
            df.render.add_blocks(frame.block_size, &frame.blocks, &frame.pixels);
            df.num_reporting += 1;
            println!("We have parts of this frame already");
        }
        // TODO: Better way here? We get stuck because self.frames is counted as borrowed
        // from the if above, making this really awkward.
        if self.frames.get_mut(&frame_num).is_none() {
            // TODO: I need to seperate the render target from the actual image as the
            // filter here isn't used in this path and it doesn't make sense to have one
            let render = Image::new(self.img_dim);
            let mut df = DistributedFrame { frame: frame_num, num_reporting: 1, render: render, completed: false };
            df.render.add_blocks(frame.block_size, &frame.blocks, &frame.pixels);
            self.frames.insert(df.frame, df);
            println!("This is a new frame");
        }
        if let Some(df) = self.frames.get_mut(&frame_num) {
            if df.num_reporting == self.workers.len() {
                df.completed = true;
                println!("We have completed frame {}", frame_num);
                let out_file = match self.config.out_path.extension() {
                    Some(_) => self.config.out_path.clone(),
                    None => self.config.out_path.join(PathBuf::from(format!("frame{:05}.png", df.frame))),
                };
                println!("Frame {}: rendered to '{}'\n--------------------", frame_num, out_file.display());
                let img = df.render.get_srgb8();
                let dim = df.render.dimensions();
                match image::save_buffer(&out_file.as_path(), &img[..], dim.0 as u32, dim.1 as u32, image::RGB(8)) {
                    Ok(_) => {},
                    Err(e) => println!("Error saving image, {}", e),
                };
            }
        }
    }
    fn read_worker_buffer(&mut self, worker: usize) -> bool {
        let mut buf = &mut self.worker_buffers[worker];
        // If we haven't read the size of data being sent, read that now
        if buf.currently_read < 8 {
            // First 8 bytes are a u64 specifying the number of bytes being sent
            buf.buf.extend(iter::repeat(0u8).take(8));
            println!("Reading size header...");
            match self.connections[worker].read(&mut buf.buf[buf.currently_read..]) {
                Ok(n) => buf.currently_read += n,
                Err(e) => println!("Error reading results from worker {}: {}", self.workers[worker], e),
            }
            println!("Read {} bytes so far, must read {} in total", buf.currently_read, buf.expected_size);
            if buf.currently_read == buf.expected_size {
                println!("Got size header");
                // How many bytes we expect to get from the worker for a frame
                buf.expected_size = decode(&buf.buf[..]).unwrap();
                println!("Expecting {} bytes", buf.expected_size);
                // Extend the Vec so we've got enough room for the remaning bytes, minus the 8 for the
                // encoded size header
                buf.buf.extend(iter::repeat(0u8).take(buf.expected_size - 8));
            }
        }
        // If we've finished reading the size header we can now start reading the frame data
        if buf.currently_read >= 8 {
            println!("Reading from worker");
            match self.connections[worker].read(&mut buf.buf[buf.currently_read..]) {
                Ok(n) => buf.currently_read += n,
                Err(e) => println!("Error reading results from worker {}: {}", self.workers[worker], e),
            }
            println!("Read {} bytes so far, must read {} in total", buf.currently_read, buf.expected_size);
        }
        buf.currently_read == buf.expected_size
    }
}

impl Handler for Master {
    type Timeout = ();
    type Message = ();

    fn ready(&mut self, event_loop: &mut EventLoop<Master>, token: Token, event: EventSet) {
        let worker = token.as_usize();
        println!("Event from {}: {:?}", worker, event);
        if event.is_error() {
            println!("Error connecting too {}", self.workers[worker]);
            match self.connections[worker].shutdown(Shutdown::Both) {
                Err(e) => println!("Error shutting down worker {}: {}", worker, e),
                Ok(_) => {},
            }
            // Remove the connection from the event loop
            match event_loop.deregister(&self.connections[worker]) {
                Err(e) => println!("Error deregistering worker {}: {}", worker, e),
                Ok(_) => {},
            }
        }
        // A worker is ready to receive instructions from us
        if event.is_writable() {
            let b_start = worker * self.blocks_per_worker;
            let b_count =
                if worker == self.workers.len() - 1 {
                    self.blocks_per_worker + self.blocks_remainder
                } else {
                    self.blocks_per_worker
                };
            let instr = Instructions::new(&self.config.scene_file, self.selected_frames, b_start, b_count);
            println!("Sending instructions {:?} to {}", instr, self.workers[worker]);
            let bytes = instr.to_json().into_bytes();
            // Loop until we successfully write the byes
            match self.connections[worker].write_all(&bytes[..]) {
                Err(e) => {
                    println!("Failed to send instructions to {}: {:?}", self.workers[worker], e);
                    return;
                },
                Ok(_) => println!("Instructions sent"),
            }
            // Register that we no longer care about writable events on this connection
            event_loop.reregister(&self.connections[worker], token,
                                  EventSet::readable() | EventSet::error() | EventSet::hup(),
                                  PollOpt::level()).expect("Re-registering failed");
            // We no longer need to write anything, so close the write end
            match self.connections[worker].shutdown(Shutdown::Write) {
                Err(e) => panic!(format!("Failed to shutdown write end to worker {}: {}",
                                         self.workers[worker], e)),
                Ok(_) => {},
            }
        }
        // Some results are available from a worker
        if event.is_readable() {
            println!("Readable event from worker {}", worker);
            if self.read_worker_buffer(worker) {
                println!("Got all frame data, now parsing");
                let frame: worker::Frame = decode(&self.worker_buffers[worker].buf[..]).unwrap();
                println!("Saving frame results");
                self.save_results(frame);
                // Clean up our read buffer
                self.worker_buffers[worker].buf.clear();
                self.worker_buffers[worker].expected_size = 8;
                self.worker_buffers[worker].currently_read = 0;
            }
        }
        // If the worker has terminated, shutdown the read end of the connection
        if event.is_hup() {
            println!("Worker {} has hung up", worker);
            match self.connections[worker].shutdown(Shutdown::Read) {
                Err(e) => println!("Error shutting down worker {}: {}", worker, e),
                Ok(_) => {},
            }
            // Remove the connection from the event loop
            match event_loop.deregister(&self.connections[worker]) {
                Err(e) => println!("Error deregistering worker {}: {}", worker, e),
                Ok(_) => {},
            }
        }
        // After getting results from the worker we check if we've completed all our frames
        // and exit if so
        if self.frames.len() == self.config.frame_info.frames
            && self.frames.values().fold(true, |all, v| all && v.completed) {
                println!("All frames have been rendered, master exiting");
                event_loop.shutdown();
        }
    }
}

<|endoftext|>


<commit_before><commit_msg>init color eyre<commit_after><|endoftext|>


<commit_before><commit_msg>Introduced context for game loop execution.<commit_after><|endoftext|>


<commit_before><commit_msg>Colorize diffs<commit_after><|endoftext|>


<commit_before><commit_msg>Remove useless comments<commit_after><|endoftext|>