<commit_before>use std::cell::{Ref, RefCell};
use std::cmp;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::rc::{Rc, Weak};
use log::LogLevel;
use mio::{Event, PollOpt, Ready, Token};
use mio::net::TcpStream;

use super::binary;
use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::ipv4_header::IPv4Header;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packet_storage::PacketStorage;
use super::packetizer::Packetizer;
use super::selector::Selector;
use super::stream_buffer::StreamBuffer;
use super::tcp_header::{self, TCPHeader};
use super::transport_header::{TransportHeader, TransportHeaderMut};

const TAG: &'static str = "TCPConnection";

const MAX_PAYLOAD_LENGTH: u16 = 1400;

pub struct TCPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    stream: TcpStream,
    token: Token,
    client_to_network: StreamBuffer,
    network_to_client: Packetizer,
    pending_packet_for_client: PacketStorage,
    packet_for_client_length: Option<u16>,
    closed: bool,
    tcb: TCB,
}

struct TCB {
    state: TCPState,
    syn_sequence_number: u32,
    sequence_number: u32,
    acknowledgement_number: u32,
    their_acknowledgement_number: u32,
    client_window: u16,
    remote_closed: bool,
}

#[derive(PartialEq, Eq)]
enum TCPState {
    Init,
    SynSent,
//    SynReceived,
    Established,
    CloseWait,
    LastAck,
}

impl TCB {
    fn new() -> Self {
        Self {
            state: TCPState::Init,
            syn_sequence_number: 0,
            sequence_number: 0,
            acknowledgement_number: 0,
            their_acknowledgement_number: 0,
            client_window: 0,
            remote_closed: false,
        }
    }

    fn numbers(&self) -> String {
        format!("(seq={}, ack={})", self.sequence_number, self.acknowledgement_number)
    }
}

impl TCPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, ipv4_header: &IPv4Header, transport_header: &TransportHeader) -> io::Result<Rc<RefCell<Self>>> {
        let stream = Self::create_stream(&id)?;

        if let TransportHeader::TCP(ref tcp_header) = *transport_header {
            // shrink the TCP options to pass a minimal refrence header to the packetizer
            let mut shrinked_tcp_header_raw = [0u8; 20];
            shrinked_tcp_header_raw.copy_from_slice(&tcp_header.raw()[..20]);
            let mut shrinked_tcp_header_data = tcp_header.data().clone();
            {
                let mut shrinked_tcp_header = shrinked_tcp_header_data.bind_mut(&mut shrinked_tcp_header_raw);
                shrinked_tcp_header.shrink_options();
                assert_eq!(20, shrinked_tcp_header.header_length());
            }

            let shrinked_transport_header = shrinked_tcp_header_data.bind(&shrinked_tcp_header_raw).into();

            let packetizer = Packetizer::new(&ipv4_header, &shrinked_transport_header);

            let rc = Rc::new(RefCell::new(Self {
                id: id,
                client: client,
                stream: stream,
                token: Token(0), // default value, will be set afterwards
                client_to_network: StreamBuffer::new(4 * MAX_PACKET_LENGTH),
                network_to_client: packetizer,
                pending_packet_for_client: PacketStorage::new(),
                packet_for_client_length: None,
                closed: false,
                tcb: TCB::new(),
            }));

            {
                let rc_clone = rc.clone();
                let handler = Box::new(move |selector: &mut Selector, ready| {
                    let mut self_ref = rc_clone.borrow_mut();
                    self_ref.on_ready(selector, ready);
                });
                let mut self_ref = rc.borrow_mut();
                let token = selector.register(&self_ref.stream, handler, Ready::readable(), PollOpt::level())?;
                self_ref.token = token;
            }
            Ok(rc)
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_stream(id: &ConnectionId) -> io::Result<TcpStream> {
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        TcpStream::connect(&rewritten_destination)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        self.disconnect(selector);

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        match self.client_to_network.write_to(&mut self.stream) {
            Ok(w) => if w == 0 {
                self.close(selector);
            },
            Err(err) => {
                error!(target: TAG, "{} Cannot write: {}", self.id, err);
                self.reset_connection(selector);
            },
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        assert!(self.packet_for_client_length.is_none(), "A pending packet was not sent");
        let remaining_client_window = self.get_remaining_client_window();
        assert!(remaining_client_window > 0, "process_received() must not be called when window == 0");
        let max_payload_length = cmp::min(remaining_client_window, MAX_PAYLOAD_LENGTH) as usize;
        Self::update_headers(&mut self.network_to_client, &self.tcb, tcp_header::FLAG_ACK | tcp_header::FLAG_PSH);
        // the packet is bound to the lifetime of self, so we cannot borrow self to call methods
        // defer the other branches in a separate match-block
        let non_lexical_lifetime_workaround = match self.network_to_client.packetize_read(&mut self.stream, Some(max_payload_length)) {
            Ok(Some(ipv4_packet)) => {
                // TODO packet source pull
                Self::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
                Ok(Some(()))
            },
            Ok(None) => Ok(None),
            Err(err) => Err(err),
        };
        match non_lexical_lifetime_workaround {
            Ok(None) => self.eof(selector),
            Err(err) => {
                error!(target: TAG, "{} Cannot read: {}", self.id, err);
                self.reset_connection(selector);
            },
            Ok(Some(_)) => (), // already handled
        }
    }

    fn send_to_client(client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        let client_rc = client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.send_to_client(selector, &ipv4_packet)
    }

    fn send_to_client_ignore(id: &ConnectionId, client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        if let Err(err) = Self::send_to_client(client, selector, ipv4_packet) {
            warn!(target: TAG, "{} Cannot send packet to client: {}", id, err);
        }
    }

    fn eof(&mut self, selector: &mut Selector) {
        self.tcb.remote_closed = true;
        if self.tcb.state == TCPState::CloseWait {
            let ipv4_packet = Self::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_FIN);
            self.tcb.sequence_number += 1; // FIN counts for 1 byte
            Self::send_to_client_ignore(&self. id, &self.client, selector, &ipv4_packet);
        }
    }

    fn update_headers(packetizer: &mut Packetizer, tcb: &TCB, flags: u16) {
        if let TransportHeaderMut::TCP(ref mut tcp_header) = packetizer.transport_header_mut() {
            tcp_header.set_sequence_number(tcb.sequence_number);
            tcp_header.set_acknowledgement_number(tcb.acknowledgement_number);
            tcp_header.set_flags(flags);
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_empty_response_packet<'a>(id: &ConnectionId, packetizer: &'a mut Packetizer, tcb: &TCB, flags: u16) -> IPv4Packet<'a> {
        Self::update_headers(packetizer, tcb, flags);
        debug!(target: TAG, "{} Forging empty response (flags={}) {}", id, flags, tcb.numbers());
        if (flags & tcp_header::FLAG_ACK) != 0 {
            debug!(target: TAG, "{} Acking {}", id, tcb.numbers());
        }
        let ipv4_packet = packetizer.packetize_empty_payload();
        if log_enabled!(target: TAG, LogLevel::Trace) {
            binary::to_string(ipv4_packet.raw());
        }
        ipv4_packet
    }

    fn reset_connection(&mut self, selector: &mut Selector) {
        {
            let ipv4_packet = Self::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_RST);
            Self::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
        }
        self.close(selector);
    }

    fn update_interests(&mut self, selector: &mut Selector) {
        if !self.closed {
            let mut ready = Ready::empty();
            if self.may_read() {
                ready = ready | Ready::readable()
            }
            if self.may_write() {
                ready = ready | Ready::writable()
            }
            selector.reregister(&self.stream, self.token, ready, PollOpt::level()).expect("Cannot register on poll");
        }
    }

    fn may_read(&self) -> bool {
        !self.tcb.remote_closed &&
                self.packet_for_client_length.is_some() &&
                self.get_remaining_client_window() > 0
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }

    fn get_remaining_client_window(&self) -> u16 {
        // TODO
        42
    }
}

impl Connection for TCPConnection {
    fn id(&self) -> &ConnectionId {
        &self.id
    }

    fn send_to_network(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        // TODO
    }

    fn disconnect(&mut self, selector: &mut Selector) {
        info!(target: TAG, "{} Close", self.id);
        selector.deregister(&self.stream, self.token).unwrap();
        // socket will be closed by RAII
    }

    fn is_expired(&self) -> bool {
        // no external timeout expiration
        false
    }
}
<commit_msg>Implement EventHandler for TCPConnection<commit_after>use std::cell::{Ref, RefCell};
use std::cmp;
use std::io;
use std::net::{Ipv4Addr, SocketAddr};
use std::rc::{Rc, Weak};
use log::LogLevel;
use mio::{Event, PollOpt, Ready, Token};
use mio::net::TcpStream;

use super::binary;
use super::client::Client;
use super::connection::{self, Connection, ConnectionId};
use super::ipv4_header::IPv4Header;
use super::ipv4_packet::{IPv4Packet, MAX_PACKET_LENGTH};
use super::packet_storage::PacketStorage;
use super::packetizer::Packetizer;
use super::selector::{EventHandler, Selector};
use super::stream_buffer::StreamBuffer;
use super::tcp_header::{self, TCPHeader};
use super::transport_header::{TransportHeader, TransportHeaderMut};

const TAG: &'static str = "TCPConnection";

const MAX_PAYLOAD_LENGTH: u16 = 1400;

