<commit_before>use crate::alloc::string::{String, ToString};
use scroll::{ctx, Pread, Pwrite};
use crate::error::{self, Error};
use crate::pe::relocation;
use std::io::Write;

#[repr(C)]
#[derive(Debug, PartialEq, Clone, Default)]
pub struct SectionTable {
    pub name: [u8; 8],
    pub real_name: Option<String>,
    pub virtual_size: u32,
    pub virtual_address: u32,
    pub size_of_raw_data: u32,
    pub pointer_to_raw_data: u32,
    pub pointer_to_relocations: u32,
    pub pointer_to_linenumbers: u32,
    pub number_of_relocations: u16,
    pub number_of_linenumbers: u16,
    pub characteristics: u32,
}

pub const SIZEOF_SECTION_TABLE: usize = 8 * 5;

// Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L70
// Decodes a string table entry in base 64 (//AAAAAA). Expects string without
// prefixed slashes.
fn base64_decode_string_entry(s: &str) -> Result<usize, ()> {
    assert!(s.len() <= 6, "String too long, possible overflow.");

    let mut val = 0;
    for c in s.bytes() {
        let v = if b'A' <= c && c <= b'Z' { // 00..=25
            c - b'A'
        } else if b'a' <= c && c <= b'z' { // 26..=51
            c - b'a' + 26
        } else if b'0' <= c && c <= b'9' { // 52..=61
            c - b'0' + 52
        } else if c == b'+' { // 62
            62
        } else if c == b'/' { // 63
            63
        } else {
            return Err(())
        };
        val = val * 64 + v as usize;
    }
    Ok(val)
}

impl SectionTable {
    pub fn parse(bytes: &[u8], offset: &mut usize, string_table_offset: usize) -> error::Result<Self> {
        let mut table = SectionTable::default();
        let mut name = [0u8; 8];
        name.copy_from_slice(bytes.gread_with(offset, 8)?);

        table.name = name;
        table.virtual_size = bytes.gread_with(offset, scroll::LE)?;
        table.virtual_address = bytes.gread_with(offset, scroll::LE)?;
        table.size_of_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.characteristics = bytes.gread_with(offset, scroll::LE)?;

        if let Some(idx) = table.name_offset()? {
            table.real_name = Some(bytes.pread::<&str>(string_table_offset + idx)?.to_string());
        }
        Ok(table)
    }

    pub fn name_offset(&self) -> error::Result<Option<usize>> {
        // Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L1054
        if self.name[0] == b'/' {
            let idx: usize = if self.name[1] == b'/' {
                let b64idx = self.name.pread::<&str>(2)?;
                base64_decode_string_entry(b64idx).map_err(|_|
                    Error::Malformed(format!("Invalid indirect section name //{}: base64 decoding failed", b64idx)))?
            } else {
                let name = self.name.pread::<&str>(1)?;
                name.parse().map_err(|err|
                    Error::Malformed(format!("Invalid indirect section name /{}: {}", name, err)))?
            };
            Ok(Some(idx))
        } else {
            Ok(None)
        }
    }

    pub fn set_name_offset(&mut self, mut idx: usize) -> error::Result<()> {
        if idx <= 9_999_999 { // 10^7 - 1
            self.name = [0; 8];
            self.name[0] = b'/';
            write!(&mut self.name[1..], "{}", idx).unwrap();
            Ok(())
        } else if idx as u64 <= 0xfff_fff_fff { // 64^6 - 1
            self.name[0] = b'/';
            self.name[1] = b'/';
            for i in 0..6 {
                let rem = (idx % 64) as u8;
                idx = idx / 64;
                let c = match rem {
                    0..=25 => b'A' + rem,
                    26..=51 => b'a' + rem - 26,
                    52..=61 => b'0' + rem - 52,
                    62 => b'+',
                    63 => b'/',
                    _ => unreachable!(),
                };
                self.name[7 - i] = c;
            }
            Ok(())
        } else {
            Err(Error::Malformed(format!("Invalid section name offset: {}", idx)))
        }
    }

    pub fn name(&self) -> error::Result<&str> {
        match self.real_name.as_ref() {
            Some(s) => Ok(s),
            None => Ok(self.name.pread(0)?)
        }
    }

    pub fn relocations<'a>(&self, bytes: &'a[u8]) -> error::Result<relocation::Relocations<'a>> {
        let offset = self.pointer_to_relocations as usize;
        let number = self.number_of_relocations as usize;
        relocation::Relocations::parse(bytes, offset, number)
    }
}

impl ctx::SizeWith<scroll::Endian> for SectionTable {
    type Units = usize;
    fn size_with(_ctx: &scroll::Endian) -> usize {
        SIZEOF_SECTION_TABLE
    }
}

impl ctx::TryIntoCtx<scroll::Endian> for SectionTable {
    type Error = error::Error;
    type Size = usize;
    fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result<Self::Size, Self::Error> {
        let offset = &mut 0;
        bytes.gwrite(&self.name[..], offset)?;
        bytes.gwrite_with(self.virtual_size, offset, ctx)?;
        bytes.gwrite_with(self.virtual_address, offset, ctx)?;
        bytes.gwrite_with(self.size_of_raw_data, offset, ctx)?;
        bytes.gwrite_with(self.pointer_to_raw_data, offset, ctx)?;
        bytes.gwrite_with(self.pointer_to_relocations, offset, ctx)?;
        bytes.gwrite_with(self.pointer_to_linenumbers, offset, ctx)?;
        bytes.gwrite_with(self.number_of_relocations, offset, ctx)?;
        bytes.gwrite_with(self.number_of_linenumbers, offset, ctx)?;
        bytes.gwrite_with(self.characteristics, offset, ctx)?;
        Ok(SIZEOF_SECTION_TABLE)
    }
}

impl ctx::IntoCtx<scroll::Endian> for SectionTable {
    fn into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) {
        bytes.pwrite_with(self, 0, ctx).unwrap();
    }
}

/// The section should not be padded to the next boundary. This flag is obsolete and is replaced
/// by `IMAGE_SCN_ALIGN_1BYTES`. This is valid only for object files.
pub const IMAGE_SCN_TYPE_NO_PAD: u32 = 0x0000_0008;
/// The section contains executable code.
pub const IMAGE_SCN_CNT_CODE: u32 = 0x0000_0020;
/// The section contains initialized data.
pub const IMAGE_SCN_CNT_INITIALIZED_DATA: u32 = 0x0000_0040;
///  The section contains uninitialized data.
pub const IMAGE_SCN_CNT_UNINITIALIZED_DATA: u32 = 0x0000_0080;
pub const IMAGE_SCN_LNK_OTHER: u32 = 0x0000_0100;
/// The section contains comments or other information. The .drectve section has this type.
/// This is valid for object files only.
pub const IMAGE_SCN_LNK_INFO: u32 = 0x0000_0200;
/// The section will not become part of the image. This is valid only for object files.
pub const IMAGE_SCN_LNK_REMOVE: u32 = 0x0000_0800;
/// The section contains COMDAT data. This is valid only for object files.
pub const IMAGE_SCN_LNK_COMDAT: u32 = 0x0000_1000;
/// The section contains data referenced through the global pointer (GP).
pub const IMAGE_SCN_GPREL: u32 = 0x0000_8000;
pub const IMAGE_SCN_MEM_PURGEABLE: u32 = 0x0002_0000;
pub const IMAGE_SCN_MEM_16BIT: u32 = 0x0002_0000;
pub const IMAGE_SCN_MEM_LOCKED: u32 = 0x0004_0000;
pub const IMAGE_SCN_MEM_PRELOAD: u32 = 0x0008_0000;

pub const IMAGE_SCN_ALIGN_1BYTES: u32 = 0x0010_0000;
pub const IMAGE_SCN_ALIGN_2BYTES: u32 = 0x0020_0000;
pub const IMAGE_SCN_ALIGN_4BYTES: u32 = 0x0030_0000;
pub const IMAGE_SCN_ALIGN_8BYTES: u32 = 0x0040_0000;
pub const IMAGE_SCN_ALIGN_16BYTES: u32 = 0x0050_0000;
pub const IMAGE_SCN_ALIGN_32BYTES: u32 = 0x0060_0000;
pub const IMAGE_SCN_ALIGN_64BYTES: u32 = 0x0070_0000;
pub const IMAGE_SCN_ALIGN_128BYTES: u32 = 0x0080_0000;
pub const IMAGE_SCN_ALIGN_256BYTES: u32 = 0x0090_0000;
pub const IMAGE_SCN_ALIGN_512BYTES: u32 = 0x00A0_0000;
pub const IMAGE_SCN_ALIGN_1024BYTES: u32 = 0x00B0_0000;
pub const IMAGE_SCN_ALIGN_2048BYTES: u32 = 0x00C0_0000;
pub const IMAGE_SCN_ALIGN_4096BYTES: u32 = 0x00D0_0000;
pub const IMAGE_SCN_ALIGN_8192BYTES: u32 = 0x00E0_0000;
pub const IMAGE_SCN_ALIGN_MASK: u32 = 0x00F0_0000;

/// The section contains extended relocations.
pub const IMAGE_SCN_LNK_NRELOC_OVFL: u32 = 0x0100_0000;
/// The section can be discarded as needed.
pub const IMAGE_SCN_MEM_DISCARDABLE: u32 = 0x0200_0000;
/// The section cannot be cached.
pub const IMAGE_SCN_MEM_NOT_CACHED: u32 = 0x0400_0000;
/// The section is not pageable.
pub const IMAGE_SCN_MEM_NOT_PAGED: u32 = 0x0800_0000;
/// The section can be shared in memory.
pub const IMAGE_SCN_MEM_SHARED: u32 = 0x1000_0000;
/// The section can be executed as code.
pub const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;
/// The section can be read.
pub const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;
/// The section can be written to.
pub const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;

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

    #[test]
    fn set_name_offset() {
        let mut section = SectionTable::default();
        for (offset, name) in vec![
            (0, b"/0\0\0\0\0\0\0"),
            (1, b"/1\0\0\0\0\0\0"),
            (9_999_999, b"/9999999"),
            (10_000_000, b"//AAmJaA"),
            #[cfg(target_pointer_width = "64")]
            (0xfff_fff_fff, b"////////"),
        ] {
            section.set_name_offset(offset).unwrap();
            assert_eq!(&section.name, name);
            assert_eq!(section.name_offset().unwrap(), Some(offset));
        }
        #[cfg(target_pointer_width = "64")]
        assert!(section.set_name_offset(0x1_000_000_000).is_err());
    }
}
<commit_msg>Fix minor clippy warnings<commit_after>use crate::alloc::string::{String, ToString};
use scroll::{ctx, Pread, Pwrite};
use crate::error::{self, Error};
use crate::pe::relocation;
use std::io::Write;

#[repr(C)]
#[derive(Debug, PartialEq, Clone, Default)]
pub struct SectionTable {
    pub name: [u8; 8],
    pub real_name: Option<String>,
    pub virtual_size: u32,
    pub virtual_address: u32,
    pub size_of_raw_data: u32,
    pub pointer_to_raw_data: u32,
    pub pointer_to_relocations: u32,
    pub pointer_to_linenumbers: u32,
    pub number_of_relocations: u16,
    pub number_of_linenumbers: u16,
    pub characteristics: u32,
}

