<commit_before>use core::marker::PhantomData;
use ffi::{self, CPtr};
use types::{c_uint, c_void};
use Error;
use Secp256k1;

#[cfg(feature = "std")]
pub use self::std_only::*;

/// A trait for all kinds of Context's that Lets you define the exact flags and a function to deallocate memory.
/// * DO NOT * implement it for your own types.
pub unsafe trait Context {
    /// Flags for the ffi.
    const FLAGS: c_uint;
    /// A constant description of the context.
    const DESCRIPTION: &'static str;
    /// A function to deallocate the memory when the context is dropped.
    fn deallocate(ptr: *mut [u8]);
}

/// Marker trait for indicating that an instance of `Secp256k1` can be used for signing.
pub trait Signing: Context {}

/// Marker trait for indicating that an instance of `Secp256k1` can be used for verification.
pub trait Verification: Context {}

/// Represents the set of capabilities needed for signing with a user preallocated memory.
pub struct SignOnlyPreallocated<'buf> {
    phantom: PhantomData<&'buf ()>,
}

/// Represents the set of capabilities needed for verification with a user preallocated memory.
pub struct VerifyOnlyPreallocated<'buf> {
    phantom: PhantomData<&'buf ()>,
}

/// Represents the set of all capabilities with a user preallocated memory.
pub struct AllPreallocated<'buf> {
    phantom: PhantomData<&'buf ()>,
}

#[cfg(feature = "std")]
mod std_only {
    use super::*;

    /// Represents the set of capabilities needed for signing.
    pub enum SignOnly {}

    /// Represents the set of capabilities needed for verification.
    pub enum VerifyOnly {}

    /// Represents the set of all capabilities.
    pub enum All {}

    impl Signing for SignOnly {}
    impl Signing for All {}

    impl Verification for VerifyOnly {}
    impl Verification for All {}

    unsafe impl Context for SignOnly {
        const FLAGS: c_uint = ffi::SECP256K1_START_SIGN;
        const DESCRIPTION: &'static str = "signing only";

        fn deallocate(ptr: *mut [u8]) {
            let _ = unsafe { Box::from_raw(ptr) };
        }
    }

    unsafe impl Context for VerifyOnly {
        const FLAGS: c_uint = ffi::SECP256K1_START_VERIFY;
        const DESCRIPTION: &'static str = "verification only";

        fn deallocate(ptr: *mut [u8]) {
            let _ = unsafe { Box::from_raw(ptr) };
        }
    }

    unsafe impl Context for All {
        const FLAGS: c_uint = VerifyOnly::FLAGS | SignOnly::FLAGS;
        const DESCRIPTION: &'static str = "all capabilities";

        fn deallocate(ptr: *mut [u8]) {
            let _ = unsafe { Box::from_raw(ptr) };
        }
    }

    impl<C: Context> Secp256k1<C> {
        /// Lets you create a context in a generic manner(sign/verify/all)
        pub fn gen_new() -> Secp256k1<C> {
            let buf = vec![0u8; Self::preallocate_size_gen()].into_boxed_slice();
            let ptr = Box::into_raw(buf);
            Secp256k1 {
                ctx: unsafe { ffi::secp256k1_context_preallocated_create(ptr as *mut c_void, C::FLAGS) },
                phantom: PhantomData,
                buf: ptr,
            }
        }
    }

    impl Secp256k1<All> {
        /// Creates a new Secp256k1 context with all capabilities
        pub fn new() -> Secp256k1<All> {
            Secp256k1::gen_new()
        }
    }

    impl Secp256k1<SignOnly> {
        /// Creates a new Secp256k1 context that can only be used for signing
        pub fn signing_only() -> Secp256k1<SignOnly> {
            Secp256k1::gen_new()
        }
    }

    impl Secp256k1<VerifyOnly> {
        /// Creates a new Secp256k1 context that can only be used for verification
        pub fn verification_only() -> Secp256k1<VerifyOnly> {
            Secp256k1::gen_new()
        }
    }

    impl Default for Secp256k1<All> {
        fn default() -> Self {
            Self::new()
        }
    }

    impl<C: Context> Clone for Secp256k1<C> {
        fn clone(&self) -> Secp256k1<C> {
            let clone_size = unsafe {ffi::secp256k1_context_preallocated_clone_size(self.ctx)};
            let ptr_buf = Box::into_raw(vec![0u8; clone_size].into_boxed_slice());
            Secp256k1 {
                ctx: unsafe { ffi::secp256k1_context_preallocated_clone(self.ctx, ptr_buf as *mut c_void) },
                phantom: PhantomData,
                buf: ptr_buf,
            }
        }
    }

}

impl<'buf> Signing for SignOnlyPreallocated<'buf> {}
impl<'buf> Signing for AllPreallocated<'buf> {}

impl<'buf> Verification for VerifyOnlyPreallocated<'buf> {}
impl<'buf> Verification for AllPreallocated<'buf> {}

unsafe impl<'buf> Context for SignOnlyPreallocated<'buf> {
    const FLAGS: c_uint = ffi::SECP256K1_START_SIGN;
    const DESCRIPTION: &'static str = "signing only";

    fn deallocate(_ptr: *mut [u8]) {
        // Allocated by the user
    }
}

unsafe impl<'buf> Context for VerifyOnlyPreallocated<'buf> {
    const FLAGS: c_uint = ffi::SECP256K1_START_VERIFY;
    const DESCRIPTION: &'static str = "verification only";

    fn deallocate(_ptr: *mut [u8]) {
        // Allocated by the user
    }
}

unsafe impl<'buf> Context for AllPreallocated<'buf> {
    const FLAGS: c_uint = SignOnlyPreallocated::FLAGS | VerifyOnlyPreallocated::FLAGS;
    const DESCRIPTION: &'static str = "all capabilities";

    fn deallocate(_ptr: *mut [u8]) {
        // Allocated by the user
    }
}

impl<'buf, C: Context + 'buf> Secp256k1<C> {
    /// Lets you create a context with preallocated buffer in a generic manner(sign/verify/all)
    pub fn preallocated_gen_new(buf: &'buf mut [u8]) -> Result<Secp256k1<C>, Error> {
        if buf.len() < Self::preallocate_size_gen() {
            return Err(Error::NotEnoughMemory);
        }
        Ok(Secp256k1 {
            ctx: unsafe {
                ffi::secp256k1_context_preallocated_create(
                    buf.as_mut_c_ptr() as *mut c_void,
                    C::FLAGS)
            },
            phantom: PhantomData,
            buf: buf as *mut [u8],
        })
    }
}

impl<'buf> Secp256k1<AllPreallocated<'buf>> {
    /// Creates a new Secp256k1 context with all capabilities
    pub fn preallocated_new(buf: &'buf mut [u8]) -> Result<Secp256k1<AllPreallocated<'buf>>, Error> {
        Secp256k1::preallocated_gen_new(buf)
    }
    /// Uses the ffi `secp256k1_context_preallocated_size` to check the memory size needed for a context
    pub fn preallocate_size() -> usize {
        Self::preallocate_size_gen()
    }
}

impl<'buf> Secp256k1<SignOnlyPreallocated<'buf>> {
    /// Creates a new Secp256k1 context that can only be used for signing
    pub fn preallocated_signing_only(buf: &'buf mut [u8]) -> Result<Secp256k1<SignOnlyPreallocated<'buf>>, Error> {
        Secp256k1::preallocated_gen_new(buf)
    }

    /// Uses the ffi `secp256k1_context_preallocated_size` to check the memory size needed for the context
    #[inline]
    pub fn preallocate_signing_size() -> usize {
        Self::preallocate_size_gen()
    }
}

impl<'buf> Secp256k1<VerifyOnlyPreallocated<'buf>> {
    /// Creates a new Secp256k1 context that can only be used for verification
    pub fn preallocated_verification_only(buf: &'buf mut [u8]) -> Result<Secp256k1<VerifyOnlyPreallocated<'buf>>, Error> {
        Secp256k1::preallocated_gen_new(buf)
    }

    /// Uses the ffi `secp256k1_context_preallocated_size` to check the memory size needed for the context
    #[inline]
    pub fn preallocate_verification_size() -> usize {
        Self::preallocate_size_gen()
    }
}
<commit_msg>Make the Context trait unimplementable<commit_after>use core::marker::PhantomData;
use ffi::{self, CPtr};
use types::{c_uint, c_void};
use Error;
use Secp256k1;

#[cfg(feature = "std")]
pub use self::std_only::*;

/// A trait for all kinds of Context's that Lets you define the exact flags and a function to deallocate memory.
/// It shouldn't be possible to implement this for types outside this crate.
pub unsafe trait Context : private::Sealed {
    /// Flags for the ffi.
    const FLAGS: c_uint;
    /// A constant description of the context.
    const DESCRIPTION: &'static str;
    /// A function to deallocate the memory when the context is dropped.
    fn deallocate(ptr: *mut [u8]);
}

/// Marker trait for indicating that an instance of `Secp256k1` can be used for signing.
pub trait Signing: Context {}

/// Marker trait for indicating that an instance of `Secp256k1` can be used for verification.
pub trait Verification: Context {}

/// Represents the set of capabilities needed for signing with a user preallocated memory.
pub struct SignOnlyPreallocated<'buf> {
    phantom: PhantomData<&'buf ()>,
}

/// Represents the set of capabilities needed for verification with a user preallocated memory.
pub struct VerifyOnlyPreallocated<'buf> {
    phantom: PhantomData<&'buf ()>,
}

/// Represents the set of all capabilities with a user preallocated memory.
pub struct AllPreallocated<'buf> {
    phantom: PhantomData<&'buf ()>,
}

mod private {
    use super::*;
    // A trick to prevent users from implementing a trait.
    // on one hand this trait is public, on the other it's in a private module
    // so it's not visible to anyone besides it's parent (the context module)
    pub trait Sealed {}

    impl<'buf> Sealed for AllPreallocated<'buf> {}
    impl<'buf> Sealed for VerifyOnlyPreallocated<'buf> {}
    impl<'buf> Sealed for SignOnlyPreallocated<'buf> {}
}

#[cfg(feature = "std")]
mod std_only {
    impl private::Sealed for SignOnly {}
    impl private::Sealed for All {}
    impl private::Sealed for VerifyOnly {}

    use super::*;

    /// Represents the set of capabilities needed for signing.
    pub enum SignOnly {}