pub struct TCPConnection {
    id: ConnectionId,
    client: Weak<RefCell<Client>>,
    stream: TcpStream,
    token: Token,
    client_to_network: StreamBuffer,
    network_to_client: Packetizer,
    pending_packet_for_client: PacketStorage,
    packet_for_client_length: Option<u16>,
    closed: bool,
    tcb: TCB,
}

struct TCB {
    state: TCPState,
    syn_sequence_number: u32,
    sequence_number: u32,
    acknowledgement_number: u32,
    their_acknowledgement_number: u32,
    client_window: u16,
    remote_closed: bool,
}

#[derive(PartialEq, Eq)]
enum TCPState {
    Init,
    SynSent,
//    SynReceived,
    Established,
    CloseWait,
    LastAck,
}

impl TCB {
    fn new() -> Self {
        Self {
            state: TCPState::Init,
            syn_sequence_number: 0,
            sequence_number: 0,
            acknowledgement_number: 0,
            their_acknowledgement_number: 0,
            client_window: 0,
            remote_closed: false,
        }
    }

    fn numbers(&self) -> String {
        format!("(seq={}, ack={})", self.sequence_number, self.acknowledgement_number)
    }
}

impl TCPConnection {
    pub fn new(selector: &mut Selector, id: ConnectionId, client: Weak<RefCell<Client>>, ipv4_header: &IPv4Header, transport_header: &TransportHeader) -> io::Result<Rc<RefCell<Self>>> {
        let stream = Self::create_stream(&id)?;

        if let TransportHeader::TCP(ref tcp_header) = *transport_header {
            // shrink the TCP options to pass a minimal refrence header to the packetizer
            let mut shrinked_tcp_header_raw = [0u8; 20];
            shrinked_tcp_header_raw.copy_from_slice(&tcp_header.raw()[..20]);
            let mut shrinked_tcp_header_data = tcp_header.data().clone();
            {
                let mut shrinked_tcp_header = shrinked_tcp_header_data.bind_mut(&mut shrinked_tcp_header_raw);
                shrinked_tcp_header.shrink_options();
                assert_eq!(20, shrinked_tcp_header.header_length());
            }

            let shrinked_transport_header = shrinked_tcp_header_data.bind(&shrinked_tcp_header_raw).into();

            let packetizer = Packetizer::new(&ipv4_header, &shrinked_transport_header);

            let rc = Rc::new(RefCell::new(Self {
                id: id,
                client: client,
                stream: stream,
                token: Token(0), // default value, will be set afterwards
                client_to_network: StreamBuffer::new(4 * MAX_PACKET_LENGTH),
                network_to_client: packetizer,
                pending_packet_for_client: PacketStorage::new(),
                packet_for_client_length: None,
                closed: false,
                tcb: TCB::new(),
            }));

            {
                let mut self_ref = rc.borrow_mut();
                // rc is an EventHandler, register() expects a Box<EventHandler>
                let handler = Box::new(rc.clone());
                let token = selector.register(&self_ref.stream, handler, Ready::readable(), PollOpt::level())?;
                self_ref.token = token;
            }
            Ok(rc)
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_stream(id: &ConnectionId) -> io::Result<TcpStream> {
        let rewritten_destination = connection::rewritten_destination(id.destination_ip(), id.destination_port()).into();
        TcpStream::connect(&rewritten_destination)
    }

    fn close(&mut self, selector: &mut Selector) {
        self.closed = true;
        self.disconnect(selector);

        // route is embedded in router which is embedded in client: the client necessarily exists
        let client_rc = self.client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.router().remove(&self.id);
    }

    fn process_send(&mut self, selector: &mut Selector) {
        match self.client_to_network.write_to(&mut self.stream) {
            Ok(w) => if w == 0 {
                self.close(selector);
            },
            Err(err) => {
                error!(target: TAG, "{} Cannot write: {}", self.id, err);
                self.reset_connection(selector);
            },
        }
    }

    fn process_receive(&mut self, selector: &mut Selector) {
        assert!(self.packet_for_client_length.is_none(), "A pending packet was not sent");
        let remaining_client_window = self.get_remaining_client_window();
        assert!(remaining_client_window > 0, "process_received() must not be called when window == 0");
        let max_payload_length = cmp::min(remaining_client_window, MAX_PAYLOAD_LENGTH) as usize;
        Self::update_headers(&mut self.network_to_client, &self.tcb, tcp_header::FLAG_ACK | tcp_header::FLAG_PSH);
        // the packet is bound to the lifetime of self, so we cannot borrow self to call methods
        // defer the other branches in a separate match-block
        let non_lexical_lifetime_workaround = match self.network_to_client.packetize_read(&mut self.stream, Some(max_payload_length)) {
            Ok(Some(ipv4_packet)) => {
                // TODO packet source pull
                Self::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
                Ok(Some(()))
            },
            Ok(None) => Ok(None),
            Err(err) => Err(err),
        };
        match non_lexical_lifetime_workaround {
            Ok(None) => self.eof(selector),
            Err(err) => {
                error!(target: TAG, "{} Cannot read: {}", self.id, err);
                self.reset_connection(selector);
            },
            Ok(Some(_)) => (), // already handled
        }
    }

    fn send_to_client(client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) -> io::Result<()> {
        let client_rc = client.upgrade().expect("expected client not found");
        let mut client = client_rc.borrow_mut();
        client.send_to_client(selector, &ipv4_packet)
    }

    fn send_to_client_ignore(id: &ConnectionId, client: &Weak<RefCell<Client>>, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        if let Err(err) = Self::send_to_client(client, selector, ipv4_packet) {
            warn!(target: TAG, "{} Cannot send packet to client: {}", id, err);
        }
    }

    fn eof(&mut self, selector: &mut Selector) {
        self.tcb.remote_closed = true;
        if self.tcb.state == TCPState::CloseWait {
            let ipv4_packet = Self::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_FIN);
            self.tcb.sequence_number += 1; // FIN counts for 1 byte
            Self::send_to_client_ignore(&self. id, &self.client, selector, &ipv4_packet);
        }
    }

    fn update_headers(packetizer: &mut Packetizer, tcb: &TCB, flags: u16) {
        if let TransportHeaderMut::TCP(ref mut tcp_header) = packetizer.transport_header_mut() {
            tcp_header.set_sequence_number(tcb.sequence_number);
            tcp_header.set_acknowledgement_number(tcb.acknowledgement_number);
            tcp_header.set_flags(flags);
        } else {
            panic!("Not a TCP header");
        }
    }

    fn create_empty_response_packet<'a>(id: &ConnectionId, packetizer: &'a mut Packetizer, tcb: &TCB, flags: u16) -> IPv4Packet<'a> {
        Self::update_headers(packetizer, tcb, flags);
        debug!(target: TAG, "{} Forging empty response (flags={}) {}", id, flags, tcb.numbers());
        if (flags & tcp_header::FLAG_ACK) != 0 {
            debug!(target: TAG, "{} Acking {}", id, tcb.numbers());
        }
        let ipv4_packet = packetizer.packetize_empty_payload();
        if log_enabled!(target: TAG, LogLevel::Trace) {
            binary::to_string(ipv4_packet.raw());
        }
        ipv4_packet
    }

    fn reset_connection(&mut self, selector: &mut Selector) {
        {
            let ipv4_packet = Self::create_empty_response_packet(&self.id, &mut self.network_to_client, &self.tcb, tcp_header::FLAG_RST);
            Self::send_to_client_ignore(&self.id, &self.client, selector, &ipv4_packet);
        }
        self.close(selector);
    }

    fn update_interests(&mut self, selector: &mut Selector) {
        if !self.closed {
            let mut ready = Ready::empty();
            if self.may_read() {
                ready = ready | Ready::readable()
            }
            if self.may_write() {
                ready = ready | Ready::writable()
            }
            selector.reregister(&self.stream, self.token, ready, PollOpt::level()).expect("Cannot register on poll");
        }
    }

    fn may_read(&self) -> bool {
        !self.tcb.remote_closed &&
                self.packet_for_client_length.is_some() &&
                self.get_remaining_client_window() > 0
    }

    fn may_write(&self) -> bool {
        !self.client_to_network.is_empty()
    }

    fn get_remaining_client_window(&self) -> u16 {
        // TODO
        42
    }
}

impl Connection for TCPConnection {
    fn id(&self) -> &ConnectionId {
        &self.id
    }

    fn send_to_network(&mut self, selector: &mut Selector, ipv4_packet: &IPv4Packet) {
        // TODO
    }

    fn disconnect(&mut self, selector: &mut Selector) {
        info!(target: TAG, "{} Close", self.id);
        selector.deregister(&self.stream, self.token).unwrap();
        // socket will be closed by RAII
    }

    fn is_expired(&self) -> bool {
        // no external timeout expiration
        false
    }
}

impl EventHandler for TCPConnection {
    fn on_ready(&mut self, selector: &mut Selector, event: Event) {
        assert!(!self.closed);
        let ready = event.readiness();
        if ready.is_writable() {
            self.process_send(selector);
        }
        if !self.closed && ready.is_readable() {
            self.process_receive(selector);
        }
        if !self.closed {
            self.update_interests(selector);
        }
    }
}
<|endoftext|>


<commit_before>use std::iter;
use std::cmp::min;

use unicode_width::UnicodeWidthStr;

use buffer::Buffer;
use widgets::{Widget, Block};
use layout::Rect;
use style::Style;

pub struct List<'a, T>
    where T: AsRef<str> + 'a
{
    block: Option<Block<'a>>,
    items: &'a [(T, &'a Style)],
    style: Style,
}