pub const SIZEOF_SECTION_TABLE: usize = 8 * 5;

// Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L70
// Decodes a string table entry in base 64 (//AAAAAA). Expects string without
// prefixed slashes.
fn base64_decode_string_entry(s: &str) -> Result<usize, ()> {
    assert!(s.len() <= 6, "String too long, possible overflow.");

    let mut val = 0;
    for c in s.bytes() {
        let v = if b'A' <= c && c <= b'Z' { // 00..=25
            c - b'A'
        } else if b'a' <= c && c <= b'z' { // 26..=51
            c - b'a' + 26
        } else if b'0' <= c && c <= b'9' { // 52..=61
            c - b'0' + 52
        } else if c == b'+' { // 62
            62
        } else if c == b'/' { // 63
            63
        } else {
            return Err(())
        };
        val = val * 64 + v as usize;
    }
    Ok(val)
}

impl SectionTable {
    pub fn parse(bytes: &[u8], offset: &mut usize, string_table_offset: usize) -> error::Result<Self> {
        let mut table = SectionTable::default();
        let mut name = [0u8; 8];
        name.copy_from_slice(bytes.gread_with(offset, 8)?);

        table.name = name;
        table.virtual_size = bytes.gread_with(offset, scroll::LE)?;
        table.virtual_address = bytes.gread_with(offset, scroll::LE)?;
        table.size_of_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.characteristics = bytes.gread_with(offset, scroll::LE)?;

        if let Some(idx) = table.name_offset()? {
            table.real_name = Some(bytes.pread::<&str>(string_table_offset + idx)?.to_string());
        }
        Ok(table)
    }

    pub fn name_offset(&self) -> error::Result<Option<usize>> {
        // Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L1054
        if self.name[0] == b'/' {
            let idx: usize = if self.name[1] == b'/' {
                let b64idx = self.name.pread::<&str>(2)?;
                base64_decode_string_entry(b64idx).map_err(|_|
                    Error::Malformed(format!("Invalid indirect section name //{}: base64 decoding failed", b64idx)))?
            } else {
                let name = self.name.pread::<&str>(1)?;
                name.parse().map_err(|err|
                    Error::Malformed(format!("Invalid indirect section name /{}: {}", name, err)))?
            };
            Ok(Some(idx))
        } else {
            Ok(None)
        }
    }

    pub fn set_name_offset(&mut self, mut idx: usize) -> error::Result<()> {
        if idx <= 9_999_999 { // 10^7 - 1
            self.name = [0; 8];
            self.name[0] = b'/';
            write!(&mut self.name[1..], "{}", idx).unwrap();
            Ok(())
        } else if idx as u64 <= 0xfff_fff_fff { // 64^6 - 1
            self.name[0] = b'/';
            self.name[1] = b'/';
            for i in 0..6 {
                let rem = (idx % 64) as u8;
                idx /= 64;
                let c = match rem {
                    0..=25 => b'A' + rem,
                    26..=51 => b'a' + rem - 26,
                    52..=61 => b'0' + rem - 52,
                    62 => b'+',
                    63 => b'/',
                    _ => unreachable!(),
                };
                self.name[7 - i] = c;
            }
            Ok(())
        } else {
            Err(Error::Malformed(format!("Invalid section name offset: {}", idx)))
        }
    }

    pub fn name(&self) -> error::Result<&str> {
        match self.real_name.as_ref() {
            Some(s) => Ok(s),
            None => Ok(self.name.pread(0)?)
        }
    }

    pub fn relocations<'a>(&self, bytes: &'a[u8]) -> error::Result<relocation::Relocations<'a>> {
        let offset = self.pointer_to_relocations as usize;
        let number = self.number_of_relocations as usize;
        relocation::Relocations::parse(bytes, offset, number)
    }
}

impl ctx::SizeWith<scroll::Endian> for SectionTable {
    type Units = usize;
    fn size_with(_ctx: &scroll::Endian) -> usize {
        SIZEOF_SECTION_TABLE
    }
}

impl ctx::TryIntoCtx<scroll::Endian> for SectionTable {
    type Error = error::Error;
    type Size = usize;
    fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result<Self::Size, Self::Error> {
        let offset = &mut 0;
        bytes.gwrite(&self.name[..], offset)?;
        bytes.gwrite_with(self.virtual_size, offset, ctx)?;
        bytes.gwrite_with(self.virtual_address, offset, ctx)?;
        bytes.gwrite_with(self.size_of_raw_data, offset, ctx)?;
        bytes.gwrite_with(self.pointer_to_raw_data, offset, ctx)?;
        bytes.gwrite_with(self.pointer_to_relocations, offset, ctx)?;
        bytes.gwrite_with(self.pointer_to_linenumbers, offset, ctx)?;
        bytes.gwrite_with(self.number_of_relocations, offset, ctx)?;
        bytes.gwrite_with(self.number_of_linenumbers, offset, ctx)?;
        bytes.gwrite_with(self.characteristics, offset, ctx)?;
        Ok(SIZEOF_SECTION_TABLE)
    }
}

impl ctx::IntoCtx<scroll::Endian> for SectionTable {
    fn into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) {
        bytes.pwrite_with(self, 0, ctx).unwrap();
    }
}

/// The section should not be padded to the next boundary. This flag is obsolete and is replaced
/// by `IMAGE_SCN_ALIGN_1BYTES`. This is valid only for object files.
pub const IMAGE_SCN_TYPE_NO_PAD: u32 = 0x0000_0008;
/// The section contains executable code.
pub const IMAGE_SCN_CNT_CODE: u32 = 0x0000_0020;
/// The section contains initialized data.
pub const IMAGE_SCN_CNT_INITIALIZED_DATA: u32 = 0x0000_0040;
///  The section contains uninitialized data.
pub const IMAGE_SCN_CNT_UNINITIALIZED_DATA: u32 = 0x0000_0080;
pub const IMAGE_SCN_LNK_OTHER: u32 = 0x0000_0100;
/// The section contains comments or other information. The .drectve section has this type.
/// This is valid for object files only.
pub const IMAGE_SCN_LNK_INFO: u32 = 0x0000_0200;
/// The section will not become part of the image. This is valid only for object files.
pub const IMAGE_SCN_LNK_REMOVE: u32 = 0x0000_0800;
/// The section contains COMDAT data. This is valid only for object files.
pub const IMAGE_SCN_LNK_COMDAT: u32 = 0x0000_1000;
/// The section contains data referenced through the global pointer (GP).
pub const IMAGE_SCN_GPREL: u32 = 0x0000_8000;
pub const IMAGE_SCN_MEM_PURGEABLE: u32 = 0x0002_0000;
pub const IMAGE_SCN_MEM_16BIT: u32 = 0x0002_0000;
pub const IMAGE_SCN_MEM_LOCKED: u32 = 0x0004_0000;
pub const IMAGE_SCN_MEM_PRELOAD: u32 = 0x0008_0000;

pub const IMAGE_SCN_ALIGN_1BYTES: u32 = 0x0010_0000;
pub const IMAGE_SCN_ALIGN_2BYTES: u32 = 0x0020_0000;
pub const IMAGE_SCN_ALIGN_4BYTES: u32 = 0x0030_0000;
pub const IMAGE_SCN_ALIGN_8BYTES: u32 = 0x0040_0000;
pub const IMAGE_SCN_ALIGN_16BYTES: u32 = 0x0050_0000;
pub const IMAGE_SCN_ALIGN_32BYTES: u32 = 0x0060_0000;
pub const IMAGE_SCN_ALIGN_64BYTES: u32 = 0x0070_0000;
pub const IMAGE_SCN_ALIGN_128BYTES: u32 = 0x0080_0000;
pub const IMAGE_SCN_ALIGN_256BYTES: u32 = 0x0090_0000;
pub const IMAGE_SCN_ALIGN_512BYTES: u32 = 0x00A0_0000;
pub const IMAGE_SCN_ALIGN_1024BYTES: u32 = 0x00B0_0000;
pub const IMAGE_SCN_ALIGN_2048BYTES: u32 = 0x00C0_0000;
pub const IMAGE_SCN_ALIGN_4096BYTES: u32 = 0x00D0_0000;
pub const IMAGE_SCN_ALIGN_8192BYTES: u32 = 0x00E0_0000;
pub const IMAGE_SCN_ALIGN_MASK: u32 = 0x00F0_0000;

/// The section contains extended relocations.
pub const IMAGE_SCN_LNK_NRELOC_OVFL: u32 = 0x0100_0000;
/// The section can be discarded as needed.
pub const IMAGE_SCN_MEM_DISCARDABLE: u32 = 0x0200_0000;
/// The section cannot be cached.
pub const IMAGE_SCN_MEM_NOT_CACHED: u32 = 0x0400_0000;
/// The section is not pageable.
pub const IMAGE_SCN_MEM_NOT_PAGED: u32 = 0x0800_0000;
/// The section can be shared in memory.
pub const IMAGE_SCN_MEM_SHARED: u32 = 0x1000_0000;
/// The section can be executed as code.
pub const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;
/// The section can be read.
pub const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;
/// The section can be written to.
pub const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;

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

    #[test]
    fn set_name_offset() {
        let mut section = SectionTable::default();
        for &(offset, name) in [
            (0usize, b"/0\0\0\0\0\0\0"),
            (1, b"/1\0\0\0\0\0\0"),
            (9_999_999, b"/9999999"),
            (10_000_000, b"//AAmJaA"),
            #[cfg(target_pointer_width = "64")]
            (0xfff_fff_fff, b"////////"),
        ].iter() {
            section.set_name_offset(offset).unwrap();
            assert_eq!(&section.name, name);
            assert_eq!(section.name_offset().unwrap(), Some(offset));
        }
        #[cfg(target_pointer_width = "64")]
        assert!(section.set_name_offset(0x1_000_000_000).is_err());
    }
}
<|endoftext|>


<commit_before>#![crate_type = "bin"]
#![feature(collections)]

extern crate gtk;
extern crate glib;
extern crate sysinfo;

use gtk::{WindowTrait, ContainerTrait, WidgetTrait, ButtonSignals, BoxTrait};
use gtk::{IconSize, Orientation, ReliefStyle};
use sysinfo::*;
use gtk::signal::Inhibit;
use gtk::signal::{WidgetSignals, TreeViewSignals};
use std::collections::VecMap;
use std::str::FromStr;
use std::rc::Rc;
use std::cell::{RefCell, RefMut};
use std::ops::DerefMut;

struct NoteBook {
    notebook: gtk::NoteBook,
    tabs: Vec<gtk::Box>
}

impl NoteBook {
    fn new() -> NoteBook {
        NoteBook {
            notebook: gtk::NoteBook::new().unwrap(),
            tabs: Vec::new()
        }
    }

    fn create_tab<'a, Widget: gtk::WidgetTrait>(&mut self, title: &'a str, widget: &Widget) -> Option<u32> {
        let label = gtk::Label::new(title).unwrap();
        let tab = gtk::Box::new(Orientation::Horizontal, 0).unwrap();

        tab.pack_start(&label, false, false, 0);
        tab.show_all();

        let index = match self.notebook.append_page(widget, Some(&tab)) {
            Some(index) => index,
            _ => return None
        };

        self.tabs.push(tab);

        Some(index)
    }
}