    /// Represents the set of capabilities needed for verification.
    pub enum VerifyOnly {}

    /// Represents the set of all capabilities.
    pub enum All {}

    impl Signing for SignOnly {}
    impl Signing for All {}

    impl Verification for VerifyOnly {}
    impl Verification for All {}

    unsafe impl Context for SignOnly {
        const FLAGS: c_uint = ffi::SECP256K1_START_SIGN;
        const DESCRIPTION: &'static str = "signing only";

        fn deallocate(ptr: *mut [u8]) {
            let _ = unsafe { Box::from_raw(ptr) };
        }
    }

    unsafe impl Context for VerifyOnly {
        const FLAGS: c_uint = ffi::SECP256K1_START_VERIFY;
        const DESCRIPTION: &'static str = "verification only";

        fn deallocate(ptr: *mut [u8]) {
            let _ = unsafe { Box::from_raw(ptr) };
        }
    }

    unsafe impl Context for All {
        const FLAGS: c_uint = VerifyOnly::FLAGS | SignOnly::FLAGS;
        const DESCRIPTION: &'static str = "all capabilities";

        fn deallocate(ptr: *mut [u8]) {
            let _ = unsafe { Box::from_raw(ptr) };
        }
    }

    impl<C: Context> Secp256k1<C> {
        /// Lets you create a context in a generic manner(sign/verify/all)
        pub fn gen_new() -> Secp256k1<C> {
            let buf = vec![0u8; Self::preallocate_size_gen()].into_boxed_slice();
            let ptr = Box::into_raw(buf);
            Secp256k1 {
                ctx: unsafe { ffi::secp256k1_context_preallocated_create(ptr as *mut c_void, C::FLAGS) },
                phantom: PhantomData,
                buf: ptr,
            }
        }
    }

    impl Secp256k1<All> {
        /// Creates a new Secp256k1 context with all capabilities
        pub fn new() -> Secp256k1<All> {
            Secp256k1::gen_new()
        }
    }

    impl Secp256k1<SignOnly> {
        /// Creates a new Secp256k1 context that can only be used for signing
        pub fn signing_only() -> Secp256k1<SignOnly> {
            Secp256k1::gen_new()
        }
    }

    impl Secp256k1<VerifyOnly> {
        /// Creates a new Secp256k1 context that can only be used for verification
        pub fn verification_only() -> Secp256k1<VerifyOnly> {
            Secp256k1::gen_new()
        }
    }

    impl Default for Secp256k1<All> {
        fn default() -> Self {
            Self::new()
        }
    }

    impl<C: Context> Clone for Secp256k1<C> {
        fn clone(&self) -> Secp256k1<C> {
            let clone_size = unsafe {ffi::secp256k1_context_preallocated_clone_size(self.ctx)};
            let ptr_buf = Box::into_raw(vec![0u8; clone_size].into_boxed_slice());
            Secp256k1 {
                ctx: unsafe { ffi::secp256k1_context_preallocated_clone(self.ctx, ptr_buf as *mut c_void) },
                phantom: PhantomData,
                buf: ptr_buf,
            }
        }
    }

}

impl<'buf> Signing for SignOnlyPreallocated<'buf> {}
impl<'buf> Signing for AllPreallocated<'buf> {}

impl<'buf> Verification for VerifyOnlyPreallocated<'buf> {}
impl<'buf> Verification for AllPreallocated<'buf> {}

unsafe impl<'buf> Context for SignOnlyPreallocated<'buf> {
    const FLAGS: c_uint = ffi::SECP256K1_START_SIGN;
    const DESCRIPTION: &'static str = "signing only";

    fn deallocate(_ptr: *mut [u8]) {
        // Allocated by the user
    }
}

unsafe impl<'buf> Context for VerifyOnlyPreallocated<'buf> {
    const FLAGS: c_uint = ffi::SECP256K1_START_VERIFY;
    const DESCRIPTION: &'static str = "verification only";

    fn deallocate(_ptr: *mut [u8]) {
        // Allocated by the user
    }
}

unsafe impl<'buf> Context for AllPreallocated<'buf> {
    const FLAGS: c_uint = SignOnlyPreallocated::FLAGS | VerifyOnlyPreallocated::FLAGS;
    const DESCRIPTION: &'static str = "all capabilities";

    fn deallocate(_ptr: *mut [u8]) {
        // Allocated by the user
    }
}

impl<'buf, C: Context + 'buf> Secp256k1<C> {
    /// Lets you create a context with preallocated buffer in a generic manner(sign/verify/all)
    pub fn preallocated_gen_new(buf: &'buf mut [u8]) -> Result<Secp256k1<C>, Error> {
        if buf.len() < Self::preallocate_size_gen() {
            return Err(Error::NotEnoughMemory);
        }
        Ok(Secp256k1 {
            ctx: unsafe {
                ffi::secp256k1_context_preallocated_create(
                    buf.as_mut_c_ptr() as *mut c_void,
                    C::FLAGS)
            },
            phantom: PhantomData,
            buf: buf as *mut [u8],
        })
    }
}

impl<'buf> Secp256k1<AllPreallocated<'buf>> {
    /// Creates a new Secp256k1 context with all capabilities
    pub fn preallocated_new(buf: &'buf mut [u8]) -> Result<Secp256k1<AllPreallocated<'buf>>, Error> {
        Secp256k1::preallocated_gen_new(buf)
    }
    /// Uses the ffi `secp256k1_context_preallocated_size` to check the memory size needed for a context
    pub fn preallocate_size() -> usize {
        Self::preallocate_size_gen()
    }
}

impl<'buf> Secp256k1<SignOnlyPreallocated<'buf>> {
    /// Creates a new Secp256k1 context that can only be used for signing
    pub fn preallocated_signing_only(buf: &'buf mut [u8]) -> Result<Secp256k1<SignOnlyPreallocated<'buf>>, Error> {
        Secp256k1::preallocated_gen_new(buf)
    }

    /// Uses the ffi `secp256k1_context_preallocated_size` to check the memory size needed for the context
    #[inline]
    pub fn preallocate_signing_size() -> usize {
        Self::preallocate_size_gen()
    }
}

impl<'buf> Secp256k1<VerifyOnlyPreallocated<'buf>> {
    /// Creates a new Secp256k1 context that can only be used for verification
    pub fn preallocated_verification_only(buf: &'buf mut [u8]) -> Result<Secp256k1<VerifyOnlyPreallocated<'buf>>, Error> {
        Secp256k1::preallocated_gen_new(buf)
    }

    /// Uses the ffi `secp256k1_context_preallocated_size` to check the memory size needed for the context
    #[inline]
    pub fn preallocate_verification_size() -> usize {
        Self::preallocate_size_gen()
    }
}
<|endoftext|>


<commit_before>//! The initializer module parses arguments and verfies commands are valid
//! The creator module handles creating the muxed directory and project files.
#![allow(experimental)]

use std::io::File;
use std::path::posix::Path;
use editor;
use root;

#[cfg(test)] use test_helper::{random_name,cleanup_file,cleanup_dir};

static TEMPLATE: &'static str = include_str!("creator/template.toml");

pub fn new(name: &str) {
    let muxed_dir = root::muxed_dir();

    let path = &Path::new(format!("{}/{}.toml", muxed_dir.display(), name));
    if !path.exists() {
        create_project_file(path);

        match editor::default_editor_set() {
            true  => editor::open_project_file(path),
            false => println!("Default editor is not set. Please define $EDITOR in your ~/.bashrc or similar file.")
        }
    } else {
        println!("Project already exists.");
    }
}

/// Copy and create the new project file from a template.
fn create_project_file(path: &Path) {
    let filename = project_filename(path);
    match File::create(path).write(modified_template(TEMPLATE, filename.as_slice()).as_bytes()) {
        Ok(())  => (),
        Err(_e) => println!("Failed to create project {}", filename),
    }
}

fn modified_template(template: &str, project_name: &str) -> String {
    template.replace("{file_name}", project_name)
}

fn project_filename(path: &Path) -> String {
    String::from_utf8(path.filename().unwrap().to_vec()).unwrap().replace(".toml", "")
}

#[test]
fn returns_filename_as_string() {
  let path = &Path::new(format!("{}/{}", "luke", "skywalker"));
  assert_eq!(project_filename(path), String::from_str("skywalker"))
}

#[test]
fn accepts_periods_in_filenames() {
  let path = &Path::new(format!("{}/{}", "luke", "skywalker1.2"));
  assert_eq!(project_filename(path), String::from_str("skywalker1.2"))
}

#[test]
fn accepts_dashes_in_filenames() {
  let path = &Path::new(format!("{}/{}", "luke", "skywalker1-2"));
  assert_eq!(project_filename(path), String::from_str("skywalker1-2"))
}

#[test]
fn populates_template_values() {
    let value = modified_template(TEMPLATE, "muxed project");
    let result = value.as_slice().contains("muxed project");
    assert!(result);
}

#[test]
fn removes_template_placeholders() {
    let value = modified_template(TEMPLATE, "muxed project");
    let result = !value.as_slice().contains("{file_name}");
    assert!(result);
}

#[test]
fn creates_project_file() {
    let muxed_dir = &root::muxed_dir();
    let path = &Path::new(format!("{}/{}.toml", muxed_dir.display(), random_name()));
    create_project_file(path);
    assert!(path.exists());

    cleanup_dir(muxed_dir);
}

#[test]
fn errors_when_creating_project_file() {
    //assert!(false);
}

#[test]
fn create_copies_the_template_file() {
    let muxed_dir = &root::muxed_dir();
    let path = &Path::new(format!("{}/{}.toml", muxed_dir.display(), random_name()));
    let filename = project_filename(path);
    create_project_file(path);
    let data = File::open(path).read_to_end().unwrap();
    let template_expectation = modified_template(TEMPLATE, filename.as_slice());
    assert_eq!(data.as_slice(), template_expectation.as_bytes());

    cleanup_dir(muxed_dir);
}

#[test]
fn new_writes_file_to_muxed_dir() {
    let name = random_name();
    let muxed_dir = &root::muxed_dir();
    let path = &Path::new(format!("{}/{}.toml", muxed_dir.display(), name));
    println!("{}", path.display())
    new(name.as_slice());
    assert!(path.exists());

    cleanup_dir(path);
}
<commit_msg>comment out unwritten test.<commit_after>//! The initializer module parses arguments and verfies commands are valid
//! The creator module handles creating the muxed directory and project files.
#![allow(experimental)]

use std::io::File;
use std::path::posix::Path;
use editor;
use root;

#[cfg(test)] use test_helper::{random_name,cleanup_file,cleanup_dir};

static TEMPLATE: &'static str = include_str!("creator/template.toml");

pub fn new(name: &str) {
    let muxed_dir = root::muxed_dir();

    let path = &Path::new(format!("{}/{}.toml", muxed_dir.display(), name));
    if !path.exists() {
        create_project_file(path);

        match editor::default_editor_set() {
            true  => editor::open_project_file(path),
            false => println!("Default editor is not set. Please define $EDITOR in your ~/.bashrc or similar file.")
        }
    } else {
        println!("Project already exists.");
    }
}

/// Copy and create the new project file from a template.
fn create_project_file(path: &Path) {
    let filename = project_filename(path);
    match File::create(path).write(modified_template(TEMPLATE, filename.as_slice()).as_bytes()) {
        Ok(())  => (),
        Err(_e) => println!("Failed to create project {}", filename),
    }
}

fn modified_template(template: &str, project_name: &str) -> String {
    template.replace("{file_name}", project_name)
}

fn project_filename(path: &Path) -> String {
    String::from_utf8(path.filename().unwrap().to_vec()).unwrap().replace(".toml", "")
}

#[test]
fn returns_filename_as_string() {
  let path = &Path::new(format!("{}/{}", "luke", "skywalker"));
  assert_eq!(project_filename(path), String::from_str("skywalker"))
}

#[test]
fn accepts_periods_in_filenames() {
  let path = &Path::new(format!("{}/{}", "luke", "skywalker1.2"));
  assert_eq!(project_filename(path), String::from_str("skywalker1.2"))
}

#[test]
fn accepts_dashes_in_filenames() {
  let path = &Path::new(format!("{}/{}", "luke", "skywalker1-2"));
  assert_eq!(project_filename(path), String::from_str("skywalker1-2"))
}

#[test]
fn populates_template_values() {
    let value = modified_template(TEMPLATE, "muxed project");
    let result = value.as_slice().contains("muxed project");
    assert!(result);
}

#[test]
fn removes_template_placeholders() {
    let value = modified_template(TEMPLATE, "muxed project");
    let result = !value.as_slice().contains("{file_name}");
    assert!(result);
}

#[test]
fn creates_project_file() {
    let muxed_dir = &root::muxed_dir();
    let path = &Path::new(format!("{}/{}.toml", muxed_dir.display(), random_name()));
    create_project_file(path);
    assert!(path.exists());

    cleanup_dir(muxed_dir);
}

//#[test]
//fn errors_when_creating_project_file() {
//    //assert!(false);
//}

#[test]
fn create_copies_the_template_file() {
    let muxed_dir = &root::muxed_dir();
    let path = &Path::new(format!("{}/{}.toml", muxed_dir.display(), random_name()));
    let filename = project_filename(path);
    create_project_file(path);
    let data = File::open(path).read_to_end().unwrap();
    let template_expectation = modified_template(TEMPLATE, filename.as_slice());
    assert_eq!(data.as_slice(), template_expectation.as_bytes());

    cleanup_dir(muxed_dir);
}

#[test]
fn new_writes_file_to_muxed_dir() {
    let name = random_name();
    let muxed_dir = &root::muxed_dir();
    let path = &Path::new(format!("{}/{}.toml", muxed_dir.display(), name));
    println!("{}", path.display())
    new(name.as_slice());
    assert!(path.exists());

    cleanup_dir(path);
}
<|endoftext|>


<commit_before>//! Partition-related types and helper functions.
//!
//! This module provides access to low-level primitives
//! to work with GPT partitions.

use bitflags::*;
use crc::crc32;
use log::*;
use std::collections::BTreeMap;
use std::fmt;
use std::fs::{File, OpenOptions};
use std::io::{Cursor, Error, ErrorKind, Read, Result, Seek, SeekFrom, Write};
use std::path::Path;
use std::str::FromStr;
use uuid;

use crate::disk;
use crate::header::{parse_uuid, Header};
use crate::partition_types::Type;

bitflags! {
    /// Partition entry attributes, defined for UEFI.
    pub struct PartitionAttributes: u64 {
        /// Required platform partition.
        const PLATFORM   = 1;
        /// No Block-IO protocol.
        const EFI        = (1 << 1);
        /// Legacy-BIOS bootable partition.
        const BOOTABLE   = (1 << 2);
    }
}

/// A partition entry in a GPT partition table.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Partition {
    /// GUID of the partition type.
    pub part_type_guid: Type,
    /// UUID of the partition.
    pub part_guid: uuid::Uuid,
    /// First LBA of the partition.
    pub first_lba: u64,
    /// Last LBA of the partition.
    pub last_lba: u64,
    /// Partition flags.
    pub flags: u64,
    /// Partition name.
    pub name: String,
}

impl Partition {
    /// Create a partition entry of type "unused", whose bytes are all 0s.
    pub fn zero() -> Self {
        Self {
            part_type_guid: crate::partition_types::UNUSED,
            part_guid: uuid::Uuid::nil(),
            first_lba: 0,
            last_lba: 0,
            flags: 0,
            name: "".to_string(),
        }
    }