impl<'a, T> Default for List<'a, T>
    where T: AsRef<str> + 'a
{
    fn default() -> List<'a, T> {
        List {
            block: None,
            items: &[],
            style: Default::default(),
        }
    }
}

impl<'a, T> List<'a, T>
    where T: AsRef<str> + 'a
{
    pub fn block(&'a mut self, block: Block<'a>) -> &mut List<'a, T> {
        self.block = Some(block);
        self
    }

    pub fn items(&'a mut self, items: &'a [(T, &'a Style)]) -> &mut List<'a, T> {
        self.items = items;
        self
    }

    pub fn style(&'a mut self, style: Style) -> &mut List<'a, T> {
        self.style = style;
        self
    }
}

impl<'a, T> Widget for List<'a, T>
    where T: AsRef<str> + 'a
{
    fn draw(&self, area: &Rect, buf: &mut Buffer) {
        let list_area = match self.block {
            Some(ref b) => {
                b.draw(area, buf);
                b.inner(area)
            }
            None => *area,
        };

        if list_area.width < 1 || list_area.height < 1 {
            return;
        }

        self.background(&list_area, buf, self.style.bg);

        let max_index = min(self.items.len(), list_area.height as usize);
        for (i, &(ref item, style)) in self.items.iter().enumerate().take(max_index) {
            buf.set_stringn(list_area.left(),
                            list_area.top() + i as u16,
                            item.as_ref(),
                            list_area.width as usize,
                            style);
        }
    }
}

/// A widget to display several items among which one can be selected (optional)
///
/// # Examples
///
/// ```
/// # extern crate tui;
/// # use tui::widgets::{Block, border, SelectableList};
/// # use tui::style::{Style, Color, Modifier};
/// # fn main() {
/// SelectableList::default()
///     .block(Block::default().title("SelectableList").borders(border::ALL))
///     .items(&["Item 1", "Item 2", "Item 3"])
///     .select(1)
///     .style(Style::default().fg(Color::White))
///     .highlight_style(Style::default().modifier(Modifier::Italic))
///     .highlight_symbol(">>");
/// # }
/// ```
pub struct SelectableList<'a> {
    block: Option<Block<'a>>,
    /// Items to be displayed
    items: &'a [&'a str],
    /// Index of the one selected
    selected: Option<usize>,
    /// Base style of the widget
    style: Style,
    /// Style used to render selected item
    highlight_style: Style,
    /// Symbol in front of the selected item (Shift all items to the right)
    highlight_symbol: Option<&'a str>,
}

impl<'a> Default for SelectableList<'a> {
    fn default() -> SelectableList<'a> {
        SelectableList {
            block: None,
            items: &[],
            selected: None,
            style: Default::default(),
            highlight_style: Default::default(),
            highlight_symbol: None,
        }
    }
}

impl<'a> SelectableList<'a> {
    pub fn block(&'a mut self, block: Block<'a>) -> &mut SelectableList<'a> {
        self.block = Some(block);
        self
    }

    pub fn items(&'a mut self, items: &'a [&'a str]) -> &mut SelectableList<'a> {
        self.items = items;
        self
    }

    pub fn style(&'a mut self, style: Style) -> &mut SelectableList<'a> {
        self.style = style;
        self
    }

    pub fn highlight_symbol(&'a mut self, highlight_symbol: &'a str) -> &mut SelectableList<'a> {
        self.highlight_symbol = Some(highlight_symbol);
        self
    }

    pub fn highlight_style(&'a mut self, highlight_style: Style) -> &mut SelectableList<'a> {
        self.highlight_style = highlight_style;
        self
    }

    pub fn select(&'a mut self, index: usize) -> &'a mut SelectableList<'a> {
        self.selected = Some(index);
        self
    }
}

impl<'a> Widget for SelectableList<'a> {
    fn draw(&self, area: &Rect, buf: &mut Buffer) {

        let list_area = match self.block {
            Some(ref b) => b.inner(area),
            None => *area,
        };

        let list_height = list_area.height as usize;

        // Use highlight_style only if something is selected
        let (selected, highlight_style) = match self.selected {
            Some(i) => (i, &self.highlight_style),
            None => (0, &self.style),
        };
        let highlight_symbol = self.highlight_symbol.unwrap_or("");
        let blank_symbol = iter::repeat(" ").take(highlight_symbol.width()).collect::<String>();
        // Make sure the list show the selected item
        let offset = if selected >= list_height {
            selected - list_height + 1
        } else {
            0
        };
        let items = self.items
            .iter()
            .cloned()
            .enumerate()
            .map(|(i, item)| if i == selected {
                (format!("{} {}", highlight_symbol, item), highlight_style)
            } else {
                (format!("{} {}", blank_symbol, item), &self.style)
            })
            .skip(offset as usize)
            .collect::<Vec<(String, &Style)>>();

        // Render items
        List::default()
            .block(self.block.unwrap_or_default())
            .items(&items)
            .draw(area, buf);
    }
}
<commit_msg>Make List more user-friendly<commit_after>use std::iter;
use std::cmp::min;

use unicode_width::UnicodeWidthStr;

use buffer::Buffer;
use widgets::{Widget, Block};
use layout::Rect;
use style::Style;

pub struct List<'a> {
    block: Option<Block<'a>>,
    items: Vec<(&'a str, &'a Style)>,
    style: Style,
}

impl<'a> Default for List<'a> {
    fn default() -> List<'a> {
        List {
            block: None,
            items: Vec::new(),
            style: Default::default(),
        }
    }
}

impl<'a> List<'a> {
    pub fn block(&'a mut self, block: Block<'a>) -> &mut List<'a> {
        self.block = Some(block);
        self
    }

    pub fn items<I>(&'a mut self, items: &'a [(I, &'a Style)]) -> &mut List<'a>
        where I: AsRef<str> + 'a
    {
        self.items =
            items.iter().map(|&(ref i, s)| (i.as_ref(), s)).collect::<Vec<(&'a str, &'a Style)>>();
        self
    }

    pub fn style(&'a mut self, style: Style) -> &mut List<'a> {
        self.style = style;
        self
    }
}

impl<'a> Widget for List<'a> {
    fn draw(&self, area: &Rect, buf: &mut Buffer) {
        let list_area = match self.block {
            Some(ref b) => {
                b.draw(area, buf);
                b.inner(area)
            }
            None => *area,
        };

        if list_area.width < 1 || list_area.height < 1 {
            return;
        }

        self.background(&list_area, buf, self.style.bg);

        let max_index = min(self.items.len(), list_area.height as usize);
        for (i, &(ref item, style)) in self.items.iter().enumerate().take(max_index) {
            buf.set_stringn(list_area.left(),
                            list_area.top() + i as u16,
                            item.as_ref(),
                            list_area.width as usize,
                            style);
        }
    }
}

/// A widget to display several items among which one can be selected (optional)
///
/// # Examples
///
/// ```
/// # extern crate tui;
/// # use tui::widgets::{Block, border, SelectableList};
/// # use tui::style::{Style, Color, Modifier};
/// # fn main() {
/// SelectableList::default()
///     .block(Block::default().title("SelectableList").borders(border::ALL))
///     .items(&["Item 1", "Item 2", "Item 3"])
///     .select(1)
///     .style(Style::default().fg(Color::White))
///     .highlight_style(Style::default().modifier(Modifier::Italic))
///     .highlight_symbol(">>");
/// # }
/// ```
pub struct SelectableList<'a> {
    block: Option<Block<'a>>,
    /// Items to be displayed
    items: Vec<&'a str>,
    /// Index of the one selected
    selected: Option<usize>,
    /// Base style of the widget
    style: Style,
    /// Style used to render selected item
    highlight_style: Style,
    /// Symbol in front of the selected item (Shift all items to the right)
    highlight_symbol: Option<&'a str>,
}

impl<'a> Default for SelectableList<'a> {
    fn default() -> SelectableList<'a> {
        SelectableList {
            block: None,
            items: Vec::new(),
            selected: None,
            style: Default::default(),
            highlight_style: Default::default(),
            highlight_symbol: None,
        }
    }
}

impl<'a> SelectableList<'a> {
    pub fn block(&'a mut self, block: Block<'a>) -> &mut SelectableList<'a> {
        self.block = Some(block);
        self
    }

    pub fn items<I>(&'a mut self, items: &'a [I]) -> &mut SelectableList<'a>
        where I: AsRef<str> + 'a
    {
        self.items = items.iter().map(|i| i.as_ref()).collect::<Vec<&str>>();
        self
    }

    pub fn style(&'a mut self, style: Style) -> &mut SelectableList<'a> {
        self.style = style;
        self
    }

    pub fn highlight_symbol(&'a mut self, highlight_symbol: &'a str) -> &mut SelectableList<'a> {
        self.highlight_symbol = Some(highlight_symbol);
        self
    }

    pub fn highlight_style(&'a mut self, highlight_style: Style) -> &mut SelectableList<'a> {
        self.highlight_style = highlight_style;
        self
    }

    pub fn select(&'a mut self, index: usize) -> &'a mut SelectableList<'a> {
        self.selected = Some(index);
        self
    }
}

impl<'a> Widget for SelectableList<'a> {
    fn draw(&self, area: &Rect, buf: &mut Buffer) {

        let list_area = match self.block {
            Some(ref b) => b.inner(area),
            None => *area,
        };

        let list_height = list_area.height as usize;

        // Use highlight_style only if something is selected
        let (selected, highlight_style) = match self.selected {
            Some(i) => (i, &self.highlight_style),
            None => (0, &self.style),
        };
        let highlight_symbol = self.highlight_symbol.unwrap_or("");
        let blank_symbol = iter::repeat(" ").take(highlight_symbol.width()).collect::<String>();
        // Make sure the list show the selected item
        let offset = if selected >= list_height {
            selected - list_height + 1
        } else {
            0
        };
        let items = self.items
            .iter()
            .cloned()
            .enumerate()
            .map(|(i, item)| if i == selected {
                (format!("{} {}", highlight_symbol, item), highlight_style)
            } else {
                (format!("{} {}", blank_symbol, item), &self.style)
            })
            .skip(offset as usize)
            .collect::<Vec<(String, &Style)>>();

        // Render items
        List::default()
            .block(self.block.unwrap_or_default())
            .items(&items)
            .draw(area, buf);
    }
}
<|endoftext|>


<commit_before>#![allow(unused_must_use)] // because we of writeln!

//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Inokentiy Babushkin <inokentiy.babushkin@googlemail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (ToDO) lstat

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg};
use std::fs;
use std::io::{ErrorKind, Write};
use uucore::InvalidEncodingHandling;

