<commit_before>// Copyright 2017 The Exonum Team
//
// 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::net::SocketAddr;
use std::thread;
use std::time::{self, Duration};

use futures::{Future, Sink, Stream};
use futures::stream::Wait;
use futures::sync::mpsc;
use tokio_core::reactor::Core;
use tokio_timer::{TimeoutStream, Timer};

use crypto::{gen_keypair, PublicKey, Signature};
use messages::{Connect, Message, MessageWriter, RawMessage};
use events::{NetworkEvent, NetworkRequest};
use events::network::{NetworkConfiguration, NetworkPart};
use events::error::log_error;
use node::{EventsPoolCapacity, NodeChannel};

#[derive(Debug)]
pub struct TestHandler {
    handle: Option<thread::JoinHandle<()>>,
    listen_address: SocketAddr,
    network_events_rx: Wait<TimeoutStream<mpsc::Receiver<NetworkEvent>>>,
    network_requests_tx: mpsc::Sender<NetworkRequest>,
}

impl TestHandler {
    pub fn new(
        listen_address: SocketAddr,
        network_requests_tx: mpsc::Sender<NetworkRequest>,
        network_events_rx: mpsc::Receiver<NetworkEvent>,
    ) -> TestHandler {
        let timer = Timer::default();
        let receiver = timer.timeout_stream(network_events_rx, Duration::from_secs(30));
        TestHandler {
            handle: None,
            listen_address,
            network_requests_tx,
            network_events_rx: receiver.wait(),
        }
    }

    pub fn wait_for_event(&mut self) -> Result<NetworkEvent, ()> {
        let event = self.network_events_rx.next().unwrap()?;
        Ok(event)
    }

    pub fn connect_with(&self, addr: SocketAddr) {
        let connect = connect_message(self.listen_address);
        self.network_requests_tx
            .clone()
            .send(NetworkRequest::SendMessage(addr, connect.raw().clone()))
            .wait()
            .unwrap();
    }

    pub fn disconnect_with(&self, addr: SocketAddr) {
        self.network_requests_tx
            .clone()
            .send(NetworkRequest::DisconnectWithPeer(addr))
            .wait()
            .unwrap();
    }

    pub fn send_to(&self, addr: SocketAddr, raw: RawMessage) {
        self.network_requests_tx
            .clone()
            .send(NetworkRequest::SendMessage(addr, raw))
            .wait()
            .unwrap();
    }

    pub fn wait_for_connect(&mut self) -> Connect {
        match self.wait_for_event() {
            Ok(NetworkEvent::PeerConnected(_addr, connect)) => connect,
            Ok(other) => panic!("Unexpected connect received, {:?}", other),
            Err(e) => panic!("An error during wait for connect occurred, {:?}", e),
        }
    }

    pub fn wait_for_disconnect(&mut self) -> SocketAddr {
        match self.wait_for_event() {
            Ok(NetworkEvent::PeerDisconnected(addr)) => addr,
            Ok(other) => panic!("Unexpected disconnect received, {:?}", other),
            Err(e) => panic!("An error during wait for disconnect occurred, {:?}", e),
        }
    }

    pub fn wait_for_message(&mut self) -> RawMessage {
        match self.wait_for_event() {
            Ok(NetworkEvent::MessageReceived(_addr, msg)) => msg,
            Ok(other) => panic!("Unexpected message received, {:?}", other),
            Err(e) => panic!("An error during wait for message occurred, {:?}", e),
        }
    }

    pub fn shutdown(&mut self) {
        self.network_requests_tx
            .clone()
            .send(NetworkRequest::Shutdown)
            .wait()
            .unwrap();
        self.handle.take().expect("shutdown twice").join().unwrap();
    }
}

impl Drop for TestHandler {
    fn drop(&mut self) {
        if !::std::thread::panicking() {
            self.shutdown();
        }
    }
}

#[derive(Debug)]
pub struct TestEvents {
    pub listen_address: SocketAddr,
    pub network_config: NetworkConfiguration,
    pub events_config: EventsPoolCapacity,
}

impl TestEvents {
    pub fn with_addr(listen_address: SocketAddr) -> TestEvents {
        TestEvents {
            listen_address,
            network_config: NetworkConfiguration::default(),
            events_config: EventsPoolCapacity::default(),
        }
    }

    pub fn spawn(self) -> TestHandler {
        let (mut handler_part, network_part) = self.into_reactor();
        let handle = thread::spawn(move || {
            let mut core = Core::new().unwrap();
            let fut = network_part.run(&core.handle());
            core.run(fut).map_err(log_error).unwrap();
        });
        handler_part.handle = Some(handle);
        handler_part
    }

    fn into_reactor(self) -> (TestHandler, NetworkPart) {
        let channel = NodeChannel::new(&self.events_config);
        let network_config = self.network_config;
        let (network_tx, network_rx) = channel.network_events;
        let network_requests_tx = channel.network_requests.0.clone();

        let network_part = NetworkPart {
            our_connect_message: connect_message(self.listen_address),
            listen_address: self.listen_address,
            network_config,
            max_message_len: ::blockchain::ConsensusConfig::DEFAULT_MESSAGE_MAX_LEN,
            network_requests: channel.network_requests,
            network_tx: network_tx.clone(),
        };

        let handler_part = TestHandler::new(self.listen_address, network_requests_tx, network_rx);
        (handler_part, network_part)
    }
}

pub fn connect_message(addr: SocketAddr) -> Connect {
    let time = time::UNIX_EPOCH;
    Connect::new_with_signature(&PublicKey::zero(), addr, time, &Signature::zero())
}

pub fn raw_message(id: u16, len: usize) -> RawMessage {
    let writer = MessageWriter::new(
        ::messages::PROTOCOL_MAJOR_VERSION,
        ::messages::TEST_NETWORK_ID,
        0,
        id,
        len,
    );
    RawMessage::new(writer.sign(&gen_keypair().1))
}

#[test]
fn test_network_handshake() {
    let first = "127.0.0.1:17230".parse().unwrap();
    let second = "127.0.0.1:17231".parse().unwrap();

    let e1 = TestEvents::with_addr(first);
    let e2 = TestEvents::with_addr(second);

    let c1 = connect_message(first);
    let c2 = connect_message(second);

    let mut e1 = e1.spawn();
    let mut e2 = e2.spawn();

    e1.connect_with(second);
    assert_eq!(e2.wait_for_connect(), c1);

    e2.connect_with(first);
    assert_eq!(e1.wait_for_connect(), c2);

    e1.disconnect_with(second);
    assert_eq!(e1.wait_for_disconnect(), second);

    e2.disconnect_with(first);
    assert_eq!(e2.wait_for_disconnect(), first);
}

#[test]
fn test_network_big_message() {
    let first = "127.0.0.1:17200".parse().unwrap();
    let second = "127.0.0.1:17201".parse().unwrap();

    let m1 = raw_message(15, 100000);
    let m2 = raw_message(16, 400);

    let e1 = TestEvents::with_addr(first);
    let e2 = TestEvents::with_addr(second);

    let mut e1 = e1.spawn();
    let mut e2 = e2.spawn();

    e1.connect_with(second);
    e2.wait_for_connect();

    e2.connect_with(first);
    e1.wait_for_connect();

    e1.send_to(second, m1.clone());
    assert_eq!(e2.wait_for_message(), m1);

    e1.send_to(second, m2.clone());
    assert_eq!(e2.wait_for_message(), m2);

    e1.send_to(second, m1.clone());
    assert_eq!(e2.wait_for_message(), m1);

    e2.send_to(first, m2.clone());
    assert_eq!(e1.wait_for_message(), m2);

    e2.send_to(first, m1.clone());
    assert_eq!(e1.wait_for_message(), m1);

    e2.send_to(first, m2.clone());
    assert_eq!(e1.wait_for_message(), m2);

    e1.disconnect_with(second);
    assert_eq!(e1.wait_for_disconnect(), second);

    e2.disconnect_with(first);
    assert_eq!(e2.wait_for_disconnect(), first);
}

#[test]
fn test_network_reconnect() {
    let first = "127.0.0.1:19100".parse().unwrap();
    let second = "127.0.0.1:19101".parse().unwrap();

    let msg = raw_message(11, 1000);
    let c1 = connect_message(first);

    let mut t1 = TestEvents::with_addr(first).spawn();

    // First connect attempt.
    let mut t2 = TestEvents::with_addr(second).spawn();

    // Handle first attempt.
    t1.connect_with(second);
    assert_eq!(t2.wait_for_connect(), c1);

    t1.send_to(second, msg.clone());
    assert_eq!(t2.wait_for_message(), msg);

    drop(t2);
    assert_eq!(t1.wait_for_disconnect(), second);

    // Handle second attempt.
    let mut t2 = TestEvents::with_addr(second).spawn();

    t1.connect_with(second);
    assert_eq!(t2.wait_for_connect(), c1);

    t1.send_to(second, msg.clone());
    assert_eq!(t2.wait_for_message(), msg);

    t1.disconnect_with(second);
    assert_eq!(t1.wait_for_disconnect(), second);
}

#[test]
fn test_network_multiple_connect() {
    let main = "127.0.0.1:19600".parse().unwrap();

    let nodes = [
        "127.0.0.1:19601".parse().unwrap(),
        "127.0.0.1:19602".parse().unwrap(),
        "127.0.0.1:19603".parse().unwrap(),
    ];

    let mut node = TestEvents::with_addr(main).spawn();

    let connect_messages: Vec<_> = nodes.iter().cloned().map(connect_message).collect();
    let connectors: Vec<_> = nodes
        .iter()
        .map(|addr| TestEvents::with_addr(*addr).spawn())
        .collect();

    connectors[0].connect_with(main);
    assert_eq!(node.wait_for_connect(), connect_messages[0]);
    connectors[1].connect_with(main);
    assert_eq!(node.wait_for_connect(), connect_messages[1]);
    connectors[2].connect_with(main);
    assert_eq!(node.wait_for_connect(), connect_messages[2]);
}

#[test]
fn test_send_first_not_connect() {
    let main = "127.0.0.1:19500".parse().unwrap();
    let other = "127.0.0.1:19501".parse().unwrap();

    let mut node = TestEvents::with_addr(main).spawn();
    let other_node = TestEvents::with_addr(other).spawn();

    let message = raw_message(11, 1000);
    other_node.send_to(main, message.clone()); // should connect before send message

    assert_eq!(node.wait_for_connect(), connect_message(other));
    assert_eq!(node.wait_for_message(), message);
}
<commit_msg>Add tests checking max_message_len limit works<commit_after>// Copyright 2017 The Exonum Team
//
// 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::net::SocketAddr;
use std::thread;
use std::time::{self, Duration};

use futures::{Future, Sink, Stream};
use futures::stream::Wait;
use futures::sync::mpsc;
use tokio_core::reactor::Core;
use tokio_timer::{TimeoutStream, Timer};