    /// Serialize this partition entry to its bytes representation.
    fn as_bytes(&self, entry_size: u16) -> Result<Vec<u8>> {
        let mut buf: Vec<u8> = Vec::with_capacity(entry_size as usize);

        // Type GUID.
        let tyguid = uuid::Uuid::from_str(self.part_type_guid.guid).map_err(|e| {
            Error::new(ErrorKind::Other, format!("Invalid guid: {}", e.to_string()))
        })?;
        let tyguid = tyguid.as_fields();
        buf.write_all(&tyguid.0.to_le_bytes())?;
        buf.write_all(&tyguid.1.to_le_bytes())?;
        buf.write_all(&tyguid.2.to_le_bytes())?;
        buf.write_all(tyguid.3)?;

        // Partition GUID.
        let pguid = self.part_guid.as_fields();
        buf.write_all(&pguid.0.to_le_bytes())?;
        buf.write_all(&pguid.1.to_le_bytes())?;
        buf.write_all(&pguid.2.to_le_bytes())?;
        buf.write_all(pguid.3)?;

        // LBAs and flags.
        buf.write_all(&self.first_lba.to_le_bytes())?;
        buf.write_all(&self.last_lba.to_le_bytes())?;
        buf.write_all(&self.flags.to_le_bytes())?;

        // Partition name as UTF16-LE.
        for utf16_char in self.name.encode_utf16().take(36) {
            buf.write_all(&utf16_char.to_le_bytes())?; // TODO: Check this
        }

        // Resize buffer to exact entry size.
        buf.resize(entry_size as usize, 0x00);

        Ok(buf)
    }

    /// Write the partition entry to the partitions area and update crc32 for the Header.
    pub fn write(
        &self,
        p: &Path,
        partition_id: u64,
        start_lba: u64,
        lb_size: disk::LogicalBlockSize,
    ) -> Result<()> {
        debug!("writing partition to: {}", p.display());
        let pstart = start_lba
            .checked_mul(lb_size.into())
            .ok_or_else(|| Error::new(ErrorKind::Other, "partition overflow - start offset"))?;
        let mut file = OpenOptions::new().write(true).read(true).open(p)?;
        // The offset is 128 * partition_id
        let offset = partition_id
            .checked_mul(128)
            .ok_or_else(|| Error::new(ErrorKind::Other, "partition overflow"))?;
        trace!("seeking to partition start: {}", pstart + offset);
        file.seek(SeekFrom::Start(pstart + offset))?;
        trace!("writing {:?}", &self.as_bytes(128));
        file.write_all(&self.as_bytes(128)?)?;

        Ok(())
    }

    /// Return the length (in bytes) of this partition.
    pub fn bytes_len(&self, lb_size: disk::LogicalBlockSize) -> Result<u64> {
        let len = self
            .last_lba
            .checked_sub(self.first_lba)
            .ok_or_else(|| Error::new(ErrorKind::Other, "partition length underflow - sectors"))?
            .checked_mul(lb_size.into())
            .ok_or_else(|| Error::new(ErrorKind::Other, "partition length overflow - bytes"))?;
        Ok(len)
    }

    /// Return the starting offset (in bytes) of this partition.
    pub fn bytes_start(&self, lb_size: disk::LogicalBlockSize) -> Result<u64> {
        let len = self
            .first_lba
            .checked_mul(lb_size.into())
            .ok_or_else(|| Error::new(ErrorKind::Other, "partition start overflow - bytes"))?;
        Ok(len)
    }

    /// Check whether this partition is in use.
    pub fn is_used(&self) -> bool {
        self.part_type_guid.guid != crate::partition_types::UNUSED.guid
    }

    /// Return the number of sectors in the partition.
    pub fn size(&self) -> Result<u64> {
        match self.last_lba.checked_sub(self.first_lba) {
            Some(size) => Ok(size),
            None => Err(Error::new(
                ErrorKind::Other,
                "Invalid partition.  last_lba < first_lba",
            )),
        }
    }
}

impl fmt::Display for Partition {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Partition:\t\t{}\nPartition GUID:\t\t{}\nPartition Type:\t\t{}\n\
             Span:\t\t\t{} - {}\nFlags:\t\t\t{}",
            self.name,
            self.part_guid,
            self.part_type_guid.guid,
            self.first_lba,
            self.last_lba,
            self.flags,
        )
    }
}

fn read_part_name(rdr: &mut Cursor<&[u8]>) -> Result<String> {
    trace!("Reading partition name");
    let mut namebytes: Vec<u16> = Vec::new();
    for _ in 0..36 {
        let b = u16::from_le_bytes(read_exact_buff!(bbuff, rdr, 2));
        if b != 0 {
            namebytes.push(b);
        }
    }

    Ok(String::from_utf16_lossy(&namebytes))
}