struct Procs {
    left_tree: gtk::TreeView,
    scroll: gtk::ScrolledWindow,
    current_pid: Rc<RefCell<Option<String>>>,
    kill_button: Rc<RefCell<gtk::Button>>,
    vertical_layout: gtk::Box,
    list_store: gtk::ListStore
}

impl Procs {
    pub fn new<'a>(proc_list: &VecMap<Processus>) -> Procs {
        let mut left_tree = gtk::TreeView::new().unwrap();
        let mut scroll = gtk::ScrolledWindow::new(None, None).unwrap();
        let mut current_pid = Rc::new(RefCell::new(None));
        let mut kill_button = Rc::new(RefCell::new(gtk::Button::new_with_label("End task").unwrap()));
        let mut current_pid1 = current_pid.clone();
        let mut current_pid2 = current_pid.clone();
        let mut kill_button1 = kill_button.clone();

        scroll.set_min_content_height(800);
        scroll.set_min_content_width(600);

        let mut columns : Vec<gtk::TreeViewColumn> = Vec::new();

        append_column("pid", &mut columns);
        append_column("process name", &mut columns);
        append_column("cpu usage", &mut columns);
        append_column("memory usage", &mut columns);

        for i in columns {
            left_tree.append_column(&i);
        }

        let mut list_store = gtk::ListStore::new(&[glib::Type::String, glib::Type::String, glib::Type::String, glib::Type::String]).unwrap();
        for (_, pro) in proc_list {
            create_and_fill_model(&mut list_store, pro.pid, &pro.name, pro.cpu_usage, pro.memory);
        }

        left_tree.set_model(&list_store.get_model().unwrap());
        left_tree.set_headers_visible(true);
        scroll.add(&left_tree);
        let mut vertical_layout = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();

        left_tree.connect_cursor_changed(move |tree_view| {
            match tree_view.get_selection() {
                Some(selection) => {
                    let model = tree_view.get_model().unwrap();
                    let mut iter = gtk::TreeIter::new();

                    if selection.get_selected(&model, &mut iter) {
                        let pid = model.get_value(&iter, 0).get_string();
                        let mut tmp = current_pid1.borrow_mut();
                        *tmp = pid;
                    }
                    (*kill_button1.borrow_mut()).set_sensitive((*current_pid.borrow()).is_some());
                }
                None => {
                    let mut tmp = current_pid1.borrow_mut();
                    *tmp = None;
                }
            }
        });
        (*kill_button.borrow_mut()).set_sensitive(false);

        vertical_layout.add(&scroll);
        vertical_layout.add(&(*kill_button.borrow_mut()));
        Procs {
            left_tree: left_tree,
            scroll: scroll,
            current_pid: current_pid2.clone(),
            kill_button: kill_button,
            vertical_layout: vertical_layout,
            list_store: list_store
        }
    }
}

fn append_column(title: &str, v: &mut Vec<gtk::TreeViewColumn>) {
    let l = v.len();
    let renderer = gtk::CellRendererText::new().unwrap();

    v.push(gtk::TreeViewColumn::new().unwrap());
    let tmp = v.get_mut(l).unwrap();

    tmp.set_title(title);
    tmp.pack_start(&renderer, true);
    tmp.add_attribute(&renderer, "text", l as i32);
}

fn create_and_fill_model(list_store: &mut gtk::ListStore, pid: i64, name: &str, cpu: f32, memory: u64) {
    if name.len() < 1 {
        return;
    }
    let mut top_level = gtk::TreeIter::new();

    list_store.append(&mut top_level);
    list_store.set_string(&top_level, 0, &format!("{}", pid));
    list_store.set_string(&top_level, 1, name);
    list_store.set_string(&top_level, 2, &format!("{}", cpu));
    list_store.set_string(&top_level, 3, &format!("{}", memory));
}

fn update_window(w: &mut (&mut gtk::ListStore, Rc<RefCell<sysinfo::System>>, Rc<RefCell<DisplaySysInfo>>)) -> i32 {
    let mut model = w.0.get_model().unwrap();
    let mut iter = gtk::TreeIter::new();

    (*w.1.borrow_mut()).refresh();
    let mut entries : VecMap<Processus> = ((*w.1.borrow()).get_processus_list()).clone();
    let mut nb = model.iter_n_children(None);

    let sys = w.1.clone();
    (*w.2.borrow_mut()).update_ram_display(sys.clone());
    (*w.2.borrow_mut()).update_process_display(sys);

    let mut i = 0;
    while i < nb {
        if model.iter_nth_child(&mut iter, None, i) {
            let pid : i64 = i64::from_str(&model.get_value(&iter, 0).get_string().unwrap()).unwrap();
            let mut to_delete = false;

            match entries.get(&(pid as usize)) {
                Some(p) => {
                    //w.0.set_string(&iter, 1, &p.name);
                    w.0.set_string(&iter, 2, &format!("{}", p.cpu_usage));
                    w.0.set_string(&iter, 3, &format!("{}", p.memory));
                    to_delete = true;
                }
                None => {
                    w.0.remove(&iter);
                }
            }
            if to_delete {
                entries.remove(&(pid as usize));
                nb = model.iter_n_children(None);
                i += 1;
            }
        } else {
            i += 1;
        }
    }
    for (_, pro) in entries {
        create_and_fill_model(&mut w.0, pro.pid, &pro.name, pro.cpu_usage, pro.memory);
    }
    1
}

struct DisplaySysInfo {
    procs : Rc<RefCell<Vec<gtk::ProgressBar>>>,
    ram : Rc<RefCell<gtk::ProgressBar>>,
    swap : Rc<RefCell<gtk::ProgressBar>>,
    vertical_layout : gtk::Box
}

impl DisplaySysInfo {
    pub fn new(sys1: Rc<RefCell<sysinfo::System>>) -> DisplaySysInfo {
        let mut tmp = DisplaySysInfo {
            procs: Rc::new(RefCell::new(Vec::new())),
            ram: Rc::new(RefCell::new(gtk::ProgressBar::new().unwrap())),
            swap: Rc::new(RefCell::new(gtk::ProgressBar::new().unwrap())),
            vertical_layout: gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap()
        };

        let sys2 = sys1.clone();
        let mut v = tmp.procs.clone();

        (*tmp.ram.borrow_mut()).set_show_text(true);
        (*tmp.swap.borrow_mut()).set_show_text(true);

        tmp.vertical_layout.set_spacing(5);
        let mut i = 1;
        let mut total = false;

        tmp.vertical_layout.pack_start(&gtk::Label::new("Memory usage").unwrap(), false, false, 15);
        tmp.vertical_layout.add(&(*tmp.ram.borrow()));
        tmp.vertical_layout.pack_start(&gtk::Label::new("Swap usage").unwrap(), false, false, 15);
        tmp.vertical_layout.add(&(*tmp.swap.borrow()));
        tmp.vertical_layout.pack_start(&gtk::Label::new("Process usage").unwrap(), false, false, 15);
        for pro in (*sys1.borrow()).get_process_list() {
            if total {
                (*v.borrow_mut()).push(gtk::ProgressBar::new().unwrap());
                let p : &gtk::ProgressBar = &(*v.borrow())[i - 1];
                let l = gtk::Label::new(&format!("{}", i)).unwrap();
                let horizontal_layout = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();

                p.set_text(&format!("{:.2} %", pro.get_cpu_usage() * 100.));
                p.set_show_text(true);
                p.set_fraction(pro.get_cpu_usage() as f64);
                horizontal_layout.pack_start(&l, false, false, 5);
                horizontal_layout.pack_start(p, true, true, 5);
                tmp.vertical_layout.add(&horizontal_layout);
                i += 1;
            } else {
                total = true;
            }
        }
        tmp.update_ram_display(sys2);
        tmp
    }

    pub fn update_ram_display(&mut self, sys: Rc<RefCell<sysinfo::System>>) {
        let total = (*sys.borrow()).get_total_memory();
        let used = (*sys.borrow()).get_free_memory();
        let disp = if total < 100000 {
            format!("{} / {}B", used, total)
        } else if total < 10000000 {
            format!("{} / {}KB", used / 1000, total / 1000)
        } else if total < 10000000000 {
            format!("{} / {}MB", used / 1000000, total / 1000000)
        } else {
            format!("{} / {}GB", used / 1000000000, total / 1000000000)
        };

        (*self.ram.borrow_mut()).set_text(&disp);
        (*self.ram.borrow_mut()).set_fraction(used as f64 / total as f64);

        let total = (*sys.borrow()).get_total_swap();
        let used = total - (*sys.borrow()).get_free_swap();
        let disp = if total < 100000 {
            format!("{} / {}B", used, total)
        } else if total < 10000000 {
            format!("{} / {}KB", used / 1000, total / 1000)
        } else if total < 10000000000 {
            format!("{} / {}MB", used / 1000000, total / 1000000)
        } else {
            format!("{} / {}GB", used / 1000000000, total / 1000000000)
        };

        (*self.swap.borrow_mut()).set_text(&disp);
        (*self.swap.borrow_mut()).set_fraction(used as f64 / total as f64);
    }

    pub fn update_process_display(&mut self, sys: Rc<RefCell<sysinfo::System>>) {
        let v = &*self.procs.borrow_mut();
        let mut total = false;
        let mut i = 0;
        for pro in (*sys.borrow()).get_process_list() {
            if total {
                v[i].set_text(&format!("{:.1} %", pro.get_cpu_usage() * 100.));
                v[i].set_show_text(true);
                v[i].set_fraction(pro.get_cpu_usage() as f64);
                i += 1;
            } else {
                total = true;
            }
        }
    }
}

fn main() {
    gtk::init();

    let mut window = gtk::Window::new(gtk::WindowType::TopLevel).unwrap();
    let mut sys : Rc<RefCell<sysinfo::System>> = Rc::new(RefCell::new(sysinfo::System::new()));
    let mut note = NoteBook::new();
    (*sys.borrow_mut()).refresh();
    let mut procs = Procs::new((*sys.borrow()).get_processus_list());
    let mut current_pid2 = procs.current_pid.clone();
    let mut sys1 = sys.clone();
    let mut sys2 = sys.clone();

    window.set_title("Processus viewer");
    window.set_window_position(gtk::WindowPosition::Center);

    window.connect_delete_event(|_, _| {
        gtk::main_quit();
        Inhibit(true)
    });

    (*procs.kill_button.borrow_mut()).connect_clicked(move |_| {
        let tmp = (*current_pid2.borrow_mut()).is_some() ;

        if tmp {
            let s = (*current_pid2.borrow()).clone();
            match (*sys.borrow_mut()).get_processus(i64::from_str(&s.unwrap()).unwrap()) {
                Some(p) => {
                    p.kill(Signal::Kill);
                },
                None => {}
            };
        }
    });

    let mut display_tab = Rc::new(RefCell::new(DisplaySysInfo::new(sys2)));

    glib::timeout::add(1500, update_window, &mut (&mut procs.list_store, sys1.clone(), display_tab.clone()));
    //window.add(&vertical_layout);
    note.create_tab("Processus list", &procs.vertical_layout);
    //let t = gtk::Button::new_with_label("test").unwrap();
    note.create_tab("System usage", &(*display_tab.borrow()).vertical_layout);
    window.add(&note.notebook);
    window.show_all();
    gtk::main();
}
<commit_msg>Add total cpu usage<commit_after>#![crate_type = "bin"]
#![feature(collections)]