use crypto::{gen_keypair, PublicKey, Signature};
use messages::{Connect, Message, MessageWriter, RawMessage};
use events::{NetworkEvent, NetworkRequest};
use events::network::{NetworkConfiguration, NetworkPart};
use events::error::log_error;
use node::{EventsPoolCapacity, NodeChannel};
use blockchain::ConsensusConfig;

#[derive(Debug)]
pub struct TestHandler {
    handle: Option<thread::JoinHandle<()>>,
    listen_address: SocketAddr,
    network_events_rx: Wait<TimeoutStream<mpsc::Receiver<NetworkEvent>>>,
    network_requests_tx: mpsc::Sender<NetworkRequest>,
}

impl TestHandler {
    pub fn new(
        listen_address: SocketAddr,
        network_requests_tx: mpsc::Sender<NetworkRequest>,
        network_events_rx: mpsc::Receiver<NetworkEvent>,
    ) -> TestHandler {
        let timer = Timer::default();
        let receiver = timer.timeout_stream(network_events_rx, Duration::from_secs(30));
        TestHandler {
            handle: None,
            listen_address,
            network_requests_tx,
            network_events_rx: receiver.wait(),
        }
    }

    pub fn wait_for_event(&mut self) -> Result<NetworkEvent, ()> {
        let event = self.network_events_rx.next().unwrap()?;
        Ok(event)
    }

    pub fn connect_with(&self, addr: SocketAddr) {
        let connect = connect_message(self.listen_address);
        self.network_requests_tx
            .clone()
            .send(NetworkRequest::SendMessage(addr, connect.raw().clone()))
            .wait()
            .unwrap();
    }

    pub fn disconnect_with(&self, addr: SocketAddr) {
        self.network_requests_tx
            .clone()
            .send(NetworkRequest::DisconnectWithPeer(addr))
            .wait()
            .unwrap();
    }

    pub fn send_to(&self, addr: SocketAddr, raw: RawMessage) {
        self.network_requests_tx
            .clone()
            .send(NetworkRequest::SendMessage(addr, raw))
            .wait()
            .unwrap();
    }

    pub fn wait_for_connect(&mut self) -> Connect {
        match self.wait_for_event() {
            Ok(NetworkEvent::PeerConnected(_addr, connect)) => connect,
            Ok(other) => panic!("Unexpected connect received, {:?}", other),
            Err(e) => panic!("An error during wait for connect occurred, {:?}", e),
        }
    }

    pub fn wait_for_disconnect(&mut self) -> SocketAddr {
        match self.wait_for_event() {
            Ok(NetworkEvent::PeerDisconnected(addr)) => addr,
            Ok(other) => panic!("Unexpected disconnect received, {:?}", other),
            Err(e) => panic!("An error during wait for disconnect occurred, {:?}", e),
        }
    }

    pub fn wait_for_message(&mut self) -> RawMessage {
        match self.wait_for_event() {
            Ok(NetworkEvent::MessageReceived(_addr, msg)) => msg,
            Ok(other) => panic!("Unexpected message received, {:?}", other),
            Err(e) => panic!("An error during wait for message occurred, {:?}", e),
        }
    }

    pub fn shutdown(&mut self) {
        self.network_requests_tx
            .clone()
            .send(NetworkRequest::Shutdown)
            .wait()
            .unwrap();
        self.handle.take().expect("shutdown twice").join().unwrap();
    }
}

impl Drop for TestHandler {
    fn drop(&mut self) {
        if !::std::thread::panicking() {
            self.shutdown();
        }
    }
}

#[derive(Debug)]
pub struct TestEvents {
    pub listen_address: SocketAddr,
    pub network_config: NetworkConfiguration,
    pub events_config: EventsPoolCapacity,
}

impl TestEvents {
    pub fn with_addr(listen_address: SocketAddr) -> TestEvents {
        TestEvents {
            listen_address,
            network_config: NetworkConfiguration::default(),
            events_config: EventsPoolCapacity::default(),
        }
    }

    pub fn spawn(self) -> TestHandler {
        let (mut handler_part, network_part) = self.into_reactor();
        let handle = thread::spawn(move || {
            let mut core = Core::new().unwrap();
            let fut = network_part.run(&core.handle());
            core.run(fut).map_err(log_error).unwrap();
        });
        handler_part.handle = Some(handle);
        handler_part
    }

    fn into_reactor(self) -> (TestHandler, NetworkPart) {
        let channel = NodeChannel::new(&self.events_config);
        let network_config = self.network_config;
        let (network_tx, network_rx) = channel.network_events;
        let network_requests_tx = channel.network_requests.0.clone();

        let network_part = NetworkPart {
            our_connect_message: connect_message(self.listen_address),
            listen_address: self.listen_address,
            network_config,
            max_message_len: ConsensusConfig::DEFAULT_MESSAGE_MAX_LEN,
            network_requests: channel.network_requests,
            network_tx: network_tx.clone(),
        };

        let handler_part = TestHandler::new(self.listen_address, network_requests_tx, network_rx);
        (handler_part, network_part)
    }
}

pub fn connect_message(addr: SocketAddr) -> Connect {
    let time = time::UNIX_EPOCH;
    Connect::new_with_signature(&PublicKey::zero(), addr, time, &Signature::zero())
}

pub fn raw_message(id: u16, len: usize) -> RawMessage {
    let writer = MessageWriter::new(
        ::messages::PROTOCOL_MAJOR_VERSION,
        ::messages::TEST_NETWORK_ID,
        0,
        id,
        len,
    );
    RawMessage::new(writer.sign(&gen_keypair().1))
}

#[test]
fn test_network_handshake() {
    let first = "127.0.0.1:17230".parse().unwrap();
    let second = "127.0.0.1:17231".parse().unwrap();

    let e1 = TestEvents::with_addr(first);
    let e2 = TestEvents::with_addr(second);

    let c1 = connect_message(first);
    let c2 = connect_message(second);

    let mut e1 = e1.spawn();
    let mut e2 = e2.spawn();

    e1.connect_with(second);
    assert_eq!(e2.wait_for_connect(), c1);

    e2.connect_with(first);
    assert_eq!(e1.wait_for_connect(), c2);

    e1.disconnect_with(second);
    assert_eq!(e1.wait_for_disconnect(), second);

    e2.disconnect_with(first);
    assert_eq!(e2.wait_for_disconnect(), first);
}

#[test]
fn test_network_big_message() {
    let first = "127.0.0.1:17200".parse().unwrap();
    let second = "127.0.0.1:17201".parse().unwrap();

    let m1 = raw_message(15, 100000);
    let m2 = raw_message(16, 400);

    let e1 = TestEvents::with_addr(first);
    let e2 = TestEvents::with_addr(second);

    let mut e1 = e1.spawn();
    let mut e2 = e2.spawn();

    e1.connect_with(second);
    e2.wait_for_connect();

    e2.connect_with(first);
    e1.wait_for_connect();

    e1.send_to(second, m1.clone());
    assert_eq!(e2.wait_for_message(), m1);

    e1.send_to(second, m2.clone());
    assert_eq!(e2.wait_for_message(), m2);

    e1.send_to(second, m1.clone());
    assert_eq!(e2.wait_for_message(), m1);

    e2.send_to(first, m2.clone());
    assert_eq!(e1.wait_for_message(), m2);

    e2.send_to(first, m1.clone());
    assert_eq!(e1.wait_for_message(), m1);

    e2.send_to(first, m2.clone());
    assert_eq!(e1.wait_for_message(), m2);

    e1.disconnect_with(second);
    assert_eq!(e1.wait_for_disconnect(), second);

    e2.disconnect_with(first);
    assert_eq!(e2.wait_for_disconnect(), first);
}

#[test]
fn test_network_max_message_len() {
    let first = "127.0.0.1:17202".parse().unwrap();
    let second = "127.0.0.1:17303".parse().unwrap();

    let max_message_length = ConsensusConfig::DEFAULT_MESSAGE_MAX_LEN as usize;
    let max_payload_length = max_message_length - ::messages::HEADER_LENGTH -
        ::crypto::SIGNATURE_LENGTH;
    let acceptable_message = raw_message(15, max_payload_length);
    let too_big_message = raw_message(16, max_payload_length + 1000);

    let e1 = TestEvents::with_addr(first);
    let e2 = TestEvents::with_addr(second);

    let mut e1 = e1.spawn();
    let mut e2 = e2.spawn();

    e1.connect_with(second);
    e2.wait_for_connect();

    e2.connect_with(first);
    e1.wait_for_connect();

    e1.send_to(second, acceptable_message.clone());
    assert_eq!(e2.wait_for_message(), acceptable_message);

    e2.send_to(first, too_big_message.clone());
    assert!(e1.wait_for_event().is_err());
}

#[test]
fn test_network_reconnect() {
    let first = "127.0.0.1:19100".parse().unwrap();
    let second = "127.0.0.1:19101".parse().unwrap();

    let msg = raw_message(11, 1000);
    let c1 = connect_message(first);

    let mut t1 = TestEvents::with_addr(first).spawn();

    // First connect attempt.
    let mut t2 = TestEvents::with_addr(second).spawn();

    // Handle first attempt.
    t1.connect_with(second);
    assert_eq!(t2.wait_for_connect(), c1);

    t1.send_to(second, msg.clone());
    assert_eq!(t2.wait_for_message(), msg);

    drop(t2);
    assert_eq!(t1.wait_for_disconnect(), second);

    // Handle second attempt.
    let mut t2 = TestEvents::with_addr(second).spawn();

    t1.connect_with(second);
    assert_eq!(t2.wait_for_connect(), c1);

    t1.send_to(second, msg.clone());
    assert_eq!(t2.wait_for_message(), msg);

    t1.disconnect_with(second);
    assert_eq!(t1.wait_for_disconnect(), second);
}

#[test]
fn test_network_multiple_connect() {
    let main = "127.0.0.1:19600".parse().unwrap();

    let nodes = [
        "127.0.0.1:19601".parse().unwrap(),
        "127.0.0.1:19602".parse().unwrap(),
        "127.0.0.1:19603".parse().unwrap(),
    ];

    let mut node = TestEvents::with_addr(main).spawn();

    let connect_messages: Vec<_> = nodes.iter().cloned().map(connect_message).collect();
    let connectors: Vec<_> = nodes
        .iter()
        .map(|addr| TestEvents::with_addr(*addr).spawn())
        .collect();

    connectors[0].connect_with(main);
    assert_eq!(node.wait_for_connect(), connect_messages[0]);
    connectors[1].connect_with(main);
    assert_eq!(node.wait_for_connect(), connect_messages[1]);
    connectors[2].connect_with(main);
    assert_eq!(node.wait_for_connect(), connect_messages[2]);
}

#[test]
fn test_send_first_not_connect() {
    let main = "127.0.0.1:19500".parse().unwrap();
    let other = "127.0.0.1:19501".parse().unwrap();

    let mut node = TestEvents::with_addr(main).spawn();
    let other_node = TestEvents::with_addr(other).spawn();

    let message = raw_message(11, 1000);
    other_node.send_to(main, message.clone()); // should connect before send message

    assert_eq!(node.wait_for_connect(), connect_message(other));
    assert_eq!(node.wait_for_message(), message);
}
<|endoftext|>