/// Read a GPT partition table.
///
/// ## Example
///
/// ```rust,no_run
/// use gpt::{header, disk, partition};
/// use std::path::Path;
///
/// let lb_size = disk::DEFAULT_SECTOR_SIZE;
/// let diskpath = Path::new("/dev/sdz");
/// let hdr = header::read_header(diskpath, lb_size).unwrap();
/// let partitions = partition::read_partitions(diskpath, &hdr, lb_size).unwrap();
/// println!("{:#?}", partitions);
/// ```
pub fn read_partitions(
    path: &Path,
    header: &Header,
    lb_size: disk::LogicalBlockSize,
) -> Result<BTreeMap<u32, Partition>> {
    debug!("reading partitions from file: {}", path.display());
    let mut file = File::open(path)?;
    file_read_partitions(&mut file, header, lb_size)
}

/// Read a GPT partition table from an open `Read` + `Seek` object.
pub fn file_read_partitions<D: Read + Seek>(
    file: &mut D,
    header: &Header,
    lb_size: disk::LogicalBlockSize,
) -> Result<BTreeMap<u32, Partition>> {
    let pstart = header
        .part_start
        .checked_mul(lb_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "partition overflow - start offset"))?;
    trace!("seeking to partitions start: {:#x}", pstart);
    let _ = file.seek(SeekFrom::Start(pstart))?;
    let mut parts: BTreeMap<u32, Partition> = BTreeMap::new();

    trace!("scanning {} partitions", header.num_parts);
    let mut count = 0;
    for i in 0..header.num_parts {
        let mut bytes: [u8; 56] = [0; 56];
        let mut nameraw: [u8; 72] = [0; 72];

        file.read_exact(&mut bytes)?;
        file.read_exact(&mut nameraw)?;

        let test: [u8; 56] = [0; 56];
        let test2: [u8; 72] = [0; 72];
        // Note: unused partition entries are zeroed, so skip them
        if bytes.eq(&test[0..]) && nameraw.eq(&test2[0..]) {
            count += 1;
        } else {
            let mut reader = Cursor::new(&bytes[..]);
            let type_guid = parse_uuid(&mut reader)?;
            let part_guid = parse_uuid(&mut reader)?;

            let partname = read_part_name(&mut Cursor::new(&nameraw[..]))?;
            let p = Partition {
                part_type_guid: Type::from_uuid(&type_guid).map_err(|e| {
                    Error::new(ErrorKind::Other, format!("Unknown Partition Type: {}", e))
                })?,
                part_guid,
                first_lba: u64::from_le_bytes(read_exact_buff!(flba, reader, 8)),
                last_lba: u64::from_le_bytes(read_exact_buff!(llba, reader, 8)),
                flags: u64::from_le_bytes(read_exact_buff!(flagbuff, reader, 8)),
                name: partname.to_string(),
            };

            parts.insert(i + 1, p);
        }
    }
    debug!("Num Zeroed partitions {:?}\n\n", count);

    debug!("checking partition table CRC");
    let _ = file.seek(SeekFrom::Start(pstart))?;
    let pt_len = u64::from(header.num_parts)
        .checked_mul(header.part_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "partitions - size"))?;
    let mut table = vec![0; pt_len as usize];
    file.read_exact(&mut table)?;

    let comp_crc = crc32::checksum_ieee(&table);
    if comp_crc != header.crc32_parts {
        return Err(Error::new(ErrorKind::Other, "partition table CRC mismatch"));
    }

    Ok(parts)
}

#[cfg(test)]
mod tests {
    use crate::disk;
    use crate::partition;

    #[test]
    fn test_zero_part() {
        let p0 = partition::Partition::zero();

        let b128 = p0.as_bytes(128).unwrap();
        assert_eq!(b128.len(), 128);
        assert_eq!(b128, vec![0u8; 128]);

        let b256 = p0.as_bytes(256).unwrap();
        assert_eq!(b256.len(), 256);
        assert_eq!(b256, vec![0u8; 256]);
    }

    #[test]
    fn test_part_bytes_len() {
        {
            // Zero.
            let p0 = partition::Partition::zero();
            let b512len = p0.bytes_len(disk::LogicalBlockSize::Lb512).unwrap();
            let b4096len = p0.bytes_len(disk::LogicalBlockSize::Lb4096).unwrap();

            assert_eq!(b512len, 0);
            assert_eq!(b4096len, 0);
        }

        {
            // Negative length.
            let mut p1 = partition::Partition::zero();
            p1.first_lba = p1.last_lba + 1;
            p1.bytes_len(disk::LogicalBlockSize::Lb512).unwrap_err();
            p1.bytes_len(disk::LogicalBlockSize::Lb4096).unwrap_err();
        }

        {
            // Overflowing u64 length.
            let mut p2 = partition::Partition::zero();
            p2.last_lba = <u64>::max_value();
            p2.bytes_len(disk::LogicalBlockSize::Lb512).unwrap_err();
            p2.bytes_len(disk::LogicalBlockSize::Lb4096).unwrap_err();
        }

        {
            // Positive value.
            let mut p3 = partition::Partition::zero();
            p3.first_lba = 2;
            p3.last_lba = 4;
            let b512len = p3.bytes_len(disk::LogicalBlockSize::Lb512).unwrap();
            let b4096len = p3.bytes_len(disk::LogicalBlockSize::Lb4096).unwrap();

            assert_eq!(b512len, 2 * 512);
            assert_eq!(b4096len, 2 * 4096);
        }
    }

    #[test]
    fn test_part_bytes_start() {
        {
            // Zero.
            let p0 = partition::Partition::zero();
            let b512len = p0.bytes_start(disk::LogicalBlockSize::Lb512).unwrap();
            let b4096len = p0.bytes_start(disk::LogicalBlockSize::Lb4096).unwrap();

            assert_eq!(b512len, 0);
            assert_eq!(b4096len, 0);
        }

        {
            // Overflowing u64 start.
            let mut p1 = partition::Partition::zero();
            p1.first_lba = <u64>::max_value();
            p1.bytes_len(disk::LogicalBlockSize::Lb512).unwrap_err();
            p1.bytes_len(disk::LogicalBlockSize::Lb4096).unwrap_err();
        }

        {
            // Positive value.
            let mut p2 = partition::Partition::zero();
            p2.first_lba = 2;
            let b512start = p2.bytes_start(disk::LogicalBlockSize::Lb512).unwrap();
            let b4096start = p2.bytes_start(disk::LogicalBlockSize::Lb4096).unwrap();

            assert_eq!(b512start, 2 * 512);
            assert_eq!(b4096start, 2 * 4096);
        }
    }
}
<commit_msg>The partition names are null-terminated. (#56)<commit_after>//! Partition-related types and helper functions.
//!
//! This module provides access to low-level primitives
//! to work with GPT partitions.

use bitflags::*;
use crc::crc32;
use log::*;
use std::collections::BTreeMap;
use std::fmt;
use std::fs::{File, OpenOptions};
use std::io::{Cursor, Error, ErrorKind, Read, Result, Seek, SeekFrom, Write};
use std::path::Path;
use std::str::FromStr;
use uuid;

use crate::disk;
use crate::header::{parse_uuid, Header};
use crate::partition_types::Type;

bitflags! {
    /// Partition entry attributes, defined for UEFI.
    pub struct PartitionAttributes: u64 {
        /// Required platform partition.
        const PLATFORM   = 1;
        /// No Block-IO protocol.
        const EFI        = (1 << 1);
        /// Legacy-BIOS bootable partition.
        const BOOTABLE   = (1 << 2);
    }
}

/// A partition entry in a GPT partition table.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Partition {
    /// GUID of the partition type.
    pub part_type_guid: Type,
    /// UUID of the partition.
    pub part_guid: uuid::Uuid,
    /// First LBA of the partition.
    pub first_lba: u64,
    /// Last LBA of the partition.
    pub last_lba: u64,
    /// Partition flags.
    pub flags: u64,
    /// Partition name.
    pub name: String,
}

impl Partition {
    /// Create a partition entry of type "unused", whose bytes are all 0s.
    pub fn zero() -> Self {
        Self {
            part_type_guid: crate::partition_types::UNUSED,
            part_guid: uuid::Uuid::nil(),
            first_lba: 0,
            last_lba: 0,
            flags: 0,
            name: "".to_string(),
        }
    }

    /// Serialize this partition entry to its bytes representation.
    fn as_bytes(&self, entry_size: u16) -> Result<Vec<u8>> {
        let mut buf: Vec<u8> = Vec::with_capacity(entry_size as usize);

        // Type GUID.
        let tyguid = uuid::Uuid::from_str(self.part_type_guid.guid).map_err(|e| {
            Error::new(ErrorKind::Other, format!("Invalid guid: {}", e.to_string()))
        })?;
        let tyguid = tyguid.as_fields();
        buf.write_all(&tyguid.0.to_le_bytes())?;
        buf.write_all(&tyguid.1.to_le_bytes())?;
        buf.write_all(&tyguid.2.to_le_bytes())?;
        buf.write_all(tyguid.3)?;

        // Partition GUID.
        let pguid = self.part_guid.as_fields();
        buf.write_all(&pguid.0.to_le_bytes())?;
        buf.write_all(&pguid.1.to_le_bytes())?;
        buf.write_all(&pguid.2.to_le_bytes())?;
        buf.write_all(pguid.3)?;

        // LBAs and flags.
        buf.write_all(&self.first_lba.to_le_bytes())?;
        buf.write_all(&self.last_lba.to_le_bytes())?;
        buf.write_all(&self.flags.to_le_bytes())?;

        // Partition name as UTF16-LE.
        for utf16_char in self.name.encode_utf16().take(36) {
            buf.write_all(&utf16_char.to_le_bytes())?; // TODO: Check this
        }

        // Resize buffer to exact entry size.
        buf.resize(entry_size as usize, 0x00);

        Ok(buf)
    }

    /// Write the partition entry to the partitions area and update crc32 for the Header.
    pub fn write(
        &self,
        p: &Path,
        partition_id: u64,
        start_lba: u64,
        lb_size: disk::LogicalBlockSize,
    ) -> Result<()> {
        debug!("writing partition to: {}", p.display());
        let pstart = start_lba
            .checked_mul(lb_size.into())
            .ok_or_else(|| Error::new(ErrorKind::Other, "partition overflow - start offset"))?;
        let mut file = OpenOptions::new().write(true).read(true).open(p)?;
        // The offset is 128 * partition_id
        let offset = partition_id
            .checked_mul(128)
            .ok_or_else(|| Error::new(ErrorKind::Other, "partition overflow"))?;
        trace!("seeking to partition start: {}", pstart + offset);
        file.seek(SeekFrom::Start(pstart + offset))?;
        trace!("writing {:?}", &self.as_bytes(128));
        file.write_all(&self.as_bytes(128)?)?;

        Ok(())
    }