// operating mode
enum Mode {
    Default, // use filesystem to determine information and limits
    Basic,   // check basic compatibility with POSIX
    Extra,   // check for leading dashes and empty names
    Both,    // a combination of `Basic` and `Extra`
}

static NAME: &str = "pathchk";
static ABOUT: &str = "Check whether file names are valid or portable";

mod options {
    pub const POSIX: &str = "posix";
    pub const POSIX_SPECIAL: &str = "posix-special";
    pub const PORTABILITY: &str = "portability";
    pub const PATH: &str = "path";
}

// a few global constants as used in the GNU implementation
const POSIX_PATH_MAX: usize = 256;
const POSIX_NAME_MAX: usize = 14;

fn get_usage() -> String {
    format!("{0} [OPTION]... NAME...", executable!())
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = get_usage();
    let args = args
        .collect_str(InvalidEncodingHandling::ConvertLossy)
        .accept_any();

    let matches = uu_app().usage(&usage[..]).get_matches_from(args);

    // set working mode
    let is_posix = matches.values_of(options::POSIX).is_some();
    let is_posix_special = matches.values_of(options::POSIX_SPECIAL).is_some();
    let is_portability = matches.values_of(options::PORTABILITY).is_some();

    let mode = if (is_posix && is_posix_special) || is_portability {
        Mode::Both
    } else if is_posix {
        Mode::Basic
    } else if is_posix_special {
        Mode::Extra
    } else {
        Mode::Default
    };

    // take necessary actions
    let paths = matches.values_of(options::PATH);
    let mut res = if paths.is_none() {
        show_error!("missing operand\nTry {} --help for more information", NAME);
        false
    } else {
        true
    };

    if res {
        // free strings are path operands
        // FIXME: TCS, seems inefficient and overly verbose (?)
        for p in paths.unwrap() {
            let mut path = Vec::new();
            for path_segment in p.split('/') {
                path.push(path_segment.to_string());
            }
            res &= check_path(&mode, &path);
        }
    }

    // determine error code
    if res {
        0
    } else {
        1
    }
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(executable!())
        .version(crate_version!())
        .about(ABOUT)
        .arg(
            Arg::with_name(options::POSIX)
                .short("p")
                .help("check for most POSIX systems"),
        )
        .arg(
            Arg::with_name(options::POSIX_SPECIAL)
                .short("P")
                .help(r#"check for empty names and leading "-""#),
        )
        .arg(
            Arg::with_name(options::PORTABILITY)
                .long(options::PORTABILITY)
                .help("check for all POSIX systems (equivalent to -p -P)"),
        )
        .arg(Arg::with_name(options::PATH).hidden(true).multiple(true))
}

// check a path, given as a slice of it's components and an operating mode
fn check_path(mode: &Mode, path: &[String]) -> bool {
    match *mode {
        Mode::Basic => check_basic(path),
        Mode::Extra => check_default(path) && check_extra(path),
        Mode::Both => check_basic(path) && check_extra(path),
        _ => check_default(path),
    }
}

// check a path in basic compatibility mode
fn check_basic(path: &[String]) -> bool {
    let joined_path = path.join("/");
    let total_len = joined_path.len();
    // path length
    if total_len > POSIX_PATH_MAX {
        writeln!(
            &mut std::io::stderr(),
            "limit {} exceeded by length {} of file name {}",
            POSIX_PATH_MAX,
            total_len,
            joined_path
        );
        return false;
    } else if total_len == 0 {
        writeln!(&mut std::io::stderr(), "empty file name");
        return false;
    }
    // components: character portability and length
    for p in path {
        let component_len = p.len();
        if component_len > POSIX_NAME_MAX {
            writeln!(
                &mut std::io::stderr(),
                "limit {} exceeded by length {} of file name component '{}'",
                POSIX_NAME_MAX,
                component_len,
                p
            );
            return false;
        }
        if !check_portable_chars(p) {
            return false;
        }
    }
    // permission checks
    check_searchable(&joined_path)
}

// check a path in extra compatibility mode
fn check_extra(path: &[String]) -> bool {
    // components: leading hyphens
    for p in path {
        if !no_leading_hyphen(p) {
            writeln!(
                &mut std::io::stderr(),
                "leading hyphen in file name component '{}'",
                p
            );
            return false;
        }
    }
    // path length
    if path.join("/").is_empty() {
        writeln!(&mut std::io::stderr(), "empty file name");
        return false;
    }
    true
}

// check a path in default mode (using the file system)
fn check_default(path: &[String]) -> bool {
    let joined_path = path.join("/");
    let total_len = joined_path.len();
    // path length
    if total_len > libc::PATH_MAX as usize {
        writeln!(
            &mut std::io::stderr(),
            "limit {} exceeded by length {} of file name '{}'",
            libc::PATH_MAX,
            total_len,
            joined_path
        );
        return false;
    }
    // components: length
    for p in path {
        let component_len = p.len();
        if component_len > libc::FILENAME_MAX as usize {
            writeln!(
                &mut std::io::stderr(),
                "limit {} exceeded by length {} of file name component '{}'",
                libc::FILENAME_MAX,
                component_len,
                p
            );
            return false;
        }
    }
    // permission checks
    check_searchable(&joined_path)
}

// check whether a path is or if other problems arise
fn check_searchable(path: &str) -> bool {
    // we use lstat, just like the original implementation
    match fs::symlink_metadata(path) {
        Ok(_) => true,
        Err(e) => {
            if e.kind() == ErrorKind::NotFound {
                true
            } else {
                writeln!(&mut std::io::stderr(), "{}", e);
                false
            }
        }
    }
}

// check for a hyphen at the beginning of a path segment
fn no_leading_hyphen(path_segment: &str) -> bool {
    !path_segment.starts_with('-')
}

// check whether a path segment contains only valid (read: portable) characters
fn check_portable_chars(path_segment: &str) -> bool {
    const VALID_CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._-";
    for (i, ch) in path_segment.as_bytes().iter().enumerate() {
        if !VALID_CHARS.contains(ch) {
            let invalid = path_segment[i..].chars().next().unwrap();
            writeln!(
                &mut std::io::stderr(),
                "nonportable character '{}' in file name component '{}'",
                invalid,
                path_segment
            );
            return false;
        }
    }
    true
}
<commit_msg>pathchk: remove double negation<commit_after>#![allow(unused_must_use)] // because we of writeln!

//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Inokentiy Babushkin <inokentiy.babushkin@googlemail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (ToDO) lstat

#[macro_use]
extern crate uucore;

use clap::{crate_version, App, Arg};
use std::fs;
use std::io::{ErrorKind, Write};
use uucore::InvalidEncodingHandling;

// operating mode
enum Mode {
    Default, // use filesystem to determine information and limits
    Basic,   // check basic compatibility with POSIX
    Extra,   // check for leading dashes and empty names
    Both,    // a combination of `Basic` and `Extra`
}

static NAME: &str = "pathchk";
static ABOUT: &str = "Check whether file names are valid or portable";

mod options {
    pub const POSIX: &str = "posix";
    pub const POSIX_SPECIAL: &str = "posix-special";
    pub const PORTABILITY: &str = "portability";
    pub const PATH: &str = "path";
}

// a few global constants as used in the GNU implementation
const POSIX_PATH_MAX: usize = 256;
const POSIX_NAME_MAX: usize = 14;

fn get_usage() -> String {
    format!("{0} [OPTION]... NAME...", executable!())
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = get_usage();
    let args = args
        .collect_str(InvalidEncodingHandling::ConvertLossy)
        .accept_any();

    let matches = uu_app().usage(&usage[..]).get_matches_from(args);

    // set working mode
    let is_posix = matches.values_of(options::POSIX).is_some();
    let is_posix_special = matches.values_of(options::POSIX_SPECIAL).is_some();
    let is_portability = matches.values_of(options::PORTABILITY).is_some();

    let mode = if (is_posix && is_posix_special) || is_portability {
        Mode::Both
    } else if is_posix {
        Mode::Basic
    } else if is_posix_special {
        Mode::Extra
    } else {
        Mode::Default
    };

    // take necessary actions
    let paths = matches.values_of(options::PATH);
    let mut res = if paths.is_none() {
        show_error!("missing operand\nTry {} --help for more information", NAME);
        false
    } else {
        true
    };

    if res {
        // free strings are path operands
        // FIXME: TCS, seems inefficient and overly verbose (?)
        for p in paths.unwrap() {
            let mut path = Vec::new();
            for path_segment in p.split('/') {
                path.push(path_segment.to_string());
            }
            res &= check_path(&mode, &path);
        }
    }

    // determine error code
    if res {
        0
    } else {
        1
    }
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(executable!())
        .version(crate_version!())
        .about(ABOUT)
        .arg(
            Arg::with_name(options::POSIX)
                .short("p")
                .help("check for most POSIX systems"),
        )
        .arg(
            Arg::with_name(options::POSIX_SPECIAL)
                .short("P")
                .help(r#"check for empty names and leading "-""#),
        )
        .arg(
            Arg::with_name(options::PORTABILITY)
                .long(options::PORTABILITY)
                .help("check for all POSIX systems (equivalent to -p -P)"),
        )
        .arg(Arg::with_name(options::PATH).hidden(true).multiple(true))
}

// check a path, given as a slice of it's components and an operating mode
fn check_path(mode: &Mode, path: &[String]) -> bool {
    match *mode {
        Mode::Basic => check_basic(path),
        Mode::Extra => check_default(path) && check_extra(path),
        Mode::Both => check_basic(path) && check_extra(path),
        _ => check_default(path),
    }
}

// check a path in basic compatibility mode
fn check_basic(path: &[String]) -> bool {
    let joined_path = path.join("/");
    let total_len = joined_path.len();
    // path length
    if total_len > POSIX_PATH_MAX {
        writeln!(
            &mut std::io::stderr(),
            "limit {} exceeded by length {} of file name {}",
            POSIX_PATH_MAX,
            total_len,
            joined_path
        );
        return false;
    } else if total_len == 0 {
        writeln!(&mut std::io::stderr(), "empty file name");
        return false;
    }
    // components: character portability and length
    for p in path {
        let component_len = p.len();
        if component_len > POSIX_NAME_MAX {
            writeln!(
                &mut std::io::stderr(),
                "limit {} exceeded by length {} of file name component '{}'",
                POSIX_NAME_MAX,
                component_len,
                p
            );
            return false;
        }
        if !check_portable_chars(p) {
            return false;
        }
    }
    // permission checks
    check_searchable(&joined_path)
}

// check a path in extra compatibility mode
fn check_extra(path: &[String]) -> bool {
    // components: leading hyphens
    for p in path {
        if p.starts_with('-') {
            writeln!(
                &mut std::io::stderr(),
                "leading hyphen in file name component '{}'",
                p
            );
            return false;
        }
    }
    // path length
    if path.join("/").is_empty() {
        writeln!(&mut std::io::stderr(), "empty file name");
        return false;
    }
    true
}

// check a path in default mode (using the file system)
fn check_default(path: &[String]) -> bool {
    let joined_path = path.join("/");
    let total_len = joined_path.len();
    // path length
    if total_len > libc::PATH_MAX as usize {
        writeln!(
            &mut std::io::stderr(),
            "limit {} exceeded by length {} of file name '{}'",
            libc::PATH_MAX,
            total_len,
            joined_path
        );
        return false;
    }
    // components: length
    for p in path {
        let component_len = p.len();
        if component_len > libc::FILENAME_MAX as usize {
            writeln!(
                &mut std::io::stderr(),
                "limit {} exceeded by length {} of file name component '{}'",
                libc::FILENAME_MAX,
                component_len,
                p
            );
            return false;
        }
    }
    // permission checks
    check_searchable(&joined_path)
}

// check whether a path is or if other problems arise
fn check_searchable(path: &str) -> bool {
    // we use lstat, just like the original implementation
    match fs::symlink_metadata(path) {
        Ok(_) => true,
        Err(e) => {
            if e.kind() == ErrorKind::NotFound {
                true
            } else {
                writeln!(&mut std::io::stderr(), "{}", e);
                false
            }
        }
    }
}

// check whether a path segment contains only valid (read: portable) characters
fn check_portable_chars(path_segment: &str) -> bool {
    const VALID_CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._-";
    for (i, ch) in path_segment.as_bytes().iter().enumerate() {
        if !VALID_CHARS.contains(ch) {
            let invalid = path_segment[i..].chars().next().unwrap();
            writeln!(
                &mut std::io::stderr(),
                "nonportable character '{}' in file name component '{}'",
                invalid,
                path_segment
            );
            return false;
        }
    }
    true
}
<|endoftext|>


<commit_before>use std::path::{Path, PathBuf};
use std::collections::BTreeMap;

pub fn find_cmd(cmd: &str, env: &BTreeMap<String, String>)
    -> Result<PathBuf, String>
{
    if cmd.contains("/") {
        return Ok(Path::new(cmd));
    } else {
        if let Some(paths) = env.get(&"PATH".to_string()) {
            for dir in paths[..].split(':') {
                let path = Path::new(dir);
                if !path.is_absolute() {
                    warn!("All items in PATH must be absolute, not {}",
                          path.display());
                    continue;
                }
                let path = path.join(cmd);
                if path.exists() {
                    return Ok(path);
                }
            }
            return Err(format!("Command {} not found in {:?}",
                cmd, paths));
        } else {
            return Err(format!("Command {} is not absolute and no PATH set",
                cmd));
        }
    }
}
<commit_msg>Fex errors in util.rs<commit_after>use std::path::{Path, PathBuf};
use std::fs::PathExt;
use std::collections::BTreeMap;

pub fn find_cmd(cmd: &str, env: &BTreeMap<String, String>)
    -> Result<PathBuf, String>
{
    if cmd.contains("/") {
        return Ok(PathBuf::from(cmd));
    } else {
        if let Some(paths) = env.get(&"PATH".to_string()) {
            for dir in paths[..].split(':') {
                let path = Path::new(dir);
                if !path.is_absolute() {
                    warn!("All items in PATH must be absolute, not {}",
                          path.display());
                    continue;
                }
                let path = path.join(cmd);
                if path.exists() {
                    return Ok(path);
                }
            }
            return Err(format!("Command {} not found in {:?}",
                cmd, paths));
        } else {
            return Err(format!("Command {} is not absolute and no PATH set",
                cmd));
        }
    }
}
<|endoftext|>


<commit_before>use byteorder::{ByteOrder, LittleEndian};
use log::{debug};
use goblin::{Object};
use failure::{Error, Fail};

use super::super::super::arch::{RVA};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

#[derive(Debug, Fail)]
pub enum RelocAnalyzerError {
    #[fail(display = "Invalid relocation type")]
    InvalidRelocType,
}

pub struct RelocAnalyzer {}

impl RelocAnalyzer {
    pub fn new() -> RelocAnalyzer {
        RelocAnalyzer {}
    }
}

struct Section {
    start: RVA,
    end: RVA,
}

fn is_in_insn(ws: &Workspace, rva: RVA) -> bool {
    let start: usize = rva.into();
    // TODO: underflow
    // TODO: remove harded max insn length
    let end: usize = start - 0x10;

    for i in (start..end).rev() {
        let i = RVA::from(i);
        if let Some(meta) = ws.get_meta(i) {
            if !meta.is_insn() {
                continue;
            }

            if let Ok(len) = meta.get_insn_length() {
                if i + len > rva {
                    return true;
                }
            }
        }
    }
    return false;
}

fn is_ptr(ws: &Workspace, rva: RVA) -> bool {
    if let Ok(ptr) = ws.read_va(rva) {
        if let Some(ptr) = ws.rva(ptr) {
            return ws.probe(ptr, 1);
        }
    }
    return false;
}

pub enum RelocationType {
    ImageRelBasedAbsolute,
    ImageRelBasedHigh,
    ImageRelBasedLow,
    ImageRelBasedHighLow,
    ImageRelBasedHighAdj,

    // ImageRelBasedMIPS_JmpAddr,
    // ImageRelBasedARM_MOV32,
    // ImageRelBasedRiscV_High20,
    ImageRelArch1,

    ImageRelReserved,

    // ImageRelBasedTHUMB_MOV32,
    // ImageRelBasedRiscV_Low12I,
    ImageRelArch2,

    ImageRelBasedRiscVLow12S,
    ImageRelBasedMIPSJmpAddr16,
    ImageRelBasedDir64,
}

pub struct Reloc {
    pub typ: RelocationType,
    pub offset: RVA,
}

fn parse_reloc(base: RVA, entry: u16) -> Result<Reloc, Error> {
    let reloc_type   = (entry & 0b1111000000000000) >> 12;
    let reloc_offset =  entry & 0b0000111111111111;

    // TODO: this should probably be on `RelocationType` as a `From` trait.
    let reloc_type = match reloc_type {
        0 => RelocationType::ImageRelBasedAbsolute,
        1 => RelocationType::ImageRelBasedHigh,
        2 => RelocationType::ImageRelBasedLow,
        3 => RelocationType::ImageRelBasedHighLow,
        4 => RelocationType::ImageRelBasedHighAdj,
        5 => RelocationType::ImageRelArch1,
        6 => RelocationType::ImageRelReserved,
        7 => RelocationType::ImageRelArch2,
        8 => RelocationType::ImageRelBasedRiscVLow12S,
        9 => RelocationType::ImageRelBasedMIPSJmpAddr16,
        10 => RelocationType::ImageRelBasedDir64,
        _ => return Err(RelocAnalyzerError::InvalidRelocType.into()),
    };

    Ok(Reloc {
        typ: reloc_type,
        offset: base + RVA::from(reloc_offset as i64),
    })
}

fn split_u32(e: u32) -> (u16, u16) {
    let m = ((e & 0x0000FFFF) >>  0) as u16;
    let n = ((e & 0xFFFF0000) >> 16) as u16;
    (m, n)
}

/// ```
/// use lancelot::rsrc::*;
/// use lancelot::arch::*;
/// use lancelot::workspace::Workspace;
/// use lancelot::analysis::pe::relocs;
///
/// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
///    .disable_analysis()
///    .load().unwrap();
/// let relocs = relocs::get_relocs(&ws).unwrap();
/// assert_eq!(relocs[0].offset,   0x76008);
/// assert_eq!(relocs[277].offset, 0xA8000);
/// ```
pub fn get_relocs(ws: &Workspace) -> Result<Vec<Reloc>, Error> {
    let pe = match Object::parse(&ws.buf) {
        Ok(Object::PE(pe)) => pe,
        _ => return Ok(vec![]),
    };

    let opt_header = match pe.header.optional_header {
        Some(opt_header) => opt_header,
        _ => return Ok(vec![]),
    };

    let reloc_directory = match opt_header.data_directories.get_base_relocation_table() {
        Some(reloc_directory) => reloc_directory,
        _ => return Ok(vec![]),
    };

    let dir_start = RVA::from(reloc_directory.virtual_address as i64);
    let buf = ws.read_bytes(dir_start, reloc_directory.size as usize)?;

    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    let mut ret = vec![];
    let mut index = 0;
    loop {
        // parse base relocation chunk
        // until there's no data left

        let page_rva = entries.get(index);
        let block_size = entries.get(index+1);
        let (page_rva, block_size) = match (page_rva, block_size) {
            (Some(0),        _               ) => break,
            (_,              Some(0)         ) => break,
            (Some(page_rva), Some(block_size)) => (RVA::from(*page_rva as i64), *block_size),
            _ => break,
        };

        // cast from u32 to usize
        let chunk_entries_count = ((block_size) / 4) as usize;

        for i in 2..chunk_entries_count {
            if let Some(&entry) = entries.get(index + i) {
                let (m, n) = split_u32(entry);

                let reloc1 = parse_reloc(page_rva, m)?;
                if !ws.probe(reloc1.offset, 4) {
                    break
                }
                ret.push(reloc1);

                let reloc2 = parse_reloc(page_rva, n)?;
                if !ws.probe(reloc2.offset, 4) {
                    break
                }
                ret.push(reloc2);
            } else {
                break
            }
        };

        index += chunk_entries_count;
    }

    Ok(ret)
}

impl Analyzer for RelocAnalyzer {
    fn get_name(&self) -> String {
        "PE relocation analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::RelocAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// let anal = RelocAnalyzer::new();
    /// anal.analyze(&mut ws).unwrap();
    /// let meta = ws.get_meta(0xC7F0).unwrap();
    /// assert!(meta.is_insn());
    /// ```
    fn analyze(&self, ws: &mut Workspace)-> Result<(), Error> {
        let text_section = match ws.module.sections.iter()
            // currently limited to the text section.
            // TODO: accept any writable section.
            .filter(|&sec| sec.name == ".text")
            .next() {
                None => return Ok(()),
                Some(s) => s,
            };

        let text_bounds = Section {
            start: text_section.addr,
            end: text_section.addr + RVA::from(text_section.buf.len()),
        };

        // scan the relocations
        // looking for pointers into the .text section
        // to things that
        //   1. are not already in an instruction
        //   2. don't appear to be a pointer
        // and assume this is code.
        let o: Vec<RVA> = get_relocs(ws)?
            .iter()
            .map(|reloc| reloc.offset)
            .map(|rva| ws.read_va(rva))
            .filter_map(Result::ok)
            .filter_map(|va| ws.rva(va))
            .filter(|&rva| text_bounds.start <= rva)
            .filter(|&rva| rva < text_bounds.end)
            .filter(|&rva| !is_in_insn(ws, rva))
            .filter(|&rva| !is_ptr(ws, rva))
            // TODO: maybe ensure that the insn decodes.
            .collect();

        o.iter().for_each(|&rva| {
            debug!("found ptr from .text section to .text section at {:#x}", rva);

            // TODO: consume result
            ws.make_insn(rva).unwrap();
            // TODO: consume result
            ws.analyze().unwrap();
        });
        Ok(())
    }
}
<commit_msg>analysis/relocs: tests passing again<commit_after>use byteorder::{ByteOrder, LittleEndian};
use log::{debug};
use goblin::{Object};
use failure::{Error, Fail};

use super::super::super::arch::{RVA};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

#[derive(Debug, Fail)]
pub enum RelocAnalyzerError {
    #[fail(display = "Invalid relocation type")]
    InvalidRelocType,
}

pub struct RelocAnalyzer {}

impl RelocAnalyzer {
    pub fn new() -> RelocAnalyzer {
        RelocAnalyzer {}
    }
}

struct Section {
    start: RVA,
    end: RVA,
}

fn is_in_insn(ws: &Workspace, rva: RVA) -> bool {
    let start: usize = rva.into();
    // TODO: underflow
    // TODO: remove harded max insn length
    let end: usize = start - 0x10;

    for i in (start..end).rev() {
        let i = RVA::from(i);
        if let Some(meta) = ws.get_meta(i) {
            if !meta.is_insn() {
                continue;
            }

            if let Ok(len) = meta.get_insn_length() {
                if i + len > rva {
                    return true;
                }
            }
        }
    }
    return false;
}

fn is_ptr(ws: &Workspace, rva: RVA) -> bool {
    if let Ok(ptr) = ws.read_va(rva) {
        if let Some(ptr) = ws.rva(ptr) {
            return ws.probe(ptr, 1);
        }
    }
    return false;
}

pub enum RelocationType {
    ImageRelBasedAbsolute,
    ImageRelBasedHigh,
    ImageRelBasedLow,
    ImageRelBasedHighLow,
    ImageRelBasedHighAdj,

    // ImageRelBasedMIPS_JmpAddr,
    // ImageRelBasedARM_MOV32,
    // ImageRelBasedRiscV_High20,
    ImageRelArch1,

    ImageRelReserved,

    // ImageRelBasedTHUMB_MOV32,
    // ImageRelBasedRiscV_Low12I,
    ImageRelArch2,

    ImageRelBasedRiscVLow12S,
    ImageRelBasedMIPSJmpAddr16,
    ImageRelBasedDir64,
}

pub struct Reloc {
    pub typ: RelocationType,
    pub offset: RVA,
}

fn parse_reloc(base: RVA, entry: u16) -> Result<Reloc, Error> {
    let reloc_type   = (entry & 0b1111000000000000) >> 12;
    let reloc_offset =  entry & 0b0000111111111111;

    // TODO: this should probably be on `RelocationType` as a `From` trait.
    let reloc_type = match reloc_type {
        0 => RelocationType::ImageRelBasedAbsolute,
        1 => RelocationType::ImageRelBasedHigh,
        2 => RelocationType::ImageRelBasedLow,
        3 => RelocationType::ImageRelBasedHighLow,
        4 => RelocationType::ImageRelBasedHighAdj,
        5 => RelocationType::ImageRelArch1,
        6 => RelocationType::ImageRelReserved,
        7 => RelocationType::ImageRelArch2,
        8 => RelocationType::ImageRelBasedRiscVLow12S,
        9 => RelocationType::ImageRelBasedMIPSJmpAddr16,
        10 => RelocationType::ImageRelBasedDir64,
        _ => return Err(RelocAnalyzerError::InvalidRelocType.into()),
    };

    Ok(Reloc {
        typ: reloc_type,
        offset: base + RVA::from(reloc_offset as i64),
    })
}

fn split_u32(e: u32) -> (u16, u16) {
    let m = ((e & 0x0000FFFF) >>  0) as u16;
    let n = ((e & 0xFFFF0000) >> 16) as u16;
    (m, n)
}

/// ```
/// use lancelot::rsrc::*;
/// use lancelot::arch::*;
/// use lancelot::workspace::Workspace;
/// use lancelot::analysis::pe::relocs;
///
/// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
///    .disable_analysis()
///    .load().unwrap();
/// let relocs = relocs::get_relocs(&ws).unwrap();
/// assert_eq!(relocs[0].offset,   RVA(0x76008));
/// assert_eq!(relocs[277].offset, RVA(0xA8000));
/// ```
pub fn get_relocs(ws: &Workspace) -> Result<Vec<Reloc>, Error> {
    let pe = match Object::parse(&ws.buf) {
        Ok(Object::PE(pe)) => pe,
        _ => return Ok(vec![]),
    };

    let opt_header = match pe.header.optional_header {
        Some(opt_header) => opt_header,
        _ => return Ok(vec![]),
    };

    let reloc_directory = match opt_header.data_directories.get_base_relocation_table() {
        Some(reloc_directory) => reloc_directory,
        _ => return Ok(vec![]),
    };

    let dir_start = RVA::from(reloc_directory.virtual_address as i64);
    let buf = ws.read_bytes(dir_start, reloc_directory.size as usize)?;

    let entries: Vec<u32> = buf
        .chunks_exact(0x4)
        .map(|b| LittleEndian::read_u32(b))
        .collect();

    let mut ret = vec![];
    let mut index = 0;
    loop {
        // parse base relocation chunk
        // until there's no data left

        let page_rva = entries.get(index);
        let block_size = entries.get(index+1);
        let (page_rva, block_size) = match (page_rva, block_size) {
            (Some(0),        _               ) => break,
            (_,              Some(0)         ) => break,
            (Some(page_rva), Some(block_size)) => (RVA::from(*page_rva as i64), *block_size),
            _ => break,
        };

        // cast from u32 to usize
        let chunk_entries_count = ((block_size) / 4) as usize;

        for i in 2..chunk_entries_count {
            if let Some(&entry) = entries.get(index + i) {
                let (m, n) = split_u32(entry);

                let reloc1 = parse_reloc(page_rva, m)?;
                if !ws.probe(reloc1.offset, 4) {
                    break
                }
                ret.push(reloc1);

                let reloc2 = parse_reloc(page_rva, n)?;
                if !ws.probe(reloc2.offset, 4) {
                    break
                }
                ret.push(reloc2);
            } else {
                break
            }
        };

        index += chunk_entries_count;
    }

    Ok(ret)
}

impl Analyzer for RelocAnalyzer {
    fn get_name(&self) -> String {
        "PE relocation analyzer".to_string()
    }

    /// ```
    /// use lancelot::rsrc::*;
    /// use lancelot::arch::*;
    /// use lancelot::analysis::Analyzer;
    /// use lancelot::workspace::Workspace;
    /// use lancelot::analysis::pe::RelocAnalyzer;
    ///
    /// let mut ws = Workspace::from_bytes("k32.dll", &get_buf(Rsrc::K32))
    ///    .disable_analysis()
    ///    .load().unwrap();
    /// let anal = RelocAnalyzer::new();
    /// anal.analyze(&mut ws).unwrap();
    /// let meta = ws.get_meta(RVA(0xC7F0)).unwrap();
    /// assert!(meta.is_insn());
    /// ```
    fn analyze(&self, ws: &mut Workspace)-> Result<(), Error> {
        let text_section = match ws.module.sections.iter()
            // currently limited to the text section.
            // TODO: accept any writable section.
            .filter(|&sec| sec.name == ".text")
            .next() {
                None => return Ok(()),
                Some(s) => s,
            };

        let text_bounds = Section {
            start: text_section.addr,
            end: text_section.addr + RVA::from(text_section.buf.len()),
        };

        // scan the relocations
        // looking for pointers into the .text section
        // to things that
        //   1. are not already in an instruction
        //   2. don't appear to be a pointer
        // and assume this is code.
        let o: Vec<RVA> = get_relocs(ws)?
            .iter()
            .map(|reloc| reloc.offset)
            .map(|rva| ws.read_va(rva))
            .filter_map(Result::ok)
            .filter_map(|va| ws.rva(va))
            .filter(|&rva| text_bounds.start <= rva)
            .filter(|&rva| rva < text_bounds.end)
            .filter(|&rva| !is_in_insn(ws, rva))
            .filter(|&rva| !is_ptr(ws, rva))
            // TODO: maybe ensure that the insn decodes.
            .collect();

        o.iter().for_each(|&rva| {
            debug!("found ptr from .text section to .text section at {:#x}", rva);

            // TODO: consume result
            ws.make_insn(rva).unwrap();
            // TODO: consume result
            ws.analyze().unwrap();
        });
        Ok(())
    }
}
<|endoftext|>


<commit_before>use error::Result;
use rustwide::{cmd::Command, Toolchain, Workspace};
use std::collections::{HashMap, HashSet};
use std::path::Path;

pub(crate) struct CargoMetadata {
    packages: HashMap<String, Package>,
    deps_graph: HashMap<String, HashSet<String>>,
    root_id: String,
}

impl CargoMetadata {
    pub(crate) fn load(
        workspace: &Workspace,
        toolchain: &Toolchain,
        source_dir: &Path,
    ) -> Result<Self> {
        let res = Command::new(workspace, toolchain.cargo())
            .args(&["metadata", "--format-version", "1"])
            .cd(source_dir)
            .log_output(false)
            .run_capture()?;

        let mut iter = res.stdout_lines().iter();
        let metadata = if let (Some(serialized), None) = (iter.next(), iter.next()) {
            ::rustc_serialize::json::decode::<DeserializedMetadata>(serialized)?
        } else {
            return Err(::failure::err_msg(
                "invalid output returned by `cargo metadata`",
            ));
        };

        // Convert from Vecs to HashMaps and HashSets to get more efficient lookups
        Ok(CargoMetadata {
            packages: metadata
                .packages
                .into_iter()
                .map(|pkg| (pkg.id.clone(), pkg))
                .collect(),
            deps_graph: metadata
                .resolve
                .nodes
                .into_iter()
                .map(|node| (node.id, node.deps.into_iter().map(|d| d.pkg).collect()))
                .collect(),
            root_id: metadata.resolve.root,
        })
    }

    pub(crate) fn root_dependencies(&self) -> Vec<&Package> {
        let ids = &self.deps_graph[&self.root_id];
        self.packages
            .iter()
            .filter(|(id, _pkg)| ids.contains(id.as_str()))
            .map(|(_id, pkg)| pkg)
            .collect()
    }

    pub(crate) fn root(&self) -> &Package {
        &self.packages[&self.root_id]
    }
}

#[derive(RustcDecodable)]
pub(crate) struct Package {
    pub(crate) id: String,
    pub(crate) name: String,
    pub(crate) version: String,
    pub(crate) license: Option<String>,
    pub(crate) repository: Option<String>,
    pub(crate) homepage: Option<String>,
    pub(crate) description: Option<String>,
    pub(crate) documentation: Option<String>,
    pub(crate) dependencies: Vec<Dependency>,
    pub(crate) targets: Vec<Target>,
    pub(crate) readme: Option<String>,
    pub(crate) keywords: Vec<String>,
    pub(crate) authors: Vec<String>,
}

impl Package {
    fn library_target(&self) -> Option<&Target> {
        self.targets.iter().find(|target| {
            target
                .kind
                .iter()
                .any(|kind| kind == "lib" || kind == "proc-macro")
        })
    }

    pub(crate) fn is_library(&self) -> bool {
        self.library_target().is_some()
    }

    fn normalize_package_name(&self, name: &str) -> String {
        name.replace('-', "_")
    }

    pub(crate) fn package_name(&self) -> String {
        self.library_name()
            .unwrap_or_else(|| self.normalize_package_name(&self.targets[0].name))
    }

    pub(crate) fn library_name(&self) -> Option<String> {
        self.library_target()
            .map(|target| self.normalize_package_name(&target.name))
    }
}

#[derive(RustcDecodable)]
pub(crate) struct Target {
    pub(crate) name: String,
    pub(crate) kind: Vec<String>,
    pub(crate) src_path: Option<String>,
}

#[derive(RustcDecodable)]
pub(crate) struct Dependency {
    pub(crate) name: String,
    pub(crate) req: String,
    pub(crate) kind: Option<String>,
}

#[derive(RustcDecodable)]
struct DeserializedMetadata {
    packages: Vec<Package>,
    resolve: DeserializedResolve,
}

#[derive(RustcDecodable)]
struct DeserializedResolve {
    root: String,
    nodes: Vec<DeserializedResolveNode>,
}

#[derive(RustcDecodable)]
struct DeserializedResolveNode {
    id: String,
    deps: Vec<DeserializedResolveDep>,
}

#[derive(RustcDecodable)]
struct DeserializedResolveDep {
    pkg: String,
}
<commit_msg>[BROKEN] treat any crate that's not a binary as a library<commit_after>use error::Result;
use rustwide::{cmd::Command, Toolchain, Workspace};
use std::collections::{HashMap, HashSet};
use std::path::Path;

pub(crate) struct CargoMetadata {
    packages: HashMap<String, Package>,
    deps_graph: HashMap<String, HashSet<String>>,
    root_id: String,
}

impl CargoMetadata {
    pub(crate) fn load(
        workspace: &Workspace,
        toolchain: &Toolchain,
        source_dir: &Path,
    ) -> Result<Self> {
        let res = Command::new(workspace, toolchain.cargo())
            .args(&["metadata", "--format-version", "1"])
            .cd(source_dir)
            .log_output(false)
            .run_capture()?;

        let mut iter = res.stdout_lines().iter();
        let metadata = if let (Some(serialized), None) = (iter.next(), iter.next()) {
            ::rustc_serialize::json::decode::<DeserializedMetadata>(serialized)?
        } else {
            return Err(::failure::err_msg(
                "invalid output returned by `cargo metadata`",
            ));
        };

        // Convert from Vecs to HashMaps and HashSets to get more efficient lookups
        Ok(CargoMetadata {
            packages: metadata
                .packages
                .into_iter()
                .map(|pkg| (pkg.id.clone(), pkg))
                .collect(),
            deps_graph: metadata
                .resolve
                .nodes
                .into_iter()
                .map(|node| (node.id, node.deps.into_iter().map(|d| d.pkg).collect()))
                .collect(),
            root_id: metadata.resolve.root,
        })
    }

    pub(crate) fn root_dependencies(&self) -> Vec<&Package> {
        let ids = &self.deps_graph[&self.root_id];
        self.packages
            .iter()
            .filter(|(id, _pkg)| ids.contains(id.as_str()))
            .map(|(_id, pkg)| pkg)
            .collect()
    }

    pub(crate) fn root(&self) -> &Package {
        &self.packages[&self.root_id]
    }
}

#[derive(RustcDecodable)]
pub(crate) struct Package {
    pub(crate) id: String,
    pub(crate) name: String,
    pub(crate) version: String,
    pub(crate) license: Option<String>,
    pub(crate) repository: Option<String>,
    pub(crate) homepage: Option<String>,
    pub(crate) description: Option<String>,
    pub(crate) documentation: Option<String>,
    pub(crate) dependencies: Vec<Dependency>,
    pub(crate) targets: Vec<Target>,
    pub(crate) readme: Option<String>,
    pub(crate) keywords: Vec<String>,
    pub(crate) authors: Vec<String>,
}

impl Package {
    fn library_target(&self) -> Option<&Target> {
        self.targets.iter().find(|target| {
            target
                .kind
                .iter()
                .any(|kind| kind != "bin")
        })
    }

    pub(crate) fn is_library(&self) -> bool {
        self.library_target().is_some()
    }

    fn normalize_package_name(&self, name: &str) -> String {
        name.replace('-', "_")
    }

    pub(crate) fn package_name(&self) -> String {
        self.library_name()
            .unwrap_or_else(|| self.normalize_package_name(&self.targets[0].name))
    }

    pub(crate) fn library_name(&self) -> Option<String> {
        self.library_target()
            .map(|target| self.normalize_package_name(&target.name))
    }
}

#[derive(RustcDecodable)]
pub(crate) struct Target {
    pub(crate) name: String,
    pub(crate) kind: Vec<String>,
    pub(crate) src_path: Option<String>,
}

#[derive(RustcDecodable)]
pub(crate) struct Dependency {
    pub(crate) name: String,
    pub(crate) req: String,
    pub(crate) kind: Option<String>,
}

#[derive(RustcDecodable)]
struct DeserializedMetadata {
    packages: Vec<Package>,
    resolve: DeserializedResolve,
}

#[derive(RustcDecodable)]
struct DeserializedResolve {
    root: String,
    nodes: Vec<DeserializedResolveNode>,
}

#[derive(RustcDecodable)]
struct DeserializedResolveNode {
    id: String,
    deps: Vec<DeserializedResolveDep>,
}

#[derive(RustcDecodable)]
struct DeserializedResolveDep {
    pkg: String,
}
<|endoftext|>


<commit_before>// * This file is part of the uutils coreutils package.
// *
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

extern crate rand;

use std::collections::BTreeMap;
use std::fmt;

use crate::numeric::{Arithmetic, Montgomery};
use crate::{miller_rabin, rho, table};

#[derive(Clone, Debug)]
pub struct Factors {
    f: BTreeMap<u64, u8>,
}

impl Factors {
    pub fn one() -> Factors {
        Factors { f: BTreeMap::new() }
    }

    pub fn add(&mut self, prime: u64, exp: u8) {
        debug_assert!(miller_rabin::is_prime(prime));
        debug_assert!(exp > 0);
        let n = *self.f.get(&prime).unwrap_or(&0);
        self.f.insert(prime, exp + n);
    }

    pub fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.f
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (p, exp) in self.f.iter() {
            for _ in 0..*exp {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn _factor<A: Arithmetic + miller_rabin::Basis>(num: u64, f: Factors) -> Factors {
    use miller_rabin::Result::*;

    // Shadow the name, so the recursion automatically goes from “Big” arithmetic to small.
    let _factor = |n, f| {
        if n < (1 << 32) {
            _factor::<Montgomery<u32>>(n, f)
        } else {
            _factor::<A>(n, f)
        }
    };

    if num == 1 {
        return f;
    }

    let n = A::new(num);
    let divisor = match miller_rabin::test::<A>(n) {
        Prime => {
            let mut r = f;
            r.push(num);
            return r;
        }

        Composite(d) => d,
        Pseudoprime => rho::find_divisor::<A>(n),
    };

    let f = _factor(divisor, f);
    _factor(num / divisor, f)
}

pub fn factor(mut n: u64) -> Factors {
    let mut factors = Factors::one();

    if n < 2 {
        return factors;
    }

    let z = n.trailing_zeros();
    if z > 0 {
        factors.add(2, z as u8);
        n >>= z;
    }

    if n == 1 {
        return factors;
    }

    let (factors, n) = table::factor(n, factors);

    if n < (1 << 32) {
        _factor::<Montgomery<u32>>(n, factors)
    } else {
        _factor::<Montgomery<u64>>(n, factors)
    }
}

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

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_strong_pseudoprime() {
        // This is a strong pseudoprime (wrt. miller_rabin::BASIS)
        //  and triggered a bug in rho::factor's codepath handling
        //  miller_rabbin::Result::Composite
        let pseudoprime = 17179869183;
        for _ in 0..20 {
            // Repeat the test 20 times, as it only fails some fraction
            // of the time.
            assert!(factor(pseudoprime).product() == pseudoprime);
        }
    }

    quickcheck! {
        fn factor_recombines(i: u64) -> bool {
            i == 0 || factor(i).product() == i
        }
    }
}
<commit_msg>factor: Add tests against (random) numbers with known factorisations (#1573) (#1578)<commit_after>// * This file is part of the uutils coreutils package.
// *
// * (c) 2020 nicoo <nicoo@debian.org>
// *
// * For the full copyright and license information, please view the LICENSE file
// * that was distributed with this source code.

extern crate rand;

use std::collections::BTreeMap;
use std::fmt;

use crate::numeric::{Arithmetic, Montgomery};
use crate::{miller_rabin, rho, table};

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Factors {
    f: BTreeMap<u64, u8>,
}

impl Factors {
    pub fn one() -> Factors {
        Factors { f: BTreeMap::new() }
    }

    pub fn add(&mut self, prime: u64, exp: u8) {
        debug_assert!(miller_rabin::is_prime(prime));
        debug_assert!(exp > 0);
        let n = *self.f.get(&prime).unwrap_or(&0);
        self.f.insert(prime, exp + n);
    }

    pub fn push(&mut self, prime: u64) {
        self.add(prime, 1)
    }

    #[cfg(test)]
    fn product(&self) -> u64 {
        self.f
            .iter()
            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
    }
}

impl fmt::Display for Factors {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (p, exp) in self.f.iter() {
            for _ in 0..*exp {
                write!(f, " {}", p)?
            }
        }

        Ok(())
    }
}

fn _factor<A: Arithmetic + miller_rabin::Basis>(num: u64, f: Factors) -> Factors {
    use miller_rabin::Result::*;

    // Shadow the name, so the recursion automatically goes from “Big” arithmetic to small.
    let _factor = |n, f| {
        if n < (1 << 32) {
            _factor::<Montgomery<u32>>(n, f)
        } else {
            _factor::<A>(n, f)
        }
    };

    if num == 1 {
        return f;
    }

    let n = A::new(num);
    let divisor = match miller_rabin::test::<A>(n) {
        Prime => {
            let mut r = f;
            r.push(num);
            return r;
        }

        Composite(d) => d,
        Pseudoprime => rho::find_divisor::<A>(n),
    };

    let f = _factor(divisor, f);
    _factor(num / divisor, f)
}

pub fn factor(mut n: u64) -> Factors {
    let mut factors = Factors::one();

    if n < 2 {
        return factors;
    }

    let z = n.trailing_zeros();
    if z > 0 {
        factors.add(2, z as u8);
        n >>= z;
    }

    if n == 1 {
        return factors;
    }

    let (factors, n) = table::factor(n, factors);

    if n < (1 << 32) {
        _factor::<Montgomery<u32>>(n, factors)
    } else {
        _factor::<Montgomery<u64>>(n, factors)
    }
}

#[cfg(test)]
mod tests {
    use super::{factor, Factors};
    use quickcheck::quickcheck;

    #[test]
    fn factor_recombines_small() {
        assert!((1..10_000)
            .map(|i| 2 * i + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_overflowing() {
        assert!((0..250)
            .map(|i| 2 * i + 2u64.pow(32) + 1)
            .all(|i| factor(i).product() == i));
    }

    #[test]
    fn factor_recombines_strong_pseudoprime() {
        // This is a strong pseudoprime (wrt. miller_rabin::BASIS)
        //  and triggered a bug in rho::factor's codepath handling
        //  miller_rabbin::Result::Composite
        let pseudoprime = 17179869183;
        for _ in 0..20 {
            // Repeat the test 20 times, as it only fails some fraction
            // of the time.
            assert!(factor(pseudoprime).product() == pseudoprime);
        }
    }

    quickcheck! {
        fn factor_recombines(i: u64) -> bool {
            i == 0 || factor(i).product() == i
        }

        fn recombines_factors(f: Factors) -> bool {
            factor(f.product()) == f
        }
    }
}

#[cfg(test)]
impl quickcheck::Arbitrary for Factors {
    fn arbitrary<G: quickcheck::Gen>(gen: &mut G) -> Self {
        use rand::Rng;
        let mut f = Factors::one();
        let mut g = 1u64;
        let mut n = u64::MAX;

        // Adam Kalai's algorithm for generating uniformly-distributed
        // integers and their factorisation.
        //
        // See Generating Random Factored Numbers, Easily, J. Cryptology (2003)
        'attempt: loop {
            while n > 1 {
                n = gen.gen_range(1, n);
                if miller_rabin::is_prime(n) {
                    if let Some(h) = g.checked_mul(n) {
                        f.push(n);
                        g = h;
                    } else {
                        // We are overflowing u64, retry
                        continue 'attempt;
                    }
                }
            }

            return f;
        }
    }
}
<|endoftext|>