<commit_before>use alloc::{
    slice::{Iter, IterMut},
    vec::{IntoIter, Vec},
};
use core::ops::{Index, IndexMut};
use crate::utils::color_gradient::Color;

const RASTER_MAX_WIDTH: u16 = 32_767;
const RASTER_MAX_HEIGHT: u16 = 32_767;

pub struct NoiseImage {
    size: (usize, usize),
    border_color: Color,
    map: Vec<Color>,
}

impl NoiseImage {
    pub fn new(width: usize, height: usize) -> Self {
        Self::initialize().set_size(width, height)
    }

    pub fn iter(&self) -> Iter<'_, Color> {
        self.map.iter()
    }

    pub fn iter_mut(&mut self) -> IterMut<'_, Color> {
        self.map.iter_mut()
    }

    pub fn set_size(self, width: usize, height: usize) -> Self {
        // Check for invalid width or height.
        assert!(width < RASTER_MAX_WIDTH as usize);
        assert!(height < RASTER_MAX_HEIGHT as usize);

        if width == 0 || height == 0 {
            // An empty noise image was specified. Return a new blank, empty map.
            Self::initialize()
        } else {
            // New noise map size specified. Allocate a new Vec unless the current Vec is large
            // enough.
            let map_size = width * height;
            if self.map.capacity() < map_size {
                // New size is too big for the current Vec. Create a new Vec with a large enough
                // capacity now so we're not reallocating when filling the map.
                Self {
                    map: vec![[0; 4]; map_size],
                    size: (width, height),
                    ..self
                }
            } else {
                // Vec capacity is already big enough, so leave it alone and just change the set size.
                Self {
                    size: (width, height),
                    ..self
                }
            }
        }
    }

    pub fn set_border_color(self, color: Color) -> Self {
        Self {
            border_color: color,
            ..self
        }
    }

    pub fn set_value(&mut self, x: usize, y: usize, value: Color) {
        let (width, height) = self.size;

        if x < width && y < height {
            self.map[x + y * width] = value;
        } else {
            // eprintln!("input point out of bounds")
        }
    }

    pub fn size(&self) -> (usize, usize) {
        self.size
    }

    pub fn border_color(&self) -> Color {
        self.border_color
    }

    pub fn get_value(&self, x: usize, y: usize) -> Color {
        let (width, height) = self.size;

        if x < width && y < height {
            self.map[x + y * width]
        } else {
            self.border_color
        }
    }

    fn initialize() -> Self {
        Self {
            size: (0, 0),
            border_color: [0; 4],
            map: Vec::new(),
        }
    }

    #[cfg(all(feature = "std", feature = "image"))]
    pub fn write_to_file(&self, filename: &str) {
        use std::{fs, path::Path};

        // Create the output directory for the images, if it doesn't already exist
        let target_dir = Path::new("example_images/");

        if !target_dir.exists() {
            fs::create_dir(target_dir).expect("failed to create example_images directory");
        }

        //concatenate the directory to the filename string
        let directory: String = "example_images/".to_owned();
        let file_path = directory + filename;

        // collect the values from the map vector into an array
        let (width, height) = self.size;
        let mut result = Vec::with_capacity(width * height);

        for i in &self.map {
            for j in i.iter() {
                result.push(*j);
            }
        }

        let _ = image::save_buffer(
            &Path::new(&file_path),
            &*result,
            self.size.0 as u32,
            self.size.1 as u32,
            image::ColorType::Rgba8,
        );

        println!("\nFinished generating {}", filename);
    }
}

impl Default for NoiseImage {
    fn default() -> Self {
        Self::initialize()
    }
}

impl Index<(usize, usize)> for NoiseImage {
    type Output = Color;

    fn index(&self, (x, y): (usize, usize)) -> &Self::Output {
        let (width, height) = self.size;
        if x < width && y < height {
            &self.map[x + y * width]
        } else {
            &self.border_color
        }
    }
}

impl IndexMut<(usize, usize)> for NoiseImage {
    fn index_mut(&mut self, (x, y): (usize, usize)) -> &mut Self::Output {
        let (width, height) = self.size;
        if x < width && y < height {
            &mut self.map[x + y * width]
        } else {
            panic!(
                "index ({}, {}) out of bounds for NoiseImage of size ({}, {})",
                x, y, width, height
            )
        }
    }
}

impl IntoIterator for NoiseImage {
    type Item = Color;

    type IntoIter = IntoIter<Color>;

    fn into_iter(self) -> Self::IntoIter {
        self.map.into_iter()
    }
}

impl<'a> IntoIterator for &'a NoiseImage {
    type Item = &'a Color;

    type IntoIter = Iter<'a, Color>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<'a> IntoIterator for &'a mut NoiseImage {
    type Item = &'a mut Color;

    type IntoIter = IterMut<'a, Color>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter_mut()
    }
}
<commit_msg>format pass<commit_after>use crate::utils::color_gradient::Color;
use alloc::{
    slice::{Iter, IterMut},
    vec::{IntoIter, Vec},
};
use core::ops::{Index, IndexMut};

const RASTER_MAX_WIDTH: u16 = 32_767;
const RASTER_MAX_HEIGHT: u16 = 32_767;

pub struct NoiseImage {
    size: (usize, usize),
    border_color: Color,
    map: Vec<Color>,
}

impl NoiseImage {
    pub fn new(width: usize, height: usize) -> Self {
        Self::initialize().set_size(width, height)
    }

    pub fn iter(&self) -> Iter<'_, Color> {
        self.map.iter()
    }

    pub fn iter_mut(&mut self) -> IterMut<'_, Color> {
        self.map.iter_mut()
    }

    pub fn set_size(self, width: usize, height: usize) -> Self {
        // Check for invalid width or height.
        assert!(width < RASTER_MAX_WIDTH as usize);
        assert!(height < RASTER_MAX_HEIGHT as usize);

        if width == 0 || height == 0 {
            // An empty noise image was specified. Return a new blank, empty map.
            Self::initialize()
        } else {
            // New noise map size specified. Allocate a new Vec unless the current Vec is large
            // enough.
            let map_size = width * height;
            if self.map.capacity() < map_size {
                // New size is too big for the current Vec. Create a new Vec with a large enough
                // capacity now so we're not reallocating when filling the map.
                Self {
                    map: vec![[0; 4]; map_size],
                    size: (width, height),
                    ..self
                }
            } else {
                // Vec capacity is already big enough, so leave it alone and just change the set size.
                Self {
                    size: (width, height),
                    ..self
                }
            }
        }
    }

    pub fn set_border_color(self, color: Color) -> Self {
        Self {
            border_color: color,
            ..self
        }
    }

    pub fn set_value(&mut self, x: usize, y: usize, value: Color) {
        let (width, height) = self.size;

        if x < width && y < height {
            self.map[x + y * width] = value;
        } else {
            // eprintln!("input point out of bounds")
        }
    }

    pub fn size(&self) -> (usize, usize) {
        self.size
    }

    pub fn border_color(&self) -> Color {
        self.border_color
    }

    pub fn get_value(&self, x: usize, y: usize) -> Color {
        let (width, height) = self.size;

        if x < width && y < height {
            self.map[x + y * width]
        } else {
            self.border_color
        }
    }

    fn initialize() -> Self {
        Self {
            size: (0, 0),
            border_color: [0; 4],
            map: Vec::new(),
        }
    }

    #[cfg(all(feature = "std", feature = "image"))]
    pub fn write_to_file(&self, filename: &str) {
        use std::{fs, path::Path};

        // Create the output directory for the images, if it doesn't already exist
        let target_dir = Path::new("example_images/");

        if !target_dir.exists() {
            fs::create_dir(target_dir).expect("failed to create example_images directory");
        }

        //concatenate the directory to the filename string
        let directory: String = "example_images/".to_owned();
        let file_path = directory + filename;

        // collect the values from the map vector into an array
        let (width, height) = self.size;
        let mut result = Vec::with_capacity(width * height);

        for i in &self.map {
            for j in i.iter() {
                result.push(*j);
            }
        }

        let _ = image::save_buffer(
            &Path::new(&file_path),
            &*result,
            self.size.0 as u32,
            self.size.1 as u32,
            image::ColorType::Rgba8,
        );

        println!("\nFinished generating {}", filename);
    }
}

impl Default for NoiseImage {
    fn default() -> Self {
        Self::initialize()
    }
}

impl Index<(usize, usize)> for NoiseImage {
    type Output = Color;

    fn index(&self, (x, y): (usize, usize)) -> &Self::Output {
        let (width, height) = self.size;
        if x < width && y < height {
            &self.map[x + y * width]
        } else {
            &self.border_color
        }
    }
}

impl IndexMut<(usize, usize)> for NoiseImage {
    fn index_mut(&mut self, (x, y): (usize, usize)) -> &mut Self::Output {
        let (width, height) = self.size;
        if x < width && y < height {
            &mut self.map[x + y * width]
        } else {
            panic!(
                "index ({}, {}) out of bounds for NoiseImage of size ({}, {})",
                x, y, width, height
            )
        }
    }
}

impl IntoIterator for NoiseImage {
    type Item = Color;

    type IntoIter = IntoIter<Color>;

    fn into_iter(self) -> Self::IntoIter {
        self.map.into_iter()
    }
}

impl<'a> IntoIterator for &'a NoiseImage {
    type Item = &'a Color;

    type IntoIter = Iter<'a, Color>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<'a> IntoIterator for &'a mut NoiseImage {
    type Item = &'a mut Color;

    type IntoIter = IterMut<'a, Color>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter_mut()
    }
}
<|endoftext|>


<commit_before>use ggez::{Context, GameResult, graphics};
use ggez::graphics::{DrawMode, Color, Rect, Point};
use event::Assets;
use states::play_state::tetromino::{Piece, PieceType};

const WIDTH: f32 = 170.0;
const HEIGHT: f32 = 50.0;
// text view's have their own height
const UITEXTVIEW_HEIGHT: f32 = 35.0;
const CIRCLE_RADIUS: f32 = 25.0;

// Ideally, there would be a general `View` trait or type or whatnot that I
// could implement these two views in a DRY way but UI design/creation wasn't
// one of my goals for this project so I just wanted to whip up something quick
// and esy.

/// Represents a UI element that displays a String.
pub struct UITextView {
    top_left: Point,
    centre: Point,
    title_text: graphics::Text,
    value: String,
    value_text: graphics::Text,
}

impl UITextView {
    /// Creates a new `UITextView` with the element title of `title` and a
    /// default value of `value.`
    pub fn new(
        ctx: &mut Context,
        assets: &Assets,
        top_left: Point,
        title: &str,
        value: &str,
    ) -> Self {
        let title_txt = graphics::Text::new(ctx, title, assets.get_font("normal").unwrap())
            .unwrap();
        let value_txt = graphics::Text::new(ctx, value, assets.get_font("normal").unwrap())
            .unwrap();
        Self {
            // ggez draws coords at center; we want to use top left as reference
            top_left: Point::new(
                top_left.x + (WIDTH / 2.0),
                top_left.y + (UITEXTVIEW_HEIGHT / 2.0),
            ),
            centre: top_left,
            title_text: title_txt,
            value: value.to_string(),
            value_text: value_txt,
        }
    }