    /// Return the length (in bytes) of this partition.
    pub fn bytes_len(&self, lb_size: disk::LogicalBlockSize) -> Result<u64> {
        let len = self
            .last_lba
            .checked_sub(self.first_lba)
            .ok_or_else(|| Error::new(ErrorKind::Other, "partition length underflow - sectors"))?
            .checked_mul(lb_size.into())
            .ok_or_else(|| Error::new(ErrorKind::Other, "partition length overflow - bytes"))?;
        Ok(len)
    }

    /// Return the starting offset (in bytes) of this partition.
    pub fn bytes_start(&self, lb_size: disk::LogicalBlockSize) -> Result<u64> {
        let len = self
            .first_lba
            .checked_mul(lb_size.into())
            .ok_or_else(|| Error::new(ErrorKind::Other, "partition start overflow - bytes"))?;
        Ok(len)
    }

    /// Check whether this partition is in use.
    pub fn is_used(&self) -> bool {
        self.part_type_guid.guid != crate::partition_types::UNUSED.guid
    }

    /// Return the number of sectors in the partition.
    pub fn size(&self) -> Result<u64> {
        match self.last_lba.checked_sub(self.first_lba) {
            Some(size) => Ok(size),
            None => Err(Error::new(
                ErrorKind::Other,
                "Invalid partition.  last_lba < first_lba",
            )),
        }
    }
}

impl fmt::Display for Partition {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Partition:\t\t{}\nPartition GUID:\t\t{}\nPartition Type:\t\t{}\n\
             Span:\t\t\t{} - {}\nFlags:\t\t\t{}",
            self.name,
            self.part_guid,
            self.part_type_guid.guid,
            self.first_lba,
            self.last_lba,
            self.flags,
        )
    }
}

fn read_part_name(rdr: &mut Cursor<&[u8]>) -> Result<String> {
    trace!("Reading partition name");
    let mut namebytes: Vec<u16> = Vec::new();
    for _ in 0..36 {
        let b = u16::from_le_bytes(read_exact_buff!(bbuff, rdr, 2));
        if b != 0 {
            namebytes.push(b);
        } else {
	    break;
	}
    }

    Ok(String::from_utf16_lossy(&namebytes))
}

/// Read a GPT partition table.
///
/// ## Example
///
/// ```rust,no_run
/// use gpt::{header, disk, partition};
/// use std::path::Path;
///
/// let lb_size = disk::DEFAULT_SECTOR_SIZE;
/// let diskpath = Path::new("/dev/sdz");
/// let hdr = header::read_header(diskpath, lb_size).unwrap();
/// let partitions = partition::read_partitions(diskpath, &hdr, lb_size).unwrap();
/// println!("{:#?}", partitions);
/// ```
pub fn read_partitions(
    path: &Path,
    header: &Header,
    lb_size: disk::LogicalBlockSize,
) -> Result<BTreeMap<u32, Partition>> {
    debug!("reading partitions from file: {}", path.display());
    let mut file = File::open(path)?;
    file_read_partitions(&mut file, header, lb_size)
}

/// Read a GPT partition table from an open `Read` + `Seek` object.
pub fn file_read_partitions<D: Read + Seek>(
    file: &mut D,
    header: &Header,
    lb_size: disk::LogicalBlockSize,
) -> Result<BTreeMap<u32, Partition>> {
    let pstart = header
        .part_start
        .checked_mul(lb_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "partition overflow - start offset"))?;
    trace!("seeking to partitions start: {:#x}", pstart);
    let _ = file.seek(SeekFrom::Start(pstart))?;
    let mut parts: BTreeMap<u32, Partition> = BTreeMap::new();

    trace!("scanning {} partitions", header.num_parts);
    let mut count = 0;
    for i in 0..header.num_parts {
        let mut bytes: [u8; 56] = [0; 56];
        let mut nameraw: [u8; 72] = [0; 72];

        file.read_exact(&mut bytes)?;
        file.read_exact(&mut nameraw)?;

        let test: [u8; 56] = [0; 56];
        let test2: [u8; 72] = [0; 72];
        // Note: unused partition entries are zeroed, so skip them
        if bytes.eq(&test[0..]) && nameraw.eq(&test2[0..]) {
            count += 1;
        } else {
            let mut reader = Cursor::new(&bytes[..]);
            let type_guid = parse_uuid(&mut reader)?;
            let part_guid = parse_uuid(&mut reader)?;

            let partname = read_part_name(&mut Cursor::new(&nameraw[..]))?;
            let p = Partition {
                part_type_guid: Type::from_uuid(&type_guid).map_err(|e| {
                    Error::new(ErrorKind::Other, format!("Unknown Partition Type: {}", e))
                })?,
                part_guid,
                first_lba: u64::from_le_bytes(read_exact_buff!(flba, reader, 8)),
                last_lba: u64::from_le_bytes(read_exact_buff!(llba, reader, 8)),
                flags: u64::from_le_bytes(read_exact_buff!(flagbuff, reader, 8)),
                name: partname.to_string(),
            };

            parts.insert(i + 1, p);
        }
    }
    debug!("Num Zeroed partitions {:?}\n\n", count);

    debug!("checking partition table CRC");
    let _ = file.seek(SeekFrom::Start(pstart))?;
    let pt_len = u64::from(header.num_parts)
        .checked_mul(header.part_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "partitions - size"))?;
    let mut table = vec![0; pt_len as usize];
    file.read_exact(&mut table)?;

    let comp_crc = crc32::checksum_ieee(&table);
    if comp_crc != header.crc32_parts {
        return Err(Error::new(ErrorKind::Other, "partition table CRC mismatch"));
    }

    Ok(parts)
}

#[cfg(test)]
mod tests {
    use crate::disk;
    use crate::partition;

    #[test]
    fn test_zero_part() {
        let p0 = partition::Partition::zero();

        let b128 = p0.as_bytes(128).unwrap();
        assert_eq!(b128.len(), 128);
        assert_eq!(b128, vec![0u8; 128]);

        let b256 = p0.as_bytes(256).unwrap();
        assert_eq!(b256.len(), 256);
        assert_eq!(b256, vec![0u8; 256]);
    }

    #[test]
    fn test_part_bytes_len() {
        {
            // Zero.
            let p0 = partition::Partition::zero();
            let b512len = p0.bytes_len(disk::LogicalBlockSize::Lb512).unwrap();
            let b4096len = p0.bytes_len(disk::LogicalBlockSize::Lb4096).unwrap();

            assert_eq!(b512len, 0);
            assert_eq!(b4096len, 0);
        }

        {
            // Negative length.
            let mut p1 = partition::Partition::zero();
            p1.first_lba = p1.last_lba + 1;
            p1.bytes_len(disk::LogicalBlockSize::Lb512).unwrap_err();
            p1.bytes_len(disk::LogicalBlockSize::Lb4096).unwrap_err();
        }

        {
            // Overflowing u64 length.
            let mut p2 = partition::Partition::zero();
            p2.last_lba = <u64>::max_value();
            p2.bytes_len(disk::LogicalBlockSize::Lb512).unwrap_err();
            p2.bytes_len(disk::LogicalBlockSize::Lb4096).unwrap_err();
        }

        {
            // Positive value.
            let mut p3 = partition::Partition::zero();
            p3.first_lba = 2;
            p3.last_lba = 4;
            let b512len = p3.bytes_len(disk::LogicalBlockSize::Lb512).unwrap();
            let b4096len = p3.bytes_len(disk::LogicalBlockSize::Lb4096).unwrap();

            assert_eq!(b512len, 2 * 512);
            assert_eq!(b4096len, 2 * 4096);
        }
    }