extern crate gtk;
extern crate glib;
extern crate sysinfo;

use gtk::{WindowTrait, ContainerTrait, WidgetTrait, ButtonSignals, BoxTrait};
use gtk::{IconSize, Orientation, ReliefStyle};
use sysinfo::*;
use gtk::signal::Inhibit;
use gtk::signal::{WidgetSignals, TreeViewSignals};
use std::collections::VecMap;
use std::str::FromStr;
use std::rc::Rc;
use std::cell::{RefCell, RefMut};
use std::ops::DerefMut;

struct NoteBook {
    notebook: gtk::NoteBook,
    tabs: Vec<gtk::Box>
}

impl NoteBook {
    fn new() -> NoteBook {
        NoteBook {
            notebook: gtk::NoteBook::new().unwrap(),
            tabs: Vec::new()
        }
    }

    fn create_tab<'a, Widget: gtk::WidgetTrait>(&mut self, title: &'a str, widget: &Widget) -> Option<u32> {
        let label = gtk::Label::new(title).unwrap();
        let tab = gtk::Box::new(Orientation::Horizontal, 0).unwrap();

        tab.pack_start(&label, false, false, 0);
        tab.show_all();

        let index = match self.notebook.append_page(widget, Some(&tab)) {
            Some(index) => index,
            _ => return None
        };

        self.tabs.push(tab);

        Some(index)
    }
}

struct Procs {
    left_tree: gtk::TreeView,
    scroll: gtk::ScrolledWindow,
    current_pid: Rc<RefCell<Option<String>>>,
    kill_button: Rc<RefCell<gtk::Button>>,
    vertical_layout: gtk::Box,
    list_store: gtk::ListStore
}

impl Procs {
    pub fn new<'a>(proc_list: &VecMap<Processus>) -> Procs {
        let mut left_tree = gtk::TreeView::new().unwrap();
        let mut scroll = gtk::ScrolledWindow::new(None, None).unwrap();
        let mut current_pid = Rc::new(RefCell::new(None));
        let mut kill_button = Rc::new(RefCell::new(gtk::Button::new_with_label("End task").unwrap()));
        let mut current_pid1 = current_pid.clone();
        let mut current_pid2 = current_pid.clone();
        let mut kill_button1 = kill_button.clone();

        scroll.set_min_content_height(800);
        scroll.set_min_content_width(600);

        let mut columns : Vec<gtk::TreeViewColumn> = Vec::new();

        append_column("pid", &mut columns);
        append_column("process name", &mut columns);
        append_column("cpu usage", &mut columns);
        append_column("memory usage", &mut columns);

        for i in columns {
            left_tree.append_column(&i);
        }

        let mut list_store = gtk::ListStore::new(&[glib::Type::String, glib::Type::String, glib::Type::String, glib::Type::String]).unwrap();
        for (_, pro) in proc_list {
            create_and_fill_model(&mut list_store, pro.pid, &pro.name, pro.cpu_usage, pro.memory);
        }

        left_tree.set_model(&list_store.get_model().unwrap());
        left_tree.set_headers_visible(true);
        scroll.add(&left_tree);
        let mut vertical_layout = gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap();

        left_tree.connect_cursor_changed(move |tree_view| {
            match tree_view.get_selection() {
                Some(selection) => {
                    let model = tree_view.get_model().unwrap();
                    let mut iter = gtk::TreeIter::new();

                    if selection.get_selected(&model, &mut iter) {
                        let pid = model.get_value(&iter, 0).get_string();
                        let mut tmp = current_pid1.borrow_mut();
                        *tmp = pid;
                    }
                    (*kill_button1.borrow_mut()).set_sensitive((*current_pid.borrow()).is_some());
                }
                None => {
                    let mut tmp = current_pid1.borrow_mut();
                    *tmp = None;
                }
            }
        });
        (*kill_button.borrow_mut()).set_sensitive(false);

        vertical_layout.add(&scroll);
        vertical_layout.add(&(*kill_button.borrow_mut()));
        Procs {
            left_tree: left_tree,
            scroll: scroll,
            current_pid: current_pid2.clone(),
            kill_button: kill_button,
            vertical_layout: vertical_layout,
            list_store: list_store
        }
    }
}

fn append_column(title: &str, v: &mut Vec<gtk::TreeViewColumn>) {
    let l = v.len();
    let renderer = gtk::CellRendererText::new().unwrap();

    v.push(gtk::TreeViewColumn::new().unwrap());
    let tmp = v.get_mut(l).unwrap();

    tmp.set_title(title);
    tmp.pack_start(&renderer, true);
    tmp.add_attribute(&renderer, "text", l as i32);
}

fn create_and_fill_model(list_store: &mut gtk::ListStore, pid: i64, name: &str, cpu: f32, memory: u64) {
    if name.len() < 1 {
        return;
    }
    let mut top_level = gtk::TreeIter::new();

    list_store.append(&mut top_level);
    list_store.set_string(&top_level, 0, &format!("{}", pid));
    list_store.set_string(&top_level, 1, name);
    list_store.set_string(&top_level, 2, &format!("{}", cpu));
    list_store.set_string(&top_level, 3, &format!("{}", memory));
}

fn update_window(w: &mut (&mut gtk::ListStore, Rc<RefCell<sysinfo::System>>, Rc<RefCell<DisplaySysInfo>>)) -> i32 {
    let mut model = w.0.get_model().unwrap();
    let mut iter = gtk::TreeIter::new();

    (*w.1.borrow_mut()).refresh();
    let mut entries : VecMap<Processus> = ((*w.1.borrow()).get_processus_list()).clone();
    let mut nb = model.iter_n_children(None);

    let sys = w.1.clone();
    (*w.2.borrow_mut()).update_ram_display(sys.clone());
    (*w.2.borrow_mut()).update_process_display(sys);

    let mut i = 0;
    while i < nb {
        if model.iter_nth_child(&mut iter, None, i) {
            let pid : i64 = i64::from_str(&model.get_value(&iter, 0).get_string().unwrap()).unwrap();
            let mut to_delete = false;

            match entries.get(&(pid as usize)) {
                Some(p) => {
                    //w.0.set_string(&iter, 1, &p.name);
                    w.0.set_string(&iter, 2, &format!("{}", p.cpu_usage));
                    w.0.set_string(&iter, 3, &format!("{}", p.memory));
                    to_delete = true;
                }
                None => {
                    w.0.remove(&iter);
                }
            }
            if to_delete {
                entries.remove(&(pid as usize));
                nb = model.iter_n_children(None);
                i += 1;
            }
        } else {
            i += 1;
        }
    }
    for (_, pro) in entries {
        create_and_fill_model(&mut w.0, pro.pid, &pro.name, pro.cpu_usage, pro.memory);
    }
    1
}

struct DisplaySysInfo {
    procs : Rc<RefCell<Vec<gtk::ProgressBar>>>,
    ram : Rc<RefCell<gtk::ProgressBar>>,
    swap : Rc<RefCell<gtk::ProgressBar>>,
    vertical_layout : gtk::Box
}

impl DisplaySysInfo {
    pub fn new(sys1: Rc<RefCell<sysinfo::System>>) -> DisplaySysInfo {
        let mut tmp = DisplaySysInfo {
            procs: Rc::new(RefCell::new(Vec::new())),
            ram: Rc::new(RefCell::new(gtk::ProgressBar::new().unwrap())),
            swap: Rc::new(RefCell::new(gtk::ProgressBar::new().unwrap())),
            vertical_layout: gtk::Box::new(gtk::Orientation::Vertical, 0).unwrap()
        };

        let sys2 = sys1.clone();
        let mut v = tmp.procs.clone();

        (*tmp.ram.borrow_mut()).set_show_text(true);
        (*tmp.swap.borrow_mut()).set_show_text(true);

        tmp.vertical_layout.set_spacing(5);
        let mut i = 0;
        let mut total = false;

        tmp.vertical_layout.pack_start(&gtk::Label::new("Memory usage").unwrap(), false, false, 15);
        tmp.vertical_layout.add(&(*tmp.ram.borrow()));
        tmp.vertical_layout.pack_start(&gtk::Label::new("Swap usage").unwrap(), false, false, 15);
        tmp.vertical_layout.add(&(*tmp.swap.borrow()));
        tmp.vertical_layout.pack_start(&gtk::Label::new("Total CPU usage").unwrap(), false, false, 15);
        for pro in (*sys1.borrow()).get_process_list() {
            if total {
                (*v.borrow_mut()).push(gtk::ProgressBar::new().unwrap());
                let p : &gtk::ProgressBar = &(*v.borrow())[i];
                let l = gtk::Label::new(&format!("{}", i)).unwrap();
                let horizontal_layout = gtk::Box::new(gtk::Orientation::Horizontal, 0).unwrap();

                p.set_text(&format!("{:.2} %", pro.get_cpu_usage() * 100.));
                p.set_show_text(true);
                p.set_fraction(pro.get_cpu_usage() as f64);
                horizontal_layout.pack_start(&l, false, false, 5);
                horizontal_layout.pack_start(p, true, true, 5);
                tmp.vertical_layout.add(&horizontal_layout);
            } else {
                (*v.borrow_mut()).push(gtk::ProgressBar::new().unwrap());
                let p : &gtk::ProgressBar = &(*v.borrow())[i];

                p.set_text(&format!("{:.2} %", pro.get_cpu_usage() * 100.));
                p.set_show_text(true);
                p.set_fraction(pro.get_cpu_usage() as f64);

                tmp.vertical_layout.add(p);
                tmp.vertical_layout.pack_start(&gtk::Label::new("Process usage").unwrap(), false, false, 15);
                total = true;
            }
            i += 1;
        }
        tmp.update_ram_display(sys2);
        tmp
    }

    pub fn update_ram_display(&mut self, sys: Rc<RefCell<sysinfo::System>>) {
        let total = (*sys.borrow()).get_total_memory();
        let used = (*sys.borrow()).get_free_memory();
        let disp = if total < 100000 {
            format!("{} / {}B", used, total)
        } else if total < 10000000 {
            format!("{} / {}KB", used / 1000, total / 1000)
        } else if total < 10000000000 {
            format!("{} / {}MB", used / 1000000, total / 1000000)
        } else {
            format!("{} / {}GB", used / 1000000000, total / 1000000000)
        };

        (*self.ram.borrow_mut()).set_text(&disp);
        (*self.ram.borrow_mut()).set_fraction(used as f64 / total as f64);

        let total = (*sys.borrow()).get_total_swap();
        let used = total - (*sys.borrow()).get_free_swap();
        let disp = if total < 100000 {
            format!("{} / {}B", used, total)
        } else if total < 10000000 {
            format!("{} / {}KB", used / 1000, total / 1000)
        } else if total < 10000000000 {
            format!("{} / {}MB", used / 1000000, total / 1000000)
        } else {
            format!("{} / {}GB", used / 1000000000, total / 1000000000)
        };

        (*self.swap.borrow_mut()).set_text(&disp);
        (*self.swap.borrow_mut()).set_fraction(used as f64 / total as f64);
    }