    /// Currently updates the current value displayed. Only updates the value if
    /// it is different from the current value.
    pub fn update(&mut self, ctx: &mut Context, assets: &Assets, new_value: &str) {
        if new_value != self.value {
            self.value.clear(); // is this necessary?
            self.value = new_value.to_string();
            self.value_text =
                graphics::Text::new(ctx, new_value, assets.get_font("normal").unwrap()).unwrap();
        }
    }

    pub fn draw(&mut self, ctx: &mut Context) -> GameResult<()> {
        // title portion
        let title_rect = Rect::new(self.top_left.x, self.top_left.y, WIDTH, UITEXTVIEW_HEIGHT);
        let title_point = Point::new(self.top_left.x as f32, self.top_left.y as f32);
        graphics::set_color(ctx, Color::new(1.0, 0.0, 0.0, 1.0))?;
        graphics::rectangle(ctx, DrawMode::Fill, title_rect)?;
        graphics::circle(
            ctx,
            DrawMode::Fill,
            Point::new(self.centre.x as f32 + WIDTH as f32, self.top_left.y as f32),
            CIRCLE_RADIUS,
            32,
        )?;
        graphics::set_color(ctx, Color::new(1.0, 1.0, 1.0, 1.0))?;
        graphics::draw(ctx, &self.title_text, title_point, 0.0)?;

        // value portion
        let value_rect = Rect::new(
            self.top_left.x - 20.0,
            self.top_left.y + UITEXTVIEW_HEIGHT,
            WIDTH - 40.0,
            UITEXTVIEW_HEIGHT,
        );
        let value_point = Point::new(
            self.top_left.x as f32,
            self.top_left.y as f32 + UITEXTVIEW_HEIGHT as f32,
        );
        graphics::set_color(ctx, Color::new(1.0, 0.0, 0.0, 1.0))?;
        graphics::rectangle(ctx, DrawMode::Fill, value_rect)?;
        graphics::set_color(ctx, Color::new(1.0, 1.0, 1.0, 1.0))?;
        graphics::draw(ctx, &self.value_text, value_point, 0.0)?;
        Ok(())
    }
}

/// A UI element that renders `Pieces`. Used for elements like 'Next' or 'Hold'.
pub struct UIBlockView {
    top_left: Point,
    centre: Point,
    title_text: graphics::Text,
    shape: Option<PieceType>,
}

impl UIBlockView {
    pub fn new(
        ctx: &mut Context,
        assets: &Assets,
        top_left: Point,
        title: &str,
        shape: Option<PieceType>,
    ) -> Self {
        let title_txt = graphics::Text::new(ctx, title, assets.get_font("normal").unwrap())
            .unwrap();
        Self {
            // ggez draws coords at center; we want to use top left as reference
            top_left: Point::new(top_left.x + (WIDTH / 2.0), top_left.y + (HEIGHT / 2.0)),
            centre: top_left,
            title_text: title_txt,
            shape: shape,
        }
    }

    pub fn update(&mut self, _: &mut Context, _: &Assets, new_value: Option<PieceType>) {
        if self.shape != new_value {
            self.shape = new_value;
        }
    }

    pub fn draw(&mut self, ctx: &mut Context, assets: &Assets) -> GameResult<()> {
        // title portion
        let title_rect = Rect::new(self.top_left.x, self.top_left.y, WIDTH, HEIGHT);
        let title_point = Point::new(self.top_left.x as f32, self.top_left.y as f32);
        graphics::set_color(ctx, Color::new(1.0, 0.0, 0.0, 1.0))?;
        graphics::rectangle(ctx, DrawMode::Fill, title_rect)?;
        graphics::circle(
            ctx,
            DrawMode::Fill,
            Point::new(self.centre.x as f32 + WIDTH as f32, self.top_left.y as f32),
            CIRCLE_RADIUS,
            32,
        )?;
        graphics::set_color(ctx, Color::new(1.0, 1.0, 1.0, 1.0))?;
        graphics::draw(ctx, &self.title_text, title_point, 0.0)?;

        // value portion
        let value_rect = Rect::new(
            self.top_left.x - 20.0,
            self.top_left.y + HEIGHT + (HEIGHT / 2.0),
            WIDTH - 40.0,
            HEIGHT * 2.0,
        );
        graphics::set_color(ctx, Color::new(1.0, 0.0, 0.0, 1.0))?;
        graphics::rectangle(ctx, DrawMode::Line, value_rect)?;
        if let Some(shape) = self.shape {
            let piece = Piece::new(shape);
            // the magic numbers below were derived by fiddling with values
            // until it looked right
            let next_piece_pos = if piece.get_type() == PieceType::I {
                Point::new(self.top_left.x - 95.0, self.top_left.y + 45.0)
            } else {
                Point::new(self.top_left.x - 110.0, self.top_left.y + 45.0)
            };
            piece.draw_at_point(
                ctx,
                assets.get_image("block")?,
                next_piece_pos,
                0.0,
            )?;
        }
        Ok(())
    }
}
<commit_msg>complete restyling of ui elements - no longer gross looking<commit_after>use ggez::{Context, GameResult, graphics};
use ggez::graphics::{DrawMode, Color, Rect, Point};
use event::Assets;
use states::play_state::tetromino::{Piece, PieceType};

const WIDTH: f32 = 150.0;
const UIBLOCK_HEIGHT: f32 = 130.0;
const UITEXT_HEIGHT: f32 = 40.0;
const UI_BG_COLOUR: Color = Color {
    r: 0.1,
    b: 0.1,
    g: 0.1,
    a: 0.1,
};

// Ideally, there would be a general `View` trait or type or whatnot that I
// could implement these two views in a DRY way but UI design/creation wasn't
// one of my goals for this project so I just wanted to whip up something quick
// and esy.

/// Represents a UI element that displays a String.
pub struct UITextView {
    top_left: Point,
    title_text: graphics::Text,
    value: String,
    value_text: graphics::Text,
}

impl UITextView {
    /// Creates a new `UITextView` with the element title of `title` and a
    /// default value of `value.`
    pub fn new(
        ctx: &mut Context,
        assets: &Assets,
        top_left: Point,
        title: &str,
        value: &str,
    ) -> Self {
        let title_txt = graphics::Text::new(ctx, title, assets.get_font("normal").unwrap())
            .unwrap();
        let value_txt = graphics::Text::new(ctx, value, assets.get_font("normal").unwrap())
            .unwrap();
        Self {
            // ggez draws coords at center; we want to use top left as reference
            top_left: Point::new(
                top_left.x + (WIDTH / 2.0),
                top_left.y + (UITEXT_HEIGHT / 2.0),
            ),
            title_text: title_txt,
            value: value.to_string(),
            value_text: value_txt,
        }
    }

    /// Currently updates the current value displayed. Only updates the value if
    /// it is different from the current value.
    pub fn update(&mut self, ctx: &mut Context, assets: &Assets, new_value: &str) {
        if new_value != self.value {
            self.value.clear(); // is this necessary?
            self.value = new_value.to_string();
            self.value_text =
                graphics::Text::new(ctx, new_value, assets.get_font("normal").unwrap()).unwrap();
        }
    }

    pub fn draw(&mut self, ctx: &mut Context) -> GameResult<()> {
        let title_rect = Rect::new(self.top_left.x, self.top_left.y, WIDTH, UITEXT_HEIGHT);
        let title_point = Point::new(
            self.top_left.x as f32 - (WIDTH / 4.0),
            self.top_left.y as f32,
        );
        graphics::set_color(ctx, UI_BG_COLOUR)?;
        graphics::rectangle(ctx, DrawMode::Fill, title_rect)?;
        graphics::set_color(ctx, Color::new(1.0, 1.0, 1.0, 1.0))?;
        graphics::draw(ctx, &self.title_text, title_point, 0.0)?;

        let value_point = Point::new(
            self.top_left.x as f32 + (WIDTH / 3.0),
            self.top_left.y as f32,
        );
        graphics::set_color(ctx, Color::new(1.0, 1.0, 0.0, 1.0))?;
        graphics::draw(ctx, &self.value_text, value_point, 0.0)?;
        Ok(())
    }
}

/// A UI element that renders `Pieces`. Used for elements like 'Next' or 'Hold'.
pub struct UIBlockView {
    top_left: Point,
    title_text: graphics::Text,
    shape: Option<PieceType>,
}

impl UIBlockView {
    pub fn new(
        ctx: &mut Context,
        assets: &Assets,
        top_left: Point,
        title: &str,
        shape: Option<PieceType>,
    ) -> Self {
        let title_txt = graphics::Text::new(ctx, title, assets.get_font("normal").unwrap())
            .unwrap();
        Self {
            // ggez draws coords at center; we want to use top left as reference
            top_left: Point::new(
                top_left.x + (WIDTH / 2.0),
                top_left.y + (UIBLOCK_HEIGHT / 2.0),
            ),
            title_text: title_txt,
            shape: shape,
        }
    }

    pub fn update(&mut self, _: &mut Context, _: &Assets, new_value: Option<PieceType>) {
        if self.shape != new_value {
            self.shape = new_value;
        }
    }