    #[test]
    fn test_part_bytes_start() {
        {
            // Zero.
            let p0 = partition::Partition::zero();
            let b512len = p0.bytes_start(disk::LogicalBlockSize::Lb512).unwrap();
            let b4096len = p0.bytes_start(disk::LogicalBlockSize::Lb4096).unwrap();

            assert_eq!(b512len, 0);
            assert_eq!(b4096len, 0);
        }

        {
            // Overflowing u64 start.
            let mut p1 = partition::Partition::zero();
            p1.first_lba = <u64>::max_value();
            p1.bytes_len(disk::LogicalBlockSize::Lb512).unwrap_err();
            p1.bytes_len(disk::LogicalBlockSize::Lb4096).unwrap_err();
        }

        {
            // Positive value.
            let mut p2 = partition::Partition::zero();
            p2.first_lba = 2;
            let b512start = p2.bytes_start(disk::LogicalBlockSize::Lb512).unwrap();
            let b4096start = p2.bytes_start(disk::LogicalBlockSize::Lb4096).unwrap();

            assert_eq!(b512start, 2 * 512);
            assert_eq!(b4096start, 2 * 4096);
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use crate::shell::Key;

/// Contains the state information for the keyboard.
///
/// The keyboard state tracks which keys are currently pressed. The
/// active state is stored in a lazy-loaded HashMap.
///
/// Beside common key activities, you may need to react on generic
/// modifier keys (`Alt`, `Ctrl`, `Hyper`, `Shift`). Helper functions
/// offer several convenience methods to handle such keyboard events.
/// A generic method comes in handy, if you don't care which
/// modifier key is down (`Shift-left` or `Shift-right` => `Shift`).
/// The example section will also tackle the case, where a combined
/// event (`Ctrl+S`) keyboard state is handled.
///
/// # Example
/// Handle a key press inside a given view.
///
/// ```
/// use orbtk {
///     widgets::behaviors::MouseBehavior,
///     prelude::*;
///     shell::event::{Key, KeyEvent};
/// };
///
/// // [KeyboardState]
///
/// #[derive(Clone)]
/// enum KeyboardAction {
///     Key(KeyEvent),
///     RequestFocus,
/// }
///
/// #[derive(AsAny, Default)]
/// struct KeyboardState {
///     action: Option<KeyboardAction>,
///     event_adapter: EventAdapter,
/// }
///
/// // Implementation of KeyboardActions for trait `State`
/// impl State for KeyboardState {
///     fn update(&mut self, _: &mut Registry, ctx: &mut Context<'_>) {
///         if let Some(action) = self.action.clone() {
///             match action {
///                 KeyboardAction::Key(key_event) => {
///                     match key_event.key {
///                         Key::Control => {
///                             // Ctrl+'a'
///                             if key_event.Key::A(false) {
///                                 self.handle_control_a(ctx);
///                             }
///                         }
///                         Key::Down => {
///                             self.handle_down_key(ctx);
///                         }
///                         Key::Enter => {
///                             self.handle_enter_key(ctx);
///                         }
///                         Key::Left => {
///                             self.handle_left_key(ctx);
///                         }
///                         Key::Right => {
///                             self.handle_right_key(ctx);
///                         }
///                         Key::Up => {
///                             self.handle_up_key(ctx);
///                         }
///                         _ => {}
///                     }
///                 }
///                 KeyboardAction::RequestFocus => {
///                     self.request_focus(ctx);
///                 }
///             }
///             self.action = None;
///         }
///
/// // associated functions
/// impl KeyboardState {
///     fn action(&mut self, action: DirectoryListAction) {
///         self.action = Some(action);
/// }
///
/// fn handle_down_key(&mut self, ctx: &mut Context<'_>) {
///     /// here goes your rust code to act on `Key::Down`
///     println!("Handle: `Key::Down`")
///     }
/// }
///
/// fn handle_enter_key(&mut self, ctx: &mut Context<'_>) {
///     /// here goes your rust code to act on `Key::Enter`
///     println!("Handle: `Key::Enter`")
///     }
/// }
///
/// fn handle_left_key(&mut self, ctx: &mut Context<'_>) {
///     /// here goes your rust code to act on `Key::Left`
///     println!("Handle: `Key::Left`")
///     }
/// }
///
/// fn handle_right_key(&mut self, ctx: &mut Context<'_>) {
///     /// here goes your rust code to act on `Key::Right`
///     println!("Handle: `Key::Right`")
///     }
/// }
///
/// fn handle_up_key(&mut self, ctx: &mut Context<'_>) {
///     /// here goes your rust code to act on `Key::Up`
///     println!("Handle: `Key::Up`")
///     }
/// }
///
/// // [KeyboardView]
/// // ... your view code
///
/// ```
#[derive(Default, Clone, Debug, PartialEq)]
pub struct KeyboardState {
    key_list: HashMap<Key, bool>,
}

impl KeyboardState {
    /// Sets whether or not the given key is currently pressed
    pub fn set_key_state(&mut self, key: Key, pressed: bool) {
        self.key_list.insert(key, pressed);
    }

    /// Returns whether or not the requested key is pressed
    pub fn is_key_down(&self, key: Key) -> bool {
        match self.key_list.get(&key) {
            // If we have the key on this list, return its state
            Some(item) => *item,
            // Otherwise, it hasn't been set as down
            None => false,
        }
    }
    /// Returns whether or not any shift key is down.
    pub fn is_shift_down(&self) -> bool {
        self.is_key_down(Key::ShiftL) || self.is_key_down(Key::ShiftR)
    }

    /// Returns whether or not any alt key is down.
    pub fn is_alt_down(&self) -> bool {
        self.is_key_down(Key::Alt)
    }

    /// Returns whether or not any control key is down.
    pub fn is_ctrl_down(&self) -> bool {
        self.is_key_down(Key::Control)
    }

    /// Returns whether or not any home key is down.
    pub fn is_home_down(&self) -> bool {
        self.is_key_down(Key::Home)
    }
}

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

    #[test]
    /// A quick test to ensure that the items are properly set.
    fn basic_test() {
        let mut state = KeyboardState::default();
        // Do a quick check beforehand
        assert!(!state.is_key_down(Key::ShiftL));
        // Set the state down and immediately check
        state.set_key_state(Key::ShiftL, true);
        assert!(state.is_key_down(Key::ShiftL));
        state.set_key_state(Key::ShiftL, false);
        assert!(!state.is_key_down(Key::ShiftL));
        // Set quit a few in a row
        state.set_key_state(Key::ShiftL, true);
        state.set_key_state(Key::ShiftR, true);
        state.set_key_state(Key::Space, true);
        state.set_key_state(Key::Control, true);
        state.set_key_state(Key::Alt, true);
        // Ensure each of these are still down
        assert!(state.is_key_down(Key::ShiftL));
        assert!(state.is_key_down(Key::ShiftR));
        assert!(state.is_key_down(Key::Space));
        assert!(state.is_key_down(Key::Control));
        assert!(state.is_key_down(Key::Alt));
    }

    #[test]
    /// Test for the convenience methods
    fn test_convenience() {
        let mut state = KeyboardState::default();
        // Check to ensure they are all false
        assert!(!state.is_alt_down());
        assert!(!state.is_ctrl_down());
        assert!(!state.is_shift_down());
        // Set ctrl and alt to true and check
        state.set_key_state(Key::Control, true);
        assert!(state.is_ctrl_down());

        state.set_key_state(Key::Alt, true);
        assert!(state.is_alt_down());

        assert!(!state.is_shift_down());
        // Set shift (via L)
        state.set_key_state(Key::ShiftL, true);
        assert!(state.is_shift_down());
        // Set shift (via R and L both set)
        state.set_key_state(Key::ShiftR, true);
        assert!(state.is_shift_down(), "{}", true);
        // Disable L Shift to ensure correct result with just R held
        state.set_key_state(Key::ShiftL, false);
        assert!(state.is_shift_down());
        // Disable both shift keys
        state.set_key_state(Key::ShiftR, false);
        assert!(!state.is_shift_down());
        // Disable alt and ctrl and check again
        state.set_key_state(Key::Control, false);
        assert!(!state.is_ctrl_down());

        state.set_key_state(Key::Alt, false);
        assert!(!state.is_alt_down());
    }
}
<commit_msg>orbtk_core: update `keyboard_state` example<commit_after>use std::collections::HashMap;

use crate::shell::Key;

/// Contains the state information for the keyboard.
///
/// The keyboard state tracks which keys are currently pressed. The
/// active state is stored in a lazy-loaded HashMap.
///
/// Beside common key activities, you may need to react on generic
/// modifier keys (`Alt`, `Ctrl`, `Hyper`, `Shift`). Helper functions
/// offer several convenience methods to handle such keyboard events.
/// A generic method comes in handy, if you don't care which
/// modifier key is down (`Shift-left` or `Shift-right` => `Shift`).
/// The example section will also tackle the case, where a combined
/// event (`Ctrl+S`) keyboard state is handled.
///
/// # Example
/// Handle a key press inside a given view.
///
/// ```
/// use orbtk:: {
///     widgets::behaviors::MouseBehavior,
///     prelude::*;
///     shell::event::{Key, KeyEvent};
/// };
///
/// // [KeyboardState]
///
/// #[derive(Clone)]
/// enum KeyboardAction {
///     Key(KeyEvent),
///     RequestFocus,
/// }
///
/// #[derive(AsAny, Default)]
/// struct KeyboardState {
///     action: Option<KeyboardAction>,
///     event_adapter: EventAdapter,
/// }
///
/// // Implementation of KeyboardActions for trait `State`
/// impl State for KeyboardState {
///     fn update(&mut self, _: &mut Registry, ctx: &mut Context<'_>) {
///         if let Some(action) = self.action.clone() {
///             match action {
///                 KeyboardAction::Key(key_event) => {
///                     match key_event.key {
///                         Key::Control => {
///                             // Ctrl+'a'
///                             if key_event.Key::A(false) {
///                                 self.handle_control_a(ctx);
///                             }
///                         }
///                         Key::Down => {
///                             self.handle_down_key(ctx);
///                         }
///                         Key::Enter => {
///                             self.handle_enter_key(ctx);
///                         }
///                         Key::Left => {
///                             self.handle_left_key(ctx);
///                         }
///                         Key::Right => {
///                             self.handle_right_key(ctx);
///                         }
///                         Key::Up => {
///                             self.handle_up_key(ctx);
///                         }
///                         _ => {}
///                     }
///                 }
///                 KeyboardAction::RequestFocus => {
///                     self.request_focus(ctx);
///                 }
///             }
///             self.action = None;
///         }
///     }
/// }
///
/// // associated functions
/// impl KeyboardState {
///     fn action(&mut self, action: DirectoryListAction) {
///         self.action = Some(action);
///     }
///
///     fn handle_down_key(&mut self, ctx: &mut Context<'_>) {
///         /// here goes your rust code to act on `Key::Down`
///         println!("Handle: `Key::Down`")
///     }
///
///     fn handle_enter_key(&mut self, ctx: &mut Context<'_>) {
///         /// here goes your rust code to act on `Key::Enter`
///         println!("Handle: `Key::Enter`")
///     }
///
///     fn handle_left_key(&mut self, ctx: &mut Context<'_>) {
///         /// here goes your rust code to act on `Key::Left`
///         println!("Handle: `Key::Left`")
///     }
///
///     fn handle_right_key(&mut self, ctx: &mut Context<'_>) {
///         /// here goes your rust code to act on `Key::Right`
///         println!("Handle: `Key::Right`")
///     }
///
///     fn handle_up_key(&mut self, ctx: &mut Context<'_>) {
///         /// here goes your rust code to act on `Key::Up`
///         println!("Handle: `Key::Up`")
///     }
/// }
///
/// // [KeyboardView]
/// // ... your view code
///
/// ```
#[derive(Default, Clone, Debug, PartialEq)]
pub struct KeyboardState {
    key_list: HashMap<Key, bool>,
}

impl KeyboardState {
    /// Sets whether or not the given key is currently pressed
    pub fn set_key_state(&mut self, key: Key, pressed: bool) {
        self.key_list.insert(key, pressed);
    }

    /// Returns whether or not the requested key is pressed
    pub fn is_key_down(&self, key: Key) -> bool {
        match self.key_list.get(&key) {
            // If we have the key on this list, return its state
            Some(item) => *item,
            // Otherwise, it hasn't been set as down
            None => false,
        }
    }
    /// Returns whether or not any shift key is down.
    pub fn is_shift_down(&self) -> bool {
        self.is_key_down(Key::ShiftL) || self.is_key_down(Key::ShiftR)
    }

    /// Returns whether or not any alt key is down.
    pub fn is_alt_down(&self) -> bool {
        self.is_key_down(Key::Alt)
    }

    /// Returns whether or not any control key is down.
    pub fn is_ctrl_down(&self) -> bool {
        self.is_key_down(Key::Control)
    }

    /// Returns whether or not any home key is down.
    pub fn is_home_down(&self) -> bool {
        self.is_key_down(Key::Home)
    }
}

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

    #[test]
    /// A quick test to ensure that the items are properly set.
    fn basic_test() {
        let mut state = KeyboardState::default();
        // Do a quick check beforehand
        assert!(!state.is_key_down(Key::ShiftL));
        // Set the state down and immediately check
        state.set_key_state(Key::ShiftL, true);
        assert!(state.is_key_down(Key::ShiftL));
        state.set_key_state(Key::ShiftL, false);
        assert!(!state.is_key_down(Key::ShiftL));
        // Set quit a few in a row
        state.set_key_state(Key::ShiftL, true);
        state.set_key_state(Key::ShiftR, true);
        state.set_key_state(Key::Space, true);
        state.set_key_state(Key::Control, true);
        state.set_key_state(Key::Alt, true);
        // Ensure each of these are still down
        assert!(state.is_key_down(Key::ShiftL));
        assert!(state.is_key_down(Key::ShiftR));
        assert!(state.is_key_down(Key::Space));
        assert!(state.is_key_down(Key::Control));
        assert!(state.is_key_down(Key::Alt));
    }

    #[test]
    /// Test for the convenience methods
    fn test_convenience() {
        let mut state = KeyboardState::default();
        // Check to ensure they are all false
        assert!(!state.is_alt_down());
        assert!(!state.is_ctrl_down());
        assert!(!state.is_shift_down());
        // Set ctrl and alt to true and check
        state.set_key_state(Key::Control, true);
        assert!(state.is_ctrl_down());

        state.set_key_state(Key::Alt, true);
        assert!(state.is_alt_down());

        assert!(!state.is_shift_down());
        // Set shift (via L)
        state.set_key_state(Key::ShiftL, true);
        assert!(state.is_shift_down());
        // Set shift (via R and L both set)
        state.set_key_state(Key::ShiftR, true);
        assert!(state.is_shift_down(), "{}", true);
        // Disable L Shift to ensure correct result with just R held
        state.set_key_state(Key::ShiftL, false);
        assert!(state.is_shift_down());
        // Disable both shift keys
        state.set_key_state(Key::ShiftR, false);
        assert!(!state.is_shift_down());
        // Disable alt and ctrl and check again
        state.set_key_state(Key::Control, false);
        assert!(!state.is_ctrl_down());

        state.set_key_state(Key::Alt, false);
        assert!(!state.is_alt_down());
    }
}
<|endoftext|>


<commit_before>use redox::{self, env, BMPFile};
use redox::event::{self, EventOption, MouseEvent};
use redox::fs::file::File;
use redox::io::Read;
use redox::orbital::Window;
use redox::time::{self, Duration};
use redox::vec::Vec;
use redox::string::{String, ToString};

pub struct FileManager {
    folder_icon: BMPFile,
    audio_icon: BMPFile,
    bin_icon: BMPFile,
    image_icon: BMPFile,
    source_icon: BMPFile,
    script_icon: BMPFile,
    text_icon: BMPFile,
    file_icon: BMPFile,
    files: Vec<String>,
    selected: isize,
    last_mouse_event: MouseEvent,
    click_time: Duration,
}

fn load_icon(path: &str) -> BMPFile {
    let mut vec: Vec<u8> = Vec::new();
    if let Some(mut file) = File::open(&("file:///ui/mimetypes/".to_string() + path + ".bmp")) {
        file.read_to_end(&mut vec);
    }
    BMPFile::from_data(&vec)
}

impl FileManager {
    pub fn new() -> Self {
        FileManager {
            folder_icon: load_icon("inode-directory"),
            audio_icon: load_icon("audio-x-wav"),
            bin_icon: load_icon("application-x-executable"),
            image_icon: load_icon("image-x-generic"),
            source_icon: load_icon("text-x-makefile"),
            script_icon: load_icon("text-x-script"),
            text_icon: load_icon("text-x-generic"),
            file_icon: load_icon("unknown"),
            files: Vec::new(),
            selected: -1,
            last_mouse_event: MouseEvent {
                x: 0,
                y: 0,
                left_button: false,
                middle_button: false,
                right_button: false,
            },
            click_time: Duration::new(0, 0),
        }
    }