    pub fn update_process_display(&mut self, sys: Rc<RefCell<sysinfo::System>>) {
        let v = &*self.procs.borrow_mut();
        let mut total = false;
        let mut i = 0;
        for pro in (*sys.borrow()).get_process_list() {
            v[i].set_text(&format!("{:.1} %", pro.get_cpu_usage() * 100.));
            v[i].set_show_text(true);
            v[i].set_fraction(pro.get_cpu_usage() as f64);
            i += 1;
        }
    }
}

fn main() {
    gtk::init();

    let mut window = gtk::Window::new(gtk::WindowType::TopLevel).unwrap();
    let mut sys : Rc<RefCell<sysinfo::System>> = Rc::new(RefCell::new(sysinfo::System::new()));
    let mut note = NoteBook::new();
    (*sys.borrow_mut()).refresh();
    let mut procs = Procs::new((*sys.borrow()).get_processus_list());
    let mut current_pid2 = procs.current_pid.clone();
    let mut sys1 = sys.clone();
    let mut sys2 = sys.clone();

    window.set_title("Processus viewer");
    window.set_window_position(gtk::WindowPosition::Center);

    window.connect_delete_event(|_, _| {
        gtk::main_quit();
        Inhibit(true)
    });

    (*procs.kill_button.borrow_mut()).connect_clicked(move |_| {
        let tmp = (*current_pid2.borrow_mut()).is_some() ;

        if tmp {
            let s = (*current_pid2.borrow()).clone();
            match (*sys.borrow_mut()).get_processus(i64::from_str(&s.unwrap()).unwrap()) {
                Some(p) => {
                    p.kill(Signal::Kill);
                },
                None => {}
            };
        }
    });

    let mut display_tab = Rc::new(RefCell::new(DisplaySysInfo::new(sys2)));

    glib::timeout::add(1500, update_window, &mut (&mut procs.list_store, sys1.clone(), display_tab.clone()));
    //window.add(&vertical_layout);
    note.create_tab("Processus list", &procs.vertical_layout);
    //let t = gtk::Button::new_with_label("test").unwrap();
    note.create_tab("System usage", &(*display_tab.borrow()).vertical_layout);
    window.add(&note.notebook);
    window.show_all();
    gtk::main();
}
<|endoftext|>


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

use std::thread;
use std::thread::JoinHandle;
use std::net::TcpStream;
use std::ops::DerefMut;
use std::os::unix::io::RawFd;
use std::collections::LinkedList;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{
    channel,
    Sender,
    Receiver,
    TryRecvError
};

use super::types::*;
use super::epoll;
use super::epoll::util::*;
use super::epoll::EpollEvent;
use super::socket::Socket;
use super::simple_stream::nbetstream::NbetStream;

#[allow(dead_code)]
pub struct EventLoop {
    /// Handle to event loop process
    prox: JoinHandle<()>,
    /// Sender<T> to proc
    prox_tx: Sender<TcpStream>
}

impl EventLoop {

    /// Returns a new EventLoop
    pub fn new(server_tx: Sender<EventTuple>) -> EventLoop {
        let (tx, rx): (Sender<TcpStream>, Receiver<TcpStream>) = channel();
        let prox = thread::Builder::new()
            .name("EventLoop".to_string())
            .spawn(move || {
                EventLoop::start(rx, server_tx);
            }).unwrap();

        EventLoop {
            prox: prox,
            prox_tx: tx
        }
    }

    /// Returns a clone of this EventLoop's Sender<TcpStream> channel
    pub fn sender(&self) -> Sender<TcpStream> {
        self.prox_tx.clone()
    }

    /// Main event loop
    fn start(rx: Receiver<TcpStream>, uspace_tx: Sender<EventTuple>) {
        // Master socket list
        let sockets = Arc::new(Mutex::new(LinkedList::<Socket>::new()));

        // Our error handler for this thread and it's children
        let (err_tx, err_rx): (Sender<()>, Receiver<()>) = channel();

        // Create an epoll instance
        let result = epoll::create1(0);
        if result.is_err() {
            let err = result.unwrap_err();
            panic!("Unable to create epoll instance: {}", err)
        }
        let epoll_instance = result.unwrap();

        // Start the epoll thread
        let c_sockets = sockets.clone();
        let c_epoll_instance = epoll_instance.clone();
        let epoll_prox = thread::Builder::new()
            .name("EpollLoop".to_string())
            .spawn(move || {
                EventLoop::epoll_loop(uspace_tx, c_sockets, c_epoll_instance, err_rx);
            }).unwrap();

        for new_stream in rx.iter() {
            match NbetStream::new(new_stream) {
                Ok(s_stream) => {
                    // Add to master list
                    let socket = Socket::new(s_stream);

                    // TODO - Replace with recoverable version once into_inner() is stable
                    // Unfortunately, this is the only stable way to use mutexes at the moment
                    // Hopefully recovering from poisoning will be in 1.2
                    // let mut s_guard = match sockets.lock() {
                    //     Ok(guard) => guard,
                    //     Err(poisoned) => poisoned.into_inner()
                    // };
                    let mut s_guard = sockets.lock().unwrap();
                    let s_list = s_guard.deref_mut();
                    let s_fd = socket.raw_fd();
                    s_list.push_back(socket);

                    // Add to epoll
                    let event = Box::new(EpollEvent {
                        data: s_fd as u64,
                        events: (event_type::EPOLLIN | event_type::EPOLLET | event_type::EPOLLRDHUP)
                    });
                    match epoll::ctl(epoll_instance, ctl_op::ADD, s_fd, event) {
                        Ok(()) => println!("New socket added to epoll list"),
                        Err(e) => println!("Epoll CtrlError during add: {}", e)
                    };
                }
                Err(e) => {
                    // TODO - write error to log file
                    println!("Error: {}", e);
                }
            };
        }

        // If we get here, shit is real bad. It means we have lost our
        // channel to the outside world.
        // Kill off event loop, so unwinding can begin
        println!("Rust.EventLoop.start - thread finished");
        let _ = err_tx.send(());
    }

    /// Calls epoll_wait forever waiting for events to be fired up from kernel
    fn epoll_loop(uspace_tx: Sender<EventTuple>,
                  sockets: SocketList,
                  epoll_instance: RawFd,
                  err_rx: Receiver<()>) {

        // This is the maximum number of events we want to be notified of at one time
        let mut events = Vec::<EpollEvent>::with_capacity(100);
        unsafe { events.set_len(100); }

        loop {
            // Wait for epoll events
            println!("epoll_wait...");
            match epoll::wait(epoll_instance, &mut events[..], -1) {
                Ok(num_events) => {
                    println!("{} epoll event(s) received", num_events);
                    for x in 0..num_events {
                        let tx_clone = uspace_tx.clone();
                        let s_clone = sockets.clone();
                        let epfd_clone = epoll_instance.clone();

                        EventLoop::handle_epoll_event(epfd_clone,
                            &events[x as usize], s_clone, tx_clone);
                    }
                }
                Err(e) => {
                    panic!("Error on epoll::wait(): {}", e)
                }
            }

            // If we receive anything on this end or lose communication, we need to halt,
            // because we have no where to send up events, our "user space" channel is
            // gone
            match err_rx.try_recv() {
                Ok(_) => {
                    println!("Error, terminating epoll_loop");
                    break;
                }
                Err(e) => {
                    match e {
                        TryRecvError::Empty => {}
                        _ => {
                            println!("Error, terminating epoll_loo");
                            break;
                        }
                    }
                }
            };
        }
        println!("epoll_loop - thread finished");
    }

    /// Processes a read on the socket that is ready for reading
    ///
    /// An epoll event is passed to this function, with the fd that is ready to be read
    /// A traversal through the list of sockets is then performed in order to locate the
    /// NbetStream instance that is wrapping that file descriptor.
    /// After a read is performed on that socket, it's data will be passed back to the
    /// "user-space" section of the Server to process the payload however it needs to.
    ///
    /// TODO - Create a faster way to gain a safe reference to the NbetStream wrapping
    /// the fd instead of an O(n) lookup for every event. Currently lookup time is
    /// O(n) * MAX_EPOLL_EVENTS, where MAX_EPOLL_EVENTS is user defined in this function's
    /// caller.
    fn handle_epoll_event(epoll_instance: RawFd,
                          event: &EpollEvent,
                          sockets: SocketList,
                          uspace_tx: Sender<EventTuple>) {

        println!("handle_epoll_event");

        // List of socket id's that produced an Err on read attempt
        // Cannot do it within the match block because the list will need
        // access to the mutex, while this scope has the mutex, which will result
        // in a spinlock that will spin forever. And, that's a trip to bummer town
        let mut errd_socket_ids = Vec::<u32>::new();

        // The fd we are traversing the list for
        let fd = event.data;

        // We need a handle before the list gets locked and used
        let s_list_clone = sockets.clone();

        // TODO - Replace with recoverable version once into_inner() is stable
        // Unfortunately, this is the only stable way to use mutexes at the moment
        // Hopefully recovering from poisoning will be in 1.2
        // let mut s_guard = match sockets.lock() {
        //     Ok(guard) => guard,
        //     Err(poisoned) => poisoned.into_inner()
        // };
        let mut s_guard = sockets.lock().unwrap();
        let s_list = s_guard.deref_mut();
        for socket in s_list.iter_mut() {
            if socket.raw_fd() as u64 == fd {
                // We've found the socket for the event passed from epoll
                // Lets attempt to read from it
                match socket.read() {
                    Ok(()) => {
                        println!("Rust.EventLoop.handle_epoll_event.socket.read.Ok()");
                        for msg in socket.buffer().iter() {
                            let msg_for_debug = msg.clone();
                            match String::from_utf8(msg_for_debug) {
                                Ok(msg_str) => {
                                    println!("Rust.EventLoop.handle_epoll_event.socket.read.ok.msg_as_string: {}", msg_str);
                                }
                                Err(e) => {
                                    println!("Rust.EventLoop.handle_epoll_event.socket.read.ok.msg_as_string.Err: {}", e);
                                }
                            };

                            let list_handle = s_list_clone.clone();
                            let socket_clone = socket.clone();
                            let msg_clone = msg.clone();
                            let _ = uspace_tx.send((list_handle, socket_clone, msg_clone));
                        }
                    }
                    Err(e) => {
                        println!("Rust.EventLoop.handle_epoll_event.socket.read.Err(e): {}", e);
                        // We really don't care what happened, it all results in the same
                        // outcome - remove the socket from system...
                        errd_socket_ids.push(socket.id());
                        let s_clone = socket.clone();
                        let epfd_clone = epoll_instance.clone();
                        EventLoop::remove_socket_from_epoll(s_clone, epfd_clone);
                    }
                }
            }
        }

        // Remove any errd sockets from the master list of sockets
        EventLoop::remove_socket_from_list(errd_socket_ids, &mut s_list);
    }