    pub fn draw(&mut self, ctx: &mut Context, assets: &Assets) -> GameResult<()> {
        // title portion
        let title_rect = Rect::new(self.top_left.x, self.top_left.y, WIDTH, UIBLOCK_HEIGHT);
        let title_point = Point::new(
            self.top_left.x as f32 - (WIDTH / 3.0) + 10.0,
            self.top_left.y as f32 - (UIBLOCK_HEIGHT / 3.0),
        );
        graphics::set_color(ctx, UI_BG_COLOUR)?;
        graphics::rectangle(ctx, DrawMode::Fill, title_rect)?;
        graphics::set_color(ctx, Color::new(1.0, 1.0, 1.0, 1.0))?;
        graphics::draw(ctx, &self.title_text, title_point, 0.0)?;

        if let Some(shape) = self.shape {
            let piece = Piece::new(shape);
            // the magic numbers below were derived by fiddling with values
            // until it looked right. ideally, some kind of object graph would
            // be used to make relative positioning easier
            let next_piece_pos = if piece.get_type() == PieceType::I {
                Point::new(self.top_left.x - 77.0, self.top_left.y - 10.0)
            } else {
                Point::new(self.top_left.x - 90.0, self.top_left.y - 15.0)
            };
            piece.draw_at_point(
                ctx,
                assets.get_image("block")?,
                next_piece_pos,
                0.0,
            )?;
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>use cell_gc;
use parser::parse;
use std::io::{self, Write};
use vm;

pub fn repl() -> io::Result<()> {
    cell_gc::with_heap(|hs| {
        let env = vm::Value::default_env(hs);

        loop {
            {
                let stdout = io::stdout();
                let mut stdout = stdout.lock();
                write!(&mut stdout, "lisp> ")?;
                stdout.flush()?;
            }

            // Read
            let mut source = String::new();
            io::stdin().read_line(&mut source)?;
            let expr = parse(hs, &source)
                .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;

            // Eval
            let result = vm::eval(hs, expr, &env);

            // Print
            println!("{:?}", result);

            // Loop...
        }

        // Unreachable...
    })
}
<commit_msg>Use the s-expression printing in the REPL<commit_after>use cell_gc;
use parser::parse;
use print::print;
use std::io::{self, Write};
use vm;

pub fn repl() -> io::Result<()> {
    cell_gc::with_heap(|hs| {
        let env = vm::Value::default_env(hs);

        loop {
            {
                let stdout = io::stdout();
                let mut stdout = stdout.lock();
                write!(&mut stdout, "lisp> ")?;
                stdout.flush()?;
            }

            // Read
            let mut source = String::new();
            io::stdin().read_line(&mut source)?;
            let expr = parse(hs, &source)
                .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;

            // Eval
            let result = vm::eval(hs, expr, &env);

            // Print
            print(val);
            println!();

            // Loop...
        }

        // Unreachable...
    })
}
<|endoftext|>


<commit_before>use direction::Direction;
use event::{AnyCb, Event, EventResult, Key};
use rect::Rect;
use vec::Vec2;
use view::{Selector, View};
use xy::XY;
use Printer;
use With;

/// Wraps a view in a scrollable area.
pub struct ScrollView<V> {
    inner_size: Vec2,
    inner: V,
    // Offset into the inner view.
    //
    // Our `(0,0)` will be inner's `offset`
    offset: Vec2,

    last_size: Vec2,

    // Can we scroll horizontally?
    enabled: XY<bool>,

    // Should we show scrollbars?
    show_scrollbars: bool,

    // How much padding should be between content and scrollbar?
    scrollbar_padding: Vec2,
}

impl<V> ScrollView<V> {
    /// Creates a new ScrollView around `view`.
    pub fn new(view: V) -> Self {
        ScrollView {
            inner: view,
            inner_size: Vec2::zero(),
            offset: Vec2::zero(),
            last_size: Vec2::zero(),
            enabled: XY::new(false, true),
            show_scrollbars: true,
            scrollbar_padding: Vec2::new(1, 0),
        }
    }

    /// Returns the viewport in the inner content.
    pub fn content_viewport(&self) -> Rect {
        Rect::from_size(self.offset, self.last_size)
    }

    /// Sets the scroll offset to the given value
    pub fn set_offset<S>(&mut self, offset: S)
    where
        S: Into<Vec2>,
    {
        let max_offset = self.inner_size.saturating_sub(self.last_size);
        self.offset = offset.into().or_min(max_offset);
    }

    /// Controls whether this view can scroll vertically.
    ///
    /// Defaults to `true`.
    pub fn set_scroll_y(&mut self, enabled: bool) {
        self.enabled.y = enabled;
    }

    /// Controls whether this view can scroll horizontally.
    ///
    /// Defaults to `false`.
    pub fn set_scroll_x(&mut self, enabled: bool) {
        self.enabled.x = enabled;
    }

    /// Controls whether this view can scroll vertically.
    ///
    /// Defaults to `true`.
    ///
    /// Chainable variant.
    pub fn scroll_y(self, enabled: bool) -> Self {
        self.with(|s| s.set_scroll_y(enabled))
    }

    /// Controls whether this view can scroll horizontally.
    ///
    /// Defaults to `false`.
    ///
    /// Chainable variant.
    pub fn scroll_x(self, enabled: bool) -> Self {
        self.with(|s| s.set_scroll_x(enabled))
    }
}

impl <V> ScrollView<V> where V: View {

    fn compute_size(&mut self, constraint: Vec2, scrollable: XY<bool>) -> (Vec2, XY<bool>) {
            let scrollbar_size = scrollable
                .select_or(self.scrollbar_padding + (1, 1), Vec2::zero());

            let inner_size = self.inner
                .required_size(constraint.saturating_sub(scrollbar_size));

            let size = self.enabled
                .select_or(Vec2::min(inner_size, constraint), inner_size)
                + scrollbar_size;

            let new_scrollable = inner_size.zip_map(size, |i, s| i > s);

            (size, new_scrollable)
    }
}

impl<V> View for ScrollView<V>
where
    V: View,
{
    fn draw(&self, printer: &Printer) {
        // Draw content
        let printer = printer
            .content_offset(self.offset)
            .inner_size(self.inner_size);
        self.inner.draw(&printer);

        // Draw scrollbar?
    }

    fn on_event(&mut self, event: Event) -> EventResult {
        // Relativize event accorging to the offset
        let mut relative_event = event.clone();
        if let Some(pos) = relative_event.mouse_position_mut() {
            *pos = *pos + self.offset;
        }
        match self.inner.on_event(relative_event) {
            EventResult::Ignored => {
                // If it's an arrow, try to scroll in the given direction.
                // If it's a mouse scroll, try to scroll as well.
                match event {
                    Event::Key(Key::Up)
                        if self.enabled.y && self.offset.y > 0 =>
                    {
                        self.offset.y -= 1;
                        EventResult::Consumed(None)
                    }
                    Event::Key(Key::Down)
                        if self.enabled.y
                            && (self.offset.y + self.last_size.y
                                < self.inner_size.y) =>
                    {
                        self.offset.y += 1;
                        EventResult::Consumed(None)
                    }
                    Event::Key(Key::Left)
                        if self.enabled.x && self.offset.x > 0 =>
                    {
                        self.offset.x -= 1;
                        EventResult::Consumed(None)
                    }
                    Event::Key(Key::Right)
                        if self.enabled.x
                            && (self.offset.x + self.last_size.x
                                < self.inner_size.x) =>
                    {
                        self.offset.x += 1;
                        EventResult::Consumed(None)
                    }
                    _ => EventResult::Ignored,
                }
            }
            other => other,
        }
    }

    fn layout(&mut self, size: Vec2) {
        self.last_size = size;

        // Ask one more time
        self.inner_size = self.inner.required_size(size).or_max(size);
        self.inner.layout(self.inner_size);

        // TODO: Refresh offset if needed!
    }

    fn needs_relayout(&self) -> bool {
        self.inner.needs_relayout()
    }

    fn required_size(&mut self, constraint: Vec2) -> Vec2 {
        // Attempt 1: try without scrollbars
        let (size, scrollable) = self.compute_size(constraint, XY::new(false, false));

        // Did it work?
        if scrollable.any() && self.show_scrollbars {
            // Attempt 2: he wants to scroll? Sure! Try again with some space for the scrollbar.
            let (size, new_scrollable) = self.compute_size(constraint, scrollable);
            if scrollable != new_scrollable {
                // Again? We're now scrolling in a new direction?
                // There is no end to this!
                let (size, _) = self.compute_size(constraint, new_scrollable);

                // That's enough. If the inner view changed again, ignore it!
                // That'll teach it.
                size
            } else {
                // Yup, scrolling did it. We're goot to go now.
                size
            }
        } else {
            // We're not showing any scrollbar, either because we don't scroll
            // or because scrollbars are hidden.
            size
        }
    }

    fn call_on_any<'a>(&mut self, selector: &Selector, cb: AnyCb<'a>) {
        self.inner.call_on_any(selector, cb)
    }

    fn focus_view(&mut self, selector: &Selector) -> Result<(), ()> {
        self.inner.focus_view(selector)
    }

    fn take_focus(&mut self, source: Direction) -> bool {
        let is_scrollable =
            self.enabled.any() && (self.inner_size != self.last_size);
        self.inner.take_focus(source) || is_scrollable
    }
}
<commit_msg>Scrollview: use proper size for child's layout() call<commit_after>use direction::Direction;
use event::{AnyCb, Event, EventResult, Key};
use rect::Rect;
use vec::Vec2;
use view::{Selector, View};
use xy::XY;
use Printer;
use With;

/// Wraps a view in a scrollable area.
pub struct ScrollView<V> {
    inner_size: Vec2,
    inner: V,
    // Offset into the inner view.
    //
    // Our `(0,0)` will be inner's `offset`
    offset: Vec2,

    last_size: Vec2,

    // Can we scroll horizontally?
    enabled: XY<bool>,

    // Should we show scrollbars?
    show_scrollbars: bool,

    // How much padding should be between content and scrollbar?
    scrollbar_padding: Vec2,
}

impl<V> ScrollView<V> {
    /// Creates a new ScrollView around `view`.
    pub fn new(view: V) -> Self {
        ScrollView {
            inner: view,
            inner_size: Vec2::zero(),
            offset: Vec2::zero(),
            last_size: Vec2::zero(),
            enabled: XY::new(false, true),
            show_scrollbars: true,
            scrollbar_padding: Vec2::new(1, 0),
        }
    }

    /// Returns the viewport in the inner content.
    pub fn content_viewport(&self) -> Rect {
        Rect::from_size(self.offset, self.last_size)
    }

    /// Sets the scroll offset to the given value
    pub fn set_offset<S>(&mut self, offset: S)
    where
        S: Into<Vec2>,
    {
        let max_offset = self.inner_size.saturating_sub(self.last_size);
        self.offset = offset.into().or_min(max_offset);
    }

    /// Controls whether this view can scroll vertically.
    ///
    /// Defaults to `true`.
    pub fn set_scroll_y(&mut self, enabled: bool) {
        self.enabled.y = enabled;
    }

    /// Controls whether this view can scroll horizontally.
    ///
    /// Defaults to `false`.
    pub fn set_scroll_x(&mut self, enabled: bool) {
        self.enabled.x = enabled;
    }

    /// Controls whether this view can scroll vertically.
    ///
    /// Defaults to `true`.
    ///
    /// Chainable variant.
    pub fn scroll_y(self, enabled: bool) -> Self {
        self.with(|s| s.set_scroll_y(enabled))
    }

    /// Controls whether this view can scroll horizontally.
    ///
    /// Defaults to `false`.
    ///
    /// Chainable variant.
    pub fn scroll_x(self, enabled: bool) -> Self {
        self.with(|s| s.set_scroll_x(enabled))
    }
}

impl<V> ScrollView<V>
where
    V: View,
{
    /// Compute the size we would need.
    ///
    /// Given the constraints, and the axis that need scrollbars.
    ///
    /// Returns `(inner_size, size, scrollable)`.
    fn sizes_when_scrolling(
        &mut self, constraint: Vec2, scrollable: XY<bool>,
    ) -> (Vec2, Vec2, XY<bool>) {
        // This is the size taken by the scrollbars.
        let scrollbar_size = scrollable
            .select_or(self.scrollbar_padding + (1, 1), Vec2::zero());

        let available = constraint.saturating_sub(scrollbar_size);

        let inner_size = self.inner.required_size(available);

        // Where we're "enabled", accept the constraints.
        // Where we're not, just forward inner_size.
        let size = self.enabled.select_or(
            Vec2::min(inner_size + scrollbar_size, constraint),
            inner_size + scrollbar_size,
        );

        let new_scrollable = inner_size.zip_map(size, |i, s| i > s);

        (inner_size, size, new_scrollable)
    }

    /// Computes the size we would need given the constraints.
    ///
    /// First be optimistic and try without scrollbars.
    /// Then try with scrollbars if needed.
    /// Then try again in case we now need to scroll both ways (!!!)
    ///
    /// Returns `(inner_size, size)`
    fn sizes(&mut self, constraint: Vec2) -> (Vec2, Vec2) {
        let (inner_size, size, scrollable) =
            self.sizes_when_scrolling(constraint, XY::new(false, false));

        // Did it work?
        if scrollable.any() && self.show_scrollbars {
            // Attempt 2: he wants to scroll? Sure! Try again with some space for the scrollbar.
            let (inner_size, size, new_scrollable) =
                self.sizes_when_scrolling(constraint, scrollable);
            if scrollable != new_scrollable {
                // Again? We're now scrolling in a new direction?
                // There is no end to this!
                let (inner_size, size, _) =
                    self.sizes_when_scrolling(constraint, new_scrollable);

                // That's enough. If the inner view changed again, ignore it!
                // That'll teach it.
                (inner_size, size)
            } else {
                // Yup, scrolling did it. We're goot to go now.
                (inner_size, size)
            }
        } else {
            // We're not showing any scrollbar, either because we don't scroll
            // or because scrollbars are hidden.
            (inner_size, size)
        }
    }
}

impl<V> View for ScrollView<V>
where
    V: View,
{
    fn draw(&self, printer: &Printer) {
        // Draw content
        let printer = printer
            .content_offset(self.offset)
            .inner_size(self.inner_size);
        self.inner.draw(&printer);

        // Draw scrollbar?
    }

    fn on_event(&mut self, event: Event) -> EventResult {
        // Relativize event accorging to the offset
        let mut relative_event = event.clone();
        if let Some(pos) = relative_event.mouse_position_mut() {
            *pos = *pos + self.offset;
        }
        match self.inner.on_event(relative_event) {
            EventResult::Ignored => {
                // If it's an arrow, try to scroll in the given direction.
                // If it's a mouse scroll, try to scroll as well.
                match event {
                    Event::Key(Key::Up)
                        if self.enabled.y && self.offset.y > 0 =>
                    {
                        self.offset.y -= 1;
                        EventResult::Consumed(None)
                    }
                    Event::Key(Key::Down)
                        if self.enabled.y
                            && (self.offset.y + self.last_size.y
                                < self.inner_size.y) =>
                    {
                        self.offset.y += 1;
                        EventResult::Consumed(None)
                    }
                    Event::Key(Key::Left)
                        if self.enabled.x && self.offset.x > 0 =>
                    {
                        self.offset.x -= 1;
                        EventResult::Consumed(None)
                    }
                    Event::Key(Key::Right)
                        if self.enabled.x
                            && (self.offset.x + self.last_size.x
                                < self.inner_size.x) =>
                    {
                        self.offset.x += 1;
                        EventResult::Consumed(None)
                    }
                    _ => EventResult::Ignored,
                }
            }
            other => other,
        }
    }

    fn layout(&mut self, size: Vec2) {
        // Size is final now
        self.last_size = size;

        let (inner_size, _) = self.sizes(size);

        // Ask one more time
        self.inner_size = inner_size;

        self.inner.layout(self.inner_size);

        // TODO: Refresh offset if needed!
    }

    fn needs_relayout(&self) -> bool {
        self.inner.needs_relayout()
    }

    fn required_size(&mut self, constraint: Vec2) -> Vec2 {
        // Attempt 1: try without scrollbars
        let (_, size) = self.sizes(constraint);

        size
    }

    fn call_on_any<'a>(&mut self, selector: &Selector, cb: AnyCb<'a>) {
        self.inner.call_on_any(selector, cb)
    }

    fn focus_view(&mut self, selector: &Selector) -> Result<(), ()> {
        self.inner.focus_view(selector)
    }

    fn take_focus(&mut self, source: Direction) -> bool {
        let is_scrollable =
            self.enabled.any() && (self.inner_size != self.last_size);
        self.inner.take_focus(source) || is_scrollable
    }
}
<|endoftext|>


<commit_before>use std::fmt::Write;

use anyhow::Result;

use crate::{analysis::dis::zydis, arch::Arch, workspace::PEWorkspace, VA};

#[derive(Default, Clone)]
struct OriginalHooks {
    print_address_abs: Option<zydis::Hook>,
    pre_instruction:   Option<zydis::Hook>,
}

struct UserData<'a> {
    ws:   &'a PEWorkspace,
    orig: OriginalHooks,
}