    fn draw_content(&mut self, window: &mut Window) {
        window.set([255, 255, 255, 255]);

        let mut i = 0;
        let mut row = 0;
        for string in self.files.iter() {
            if i == self.selected {
                let width = window.width();
                window.rect(0, 32 * row as isize, width, 32, [224, 224, 224, 255]);
            }

            if string.ends_with('/') {
                window.image(0,
                             32 * row as isize,
                             self.folder_icon.width(),
                             self.folder_icon.height(),
                             self.folder_icon.as_slice());
            } else if string.ends_with(".wav") {
                window.image(0,
                             32 * row as isize,
                             self.audio_icon.width(),
                             self.audio_icon.height(),
                             self.audio_icon.as_slice());
            } else if string.ends_with(".bin") {
                window.image(0,
                             32 * row as isize,
                             self.bin_icon.width(),
                             self.bin_icon.height(),
                             self.bin_icon.as_slice());
            } else if string.ends_with(".bmp") {
                window.image(0,
                             32 * row as isize,
                             self.image_icon.width(),
                             self.image_icon.height(),
                             self.image_icon.as_slice());
            } else if string.ends_with(".rs") || string.ends_with(".asm") || string.ends_with(".list") {
                window.image(0,
                             32 * row as isize,
                             self.source_icon.width(),
                             self.source_icon.height(),
                             self.source_icon.as_slice());
            } else if string.ends_with(".sh") || string.ends_with(".lua") {
                window.image(0,
                             32 * row as isize,
                             self.script_icon.width(),
                             self.script_icon.height(),
                             self.script_icon.as_slice());
            } else if string.ends_with(".md") || string.ends_with(".txt") {
                window.image(0,
                             32 * row as isize,
                             self.text_icon.width(),
                             self.text_icon.height(),
                             self.text_icon.as_slice());
            } else {
                window.image(0,
                             32 * row as isize,
                             self.file_icon.width(),
                             self.file_icon.height(),
                             self.file_icon.as_slice());
            }

            let mut col = 0;
            for c in string.chars() {
                if c == '\n' {
                    col = 0;
                    row += 1;
                } else if c == '\t' {
                    col += 8 - col % 8;
                } else {
                    if col < window.width() / 8 && row < window.height() / 32 {
                        window.char(8 * col as isize + 40,
                                    32 * row as isize + 8,
                                    c,
                                    [0, 0, 0, 255]);
                        col += 1;
                    }
                }
                if col >= window.width() / 8 {
                    col = 0;
                    row += 1;
                }
            }
            row += 1;
            i += 1;
        }

        window.sync();
    }

    fn main(&mut self, path: &str) {
        let mut width = 160;
        let mut height = 0;
        if let Some(mut file) = File::open(path) {
            let mut list = String::new();
            file.read_to_string(&mut list);

            for entry in list.split('\n') {
                if width < 40 + (entry.len() + 1) * 8 {
                    width = 40 + (entry.len() + 1) * 8;
                }
                self.files.push(entry.to_string());
            }

            if height < self.files.len() * 32 {
                height = self.files.len() * 32;
            }
        }

        let mut window = Window::new((redox::rand() % 400 + 50) as isize,
                                     (redox::rand() % 300 + 50) as isize,
                                     width,
                                     height,
                                     &path).unwrap();

        self.draw_content(&mut window);

        while let Some(event) = window.poll() {
            match event.to_option() {
                EventOption::Key(key_event) => {
                    if key_event.pressed {
                        match key_event.scancode {
                            event::K_ESC => break,
                            event::K_HOME => self.selected = 0,
                            event::K_UP => if self.selected > 0 {
                                self.selected -= 1;
                            },
                            event::K_END => self.selected = self.files.len() as isize - 1,
                            event::K_DOWN => if self.selected < self.files.len() as isize - 1 {
                                self.selected += 1;
                            },
                            _ => match key_event.character {
                                '\0' => (),
                                '\n' => {
                                    if self.selected >= 0 &&
                                       self.selected < self.files.len() as isize {
                                        match self.files.get(self.selected as usize) {
                                            Some(file) => {
                                                File::exec(&(path.to_string() + &file));
                                            },
                                            None => (),
                                        }
                                    }
                                }
                                _ => {
                                    let mut i = 0;
                                    for file in self.files.iter() {
                                        if file.starts_with(key_event.character) {
                                            self.selected = i;
                                            break;
                                        }
                                        i += 1;
                                    }
                                }
                            },
                        }

                        self.draw_content(&mut window);
                    }
                }
                EventOption::Mouse(mouse_event) => {
                    let mut redraw = false;
                    let mut i = 0;
                    let mut row = 0;
                    for file in self.files.iter() {
                        let mut col = 0;
                        for c in file.chars() {
                            if mouse_event.y >= 32 * row as isize &&
                               mouse_event.y < 32 * row as isize + 32 {
                                self.selected = i;
                                redraw = true;
                            }

                            if c == '\n' {
                                col = 0;
                                row += 1;
                            } else if c == '\t' {
                                col += 8 - col % 8;
                            } else {
                                if col < window.width() / 8 && row < window.height() / 32 {
                                    col += 1;
                                }
                            }
                            if col >= window.width() / 8 {
                                col = 0;
                                row += 1;
                            }
                        }
                        row += 1;
                        i += 1;
                    }

                    if redraw {
                        self.draw_content(&mut window);
                    }

                    //Check for double click
                    if mouse_event.left_button {
                        let click_time = Duration::realtime();

                        if click_time - self.click_time < Duration::new(0, 500 * time::NANOS_PER_MILLI)
                            && self.last_mouse_event.x == mouse_event.x
                            && self.last_mouse_event.y == mouse_event.y {
                            if self.selected >= 0 && self.selected < self.files.len() as isize {
                                match self.files.get(self.selected as usize) {
                                    Some(file) => {
                                        File::exec(&(path.to_string() + &file));
                                    },
                                    None => (),
                                }
                            }
                            self.click_time = Duration::new(0, 0);
                        } else {
                            self.click_time = click_time;
                        }
                    }
                    self.last_mouse_event = mouse_event;
                }
                _ => (),
            }
        }
    }
}

pub fn main() {
    match env::args().get(1) {
        Some(arg) => FileManager::new().main(arg),
        None => FileManager::new().main("file:///"),
    }
}
<commit_msg>Implement file size and entry count for the file manager<commit_after>use redox::{self, env, BMPFile};
use redox::event::{self, EventOption, MouseEvent};
use redox::fs::file::File;
use redox::io::{Read, Seek, SeekFrom};
use redox::orbital::Window;
use redox::time::{self, Duration};
use redox::vec::Vec;
use redox::string::{String, ToString};

pub struct FileManager {
    folder_icon: BMPFile,
    audio_icon: BMPFile,
    bin_icon: BMPFile,
    image_icon: BMPFile,
    source_icon: BMPFile,
    script_icon: BMPFile,
    text_icon: BMPFile,
    file_icon: BMPFile,
    files: Vec<String>,
    file_sizes: Vec<String>,
    selected: isize,
    last_mouse_event: MouseEvent,
    click_time: Duration,
}

fn load_icon(path: &str) -> BMPFile {
    let mut vec: Vec<u8> = Vec::new();
    if let Some(mut file) = File::open(&("file:///ui/mimetypes/".to_string() + path + ".bmp")) {
        file.read_to_end(&mut vec);
    }
    BMPFile::from_data(&vec)
}

impl FileManager {
    pub fn new() -> Self {
        FileManager {
            folder_icon: load_icon("inode-directory"),
            audio_icon: load_icon("audio-x-wav"),
            bin_icon: load_icon("application-x-executable"),
            image_icon: load_icon("image-x-generic"),
            source_icon: load_icon("text-x-makefile"),
            script_icon: load_icon("text-x-script"),
            text_icon: load_icon("text-x-generic"),
            file_icon: load_icon("unknown"),
            files: Vec::new(),
            file_sizes: Vec::new(),
            selected: -1,
            last_mouse_event: MouseEvent {
                x: 0,
                y: 0,
                left_button: false,
                middle_button: false,
                right_button: false,
            },
            click_time: Duration::new(0, 0),
        }
    }