    /// Removes socket from the epoll watch list
    fn remove_socket_from_epoll(socket: Socket, epoll_instance: RawFd) {
        println!("remove_socket_from_epoll");

        // In kernel versions before 2.6.9, the EPOLL_CTL_DEL operation required
        // a non-null pointer in event, even though this argument is ignored.
        // Since Linux 2.6.9, event can be specified as NULL when using
        // EPOLL_CTL_DEL.  Applications that need to be portable to kernels
        // before 2.6.9 should specify a non-null pointer in event.
        let event = Box::new(EpollEvent {
            data: 0 as u64,
            events: 0 as u32
        });
        let s_fd = socket.raw_fd();

        // Depending on how fd are duplicated with .clone(), this may fail
        // If the failure case is CtlError::ENOENT, we do not care, because
        // epoll will clean the up the descriptor after they are all dropped from
        // program memory
        match epoll::ctl(epoll_instance, ctl_op::DEL, s_fd, event) {
            Ok(()) => println!("Socket removed from epoll watch list"),
            Err(e) => {
                match e {
                    CtlError::ENOENT => {
                        println!("Fd not found in epoll, will remove when fd is syscall closed");
                    }
                    _ => println!("Epoll CtrlError during del: {}", e)
                }
            }
        };
    }

    /// Removes socket from master list
    fn remove_socket_from_list(socket_ids: Vec<u32>, sockets: &mut LinkedList<Socket>) {
        println!("remove_socket_from_list");

        // TODO - Replace with recoverable version once into_inner() is stable
        // Unfortunately, this is the only stable way to use mutexes at the moment
        // Hopefully recovering from poisoning will be in 1.2
        // let mut s_guard = match sockets.lock() {
        //     Ok(guard) => guard,
        //     Err(poisoned) => poisoned.into_inner()
        // };
        // let mut s_guard = sockets.lock().unwrap();
        // let s_list = s_guard.deref_mut();

        for socket_id in socket_ids.iter() {
            let mut socket_found = false;
            let mut index: usize = 1;
            for socket in sockets.iter() {
                if socket.id() == *socket_id {
                    socket_found = true;
                    break;
                }
                index += 1;
            }

            if !socket_found {
                println!("Socket not found in list...?");
                return;
            }

            if index == 1 {
                sockets.pop_front();
            } else {
                let mut split = sockets.split_off(index - 1);
                split.pop_front();
                sockets.append(&mut split);
            }

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

use std::thread;
use std::thread::JoinHandle;
use std::net::TcpStream;
use std::ops::DerefMut;
use std::os::unix::io::RawFd;
use std::collections::LinkedList;
use std::sync::{Arc, Mutex};
use std::sync::mpsc::{
    channel,
    Sender,
    Receiver,
    TryRecvError
};

use super::types::*;
use super::epoll;
use super::epoll::util::*;
use super::epoll::EpollEvent;
use super::socket::Socket;
use super::simple_stream::nbetstream::NbetStream;

#[allow(dead_code)]
pub struct EventLoop {
    /// Handle to event loop process
    prox: JoinHandle<()>,
    /// Sender<T> to proc
    prox_tx: Sender<TcpStream>
}

impl EventLoop {

    /// Returns a new EventLoop
    pub fn new(server_tx: Sender<EventTuple>) -> EventLoop {
        let (tx, rx): (Sender<TcpStream>, Receiver<TcpStream>) = channel();
        let prox = thread::Builder::new()
            .name("EventLoop".to_string())
            .spawn(move || {
                EventLoop::start(rx, server_tx);
            }).unwrap();

        EventLoop {
            prox: prox,
            prox_tx: tx
        }
    }

    /// Returns a clone of this EventLoop's Sender<TcpStream> channel
    pub fn sender(&self) -> Sender<TcpStream> {
        self.prox_tx.clone()
    }

    /// Main event loop
    fn start(rx: Receiver<TcpStream>, uspace_tx: Sender<EventTuple>) {
        // Master socket list
        let sockets = Arc::new(Mutex::new(LinkedList::<Socket>::new()));

        // Our error handler for this thread and it's children
        let (err_tx, err_rx): (Sender<()>, Receiver<()>) = channel();

        // Create an epoll instance
        let result = epoll::create1(0);
        if result.is_err() {
            let err = result.unwrap_err();
            panic!("Unable to create epoll instance: {}", err)
        }
        let epoll_instance = result.unwrap();

        // Start the epoll thread
        let c_sockets = sockets.clone();
        let c_epoll_instance = epoll_instance.clone();
        let epoll_prox = thread::Builder::new()
            .name("EpollLoop".to_string())
            .spawn(move || {
                EventLoop::epoll_loop(uspace_tx, c_sockets, c_epoll_instance, err_rx);
            }).unwrap();

        for new_stream in rx.iter() {
            match NbetStream::new(new_stream) {
                Ok(s_stream) => {
                    // Add to master list
                    let socket = Socket::new(s_stream);

                    // TODO - Replace with recoverable version once into_inner() is stable
                    // Unfortunately, this is the only stable way to use mutexes at the moment
                    // Hopefully recovering from poisoning will be in 1.2
                    // let mut s_guard = match sockets.lock() {
                    //     Ok(guard) => guard,
                    //     Err(poisoned) => poisoned.into_inner()
                    // };
                    let mut s_guard = sockets.lock().unwrap();
                    let s_list = s_guard.deref_mut();
                    let s_fd = socket.raw_fd();
                    s_list.push_back(socket);

                    // Add to epoll
                    let event = Box::new(EpollEvent {
                        data: s_fd as u64,
                        events: (event_type::EPOLLIN | event_type::EPOLLET | event_type::EPOLLRDHUP)
                    });
                    match epoll::ctl(epoll_instance, ctl_op::ADD, s_fd, event) {
                        Ok(()) => println!("New socket added to epoll list"),
                        Err(e) => println!("Epoll CtrlError during add: {}", e)
                    };
                }
                Err(e) => {
                    // TODO - write error to log file
                    println!("Error: {}", e);
                }
            };
        }

        // If we get here, shit is real bad. It means we have lost our
        // channel to the outside world.
        // Kill off event loop, so unwinding can begin
        println!("Rust.EventLoop.start - thread finished");
        let _ = err_tx.send(());
    }

    /// Calls epoll_wait forever waiting for events to be fired up from kernel
    fn epoll_loop(uspace_tx: Sender<EventTuple>,
                  sockets: SocketList,
                  epoll_instance: RawFd,
                  err_rx: Receiver<()>) {

        // This is the maximum number of events we want to be notified of at one time
        let mut events = Vec::<EpollEvent>::with_capacity(100);
        unsafe { events.set_len(100); }

        loop {
            // Wait for epoll events
            println!("epoll_wait...");
            match epoll::wait(epoll_instance, &mut events[..], -1) {
                Ok(num_events) => {
                    println!("{} epoll event(s) received", num_events);
                    for x in 0..num_events {
                        let tx_clone = uspace_tx.clone();
                        let s_clone = sockets.clone();
                        let epfd_clone = epoll_instance.clone();

                        EventLoop::handle_epoll_event(epfd_clone,
                            &events[x as usize], s_clone, tx_clone);
                    }
                }
                Err(e) => {
                    panic!("Error on epoll::wait(): {}", e)
                }
            }

            // If we receive anything on this end or lose communication, we need to halt,
            // because we have no where to send up events, our "user space" channel is
            // gone
            match err_rx.try_recv() {
                Ok(_) => {
                    println!("Error, terminating epoll_loop");
                    break;
                }
                Err(e) => {
                    match e {
                        TryRecvError::Empty => {}
                        _ => {
                            println!("Error, terminating epoll_loo");
                            break;
                        }
                    }
                }
            };
        }
        println!("epoll_loop - thread finished");
    }

    /// Processes a read on the socket that is ready for reading
    ///
    /// An epoll event is passed to this function, with the fd that is ready to be read
    /// A traversal through the list of sockets is then performed in order to locate the
    /// NbetStream instance that is wrapping that file descriptor.
    /// After a read is performed on that socket, it's data will be passed back to the
    /// "user-space" section of the Server to process the payload however it needs to.
    ///
    /// TODO - Create a faster way to gain a safe reference to the NbetStream wrapping
    /// the fd instead of an O(n) lookup for every event. Currently lookup time is
    /// O(n) * MAX_EPOLL_EVENTS, where MAX_EPOLL_EVENTS is user defined in this function's
    /// caller.
    fn handle_epoll_event(epoll_instance: RawFd,
                          event: &EpollEvent,
                          sockets: SocketList,
                          uspace_tx: Sender<EventTuple>) {

        println!("handle_epoll_event");

        // List of socket id's that produced an Err on read attempt
        // Cannot do it within the match block because the list will need
        // access to the mutex, while this scope has the mutex, which will result
        // in a spinlock that will spin forever. And, that's a trip to bummer town
        let mut errd_socket_ids = Vec::<u32>::new();

        // The fd we are traversing the list for
        let fd = event.data;

        // We need a handle before the list gets locked and used
        let s_list_clone = sockets.clone();

        // TODO - Replace with recoverable version once into_inner() is stable
        // Unfortunately, this is the only stable way to use mutexes at the moment
        // Hopefully recovering from poisoning will be in 1.2
        // let mut s_guard = match sockets.lock() {
        //     Ok(guard) => guard,
        //     Err(poisoned) => poisoned.into_inner()
        // };
        let mut s_guard = sockets.lock().unwrap();
        let mut s_list = s_guard.deref_mut();
        for socket in s_list.iter_mut() {
            if socket.raw_fd() as u64 == fd {
                // We've found the socket for the event passed from epoll
                // Lets attempt to read from it
                match socket.read() {
                    Ok(()) => {
                        println!("Rust.EventLoop.handle_epoll_event.socket.read.Ok()");
                        for msg in socket.buffer().iter() {
                            let msg_for_debug = msg.clone();
                            match String::from_utf8(msg_for_debug) {
                                Ok(msg_str) => {
                                    println!("Rust.EventLoop.handle_epoll_event.socket.read.ok.msg_as_string: {}", msg_str);
                                }
                                Err(e) => {
                                    println!("Rust.EventLoop.handle_epoll_event.socket.read.ok.msg_as_string.Err: {}", e);
                                }
                            };

                            let list_handle = s_list_clone.clone();
                            let socket_clone = socket.clone();
                            let msg_clone = msg.clone();
                            let _ = uspace_tx.send((list_handle, socket_clone, msg_clone));
                        }
                    }
                    Err(e) => {
                        println!("Rust.EventLoop.handle_epoll_event.socket.read.Err(e): {}", e);
                        // We really don't care what happened, it all results in the same
                        // outcome - remove the socket from system...
                        errd_socket_ids.push(socket.id());
                        let s_clone = socket.clone();
                        let epfd_clone = epoll_instance.clone();
                        EventLoop::remove_socket_from_epoll(s_clone, epfd_clone);
                    }
                }
            }
        }

        // Remove any errd sockets from the master list of sockets
        EventLoop::remove_socket_from_list(errd_socket_ids, &mut s_list);
    }

    /// Removes socket from the epoll watch list
    fn remove_socket_from_epoll(socket: Socket, epoll_instance: RawFd) {
        println!("remove_socket_from_epoll");

        // In kernel versions before 2.6.9, the EPOLL_CTL_DEL operation required
        // a non-null pointer in event, even though this argument is ignored.
        // Since Linux 2.6.9, event can be specified as NULL when using
        // EPOLL_CTL_DEL.  Applications that need to be portable to kernels
        // before 2.6.9 should specify a non-null pointer in event.
        let event = Box::new(EpollEvent {
            data: 0 as u64,
            events: 0 as u32
        });
        let s_fd = socket.raw_fd();

        // Depending on how fd are duplicated with .clone(), this may fail
        // If the failure case is CtlError::ENOENT, we do not care, because
        // epoll will clean the up the descriptor after they are all dropped from
        // program memory
        match epoll::ctl(epoll_instance, ctl_op::DEL, s_fd, event) {
            Ok(()) => println!("Socket removed from epoll watch list"),
            Err(e) => {
                match e {
                    CtlError::ENOENT => {
                        println!("Fd not found in epoll, will remove when fd is syscall closed");
                    }
                    _ => println!("Epoll CtrlError during del: {}", e)
                }
            }
        };
    }

    /// Removes socket from master list
    fn remove_socket_from_list(socket_ids: Vec<u32>, sockets: &mut LinkedList<Socket>) {
        println!("remove_socket_from_list");

        // TODO - Replace with recoverable version once into_inner() is stable
        // Unfortunately, this is the only stable way to use mutexes at the moment
        // Hopefully recovering from poisoning will be in 1.2
        // let mut s_guard = match sockets.lock() {
        //     Ok(guard) => guard,
        //     Err(poisoned) => poisoned.into_inner()
        // };
        // let mut s_guard = sockets.lock().unwrap();
        // let s_list = s_guard.deref_mut();

        for socket_id in socket_ids.iter() {
            let mut socket_found = false;
            let mut index: usize = 1;
            for socket in sockets.iter() {
                if socket.id() == *socket_id {
                    socket_found = true;
                    break;
                }
                index += 1;
            }

            if !socket_found {
                println!("Socket not found in list...?");
                return;
            }

            if index == 1 {
                sockets.pop_front();
            } else {
                let mut split = sockets.split_off(index - 1);
                split.pop_front();
                sockets.append(&mut split);
            }

            println!("Socket removed sucessfully");
        }
    }
}
<|endoftext|>


<commit_before>use std::{cmp, mem};

use orbclient::{Color, Event, EventOption, Window};
use orbclient::event;

fn ansi_color(value: u8) -> Color {
    match value {
        0 => Color::rgb(0x00, 0x00, 0x00),
        1 => Color::rgb(0x80, 0x00, 0x00),
        2 => Color::rgb(0x00, 0x80, 0x00),
        3 => Color::rgb(0x80, 0x80, 0x00),
        4 => Color::rgb(0x00, 0x00, 0x80),
        5 => Color::rgb(0x80, 0x00, 0x80),
        6 => Color::rgb(0x00, 0x80, 0x80),
        7 => Color::rgb(0xc0, 0xc0, 0xc0),
        8 => Color::rgb(0x80, 0x80, 0x80),
        9 => Color::rgb(0xff, 0x00, 0x00),
        10 => Color::rgb(0x00, 0xff, 0x00),
        11 => Color::rgb(0xff, 0xff, 0x00),
        12 => Color::rgb(0x00, 0x00, 0xff),
        13 => Color::rgb(0xff, 0x00, 0xff),
        14 => Color::rgb(0x00, 0xff, 0xff),
        15 => Color::rgb(0xff, 0xff, 0xff),
        16 ... 231 => {
            let convert = |value: u8| -> u8 {
                match value {
                    0 => 0,
                    _ => value * 0x28 + 0x28
                }
            };

            let r = convert((value - 16)/36 % 6);
            let g = convert((value - 16)/6 % 6);
            let b = convert((value - 16) % 6);
            Color::rgb(r, g, b)
        },
        232 ... 255 => {
            let gray = (value - 232) * 10 + 8;
            Color::rgb(gray, gray, gray)
        },
        _ => Color::rgb(0, 0, 0)
    }
}

pub struct Console {
    pub window: Box<Window>,
    pub point_x: i32,
    pub point_y: i32,
    pub foreground: Color,
    pub background: Color,
    pub redraw: bool,
    pub command: String,
    pub escape: bool,
    pub escape_sequence: bool,
    pub sequence: Vec<String>,
    pub raw_mode: bool,
}

impl Console {
    pub fn new(width: u32, height: u32) -> Console {
        let mut window = Window::new_flags(-1, -1, width, height, "Terminal", true).unwrap();
        window.sync();
        Console {
            window: window,
            point_x: 0,
            point_y: 0,
            foreground: ansi_color(7),
            background: ansi_color(0),
            redraw: true,
            command: String::new(),
            escape: false,
            escape_sequence: false,
            sequence: Vec::new(),
            raw_mode: false,
        }
    }

    pub fn code(&mut self, c: char) {
        if self.escape_sequence {
            match c {
                '0' ... '9' => {
                    // Add a number to the sequence list
                    if let Some(mut value) = self.sequence.last_mut() {
                        value.push(c);
                    }
                },
                ';' => {
                    // Split sequence into list
                    self.sequence.push(String::new());
                },
                'm' => {
                    // Display attributes
                    let mut value_iter = self.sequence.iter();
                    while let Some(value_str) = value_iter.next() {
                        let value = value_str.parse::<u8>().unwrap_or(0);
                        match value {
                            0 => {
                                self.foreground = ansi_color(7);
                                self.background = ansi_color(0);
                            },
                            30 ... 37 => self.foreground = ansi_color(value - 30),
                            38 => match value_iter.next().map_or("", |s| &s).parse::<usize>().unwrap_or(0) {
                                2 => {
                                    //True color
                                    let r = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    let g = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    let b = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    self.foreground = Color::rgb(r, g, b);
                                },
                                5 => {
                                    //256 color
                                    let color_value = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    self.foreground = ansi_color(color_value);
                                },
                                _ => {}
                            },
                            40 ... 47 => self.background = ansi_color(value - 40),
                            48 => match value_iter.next().map_or("", |s| &s).parse::<usize>().unwrap_or(0) {
                                2 => {
                                    //True color
                                    let r = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    let g = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    let b = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    self.background = Color::rgb(r, g, b);
                                },
                                5 => {
                                    //256 color
                                    let color_value = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    self.background = ansi_color(color_value);
                                },
                                _ => {}
                            },
                            _ => {},
                        }
                    }

                    self.escape_sequence = false;
                },
                'J' => {
                    match self.sequence.get(0).map_or("", |p| &p).parse::<usize>().unwrap_or(0) {
                        0 => {
                            //TODO: Erase down
                        },
                        1 => {
                            //TODO: Erase up
                        },
                        2 => {
                            // Erase all
                            self.point_x = 0;
                            self.point_y = 0;
                            self.window.set(self.background);
                            self.redraw = true;
                        },
                        _ => {}
                    }

                    self.escape_sequence = false;
                },
                'H' | 'f' => {
                    self.window.rect(self.point_x, self.point_y, 8, 16, self.background);

                    let row = self.sequence.get(0).map_or("", |p| &p).parse::<isize>().unwrap_or(1);
                    self.point_y = cmp::max(0, row - 1) as i32 * 16;

                    let col = self.sequence.get(1).map_or("", |p| &p).parse::<isize>().unwrap_or(1);
                    self.point_x = cmp::max(0, col - 1) as i32 * 8;

                    self.window.rect(self.point_x, self.point_y, 8, 16, self.foreground);
                    self.redraw = true;

                    self.escape_sequence = false;
                },
/*
@MANSTART{terminal-raw-mode}
INTRODUCTION
    Since Redox has no ioctl syscall, it uses escape codes for switching to raw mode.

ENTERING AND EXITING RAW MODE
    Entering raw mode is done using CSI-r (^[r). Unsetting raw mode is done by CSI-R (^[R).

RAW MODE
    Raw mode means that the stdin must be handled solely by the program itself. It will not automatically be printed nor will it be modified in any way (modulo escape codes).

    This means that:
        - stdin is not printed.
        - newlines are interpreted as carriage returns in stdin.
        - stdin is not buffered, meaning that the stream of bytes goes directly to the program, without the user having to press enter.
@MANEND
*/
                'r' => {
                    self.raw_mode = true;
                    self.escape_sequence = false;
                },
                'R' => {
                    self.raw_mode = false;
                    self.escape_sequence = false;
                },
                _ => self.escape_sequence = false,

            }

            if !self.escape_sequence {
                self.sequence.clear();
                self.escape = false;
            }
        } else {
            match c {
                '[' => {
                    // Control sequence initiator

                    self.escape_sequence = true;
                    self.sequence.push(String::new());
                },
                'c' => {
                    // Reset
                    self.point_x = 0;
                    self.point_y = 0;
                    self.raw_mode = false;
                    self.foreground = ansi_color(7);
                    self.background = ansi_color(0);
                    self.window.set(self.background);
                    self.redraw = true;

                    self.escape = false;
                }
                _ => self.escape = false,
            }
        }
    }

    pub fn scroll(&mut self, rows: usize) {
        if rows > 0 && rows < self.window.height() as usize {
            let offset = rows * self.window.width() as usize;
            let data = self.window.data_mut();
            for i in 0..data.len() - offset {
                let color = data[i + offset];
                data[i] = color;
            }
            for i in data.len() - offset..data.len() {
                data[i] = self.background;
            }
        }
    }

    pub fn character(&mut self, c: char) {
        self.window.rect(self.point_x, self.point_y, 8, 16, self.background);

        match c {
            '\0' => {},
            '\x1B' => self.escape = true,
            '\n' => {
                self.point_x = 0;
                self.point_y += 16;
            },
            '\t' => self.point_x = ((self.point_x / 64) + 1) * 64,
            '\r' => self.point_x = 0,
            '\x08' => {
                self.point_x -= 8;
                if self.point_x < 0 {
                    self.point_x = 0
                }
                self.window.rect(self.point_x, self.point_y, 8, 16, self.background);
            },
            _ => {
                self.window.char(self.point_x, self.point_y, c, self.foreground);
                self.point_x += 8;
            }
        }

        if self.point_x >= self.window.width() as i32 {
            self.point_x = 0;
            self.point_y += 16;
        }

        while self.point_y + 16 > self.window.height() as i32 {
            self.scroll(16);
            self.point_y -= 16;
        }
        self.window.rect(self.point_x, self.point_y, 8, 16, self.foreground);
        self.redraw = true;
    }

    pub fn event(&mut self, event: Event) -> Option<String> {
        match event.to_option() {
            EventOption::Key(key_event) => {
                if key_event.pressed {
                    if self.raw_mode {
                        match key_event.scancode {
                            event::K_BKSP => self.command.push_str("\x08"),
                            event::K_UP => self.command.push_str("\x1B[A"),
                            event::K_DOWN => self.command.push_str("\x1B[B"),
                            event::K_RIGHT => self.command.push_str("\x1B[C"),
                            event::K_LEFT => self.command.push_str("\x1B[D"),
                            _ => match key_event.character {
                                '\0' => {},
                                c => {
                                    self.command.push(c);
                                }
                            },
                        }

                        if ! self.command.is_empty() {
                            let mut command = String::new();
                            mem::swap(&mut self.command, &mut command);
                            return Some(command);
                        }
                    } else {
                        match key_event.scancode {
                            event::K_BKSP => if ! self.command.is_empty() {
                                self.write(&[8]);
                                self.command.pop();
                            },
                            _ => match key_event.character {
                                '\0' => (),
                                c => {
                                    self.write(&[c as u8]);
                                    self.command.push(c);

                                    if c == '\n' {
                                        let mut command = String::new();
                                        mem::swap(&mut self.command, &mut command);
                                        return Some(command);
                                    }
                                }
                            },
                        }
                    }
                }
            },
            _ => (),
        }

        None
    }

    pub fn write(&mut self, bytes: &[u8]) {
        for byte in bytes.iter() {
            let c = *byte as char;

            if self.escape {
                self.code(c);
            } else {
                self.character(c);
            }
        }

        if self.redraw {
            self.redraw = false;
            self.window.sync();
        }
    }
}
<commit_msg>Fixes for performance<commit_after>extern crate system;

use self::system::graphics::{fast_copy, fast_set};

use std::{cmp, mem};

use orbclient::{Color, Event, EventOption, Window};
use orbclient::event;

fn ansi_color(value: u8) -> Color {
    match value {
        0 => Color::rgb(0x00, 0x00, 0x00),
        1 => Color::rgb(0x80, 0x00, 0x00),
        2 => Color::rgb(0x00, 0x80, 0x00),
        3 => Color::rgb(0x80, 0x80, 0x00),
        4 => Color::rgb(0x00, 0x00, 0x80),
        5 => Color::rgb(0x80, 0x00, 0x80),
        6 => Color::rgb(0x00, 0x80, 0x80),
        7 => Color::rgb(0xc0, 0xc0, 0xc0),
        8 => Color::rgb(0x80, 0x80, 0x80),
        9 => Color::rgb(0xff, 0x00, 0x00),
        10 => Color::rgb(0x00, 0xff, 0x00),
        11 => Color::rgb(0xff, 0xff, 0x00),
        12 => Color::rgb(0x00, 0x00, 0xff),
        13 => Color::rgb(0xff, 0x00, 0xff),
        14 => Color::rgb(0x00, 0xff, 0xff),
        15 => Color::rgb(0xff, 0xff, 0xff),
        16 ... 231 => {
            let convert = |value: u8| -> u8 {
                match value {
                    0 => 0,
                    _ => value * 0x28 + 0x28
                }
            };

            let r = convert((value - 16)/36 % 6);
            let g = convert((value - 16)/6 % 6);
            let b = convert((value - 16) % 6);
            Color::rgb(r, g, b)
        },
        232 ... 255 => {
            let gray = (value - 232) * 10 + 8;
            Color::rgb(gray, gray, gray)
        },
        _ => Color::rgb(0, 0, 0)
    }
}

pub struct Console {
    pub window: Box<Window>,
    pub point_x: i32,
    pub point_y: i32,
    pub foreground: Color,
    pub background: Color,
    pub redraw: bool,
    pub command: String,
    pub escape: bool,
    pub escape_sequence: bool,
    pub sequence: Vec<String>,
    pub raw_mode: bool,
}

impl Console {
    pub fn new(width: u32, height: u32) -> Console {
        let mut window = Window::new_flags(-1, -1, width, height, "Terminal", true).unwrap();
        window.sync();
        Console {
            window: window,
            point_x: 0,
            point_y: 0,
            foreground: ansi_color(7),
            background: ansi_color(0),
            redraw: true,
            command: String::new(),
            escape: false,
            escape_sequence: false,
            sequence: Vec::new(),
            raw_mode: false,
        }
    }

    pub fn code(&mut self, c: char) {
        if self.escape_sequence {
            match c {
                '0' ... '9' => {
                    // Add a number to the sequence list
                    if let Some(mut value) = self.sequence.last_mut() {
                        value.push(c);
                    }
                },
                ';' => {
                    // Split sequence into list
                    self.sequence.push(String::new());
                },
                'm' => {
                    // Display attributes
                    let mut value_iter = self.sequence.iter();
                    while let Some(value_str) = value_iter.next() {
                        let value = value_str.parse::<u8>().unwrap_or(0);
                        match value {
                            0 => {
                                self.foreground = ansi_color(7);
                                self.background = ansi_color(0);
                            },
                            30 ... 37 => self.foreground = ansi_color(value - 30),
                            38 => match value_iter.next().map_or("", |s| &s).parse::<usize>().unwrap_or(0) {
                                2 => {
                                    //True color
                                    let r = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    let g = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    let b = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    self.foreground = Color::rgb(r, g, b);
                                },
                                5 => {
                                    //256 color
                                    let color_value = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    self.foreground = ansi_color(color_value);
                                },
                                _ => {}
                            },
                            40 ... 47 => self.background = ansi_color(value - 40),
                            48 => match value_iter.next().map_or("", |s| &s).parse::<usize>().unwrap_or(0) {
                                2 => {
                                    //True color
                                    let r = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    let g = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    let b = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    self.background = Color::rgb(r, g, b);
                                },
                                5 => {
                                    //256 color
                                    let color_value = value_iter.next().map_or("", |s| &s).parse::<u8>().unwrap_or(0);
                                    self.background = ansi_color(color_value);
                                },
                                _ => {}
                            },
                            _ => {},
                        }
                    }

                    self.escape_sequence = false;
                },
                'J' => {
                    match self.sequence.get(0).map_or("", |p| &p).parse::<usize>().unwrap_or(0) {
                        0 => {
                            //TODO: Erase down
                        },
                        1 => {
                            //TODO: Erase up
                        },
                        2 => {
                            // Erase all
                            self.point_x = 0;
                            self.point_y = 0;
                            self.window.set(self.background);
                            self.redraw = true;
                        },
                        _ => {}
                    }

                    self.escape_sequence = false;
                },
                'H' | 'f' => {
                    self.window.rect(self.point_x, self.point_y, 8, 16, self.background);

                    let row = self.sequence.get(0).map_or("", |p| &p).parse::<isize>().unwrap_or(1);
                    self.point_y = cmp::max(0, row - 1) as i32 * 16;

                    let col = self.sequence.get(1).map_or("", |p| &p).parse::<isize>().unwrap_or(1);
                    self.point_x = cmp::max(0, col - 1) as i32 * 8;

                    self.window.rect(self.point_x, self.point_y, 8, 16, self.foreground);
                    self.redraw = true;

                    self.escape_sequence = false;
                },
/*
@MANSTART{terminal-raw-mode}
INTRODUCTION
    Since Redox has no ioctl syscall, it uses escape codes for switching to raw mode.

ENTERING AND EXITING RAW MODE
    Entering raw mode is done using CSI-r (^[r). Unsetting raw mode is done by CSI-R (^[R).

RAW MODE
    Raw mode means that the stdin must be handled solely by the program itself. It will not automatically be printed nor will it be modified in any way (modulo escape codes).

    This means that:
        - stdin is not printed.
        - newlines are interpreted as carriage returns in stdin.
        - stdin is not buffered, meaning that the stream of bytes goes directly to the program, without the user having to press enter.
@MANEND
*/
                'r' => {
                    self.raw_mode = true;
                    self.escape_sequence = false;
                },
                'R' => {
                    self.raw_mode = false;
                    self.escape_sequence = false;
                },
                _ => self.escape_sequence = false,

            }

            if !self.escape_sequence {
                self.sequence.clear();
                self.escape = false;
            }
        } else {
            match c {
                '[' => {
                    // Control sequence initiator

                    self.escape_sequence = true;
                    self.sequence.push(String::new());
                },
                'c' => {
                    // Reset
                    self.point_x = 0;
                    self.point_y = 0;
                    self.raw_mode = false;
                    self.foreground = ansi_color(7);
                    self.background = ansi_color(0);
                    self.window.set(self.background);
                    self.redraw = true;

                    self.escape = false;
                }
                _ => self.escape = false,
            }
        }
    }

    pub fn scroll(&mut self, rows: usize) {
        if rows > 0 && rows < self.window.height() as usize {
            let offset = rows * self.window.width() as usize;
            let data = self.window.data_mut();
            unsafe {
                fast_copy(data.as_mut_ptr() as *mut u32, data.as_ptr().offset(offset as isize) as *const u32, data.len() - offset);
                fast_set(data.as_mut_ptr().offset((data.len() - offset) as isize) as *mut u32, self.background.data, offset);
            }
        }
    }

    pub fn character(&mut self, c: char) {
        self.window.rect(self.point_x, self.point_y, 8, 16, self.background);

        match c {
            '\0' => {},
            '\x1B' => self.escape = true,
            '\n' => {
                self.point_x = 0;
                self.point_y += 16;
            },
            '\t' => self.point_x = ((self.point_x / 64) + 1) * 64,
            '\r' => self.point_x = 0,
            '\x08' => {
                self.point_x -= 8;
                if self.point_x < 0 {
                    self.point_x = 0
                }
                self.window.rect(self.point_x, self.point_y, 8, 16, self.background);
            },
            _ => {
                self.window.char(self.point_x, self.point_y, c, self.foreground);
                self.point_x += 8;
            }
        }

        if self.point_x >= self.window.width() as i32 {
            self.point_x = 0;
            self.point_y += 16;
        }

        while self.point_y + 16 > self.window.height() as i32 {
            self.scroll(16);
            self.point_y -= 16;
        }
        self.window.rect(self.point_x, self.point_y, 8, 16, self.foreground);
        self.redraw = true;
    }

    pub fn event(&mut self, event: Event) -> Option<String> {
        match event.to_option() {
            EventOption::Key(key_event) => {
                if key_event.pressed {
                    if self.raw_mode {
                        match key_event.scancode {
                            event::K_BKSP => self.command.push_str("\x08"),
                            event::K_UP => self.command.push_str("\x1B[A"),
                            event::K_DOWN => self.command.push_str("\x1B[B"),
                            event::K_RIGHT => self.command.push_str("\x1B[C"),
                            event::K_LEFT => self.command.push_str("\x1B[D"),
                            _ => match key_event.character {
                                '\0' => {},
                                c => {
                                    self.command.push(c);
                                }
                            },
                        }

                        if ! self.command.is_empty() {
                            let mut command = String::new();
                            mem::swap(&mut self.command, &mut command);
                            return Some(command);
                        }
                    } else {
                        match key_event.scancode {
                            event::K_BKSP => if ! self.command.is_empty() {
                                self.write(&[8]);
                                self.command.pop();
                            },
                            _ => match key_event.character {
                                '\0' => (),
                                c => {
                                    self.write(&[c as u8]);
                                    self.command.push(c);

                                    if c == '\n' {
                                        let mut command = String::new();
                                        mem::swap(&mut self.command, &mut command);
                                        return Some(command);
                                    }
                                }
                            },
                        }
                    }
                }
            },
            _ => (),
        }

        None
    }

    pub fn write(&mut self, bytes: &[u8]) {
        for byte in bytes.iter() {
            let c = *byte as char;

            if self.escape {
                self.code(c);
            } else {
                self.character(c);
            }
        }

        if self.redraw {
            self.redraw = false;
            self.window.sync();
        }
    }
}
<|endoftext|>