pub struct FormatterOptions {
    colors: bool,
}

pub struct FormatterBuilder {
    options: FormatterOptions,
}

impl FormatterBuilder {
    pub fn build(self) -> Formatter {
        Formatter::from_options(self.options)
    }

    pub fn with_colors(mut self, colors: bool) -> FormatterBuilder {
        self.options.colors = colors;
        self
    }
}

pub struct Formatter {
    options: FormatterOptions,
    inner:   zydis::Formatter,
    orig:    OriginalHooks,
}

pub const TOKEN_USER_SYMBOLNAME: zydis::Token = zydis::Token(zydis::TOKEN_USER.0 + 1);

impl Formatter {
    // default theme
    // TODO: move this to a struct that can be configured
    const COLOR_ADDRESS_ABS: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_ADDRESS_REL: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_DECORATOR: ansi_term::Color = Formatter::GREY;
    const COLOR_DELIMITER: ansi_term::Color = Formatter::GREY;
    const COLOR_DISPLACEMENT: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_IMMEDIATE: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_INVALID: ansi_term::Color = ansi_term::Color::Red;
    const COLOR_MNEMONIC: ansi_term::Color = ansi_term::Color::Green;
    const COLOR_PARENTHESIS_CLOSE: ansi_term::Color = Formatter::GREY;
    const COLOR_PARENTHESIS_OPEN: ansi_term::Color = Formatter::GREY;
    const COLOR_PREFIX: ansi_term::Color = Formatter::GREY;
    const COLOR_REGISTER: ansi_term::Color = ansi_term::Color::Yellow;
    const COLOR_SYMBOL: ansi_term::Color = Formatter::GREY;
    const COLOR_SYMBOLNAME: ansi_term::Color = ansi_term::Color::Purple;
    const COLOR_TYPECAST: ansi_term::Color = Formatter::GREY;
    const COLOR_USER: ansi_term::Color = Formatter::GREY;
    const COLOR_WHITESPACE: ansi_term::Color = ansi_term::Color::Black;
    const GREY: ansi_term::Color = ansi_term::Color::Fixed(242);

    pub fn new() -> FormatterBuilder {
        FormatterBuilder {
            options: FormatterOptions { colors: true },
        }
    }

    pub fn from_options(options: FormatterOptions) -> Formatter {
        let mut inner = zydis::Formatter::new(zydis::FormatterStyle::INTEL).unwrap();

        let mut orig: OriginalHooks = Default::default();

        // TODO: align mnemonic

        let f = inner
            .set_pre_instruction(Box::new(
                |_formatter: &zydis::Formatter,
                 buf: &mut zydis::FormatterBuffer,
                 ctx: &mut zydis::FormatterContext,
                 userdata: Option<&mut dyn core::any::Any>|
                 -> zydis::Result<()> {
                    // programming error: userdata must be provided. this is guaranteed within
                    // Formatter.
                    let userdata = userdata.expect("no userdata");

                    // programming error: userdata must be a Box<UserData>. this is guaranteed
                    // within Formatter.
                    let userdata = userdata.downcast_ref::<UserData>().expect("incorrect userdata");

                    let va = ctx.runtime_address;

                    if let Some(sec) = userdata
                        .ws
                        .pe
                        .module
                        .sections
                        .iter()
                        .find(|&sec| sec.virtual_range.contains(&va))
                    {
                        buf.append(TOKEN_USER_SYMBOLNAME)?;
                        buf.get_string()?.append(&sec.name)?;
                    } else {
                        buf.append(zydis::TOKEN_INVALID)?;
                        buf.get_string()?.append("???")?;
                    }
                    buf.append(zydis::TOKEN_DELIMITER)?;
                    buf.get_string()?.append(":")?;

                    // TODO: insn bytes

                    buf.append(zydis::TOKEN_ADDRESS_ABS)?;
                    match userdata.ws.pe.module.arch {
                        Arch::X32 => {
                            buf.get_string()?.append(&format!("{:08x}", va))?;
                        }
                        Arch::X64 => {
                            buf.get_string()?.append(&format!("{:016x}", va))?;
                        }
                    }

                    buf.append(zydis::TOKEN_WHITESPACE)?;
                    buf.get_string()?.append("  ")?;

                    Ok(())
                },
            ))
            .unwrap();
        orig.pre_instruction = Some(f);

        let f = inner
            .set_print_address_abs(Box::new(
                |formatter: &zydis::Formatter,
                 buf: &mut zydis::FormatterBuffer,
                 ctx: &mut zydis::FormatterContext,
                 userdata: Option<&mut dyn core::any::Any>|
                 -> zydis::Result<()> {
                    // programming error: userdata must be provided. this is guaranteed within
                    // Formatter.
                    let userdata = userdata.expect("no userdata");

                    // programming error: userdata must be a Box<UserData>. this is guaranteed
                    // within Formatter.
                    let userdata = userdata.downcast_ref::<UserData>().expect("incorrect userdata");

                    let absolute_address = unsafe {
                        // safety: the insn and operands come from zydis, so we assume they contain
                        // valid data.
                        let insn: &zydis::DecodedInstruction = &*ctx.instruction;
                        let op: &zydis::DecodedOperand = &*ctx.operand;
                        insn.calc_absolute_address(ctx.runtime_address, op)
                            .expect("failed to calculate absolute address")
                    };

                    if let Some(name) = userdata.ws.analysis.names.names_by_address.get(&absolute_address) {
                        // name is found in map, use that.
                        buf.append(TOKEN_USER_SYMBOLNAME)?;
                        return buf.get_string()?.append(name);
                    } else {
                        // name is not found, use original formatter.

                        // programming error: the original hook must be recorded. this is guaranteed
                        // within Formatter.
                        let orig = userdata.orig.print_address_abs.as_ref().expect("no original hook");

                        if let zydis::Hook::PrintAddressAbs(Some(f)) = orig {
                            // safety: zydis::Formatter <-> zydis::ffi::ZydisFormatter is safe according to
                            // here: https://docs.rs/zydis/3.1.2/src/zydis/formatter.rs.html#306
                            let status =
                                unsafe { f(formatter as *const _ as *const zydis::ffi::ZydisFormatter, buf, ctx) };
                            if status.is_error() {
                                return Err(status);
                            } else {
                                return Ok(());
                            }
                        } else {
                            // I'm not sure how this could ever be the case, as zydis initializes the hook
                            // with a default. I suppose if you explicitly set
                            // the callback to NULL/None? Which we don't do here.
                            panic!("unexpected original hook");
                        }
                    }
                },
            ))
            .unwrap();
        orig.print_address_abs = Some(f);

        Formatter { options, inner, orig }
    }