    fn draw_content(&mut self, window: &mut Window) {
        window.set([255, 255, 255, 255]);

        let mut i = 0;
        let mut row = 0;
        let column1 = {
            let mut tmp = 0;
            for string in self.files.iter() {
                if tmp < string.len() {
                    tmp = string.len();
                }
            }
            tmp + 1
        };
        for (file_name, file_size) in self.files.iter().zip(self.file_sizes.iter()) {
            if i == self.selected {
                let width = window.width();
                window.rect(0, 32 * row as isize, width, 32, [224, 224, 224, 255]);
            }

            if file_name.ends_with('/') {
                window.image(0,
                             32 * row as isize,
                             self.folder_icon.width(),
                             self.folder_icon.height(),
                             self.folder_icon.as_slice());
            } else if file_name.ends_with(".wav") {
                window.image(0,
                             32 * row as isize,
                             self.audio_icon.width(),
                             self.audio_icon.height(),
                             self.audio_icon.as_slice());
            } else if file_name.ends_with(".bin") {
                window.image(0,
                             32 * row as isize,
                             self.bin_icon.width(),
                             self.bin_icon.height(),
                             self.bin_icon.as_slice());
            } else if file_name.ends_with(".bmp") {
                window.image(0,
                             32 * row as isize,
                             self.image_icon.width(),
                             self.image_icon.height(),
                             self.image_icon.as_slice());
            } else if file_name.ends_with(".rs") || file_name.ends_with(".asm") || file_name.ends_with(".list") {
                window.image(0,
                             32 * row as isize,
                             self.source_icon.width(),
                             self.source_icon.height(),
                             self.source_icon.as_slice());
            } else if file_name.ends_with(".sh") || file_name.ends_with(".lua") {
                window.image(0,
                             32 * row as isize,
                             self.script_icon.width(),
                             self.script_icon.height(),
                             self.script_icon.as_slice());
            } else if file_name.ends_with(".md") || file_name.ends_with(".txt") {
                window.image(0,
                             32 * row as isize,
                             self.text_icon.width(),
                             self.text_icon.height(),
                             self.text_icon.as_slice());
            } else {
                window.image(0,
                             32 * row as isize,
                             self.file_icon.width(),
                             self.file_icon.height(),
                             self.file_icon.as_slice());
            }

            let mut col = 0;
            for c in file_name.chars() {
                if c == '\n' {
                    col = 0;
                    row += 1;
                } else if c == '\t' {
                    col += 8 - col % 8;
                } else {
                    if col < window.width() / 8 && row < window.height() / 32 {
                        window.char(8 * col as isize + 40,
                                    32 * row as isize + 8,
                                    c,
                                    [0, 0, 0, 255]);
                        col += 1;
                    }
                }
                if col >= window.width() / 8 {
                    col = 0;
                    row += 1;
                }
            }

            col = column1;

            for c in file_size.chars() {
                if c == '\n' {
                    col = 0;
                    row += 1;
                } else if c == '\t' {
                    col += 8 - col % 8;
                } else {
                    if col < window.width() / 8 && row < window.height() / 32 {
                        window.char(8 * col as isize + 40,
                                    32 * row as isize + 8,
                                    c,
                                    [0, 0, 0, 255]);
                        col += 1;
                    }
                }
                if col >= window.width() / 8 {
                    col = 0;
                    row += 1;
                }
            }

            row += 1;
            i += 1;
        }

        window.sync();
    }

    fn main(&mut self, path: &str) {
        let mut width = 160;
        let mut height = 0;
        if let Some(mut file) = File::open(path) {
            let mut list = String::new();
            file.read_to_string(&mut list);

            for entry in list.split('\n') {
                self.files.push(entry.to_string());
                self.file_sizes.push(
                    match File::open(&(path.to_string() + entry)) {
                        Some(mut file) => {
                            // When the entry is a folder
                            if entry.ends_with('/') {
                                let mut string = String::new();
                                file.read_to_string(&mut string);

                                let count = string.split('\n').count();
                                if count == 1 {
                                    "1 entry".to_string()
                                } else {
                                    format!("{} entries", count)
                                }
                            } else {
                                match file.seek(SeekFrom::End(0)) {
                                    Some(size) => {
                                        if size >= 1_000 {
                                            format!("{:.1} KB", (size as f64)/1_000.0)
                                        } else if size >= 1_000_000 {
                                            format!("{:.1} MB", (size as f64)/1_000_000.0)
                                        } else if size >= 1_000_000_000 {
                                            format!("{:.1} GB", (size as f64)/1_000_000_000.0)
                                        } else {
                                            format!("{:.1} bytes", size)
                                        }
                                    }
                                    None => "Failed to seek".to_string()
                                }
                            }
                        },
                        None => "Failed to open".to_string(),
                    }
                );
                // Unwrapping the last file size will not panic since it has
                // been at least pushed once in the vector
                let current_width = (40 + (entry.len() + 1) * 8) +
                                    (8 + (self.file_sizes.last().unwrap().len() + 1) * 8);
                if width < current_width {
                    width = current_width;
                }
            }

            if height < self.files.len() * 32 {
                height = self.files.len() * 32;
            }
        }

        let mut window = Window::new((redox::rand() % 400 + 50) as isize,
                                     (redox::rand() % 300 + 50) as isize,
                                     width,
                                     height,
                                     &path).unwrap();

        self.draw_content(&mut window);

        while let Some(event) = window.poll() {
            match event.to_option() {
                EventOption::Key(key_event) => {
                    if key_event.pressed {
                        match key_event.scancode {
                            event::K_ESC => break,
                            event::K_HOME => self.selected = 0,
                            event::K_UP => if self.selected > 0 {
                                self.selected -= 1;
                            },
                            event::K_END => self.selected = self.files.len() as isize - 1,
                            event::K_DOWN => if self.selected < self.files.len() as isize - 1 {
                                self.selected += 1;
                            },
                            _ => match key_event.character {
                                '\0' => (),
                                '\n' => {
                                    if self.selected >= 0 &&
                                       self.selected < self.files.len() as isize {
                                        match self.files.get(self.selected as usize) {
                                            Some(file) => {
                                                File::exec(&(path.to_string() + &file));
                                            },
                                            None => (),
                                        }
                                    }
                                }
                                _ => {
                                    let mut i = 0;
                                    for file in self.files.iter() {
                                        if file.starts_with(key_event.character) {
                                            self.selected = i;
                                            break;
                                        }
                                        i += 1;
                                    }
                                }
                            },
                        }

                        self.draw_content(&mut window);
                    }
                }
                EventOption::Mouse(mouse_event) => {
                    let mut redraw = false;
                    let mut i = 0;
                    let mut row = 0;
                    for file in self.files.iter() {
                        let mut col = 0;
                        for c in file.chars() {
                            if mouse_event.y >= 32 * row as isize &&
                               mouse_event.y < 32 * row as isize + 32 {
                                self.selected = i;
                                redraw = true;
                            }

                            if c == '\n' {
                                col = 0;
                                row += 1;
                            } else if c == '\t' {
                                col += 8 - col % 8;
                            } else {
                                if col < window.width() / 8 && row < window.height() / 32 {
                                    col += 1;
                                }
                            }
                            if col >= window.width() / 8 {
                                col = 0;
                                row += 1;
                            }
                        }
                        row += 1;
                        i += 1;
                    }

                    if redraw {
                        self.draw_content(&mut window);
                    }

                    //Check for double click
                    if mouse_event.left_button {
                        let click_time = Duration::realtime();

                        if click_time - self.click_time < Duration::new(0, 500 * time::NANOS_PER_MILLI)
                            && self.last_mouse_event.x == mouse_event.x
                            && self.last_mouse_event.y == mouse_event.y {
                            if self.selected >= 0 && self.selected < self.files.len() as isize {
                                match self.files.get(self.selected as usize) {
                                    Some(file) => {
                                        File::exec(&(path.to_string() + &file));
                                    },
                                    None => (),
                                }
                            }
                            self.click_time = Duration::new(0, 0);
                        } else {
                            self.click_time = click_time;
                        }
                    }
                    self.last_mouse_event = mouse_event;
                }
                _ => (),
            }
        }
    }
}

pub fn main() {
    match env::args().get(1) {
        Some(arg) => FileManager::new().main(arg),
        None => FileManager::new().main("file:///"),
    }
}
<|endoftext|>