    // grey

    fn get_token_color(token: zydis::Token) -> ansi_term::Color {
        match token {
            zydis::TOKEN_INVALID => Formatter::COLOR_INVALID,
            zydis::TOKEN_WHITESPACE => Formatter::COLOR_WHITESPACE,
            zydis::TOKEN_DELIMITER => Formatter::COLOR_DELIMITER,
            zydis::TOKEN_PARENTHESIS_OPEN => Formatter::COLOR_PARENTHESIS_OPEN,
            zydis::TOKEN_PARENTHESIS_CLOSE => Formatter::COLOR_PARENTHESIS_CLOSE,
            zydis::TOKEN_PREFIX => Formatter::COLOR_PREFIX,
            zydis::TOKEN_MNEMONIC => Formatter::COLOR_MNEMONIC,
            zydis::TOKEN_REGISTER => Formatter::COLOR_REGISTER,
            zydis::TOKEN_ADDRESS_ABS => Formatter::COLOR_ADDRESS_ABS,
            zydis::TOKEN_ADDRESS_REL => Formatter::COLOR_ADDRESS_REL,
            zydis::TOKEN_DISPLACEMENT => Formatter::COLOR_DISPLACEMENT,
            zydis::TOKEN_IMMEDIATE => Formatter::COLOR_IMMEDIATE,
            zydis::TOKEN_TYPECAST => Formatter::COLOR_TYPECAST,
            zydis::TOKEN_DECORATOR => Formatter::COLOR_DECORATOR,
            zydis::TOKEN_SYMBOL => Formatter::COLOR_SYMBOL,
            zydis::TOKEN_USER => Formatter::COLOR_USER,
            TOKEN_USER_SYMBOLNAME => Formatter::COLOR_SYMBOLNAME,
            _ => unimplemented!("token: {}", token),
        }
    }

    fn render_token<T: Write>(&self, o: &mut T, token: zydis::Token, s: &str) -> Result<()> {
        if self.options.colors {
            o.write_str(&Formatter::get_token_color(token).paint(s).to_string())?;
        } else {
            o.write_str(s)?;
        }

        Ok(())
    }

    pub fn format_instruction(&self, ws: &PEWorkspace, insn: &zydis::DecodedInstruction, va: VA) -> Result<String> {
        let mut buffer = [0u8; 200];

        // we pass our userdata to ZydisFormatterFormatInstruction.
        // but to make it work, we have to play games with the lifetimes:
        // we need to convince the compiler that the userdata pointer lives long enough
        // to by used by the callbacks.
        //
        // we do this by extending the CFG lifetime from '_ to 'static.
        //
        // userdata is passed into ZydisFormatterFormatInstruction,
        // which passes userdata to each of the formatter callbacks.
        // those read strictly from insn/ctx/userdata and write strictly to output
        // buffer. there is no state maintained within these routines.
        // the callbacks won't be invoked beyond the call into FormatInstruction.
        //
        // therefore, i believe its safe to extend the lifetime here to work with zydis.
        let x = unsafe { std::mem::transmute::<&'_ PEWorkspace, &'static PEWorkspace>(ws) };

        let mut ud = UserData {
            orig: self.orig.clone(),
            ws:   x,
        };

        let mut out = String::new();
        for (token, s) in self
            .inner
            .tokenize_instruction(&insn, &mut buffer, Some(va), Some(&mut ud))?
        {
            self.render_token(&mut out, token, s)?;
        }

        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::{super::*, *};
    use crate::rsrc::*;
    use anyhow::Result;

    #[test]
    fn with_colors() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let ws = PEWorkspace::from_pe(pe)?;

        let fmt = Formatter::new().with_colors(true).build();

        // ```
        //     .text:00401C4E 000 68 F4 61 40 00          push    offset ModuleName ; "mscoree.dll"
        //     .text:00401C53 004 FF 15 00 60 40 00       call    ds:GetModuleHandleA
        //     .text:00401C59 000 85 C0                   test    eax, eax
        // ```
        let insn = crate::test::read_insn(&ws.pe.module, 0x401C53);
        let s = fmt.format_instruction(&ws, &insn, 0x401C53)?;
        assert!(s.contains("\u{1b}"));
        assert!(s.contains("call"));
        assert!(s.contains("GetModuleHandleA"));

        Ok(())
    }

    #[test]
    fn no_colors() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let ws = PEWorkspace::from_pe(pe)?;

        let fmt = Formatter::new().with_colors(false).build();

        // ```
        //     .text:00401C4E 000 68 F4 61 40 00          push    offset ModuleName ; "mscoree.dll"
        //     .text:00401C53 004 FF 15 00 60 40 00       call    ds:GetModuleHandleA
        //     .text:00401C59 000 85 C0                   test    eax, eax
        // ```
        let insn = crate::test::read_insn(&ws.pe.module, 0x401C53);
        assert_eq!(
            fmt.format_instruction(&ws, &insn, 0x401C53)?,
            "call [kernel32.dll!GetModuleHandleA]"
        );

        Ok(())
    }
}
<commit_msg>formatter: show instruction bytes and align mnemonics<commit_after>use std::{cmp::min, fmt::Write};

use anyhow::Result;

use crate::{analysis::dis::zydis, arch::Arch, aspace::AddressSpace, workspace::PEWorkspace, VA};

#[derive(Default, Clone)]
struct OriginalHooks {
    print_address_abs: Option<zydis::Hook>,
    print_mnemonic:    Option<zydis::Hook>,
    pre_instruction:   Option<zydis::Hook>,
}

#[derive(Clone, Copy)]
pub struct FormatterOptions {
    colors:          bool,
    /// show up to the given number of bytes for each instruction,
    /// or ... (ellipsis) when truncated.
    /// use zero to disable this column.
    ///
    /// unit: bytes
    /// default: 8
    /// max: 16
    hex_column_size: usize,

    /// pad mnemonics to at least the given width.
    /// won't truncate longer instructions, just wont be nicely aligned.
    /// you probably don't need to touch this unless you want to.
    ///
    /// unit: characters
    /// default: 7
    mnemonic_width: usize,
}

struct UserData<'a> {
    ws:      &'a PEWorkspace,
    orig:    OriginalHooks,
    options: FormatterOptions,
}

pub struct FormatterBuilder {
    options: FormatterOptions,
}

impl FormatterBuilder {
    pub fn build(self) -> Formatter {
        Formatter::from_options(self.options)
    }

    pub fn with_colors(mut self, colors: bool) -> FormatterBuilder {
        self.options.colors = colors;
        self
    }

    pub fn with_hex_column_size(mut self, hex_column_size: usize) -> FormatterBuilder {
        // 0x10: max instruction length
        self.options.hex_column_size = min(hex_column_size, 0x10);
        self
    }
}

pub struct Formatter {
    options: FormatterOptions,
    inner:   zydis::Formatter,
    orig:    OriginalHooks,
}

pub const TOKEN_USER_SYMBOLNAME: zydis::Token = zydis::Token(zydis::TOKEN_USER.0 + 1);
pub const TOKEN_USER_HEX: zydis::Token = zydis::Token(zydis::TOKEN_USER.0 + 2);

impl Formatter {
    // default theme
    // TODO: move this to a struct that can be configured
    const COLOR_ADDRESS_ABS: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_ADDRESS_REL: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_DECORATOR: ansi_term::Color = Formatter::GREY;
    const COLOR_DELIMITER: ansi_term::Color = Formatter::GREY;
    const COLOR_DISPLACEMENT: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_HEX: ansi_term::Color = ansi_term::Color::Cyan;
    const COLOR_IMMEDIATE: ansi_term::Color = ansi_term::Color::Blue;
    const COLOR_INVALID: ansi_term::Color = ansi_term::Color::Red;
    const COLOR_MNEMONIC: ansi_term::Color = ansi_term::Color::Green;
    const COLOR_PARENTHESIS_CLOSE: ansi_term::Color = Formatter::GREY;
    const COLOR_PARENTHESIS_OPEN: ansi_term::Color = Formatter::GREY;
    const COLOR_PREFIX: ansi_term::Color = Formatter::GREY;
    const COLOR_REGISTER: ansi_term::Color = ansi_term::Color::Yellow;
    const COLOR_SYMBOL: ansi_term::Color = Formatter::GREY;
    const COLOR_SYMBOLNAME: ansi_term::Color = ansi_term::Color::Purple;
    const COLOR_TYPECAST: ansi_term::Color = Formatter::GREY;
    const COLOR_USER: ansi_term::Color = Formatter::GREY;
    const COLOR_WHITESPACE: ansi_term::Color = ansi_term::Color::Black;
    const GREY: ansi_term::Color = ansi_term::Color::Fixed(242);

    pub fn new() -> FormatterBuilder {
        FormatterBuilder {
            options: FormatterOptions {
                colors:          true,
                hex_column_size: 7,
                mnemonic_width:  7,
            },
        }
    }

    pub fn from_options(options: FormatterOptions) -> Formatter {
        let mut inner = zydis::Formatter::new(zydis::FormatterStyle::INTEL).unwrap();

        let mut orig: OriginalHooks = Default::default();

        let f = inner
            .set_pre_instruction(Box::new(
                |_formatter: &zydis::Formatter,
                 buf: &mut zydis::FormatterBuffer,
                 ctx: &mut zydis::FormatterContext,
                 userdata: Option<&mut dyn core::any::Any>|
                 -> zydis::Result<()> {
                    // programming error: userdata must be provided. this is guaranteed within
                    // Formatter.
                    let userdata = userdata.expect("no userdata");

                    // programming error: userdata must be a Box<UserData>. this is guaranteed
                    // within Formatter.
                    let userdata = userdata.downcast_ref::<UserData>().expect("incorrect userdata");

                    let va = ctx.runtime_address;

                    if let Some(sec) = userdata
                        .ws
                        .pe
                        .module
                        .sections
                        .iter()
                        .find(|&sec| sec.virtual_range.contains(&va))
                    {
                        buf.append(TOKEN_USER_SYMBOLNAME)?;
                        buf.get_string()?.append(&sec.name)?;
                    } else {
                        buf.append(zydis::TOKEN_INVALID)?;
                        buf.get_string()?.append("???")?;
                    }
                    buf.append(zydis::TOKEN_DELIMITER)?;
                    buf.get_string()?.append(":")?;

                    buf.append(zydis::TOKEN_ADDRESS_ABS)?;
                    match userdata.ws.pe.module.arch {
                        Arch::X32 => {
                            buf.get_string()?.append(&format!("{:08x}", va))?;
                        }
                        Arch::X64 => {
                            buf.get_string()?.append(&format!("{:016x}", va))?;
                        }
                    }

                    buf.append(zydis::TOKEN_WHITESPACE)?;
                    buf.get_string()?.append("  ")?;

                    if userdata.options.hex_column_size > 0 {
                        let mut insn_buf = [0u8; 0x10];
                        let insn_len = (unsafe { &*ctx.instruction }).length as usize;
                        let col_count = userdata.options.hex_column_size;
                        userdata
                            .ws
                            .pe
                            .module
                            .address_space
                            .read_into(va, &mut insn_buf[..insn_len])
                            .expect("failed to read instruction");

                        let mut hex = String::new();
                        for i in 0..col_count {
                            if insn_len > col_count && i == col_count - 1 {
                                // instruction is larger than reserved space,
                                // and this is the final spot for hex,
                                // which is 3 characters wide,
                                // so show "..." instead of the last byte.
                                hex.write_str("...").unwrap();
                            } else if i < insn_len {
                                // most common case: bytes of the instruction

                                if i != 0 {
                                    hex.write_str(" ").unwrap();
                                }

                                hex.write_str(&format!("{:02X}", insn_buf[i])).unwrap();
                            } else {
                                // common case, insn is smaller than reserved space,
                                // so fill with spaces.
                                hex.write_str("   ").unwrap();
                            }
                        }

                        buf.append(TOKEN_USER_HEX)?;
                        buf.get_string()?.append(&hex)?;

                        buf.append(zydis::TOKEN_WHITESPACE)?;
                        buf.get_string()?.append("  ")?;
                    }

                    Ok(())
                },
            ))
            .unwrap();
        orig.pre_instruction = Some(f);

        let f = inner
            .set_print_address_abs(Box::new(
                |formatter: &zydis::Formatter,
                 buf: &mut zydis::FormatterBuffer,
                 ctx: &mut zydis::FormatterContext,
                 userdata: Option<&mut dyn core::any::Any>|
                 -> zydis::Result<()> {
                    // programming error: userdata must be provided. this is guaranteed within
                    // Formatter.
                    let userdata = userdata.expect("no userdata");

                    // programming error: userdata must be a Box<UserData>. this is guaranteed
                    // within Formatter.
                    let userdata = userdata.downcast_ref::<UserData>().expect("incorrect userdata");

                    let absolute_address = unsafe {
                        // safety: the insn and operands come from zydis, so we assume they contain
                        // valid data.
                        let insn: &zydis::DecodedInstruction = &*ctx.instruction;
                        let op: &zydis::DecodedOperand = &*ctx.operand;
                        insn.calc_absolute_address(ctx.runtime_address, op)
                            .expect("failed to calculate absolute address")
                    };

                    if let Some(name) = userdata.ws.analysis.names.names_by_address.get(&absolute_address) {
                        // name is found in map, use that.
                        buf.append(TOKEN_USER_SYMBOLNAME)?;
                        return buf.get_string()?.append(name);
                    } else {
                        // name is not found, use original formatter.

                        // programming error: the original hook must be recorded. this is guaranteed
                        // within Formatter.
                        let orig = userdata.orig.print_address_abs.as_ref().expect("no original hook");

                        if let zydis::Hook::PrintAddressAbs(Some(f)) = orig {
                            // safety: zydis::Formatter <-> zydis::ffi::ZydisFormatter is safe according to
                            // here: https://docs.rs/zydis/3.1.2/src/zydis/formatter.rs.html#306
                            let status =
                                unsafe { f(formatter as *const _ as *const zydis::ffi::ZydisFormatter, buf, ctx) };
                            if status.is_error() {
                                return Err(status);
                            } else {
                                return Ok(());
                            }
                        } else {
                            // I'm not sure how this could ever be the case, as zydis initializes the hook
                            // with a default. I suppose if you explicitly set
                            // the callback to NULL/None? Which we don't do here.
                            panic!("unexpected original hook");
                        }
                    }
                },
            ))
            .unwrap();
        orig.print_address_abs = Some(f);

        let f = inner
            .set_print_mnemonic(Box::new(
                |formatter: &zydis::Formatter,
                 buf: &mut zydis::FormatterBuffer,
                 ctx: &mut zydis::FormatterContext,
                 userdata: Option<&mut dyn core::any::Any>|
                 -> zydis::Result<()> {
                    // programming error: userdata must be provided. this is guaranteed within
                    // Formatter.
                    let userdata = userdata.expect("no userdata");

                    // programming error: userdata must be a Box<UserData>. this is guaranteed
                    // within Formatter.
                    let userdata = userdata.downcast_ref::<UserData>().expect("incorrect userdata");

                    let orig = userdata.orig.print_mnemonic.as_ref().expect("no original hook");

                    if let zydis::Hook::PrintMnemonic(Some(f)) = orig {
                        // safety: zydis::Formatter <-> zydis::ffi::ZydisFormatter is safe according to
                        // here: https://docs.rs/zydis/3.1.2/src/zydis/formatter.rs.html#306
                        let status = unsafe { f(formatter as *const _ as *const zydis::ffi::ZydisFormatter, buf, ctx) };
                        if status.is_error() {
                            return Err(status);
                        } else {
                            let (_, mnemonic) = buf.get_token()?.get_value()?;

                            if mnemonic.len() < userdata.options.mnemonic_width {
                                let mut padding = String::new();

                                for _ in 0..userdata.options.mnemonic_width - mnemonic.len() {
                                    padding.write_str(" ").unwrap();
                                }

                                buf.append(zydis::TOKEN_WHITESPACE)?;
                                buf.get_string()?.append(&padding)?;
                            }

                            return Ok(());
                        }
                    } else {
                        // I'm not sure how this could ever be the case, as zydis initializes the hook
                        // with a default. I suppose if you explicitly set
                        // the callback to NULL/None? Which we don't do here.
                        panic!("unexpected original hook");
                    }
                },
            ))
            .unwrap();
        orig.print_mnemonic = Some(f);

        Formatter { options, inner, orig }
    }

    fn get_token_color(token: zydis::Token) -> ansi_term::Color {
        match token {
            zydis::TOKEN_INVALID => Formatter::COLOR_INVALID,
            zydis::TOKEN_WHITESPACE => Formatter::COLOR_WHITESPACE,
            zydis::TOKEN_DELIMITER => Formatter::COLOR_DELIMITER,
            zydis::TOKEN_PARENTHESIS_OPEN => Formatter::COLOR_PARENTHESIS_OPEN,
            zydis::TOKEN_PARENTHESIS_CLOSE => Formatter::COLOR_PARENTHESIS_CLOSE,
            zydis::TOKEN_PREFIX => Formatter::COLOR_PREFIX,
            zydis::TOKEN_MNEMONIC => Formatter::COLOR_MNEMONIC,
            zydis::TOKEN_REGISTER => Formatter::COLOR_REGISTER,
            zydis::TOKEN_ADDRESS_ABS => Formatter::COLOR_ADDRESS_ABS,
            zydis::TOKEN_ADDRESS_REL => Formatter::COLOR_ADDRESS_REL,
            zydis::TOKEN_DISPLACEMENT => Formatter::COLOR_DISPLACEMENT,
            zydis::TOKEN_IMMEDIATE => Formatter::COLOR_IMMEDIATE,
            zydis::TOKEN_TYPECAST => Formatter::COLOR_TYPECAST,
            zydis::TOKEN_DECORATOR => Formatter::COLOR_DECORATOR,
            zydis::TOKEN_SYMBOL => Formatter::COLOR_SYMBOL,
            zydis::TOKEN_USER => Formatter::COLOR_USER,
            TOKEN_USER_SYMBOLNAME => Formatter::COLOR_SYMBOLNAME,
            TOKEN_USER_HEX => Formatter::COLOR_HEX,
            _ => unimplemented!("token: {}", token),
        }
    }

    fn render_token<T: Write>(&self, o: &mut T, token: zydis::Token, s: &str) -> Result<()> {
        if self.options.colors {
            o.write_str(&Formatter::get_token_color(token).paint(s).to_string())?;
        } else {
            o.write_str(s)?;
        }

        Ok(())
    }

    pub fn format_instruction(&self, ws: &PEWorkspace, insn: &zydis::DecodedInstruction, va: VA) -> Result<String> {
        let mut buffer = [0u8; 200];

        // we pass our userdata to ZydisFormatterFormatInstruction.
        // but to make it work, we have to play games with the lifetimes:
        // we need to convince the compiler that the userdata pointer lives long enough
        // to by used by the callbacks.
        //
        // we do this by extending the CFG lifetime from '_ to 'static.
        //
        // userdata is passed into ZydisFormatterFormatInstruction,
        // which passes userdata to each of the formatter callbacks.
        // those read strictly from insn/ctx/userdata and write strictly to output
        // buffer. there is no state maintained within these routines.
        // the callbacks won't be invoked beyond the call into FormatInstruction.
        //
        // therefore, i believe its safe to extend the lifetime here to work with zydis.
        let x = unsafe { std::mem::transmute::<&'_ PEWorkspace, &'static PEWorkspace>(ws) };

        let mut ud = UserData {
            orig:    self.orig.clone(),
            ws:      x,
            options: self.options,
        };

        let mut out = String::new();
        for (token, s) in self
            .inner
            .tokenize_instruction(&insn, &mut buffer, Some(va), Some(&mut ud))?
        {
            self.render_token(&mut out, token, s)?;
        }

        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::{super::*, *};
    use crate::rsrc::*;
    use anyhow::Result;

    #[test]
    fn with_colors() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let ws = PEWorkspace::from_pe(pe)?;

        let fmt = Formatter::new().with_colors(true).build();

        // ```
        //     .text:00401C4E 000 68 F4 61 40 00          push    offset ModuleName ; "mscoree.dll"
        //     .text:00401C53 004 FF 15 00 60 40 00       call    ds:GetModuleHandleA
        //     .text:00401C59 000 85 C0                   test    eax, eax
        // ```
        let insn = crate::test::read_insn(&ws.pe.module, 0x401C53);
        let s = fmt.format_instruction(&ws, &insn, 0x401C53)?;
        assert!(s.contains("\u{1b}"));
        assert!(s.contains("call"));
        assert!(s.contains("GetModuleHandleA"));

        Ok(())
    }

    #[test]
    fn no_colors() -> Result<()> {
        let buf = get_buf(Rsrc::NOP);
        let pe = crate::loader::pe::PE::from_bytes(&buf)?;

        let ws = PEWorkspace::from_pe(pe)?;

        let fmt = Formatter::new().with_colors(false).build();

        // ```
        //     .text:00401C4E 000 68 F4 61 40 00          push    offset ModuleName ; "mscoree.dll"
        //     .text:00401C53 004 FF 15 00 60 40 00       call    ds:GetModuleHandleA
        //     .text:00401C59 000 85 C0                   test    eax, eax
        // ```
        let insn = crate::test::read_insn(&ws.pe.module, 0x401C53);
        assert_eq!(
            fmt.format_instruction(&ws, &insn, 0x401C53)?,
            "call [kernel32.dll!GetModuleHandleA]"
        );

        Ok(())
    }
}
<|endoftext|>