<commit_before>use rocksdb::{ColumnFamilyDescriptor, DBWithThreadMode, MergeOperands, Options};
use std::io::Cursor;
use crate::model::PersonalEntry;

#[derive(Debug)]
pub struct Database {
    pub inner: DBWithThreadMode<rocksdb::SingleThreaded>,
    _debug_drop: DebugDrop,
}

impl Database {
    pub fn open() -> Database {
        let mut db_opts = Options::default();
        db_opts.create_if_missing(true);
        db_opts.create_missing_column_families(true);

        let mut personal_opts = Options::default();
        personal_opts.set_merge_operator_associative("personal merge", personal_merge);

        let inner = DBWithThreadMode::open_cf_descriptors(
            &db_opts,
            "_db",
            vec![ColumnFamilyDescriptor::new("personal", personal_opts)],
        )
        .expect("open db");

        Database {
            inner,
            _debug_drop: DebugDrop,
        }
    }
}

fn personal_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let mut entry = PersonalEntry::default();
    let mut size_hint = 0;
    if let Some(existing) = existing {
        let mut cursor = Cursor::new(existing);
        entry
            .extend_from_reader(&mut cursor)
            .expect("read existing personal entry");
        size_hint += existing.len();
    }
    for op in operands {
        let mut cursor = Cursor::new(op);
        entry
            .extend_from_reader(&mut cursor)
            .expect("read personal merge operand");
        size_hint += op.len();
    }
    let mut writer = Cursor::new(Vec::with_capacity(size_hint));
    entry.write(&mut writer).expect("write personal entry");
    Some(writer.into_inner())
}

#[derive(Debug)]
struct DebugDrop;

impl Drop for DebugDrop {
    fn drop(&mut self) {
        println!("dropped");
    }
}
<commit_msg>remove drop debug<commit_after>use rocksdb::{ColumnFamilyDescriptor, DBWithThreadMode, MergeOperands, Options};
use std::io::Cursor;
use crate::model::PersonalEntry;

#[derive(Debug)]
pub struct Database {
    pub inner: DBWithThreadMode<rocksdb::SingleThreaded>,
}

impl Database {
    pub fn open() -> Database {
        let mut db_opts = Options::default();
        db_opts.create_if_missing(true);
        db_opts.create_missing_column_families(true);

        let mut personal_opts = Options::default();
        personal_opts.set_merge_operator_associative("personal merge", personal_merge);

        let inner = DBWithThreadMode::open_cf_descriptors(
            &db_opts,
            "_db",
            vec![ColumnFamilyDescriptor::new("personal", personal_opts)],
        )
        .expect("open db");

        Database {
            inner,
        }
    }
}

fn personal_merge(
    _key: &[u8],
    existing: Option<&[u8]>,
    operands: &mut MergeOperands,
) -> Option<Vec<u8>> {
    let mut entry = PersonalEntry::default();
    let mut size_hint = 0;
    if let Some(existing) = existing {
        let mut cursor = Cursor::new(existing);
        entry
            .extend_from_reader(&mut cursor)
            .expect("read existing personal entry");
        size_hint += existing.len();
    }
    for op in operands {
        let mut cursor = Cursor::new(op);
        entry
            .extend_from_reader(&mut cursor)
            .expect("read personal merge operand");
        size_hint += op.len();
    }
    let mut writer = Cursor::new(Vec::with_capacity(size_hint));
    entry.write(&mut writer).expect("write personal entry");
    Some(writer.into_inner())
}
<|endoftext|>


<commit_before>use core::cell::RefCell;
use core::cmp;
use std::io::prelude::*;
use std::io::{Error, ErrorKind, SeekFrom};
use std::io;
use byteorder::{LittleEndian, ReadBytesExt};

use file::File;
use dir::{DirRawStream, Dir};
use table::{ClusterIterator, alloc_cluster};

// FAT implementation based on:
//   http://wiki.osdev.org/FAT
//   https://www.win.tue.nl/~aeb/linux/fs/fat/fat-1.html

#[derive(Debug, Copy, Clone, PartialEq)]
pub enum FatType {
    Fat12, Fat16, Fat32,
}

pub trait ReadSeek: Read + Seek {}
impl<T> ReadSeek for T where T: Read + Seek {}

pub trait ReadWriteSeek: Read + Write + Seek {}
impl<T> ReadWriteSeek for T where T: Read + Write + Seek {}

#[allow(dead_code)]
#[derive(Default, Debug, Clone)]
pub(crate) struct BiosParameterBlock {
    bytes_per_sector: u16,
    sectors_per_cluster: u8,
    reserved_sectors: u16,
    fats: u8,
    root_entries: u16,
    total_sectors_16: u16,
    media: u8,
    sectors_per_fat_16: u16,
    sectors_per_track: u16,
    heads: u16,
    hidden_sectors: u32,
    total_sectors_32: u32,

    // Extended BIOS Parameter Block
    sectors_per_fat_32: u32,
    extended_flags: u16,
    fs_version: u16,
    root_dir_first_cluster: u32,
    fs_info_sector: u16,
    backup_boot_sector: u16,
    reserved_0: [u8; 12],
    drive_num: u8,
    reserved_1: u8,
    ext_sig: u8,
    volume_id: u32,
    volume_label: [u8; 11],
    fs_type_label: [u8; 8],
}

#[allow(dead_code)]
pub(crate) struct BootRecord {
    bootjmp: [u8; 3],
    oem_name: [u8; 8],
    bpb: BiosParameterBlock,
    boot_code: [u8; 448],
    boot_sig: [u8; 2],
}

impl Default for BootRecord {
    fn default() -> BootRecord {
        BootRecord {
            bootjmp: Default::default(),
            oem_name: Default::default(),
            bpb: Default::default(),
            boot_code: [0; 448],
            boot_sig: Default::default(),
        }
    }
}

pub(crate) type FileSystemRef<'a, 'b: 'a> = &'a FileSystem<'b>;

pub struct FileSystem<'a> {
    pub(crate) disk: RefCell<&'a mut ReadWriteSeek>,
    pub(crate) fat_type: FatType,
    pub(crate) boot: BootRecord,
    pub(crate) first_data_sector: u32,
    pub(crate) root_dir_sectors: u32,
}

impl <'a> FileSystem<'a> {

    pub fn new<T: ReadWriteSeek>(disk: &'a mut T) -> io::Result<FileSystem<'a>> {
        let boot = Self::read_boot_record(disk)?;
        if boot.boot_sig != [0x55, 0xAA] {
            return Err(Error::new(ErrorKind::Other, "invalid signature"));
        }

        let total_sectors = if boot.bpb.total_sectors_16 == 0 { boot.bpb.total_sectors_32 } else { boot.bpb.total_sectors_16 as u32 };
        let sectors_per_fat = if boot.bpb.sectors_per_fat_16 == 0 { boot.bpb.sectors_per_fat_32 } else { boot.bpb.sectors_per_fat_16 as u32 };
        let root_dir_sectors = (((boot.bpb.root_entries * 32) + (boot.bpb.bytes_per_sector - 1)) / boot.bpb.bytes_per_sector) as u32;
        let first_data_sector = boot.bpb.reserved_sectors as u32 + (boot.bpb.fats as u32 * sectors_per_fat) + root_dir_sectors;
        let data_sectors = total_sectors - (boot.bpb.reserved_sectors as u32 + (boot.bpb.fats as u32 * sectors_per_fat) + root_dir_sectors as u32);
        let total_clusters = data_sectors / boot.bpb.sectors_per_cluster as u32;
        let fat_type = Self::fat_type_from_clusters(total_clusters);

        Ok(FileSystem {
            disk: RefCell::new(disk),
            fat_type,
            boot,
            first_data_sector,
            root_dir_sectors,
        })
    }

    pub fn fat_type(&self) -> FatType {
        self.fat_type
    }

    pub fn volume_id(&self) -> u32 {
        self.boot.bpb.volume_id
    }

    pub fn volume_label(&self) -> String {
        String::from_utf8_lossy(&self.boot.bpb.volume_label).trim_right().to_string()
    }

    pub fn root_dir<'b>(&'b self) -> Dir<'b, 'a> {
        let root_rdr = {
            match self.fat_type {
                FatType::Fat12 | FatType::Fat16 => DirRawStream::Root(DiskSlice::from_sectors(
                   self.first_data_sector - self.root_dir_sectors, self.root_dir_sectors, self)),
                _ => DirRawStream::File(File::new(Some(self.boot.bpb.root_dir_first_cluster), None, self)),
            }
        };
        Dir::new(root_rdr, self)
    }

    fn read_bpb(rdr: &mut Read) -> io::Result<BiosParameterBlock> {
        let mut bpb: BiosParameterBlock = Default::default();
        bpb.bytes_per_sector = rdr.read_u16::<LittleEndian>()?;
        bpb.sectors_per_cluster = rdr.read_u8()?;
        bpb.reserved_sectors = rdr.read_u16::<LittleEndian>()?;
        bpb.fats = rdr.read_u8()?;
        bpb.root_entries = rdr.read_u16::<LittleEndian>()? ;
        bpb.total_sectors_16 = rdr.read_u16::<LittleEndian>()?;
        bpb.media = rdr.read_u8()?;
        bpb.sectors_per_fat_16 = rdr.read_u16::<LittleEndian>()?;
        bpb.sectors_per_track = rdr.read_u16::<LittleEndian>()?;
        bpb.heads = rdr.read_u16::<LittleEndian>()?;
        bpb.hidden_sectors = rdr.read_u32::<LittleEndian>()?; // hidden_sector_count
        bpb.total_sectors_32 = rdr.read_u32::<LittleEndian>()?;

        if bpb.sectors_per_fat_16 == 0 {
            bpb.sectors_per_fat_32 = rdr.read_u32::<LittleEndian>()?;
            bpb.extended_flags = rdr.read_u16::<LittleEndian>()?;
            bpb.fs_version = rdr.read_u16::<LittleEndian>()?;
            bpb.root_dir_first_cluster = rdr.read_u32::<LittleEndian>()?;
            bpb.fs_info_sector = rdr.read_u16::<LittleEndian>()?;
            bpb.backup_boot_sector = rdr.read_u16::<LittleEndian>()?;
            rdr.read(&mut bpb.reserved_0)?;
            bpb.drive_num = rdr.read_u8()?;
            bpb.reserved_1 = rdr.read_u8()?;
            bpb.ext_sig = rdr.read_u8()?; // 0x29
            bpb.volume_id = rdr.read_u32::<LittleEndian>()?;
            rdr.read(&mut bpb.volume_label)?;
            rdr.read(&mut bpb.fs_type_label)?;
        } else {
            bpb.drive_num = rdr.read_u8()?;
            bpb.reserved_1 = rdr.read_u8()?;
            bpb.ext_sig = rdr.read_u8()?; // 0x29
            bpb.volume_id = rdr.read_u32::<LittleEndian>()?;
            rdr.read(&mut bpb.volume_label)?;
            rdr.read(&mut bpb.fs_type_label)?;
        }
        Ok(bpb)
    }

    fn fat_type_from_clusters(total_clusters: u32) -> FatType {
        if total_clusters < 4085 {
            FatType::Fat12
        } else if total_clusters < 65525 {
            FatType::Fat16
        } else {
            FatType::Fat32
        }
    }

    fn read_boot_record(rdr: &mut Read) -> io::Result<BootRecord> {
        let mut boot: BootRecord = Default::default();
        rdr.read(&mut boot.bootjmp)?;
        rdr.read(&mut boot.oem_name)?;
        boot.bpb = Self::read_bpb(rdr)?;

        if boot.bpb.sectors_per_fat_16 == 0 {
            rdr.read_exact(&mut boot.boot_code[0..420])?;
        } else {
            rdr.read_exact(&mut boot.boot_code[0..448])?;
        }
        rdr.read(&mut boot.boot_sig)?;
        Ok(boot)
    }

    pub(crate) fn offset_from_sector(&self, sector: u32) -> u64 {
        (sector as u64) * self.boot.bpb.bytes_per_sector as u64
    }

    pub(crate) fn sector_from_cluster(&self, cluster: u32) -> u32 {
        ((cluster - 2) * self.boot.bpb.sectors_per_cluster as u32) + self.first_data_sector
    }

    pub(crate) fn get_cluster_size(&self) -> u32 {
        self.boot.bpb.sectors_per_cluster as u32 * self.boot.bpb.bytes_per_sector as u32
    }

    pub(crate) fn offset_from_cluster(&self, cluser: u32) -> u64 {
        self.offset_from_sector(self.sector_from_cluster(cluser))
    }

    fn fat_slice<'b>(&'b self) -> DiskSlice<'b, 'a> {
        let bytes_per_sector = self.boot.bpb.bytes_per_sector as u64;
        let fat_offset = self.boot.bpb.reserved_sectors as u64 * bytes_per_sector;
        let sectors_per_fat =
            if self.boot.bpb.sectors_per_fat_16 == 0 { self.boot.bpb.sectors_per_fat_32 }
            else { self.boot.bpb.sectors_per_fat_16 as u32 };
        let fat_size = sectors_per_fat as u64 * bytes_per_sector;
        DiskSlice::new(fat_offset, fat_size, self)
    }

    pub(crate) fn cluster_iter<'b>(&'b self, cluster: u32) -> ClusterIterator<'b, 'a> {
        let disk_slice = self.fat_slice();
        ClusterIterator::new(disk_slice, self.fat_type, cluster)
    }

    pub(crate) fn alloc_cluster(&self, prev_cluster: Option<u32>) -> io::Result<u32> {
        let mut disk_slice = self.fat_slice();
        alloc_cluster(&mut disk_slice, self.fat_type, prev_cluster)
    }
}

#[derive(Clone)]
pub(crate) struct DiskSlice<'a, 'b: 'a> {
    begin: u64,
    size: u64,
    offset: u64,
    fs: &'a FileSystem<'b>,
}

impl <'a, 'b> DiskSlice<'a, 'b> {
    pub(crate) fn new(begin: u64, size: u64, fs: FileSystemRef<'a, 'b>) -> Self {
        DiskSlice { begin, size, fs, offset: 0 }
    }

    pub(crate) fn from_sectors(first_sector: u32, sectors_count: u32, fs: FileSystemRef<'a, 'b>) -> Self {
        let bytes_per_sector = fs.boot.bpb.bytes_per_sector as u64;
        Self::new(first_sector as u64 * bytes_per_sector, sectors_count as u64 * bytes_per_sector, fs)
    }

    pub(crate) fn global_pos(&self) -> u64 {
        self.begin + self.offset
    }
}

impl <'a, 'b> Read for DiskSlice<'a, 'b> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let offset = self.begin + self.offset;
        let read_size = cmp::min((self.size - self.offset) as usize, buf.len());
        let mut disk = self.fs.disk.borrow_mut();
        disk.seek(SeekFrom::Start(offset))?;
        let size = disk.read(&mut buf[..read_size])?;
        self.offset += size as u64;
        Ok(size)
    }
}

impl <'a, 'b> Write for DiskSlice<'a, 'b> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let offset = self.begin + self.offset;
        let write_size = cmp::min((self.size - self.offset) as usize, buf.len());
        let mut disk = self.fs.disk.borrow_mut();
        disk.seek(SeekFrom::Start(offset))?;
        let size = disk.write(&buf[..write_size])?;
        self.offset += size as u64;
        Ok(size)
    }

    fn flush(&mut self) -> io::Result<()> {
        let mut disk = self.fs.disk.borrow_mut();
        disk.flush()
    }
}

impl <'a, 'b> Seek for DiskSlice<'a, 'b> {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        let new_offset = match pos {
            SeekFrom::Current(x) => self.offset as i64 + x,
            SeekFrom::Start(x) => x as i64,
            SeekFrom::End(x) => self.size as i64 + x,
        };
        if new_offset < 0 || new_offset as u64 > self.size {
            Err(io::Error::new(ErrorKind::InvalidInput, "invalid seek"))
        } else {
            self.offset = new_offset as u64;
            Ok(self.offset)
        }
    }
}
<commit_msg>Support FAT mirroring.<commit_after>use core::cell::RefCell;
use core::cmp;
use std::io::prelude::*;
use std::io::{Error, ErrorKind, SeekFrom};
use std::io;
use byteorder::{LittleEndian, ReadBytesExt};

use file::File;
use dir::{DirRawStream, Dir};
use table::{ClusterIterator, alloc_cluster};

// FAT implementation based on:
//   http://wiki.osdev.org/FAT
//   https://www.win.tue.nl/~aeb/linux/fs/fat/fat-1.html

#[derive(Debug, Copy, Clone, PartialEq)]
pub enum FatType {
    Fat12, Fat16, Fat32,
}

pub trait ReadSeek: Read + Seek {}
impl<T> ReadSeek for T where T: Read + Seek {}

pub trait ReadWriteSeek: Read + Write + Seek {}
impl<T> ReadWriteSeek for T where T: Read + Write + Seek {}

#[allow(dead_code)]
#[derive(Default, Debug, Clone)]
pub(crate) struct BiosParameterBlock {
    bytes_per_sector: u16,
    sectors_per_cluster: u8,
    reserved_sectors: u16,
    fats: u8,
    root_entries: u16,
    total_sectors_16: u16,
    media: u8,
    sectors_per_fat_16: u16,
    sectors_per_track: u16,
    heads: u16,
    hidden_sectors: u32,
    total_sectors_32: u32,

    // Extended BIOS Parameter Block
    sectors_per_fat_32: u32,
    extended_flags: u16,
    fs_version: u16,
    root_dir_first_cluster: u32,
    fs_info_sector: u16,
    backup_boot_sector: u16,
    reserved_0: [u8; 12],
    drive_num: u8,
    reserved_1: u8,
    ext_sig: u8,
    volume_id: u32,
    volume_label: [u8; 11],
    fs_type_label: [u8; 8],
}

#[allow(dead_code)]
pub(crate) struct BootRecord {
    bootjmp: [u8; 3],
    oem_name: [u8; 8],
    bpb: BiosParameterBlock,
    boot_code: [u8; 448],
    boot_sig: [u8; 2],
}

impl Default for BootRecord {
    fn default() -> BootRecord {
        BootRecord {
            bootjmp: Default::default(),
            oem_name: Default::default(),
            bpb: Default::default(),
            boot_code: [0; 448],
            boot_sig: Default::default(),
        }
    }
}

pub(crate) type FileSystemRef<'a, 'b: 'a> = &'a FileSystem<'b>;

pub struct FileSystem<'a> {
    pub(crate) disk: RefCell<&'a mut ReadWriteSeek>,
    pub(crate) fat_type: FatType,
    pub(crate) boot: BootRecord,
    pub(crate) first_data_sector: u32,
    pub(crate) root_dir_sectors: u32,
}

impl <'a> FileSystem<'a> {

    pub fn new<T: ReadWriteSeek>(disk: &'a mut T) -> io::Result<FileSystem<'a>> {
        let boot = Self::read_boot_record(disk)?;
        if boot.boot_sig != [0x55, 0xAA] {
            return Err(Error::new(ErrorKind::Other, "invalid signature"));
        }

        let total_sectors = if boot.bpb.total_sectors_16 == 0 { boot.bpb.total_sectors_32 } else { boot.bpb.total_sectors_16 as u32 };
        let sectors_per_fat = if boot.bpb.sectors_per_fat_16 == 0 { boot.bpb.sectors_per_fat_32 } else { boot.bpb.sectors_per_fat_16 as u32 };
        let root_dir_sectors = (((boot.bpb.root_entries * 32) + (boot.bpb.bytes_per_sector - 1)) / boot.bpb.bytes_per_sector) as u32;
        let first_data_sector = boot.bpb.reserved_sectors as u32 + (boot.bpb.fats as u32 * sectors_per_fat) + root_dir_sectors;
        let data_sectors = total_sectors - (boot.bpb.reserved_sectors as u32 + (boot.bpb.fats as u32 * sectors_per_fat) + root_dir_sectors as u32);
        let total_clusters = data_sectors / boot.bpb.sectors_per_cluster as u32;
        let fat_type = Self::fat_type_from_clusters(total_clusters);

        Ok(FileSystem {
            disk: RefCell::new(disk),
            fat_type,
            boot,
            first_data_sector,
            root_dir_sectors,
        })
    }

    pub fn fat_type(&self) -> FatType {
        self.fat_type
    }

    pub fn volume_id(&self) -> u32 {
        self.boot.bpb.volume_id
    }

    pub fn volume_label(&self) -> String {
        String::from_utf8_lossy(&self.boot.bpb.volume_label).trim_right().to_string()
    }

    pub fn root_dir<'b>(&'b self) -> Dir<'b, 'a> {
        let root_rdr = {
            match self.fat_type {
                FatType::Fat12 | FatType::Fat16 => DirRawStream::Root(DiskSlice::from_sectors(
                   self.first_data_sector - self.root_dir_sectors, self.root_dir_sectors, 1, self)),
                _ => DirRawStream::File(File::new(Some(self.boot.bpb.root_dir_first_cluster), None, self)),
            }
        };
        Dir::new(root_rdr, self)
    }

    fn read_bpb(rdr: &mut Read) -> io::Result<BiosParameterBlock> {
        let mut bpb: BiosParameterBlock = Default::default();
        bpb.bytes_per_sector = rdr.read_u16::<LittleEndian>()?;
        bpb.sectors_per_cluster = rdr.read_u8()?;
        bpb.reserved_sectors = rdr.read_u16::<LittleEndian>()?;
        bpb.fats = rdr.read_u8()?;
        bpb.root_entries = rdr.read_u16::<LittleEndian>()? ;
        bpb.total_sectors_16 = rdr.read_u16::<LittleEndian>()?;
        bpb.media = rdr.read_u8()?;
        bpb.sectors_per_fat_16 = rdr.read_u16::<LittleEndian>()?;
        bpb.sectors_per_track = rdr.read_u16::<LittleEndian>()?;
        bpb.heads = rdr.read_u16::<LittleEndian>()?;
        bpb.hidden_sectors = rdr.read_u32::<LittleEndian>()?; // hidden_sector_count
        bpb.total_sectors_32 = rdr.read_u32::<LittleEndian>()?;

        if bpb.sectors_per_fat_16 == 0 {
            bpb.sectors_per_fat_32 = rdr.read_u32::<LittleEndian>()?;
            bpb.extended_flags = rdr.read_u16::<LittleEndian>()?;
            bpb.fs_version = rdr.read_u16::<LittleEndian>()?;
            bpb.root_dir_first_cluster = rdr.read_u32::<LittleEndian>()?;
            bpb.fs_info_sector = rdr.read_u16::<LittleEndian>()?;
            bpb.backup_boot_sector = rdr.read_u16::<LittleEndian>()?;
            rdr.read(&mut bpb.reserved_0)?;
            bpb.drive_num = rdr.read_u8()?;
            bpb.reserved_1 = rdr.read_u8()?;
            bpb.ext_sig = rdr.read_u8()?; // 0x29
            bpb.volume_id = rdr.read_u32::<LittleEndian>()?;
            rdr.read(&mut bpb.volume_label)?;
            rdr.read(&mut bpb.fs_type_label)?;
        } else {
            bpb.drive_num = rdr.read_u8()?;
            bpb.reserved_1 = rdr.read_u8()?;
            bpb.ext_sig = rdr.read_u8()?; // 0x29
            bpb.volume_id = rdr.read_u32::<LittleEndian>()?;
            rdr.read(&mut bpb.volume_label)?;
            rdr.read(&mut bpb.fs_type_label)?;
        }
        Ok(bpb)
    }

    fn fat_type_from_clusters(total_clusters: u32) -> FatType {
        if total_clusters < 4085 {
            FatType::Fat12
        } else if total_clusters < 65525 {
            FatType::Fat16
        } else {
            FatType::Fat32
        }
    }

    fn read_boot_record(rdr: &mut Read) -> io::Result<BootRecord> {
        let mut boot: BootRecord = Default::default();
        rdr.read(&mut boot.bootjmp)?;
        rdr.read(&mut boot.oem_name)?;
        boot.bpb = Self::read_bpb(rdr)?;

        if boot.bpb.sectors_per_fat_16 == 0 {
            rdr.read_exact(&mut boot.boot_code[0..420])?;
        } else {
            rdr.read_exact(&mut boot.boot_code[0..448])?;
        }
        rdr.read(&mut boot.boot_sig)?;
        Ok(boot)
    }

    pub(crate) fn offset_from_sector(&self, sector: u32) -> u64 {
        (sector as u64) * self.boot.bpb.bytes_per_sector as u64
    }

    pub(crate) fn sector_from_cluster(&self, cluster: u32) -> u32 {
        ((cluster - 2) * self.boot.bpb.sectors_per_cluster as u32) + self.first_data_sector
    }

    pub(crate) fn get_cluster_size(&self) -> u32 {
        self.boot.bpb.sectors_per_cluster as u32 * self.boot.bpb.bytes_per_sector as u32
    }

    pub(crate) fn offset_from_cluster(&self, cluser: u32) -> u64 {
        self.offset_from_sector(self.sector_from_cluster(cluser))
    }

    fn fat_slice<'b>(&'b self) -> DiskSlice<'b, 'a> {
        let sectors_per_fat =
            if self.boot.bpb.sectors_per_fat_16 == 0 { self.boot.bpb.sectors_per_fat_32 }
            else { self.boot.bpb.sectors_per_fat_16 as u32 };
        let mirroring_enabled = self.boot.bpb.extended_flags & 0x80 == 0;
        let (fat_first_sector, mirrors) = if mirroring_enabled {
            (self.boot.bpb.reserved_sectors as u32, self.boot.bpb.fats)
        } else {
            let active_fat = (self.boot.bpb.extended_flags & 0x0F) as u32;
            let fat_first_sector = (self.boot.bpb.reserved_sectors as u32) + active_fat * sectors_per_fat;
            (fat_first_sector, 1)
        };
        DiskSlice::from_sectors(fat_first_sector, sectors_per_fat, mirrors, self)
    }

    pub(crate) fn cluster_iter<'b>(&'b self, cluster: u32) -> ClusterIterator<'b, 'a> {
        let disk_slice = self.fat_slice();
        ClusterIterator::new(disk_slice, self.fat_type, cluster)
    }

    pub(crate) fn alloc_cluster(&self, prev_cluster: Option<u32>) -> io::Result<u32> {
        let mut disk_slice = self.fat_slice();
        alloc_cluster(&mut disk_slice, self.fat_type, prev_cluster)
    }
}

#[derive(Clone)]
pub(crate) struct DiskSlice<'a, 'b: 'a> {
    begin: u64,
    size: u64,
    offset: u64,
    mirrors: u8,
    fs: &'a FileSystem<'b>,
}

impl <'a, 'b> DiskSlice<'a, 'b> {
    pub(crate) fn new(begin: u64, size: u64, mirrors: u8, fs: FileSystemRef<'a, 'b>) -> Self {
        DiskSlice { begin, size, mirrors, fs, offset: 0 }
    }

    pub(crate) fn from_sectors(first_sector: u32, sector_count: u32, mirrors: u8, fs: FileSystemRef<'a, 'b>) -> Self {
        let bytes_per_sector = fs.boot.bpb.bytes_per_sector as u64;
        Self::new(first_sector as u64 * bytes_per_sector, sector_count as u64 * bytes_per_sector, mirrors, fs)
    }

    pub(crate) fn global_pos(&self) -> u64 {
        self.begin + self.offset
    }
}

impl <'a, 'b> Read for DiskSlice<'a, 'b> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let offset = self.begin + self.offset;
        let read_size = cmp::min((self.size - self.offset) as usize, buf.len());
        let mut disk = self.fs.disk.borrow_mut();
        disk.seek(SeekFrom::Start(offset))?;
        let size = disk.read(&mut buf[..read_size])?;
        self.offset += size as u64;
        Ok(size)
    }
}

impl <'a, 'b> Write for DiskSlice<'a, 'b> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let offset = self.begin + self.offset;
        let write_size = cmp::min((self.size - self.offset) as usize, buf.len());
        for i in 0..self.mirrors {
            let mut disk = self.fs.disk.borrow_mut();
            disk.seek(SeekFrom::Start(offset + i as u64 * self.size))?;
            disk.write_all(&buf[..write_size])?;
        }
        self.offset += write_size as u64;
        Ok(write_size)
    }

    fn flush(&mut self) -> io::Result<()> {
        let mut disk = self.fs.disk.borrow_mut();
        disk.flush()
    }
}

impl <'a, 'b> Seek for DiskSlice<'a, 'b> {
    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
        let new_offset = match pos {
            SeekFrom::Current(x) => self.offset as i64 + x,
            SeekFrom::Start(x) => x as i64,
            SeekFrom::End(x) => self.size as i64 + x,
        };
        if new_offset < 0 || new_offset as u64 > self.size {
            Err(io::Error::new(ErrorKind::InvalidInput, "invalid seek"))
        } else {
            self.offset = new_offset as u64;
            Ok(self.offset)
        }
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2015-2016 Torrie Fischer <tdfischer@hackerbots.net>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::io::Read;
use std::str;

#[derive(Debug)]
pub struct LineReader {
    linebuf: String
}

impl LineReader {
    pub fn new() -> Self {
        LineReader {
            linebuf: String::new()
        }
    }

    pub fn read(&mut self, stream: &mut Read) -> Option<String> {
        match self.split_next_line() {
            None => self.read_and_split(stream),
            Some(line) => Some(line)
        }
    }

    fn read_and_split(&mut self, stream: &mut Read) -> Option<String> {
        let mut buf = [0; 1024];
        let next_msg = stream.read(&mut buf);
        match next_msg {
            Ok(count) => {
                self.linebuf.push_str(str::from_utf8(&buf[0..count]).unwrap());
                self.split_next_line()
            }
            Err(_) => None
        }
    }

    fn split_next_line(&mut self) -> Option<String> {
        let new_str;
        let split;
        match self.linebuf.find("\r\n") {
            Some(idx) => {
                new_str = self.linebuf.clone();
                split = new_str .split_at(idx);
            },
            None =>
                return None

        }
        self.linebuf = split.1[2..].to_string().clone();
        Some(split.0.to_string())
    }
}

<commit_msg>irc: util: linereader: also handle the case of \n without \r<commit_after>/*
 * Copyright 2015-2016 Torrie Fischer <tdfischer@hackerbots.net>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::io::Read;
use std::str;

#[derive(Debug)]
pub struct LineReader {
    linebuf: String
}

impl LineReader {
    pub fn new() -> Self {
        LineReader {
            linebuf: String::new()
        }
    }

    pub fn read(&mut self, stream: &mut Read) -> Option<String> {
        match self.split_next_line() {
            None => self.read_and_split(stream),
            Some(line) => Some(line)
        }
    }

    fn read_and_split(&mut self, stream: &mut Read) -> Option<String> {
        let mut buf = [0; 1024];
        let next_msg = stream.read(&mut buf);
        match next_msg {
            Ok(count) => {
                self.linebuf.push_str(str::from_utf8(&buf[0..count]).unwrap());
                self.split_next_line()
            }
            Err(_) => None
        }
    }

    fn split_next_line(&mut self) -> Option<String> {
        let new_str;
        let split;
        match self.linebuf.find("\r\n") {
            Some(idx) => {
                new_str = self.linebuf.clone();
                split = new_str .split_at(idx);
            },
            None => {
                match self.linebuf.find("\n") {
                    Some(idx) => {
                        new_str = self.linebuf.clone();
                        split = new_str .split_at(idx);
                    },
                    None =>
                        return None
                }
            }
        }
        self.linebuf = split.1[1..].to_string().clone();
        Some(split.0.to_string())
    }
}

<|endoftext|>


<commit_before>use std::io::{stdin, Read};
use std::fmt;
use std::fs::File;
use std::sync::Mutex;

use futures::{Future, IntoFuture};
use futures::future::Either;

use vm::{self, ExternModule, Result, Variants};
use vm::future::FutureValue;
use vm::gc::{Gc, Traverseable};
use vm::types::*;
use vm::thread::{Thread, ThreadInternal};
use vm::api::{Array, FutureResult, Generic, Getable, OpaqueValue, OwnedFunction, PrimitiveFuture,
              TypedBytecode, Userdata, VmType, WithVM, IO};
use vm::api::generic::{A, B};
use vm::stack::StackFrame;
use vm::internal::ValuePrinter;

use vm::internal::Value;

use compiler_pipeline::*;

use super::{Compiler, Error};

fn print(s: &str) -> IO<()> {
    print!("{}", s);
    IO::Value(())
}

fn println(s: &str) -> IO<()> {
    println!("{}", s);
    IO::Value(())
}

struct GluonFile(Mutex<File>);

impl Userdata for GluonFile {}

impl fmt::Debug for GluonFile {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "File")
    }
}

impl VmType for GluonFile {
    type Type = GluonFile;
}

impl Traverseable for GluonFile {
    fn traverse(&self, _: &mut Gc) {}
}

fn open_file(s: &str) -> IO<GluonFile> {
    match File::open(s) {
        Ok(f) => IO::Value(GluonFile(Mutex::new(f))),
        Err(err) => IO::Exception(format!("{}", err)),
    }
}

fn read_file<'vm>(file: WithVM<'vm, &GluonFile>, count: usize) -> IO<Array<'vm, u8>> {
    let WithVM { vm, value: file } = file;
    let mut file = file.0.lock().unwrap();
    let mut buffer = Vec::with_capacity(count);
    unsafe {
        buffer.set_len(count);
        match file.read(&mut *buffer) {
            Ok(bytes_read) => {
                let value = {
                    let mut context = vm.context();
                    match context.alloc(&buffer[..bytes_read]) {
                        Ok(value) => value,
                        Err(err) => return IO::Exception(format!("{}", err)),
                    }
                };
                IO::Value(Getable::from_value(vm, Variants::new(&Value::Array(value))))
            }
            Err(err) => IO::Exception(format!("{}", err)),
        }
    }
}

fn read_file_to_string(s: &str) -> IO<String> {
    let mut buffer = String::new();
    match File::open(s).and_then(|mut file| file.read_to_string(&mut buffer)) {
        Ok(_) => IO::Value(buffer),
        Err(err) => {
            use std::fmt::Write;
            buffer.clear();
            let _ = write!(&mut buffer, "{}", err);
            IO::Exception(buffer)
        }
    }
}

fn read_char() -> IO<char> {
    match stdin().bytes().next() {
        Some(result) => match result {
            Ok(b) => ::std::char::from_u32(b as u32)
                .map(IO::Value)
                .unwrap_or_else(|| IO::Exception("Not a valid char".into())),
            Err(err) => IO::Exception(format!("{}", err)),
        },
        None => IO::Exception("No read".into()),
    }
}

fn read_line() -> IO<String> {
    let mut buffer = String::new();
    match stdin().read_line(&mut buffer) {
        Ok(_) => IO::Value(buffer),
        Err(err) => {
            use std::fmt::Write;
            buffer.clear();
            let _ = write!(&mut buffer, "{}", err);
            IO::Exception(buffer)
        }
    }
}

/// IO a -> (String -> IO a) -> IO a
fn catch<'vm>(
    action: OpaqueValue<&'vm Thread, IO<A>>,
    mut catch: OwnedFunction<fn(String) -> IO<Generic<A>>>,
) -> FutureResult<Box<Future<Item = IO<Generic<A>>, Error = vm::Error> + Send>> {
    let vm = action.vm().root_thread();
    let frame_level = vm.context().stack.get_frames().len();
    let mut action: OwnedFunction<fn(()) -> Generic<A>> =
        unsafe { Getable::from_value(&vm, Variants::new(&action.get_value())) };

    let future = action.call_async(()).then(move |result| match result {
        Ok(value) => Either::A(Ok(IO::Value(value)).into_future()),
        Err(err) => {
            {
                let mut context = vm.context();
                let mut stack = StackFrame::current(&mut context.stack);
                while stack.stack.get_frames().len() > frame_level {
                    if stack.exit_scope().is_err() {
                        return Either::A(Ok(IO::Exception("Unknown error".into())).into_future());
                    }
                }
            }
            Either::B(catch.call_async(format!("{}", err)).then(|result| {
                Ok(match result {
                    Ok(value) => value,
                    Err(err) => IO::Exception(format!("{}", err)),
                })
            }))
        }
    });

    FutureResult(Box::new(future))
}

fn clear_frames(err: Error, frame_level: usize, mut stack: StackFrame) -> IO<String> {
    let fmt = match err {
        Error::VM(vm::Error::Panic(_)) => {
            let trace = stack.stack.stacktrace(frame_level);
            format!("{}\n{}", err, trace)
        }
        _ => format!("{}", err),
    };
    while stack.stack.get_frames().len() > frame_level {
        if stack.exit_scope().is_err() {
            return IO::Exception(fmt);
        }
    }
    IO::Exception(fmt)
}

fn run_expr(WithVM { vm, value: expr }: WithVM<&str>) -> PrimitiveFuture<IO<String>> {
    let vm = vm.root_thread();
    let frame_level = vm.context().stack.get_frames().len();

    let vm1 = vm.clone();
    let future = expr.run_expr(&mut Compiler::new(), vm1, "<top>", expr, None)
        .then(move |run_result| {
            let mut context = vm.context();
            let stack = StackFrame::current(&mut context.stack);
            FutureValue::sync(Ok(match run_result {
                Ok(execute_value) => {
                    let env = vm.global_env().get_env();
                    let typ = execute_value.typ;
                    IO::Value(format!(
                        "{} : {}",
                        ValuePrinter::new(&*env, &typ, *execute_value.value).width(80),
                        typ
                    ))
                }
                Err(err) => clear_frames(err, frame_level, stack),
            }))
        });

    future.boxed()
}

fn load_script(
    WithVM { vm, value: name }: WithVM<&str>,
    expr: &str,
) -> PrimitiveFuture<IO<String>> {
    let frame_level = vm.context().stack.get_frames().len();

    let vm1 = vm.root_thread();
    let vm = vm.root_thread();
    let name = name.to_string();
    let future = expr.load_script(&mut Compiler::new(), vm1, &name, expr, None)
        .then(move |run_result| {
            let mut context = vm.context();
            let stack = StackFrame::current(&mut context.stack);
            let io = match run_result {
                Ok(()) => IO::Value(format!("Loaded {}", name)),
                Err(err) => clear_frames(err, frame_level, stack),
            };
            Ok(io).into()
        });
    future.boxed()
}

mod std {
    pub mod io {
        pub use io as prim;
    }
}

pub fn load(vm: &Thread) -> Result<ExternModule> {
    vm.register_type::<GluonFile>("File", &[])?;

    // flat_map f m : (a -> IO b) -> IO a -> IO b
    //     = f (m ())
    let flat_map = vec![
        // [f, m, ()]       Initial stack
        Call(1),     // [f, m_ret]       Call m ()
        PushInt(0),  // [f, m_ret, ()]   Add a dummy argument ()
        TailCall(2), /* [f_ret]          Call f m_ret () */
    ];

    type FlatMap = fn(fn(A) -> IO<B>, IO<A>) -> IO<B>;
    type Wrap = fn(A) -> IO<A>;

    let wrap = vec![Pop(1)];

    use self::std;

    // IO functions
    ExternModule::new(
        vm,
        record! {
            flat_map => TypedBytecode::<FlatMap>::new("std.io.prim.flat_map", 3, flat_map),
            wrap => TypedBytecode::<Wrap>::new("std.io.prim.wrap", 2, wrap),
            open_file => primitive!(1 std::io::prim::open_file),
            read_file => primitive!(2 std::io::prim::read_file),
            read_file_to_string => primitive!(1 std::io::prim::read_file_to_string),
            read_char => primitive!(0 std::io::prim::read_char),
            read_line => primitive!(0 std::io::prim::read_line),
            print => primitive!(1 std::io::prim::print),
            println => primitive!(1 std::io::prim::println),
            catch => primitive!(2 std::io::prim::catch),
            run_expr => primitive!(1 std::io::prim::run_expr),
            load_script => primitive!(2 std::io::prim::load_script),
        },
    )
}
<commit_msg>fix(vm): Don't rely on frame_levels to clear frames after errors<commit_after>use std::io::{stdin, Read};
use std::fmt;
use std::fs::File;
use std::sync::Mutex;

use futures::{Future, IntoFuture};
use futures::future::Either;

use vm::{self, ExternModule, Result, Variants};
use vm::future::FutureValue;
use vm::gc::{Gc, Traverseable};
use vm::types::*;
use vm::thread::{Thread, ThreadInternal};
use vm::api::{Array, FutureResult, Generic, Getable, OpaqueValue, OwnedFunction, PrimitiveFuture,
              TypedBytecode, Userdata, VmType, WithVM, IO};
use vm::api::generic::{A, B};
use vm::stack::{StackFrame, State};
use vm::internal::ValuePrinter;

use vm::internal::Value;

use compiler_pipeline::*;

use super::{Compiler, Error};

fn print(s: &str) -> IO<()> {
    print!("{}", s);
    IO::Value(())
}

fn println(s: &str) -> IO<()> {
    println!("{}", s);
    IO::Value(())
}

struct GluonFile(Mutex<File>);

impl Userdata for GluonFile {}

impl fmt::Debug for GluonFile {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "File")
    }
}

impl VmType for GluonFile {
    type Type = GluonFile;
}

impl Traverseable for GluonFile {
    fn traverse(&self, _: &mut Gc) {}
}

fn open_file(s: &str) -> IO<GluonFile> {
    match File::open(s) {
        Ok(f) => IO::Value(GluonFile(Mutex::new(f))),
        Err(err) => IO::Exception(format!("{}", err)),
    }
}

fn read_file<'vm>(file: WithVM<'vm, &GluonFile>, count: usize) -> IO<Array<'vm, u8>> {
    let WithVM { vm, value: file } = file;
    let mut file = file.0.lock().unwrap();
    let mut buffer = Vec::with_capacity(count);
    unsafe {
        buffer.set_len(count);
        match file.read(&mut *buffer) {
            Ok(bytes_read) => {
                let value = {
                    let mut context = vm.context();
                    match context.alloc(&buffer[..bytes_read]) {
                        Ok(value) => value,
                        Err(err) => return IO::Exception(format!("{}", err)),
                    }
                };
                IO::Value(Getable::from_value(vm, Variants::new(&Value::Array(value))))
            }
            Err(err) => IO::Exception(format!("{}", err)),
        }
    }
}

fn read_file_to_string(s: &str) -> IO<String> {
    let mut buffer = String::new();
    match File::open(s).and_then(|mut file| file.read_to_string(&mut buffer)) {
        Ok(_) => IO::Value(buffer),
        Err(err) => {
            use std::fmt::Write;
            buffer.clear();
            let _ = write!(&mut buffer, "{}", err);
            IO::Exception(buffer)
        }
    }
}

fn read_char() -> IO<char> {
    match stdin().bytes().next() {
        Some(result) => match result {
            Ok(b) => ::std::char::from_u32(b as u32)
                .map(IO::Value)
                .unwrap_or_else(|| IO::Exception("Not a valid char".into())),
            Err(err) => IO::Exception(format!("{}", err)),
        },
        None => IO::Exception("No read".into()),
    }
}

fn read_line() -> IO<String> {
    let mut buffer = String::new();
    match stdin().read_line(&mut buffer) {
        Ok(_) => IO::Value(buffer),
        Err(err) => {
            use std::fmt::Write;
            buffer.clear();
            let _ = write!(&mut buffer, "{}", err);
            IO::Exception(buffer)
        }
    }
}

/// IO a -> (String -> IO a) -> IO a
fn catch<'vm>(
    action: OpaqueValue<&'vm Thread, IO<A>>,
    mut catch: OwnedFunction<fn(String) -> IO<Generic<A>>>,
) -> FutureResult<Box<Future<Item = IO<Generic<A>>, Error = vm::Error> + Send>> {
    let vm = action.vm().root_thread();
    let frame_level = vm.context().stack.get_frames().len();
    let mut action: OwnedFunction<fn(()) -> Generic<A>> =
        unsafe { Getable::from_value(&vm, Variants::new(&action.get_value())) };

    let future = action.call_async(()).then(move |result| match result {
        Ok(value) => Either::A(Ok(IO::Value(value)).into_future()),
        Err(err) => {
            {
                let mut context = vm.context();
                let mut stack = StackFrame::current(&mut context.stack);
                while stack.stack.get_frames().len() > frame_level {
                    if stack.exit_scope().is_err() {
                        return Either::A(Ok(IO::Exception("Unknown error".into())).into_future());
                    }
                }
            }
            Either::B(catch.call_async(format!("{}", err)).then(|result| {
                Ok(match result {
                    Ok(value) => value,
                    Err(err) => IO::Exception(format!("{}", err)),
                })
            }))
        }
    });

    FutureResult(Box::new(future))
}

fn clear_frames(err: Error, mut stack: StackFrame) -> IO<String> {
    let frame_level = stack
        .stack
        .get_frames()
        .iter()
        .rposition(|frame| frame.state == State::Lock)
        .unwrap_or(0);

    let fmt = match err {
        Error::VM(vm::Error::Panic(_)) => {
            let trace = stack.stack.stacktrace(frame_level);
            format!("{}\n{}", err, trace)
        }
        _ => format!("{}", err),
    };
    while let Ok(_) = stack.exit_scope() {}
    IO::Exception(fmt)
}

fn run_expr(WithVM { vm, value: expr }: WithVM<&str>) -> PrimitiveFuture<IO<String>> {
    let vm = vm.root_thread();

    let vm1 = vm.clone();
    let future = expr.run_expr(&mut Compiler::new(), vm1, "<top>", expr, None)
        .then(move |run_result| {
            let mut context = vm.context();
            let stack = StackFrame::current(&mut context.stack);
            FutureValue::sync(Ok(match run_result {
                Ok(execute_value) => {
                    let env = vm.global_env().get_env();
                    let typ = execute_value.typ;
                    IO::Value(format!(
                        "{} : {}",
                        ValuePrinter::new(&*env, &typ, *execute_value.value).width(80),
                        typ
                    ))
                }
                Err(err) => clear_frames(err, stack),
            }))
        });

    future.boxed()
}

fn load_script(
    WithVM { vm, value: name }: WithVM<&str>,
    expr: &str,
) -> PrimitiveFuture<IO<String>> {

    let vm1 = vm.root_thread();
    let vm = vm.root_thread();
    let name = name.to_string();
    let future = expr.load_script(&mut Compiler::new(), vm1, &name, expr, None)
        .then(move |run_result| {
            let mut context = vm.context();
            let stack = StackFrame::current(&mut context.stack);
            let io = match run_result {
                Ok(()) => IO::Value(format!("Loaded {}", name)),
                Err(err) => clear_frames(err, stack),
            };
            Ok(io).into()
        });
    future.boxed()
}

mod std {
    pub mod io {
        pub use io as prim;
    }
}

pub fn load(vm: &Thread) -> Result<ExternModule> {
    vm.register_type::<GluonFile>("File", &[])?;

    // flat_map f m : (a -> IO b) -> IO a -> IO b
    //     = f (m ())
    let flat_map = vec![
        // [f, m, ()]       Initial stack
        Call(1),     // [f, m_ret]       Call m ()
        PushInt(0),  // [f, m_ret, ()]   Add a dummy argument ()
        TailCall(2), /* [f_ret]          Call f m_ret () */
    ];

    type FlatMap = fn(fn(A) -> IO<B>, IO<A>) -> IO<B>;
    type Wrap = fn(A) -> IO<A>;

    let wrap = vec![Pop(1)];

    use self::std;

    // IO functions
    ExternModule::new(
        vm,
        record! {
            flat_map => TypedBytecode::<FlatMap>::new("std.io.prim.flat_map", 3, flat_map),
            wrap => TypedBytecode::<Wrap>::new("std.io.prim.wrap", 2, wrap),
            open_file => primitive!(1 std::io::prim::open_file),
            read_file => primitive!(2 std::io::prim::read_file),
            read_file_to_string => primitive!(1 std::io::prim::read_file_to_string),
            read_char => primitive!(0 std::io::prim::read_char),
            read_line => primitive!(0 std::io::prim::read_line),
            print => primitive!(1 std::io::prim::print),
            println => primitive!(1 std::io::prim::println),
            catch => primitive!(2 std::io::prim::catch),
            run_expr => primitive!(1 std::io::prim::run_expr),
            load_script => primitive!(2 std::io::prim::load_script),
        },
    )
}
<|endoftext|>


<commit_before>/*!
 * Bounded package merge algorithm, based on the paper
 * "A Fast and Space-Economical Algorithm for Length-Limited Coding
 * Jyrki Katajainen, Alistair Moffat, Andrew Turpin".
 */

use std::mem;
use std::ptr::{null, null_mut};

use libc::{c_void, size_t};
use libc::funcs::c95::stdlib::{free, malloc};

/// Nodes forming chains. Also used to represent leaves.
struct Node {
    /// Total weight (symbol count) of this chain.
    weight: usize,
    /// Previous node(s) of this chain, or 0 if none.
    tail: *mut Node,
    /// Leaf symbol index, or number of leaves before this chain.
    count: i32,
    /// Tracking for garbage collection.
    in_use: bool,
}

/// Memory pool for nodes.
struct NodePool {
    /// The pool.
    nodes: *mut Node,
    /// Pointer to a possibly free node in the pool.
    next: *mut Node,
    /// Size of the memory pool.
    size: u32,
}

/// Initializes a chain node with the given values and marks it as in use.
unsafe fn init_node(weight: usize, count: i32, tail: *mut Node, node: *mut Node) {
    (*node).weight = weight;
    (*node).count = count;
    (*node).tail = tail;
    (*node).in_use = true;
}

/**
 * Finds a free location in the memory pool. Performs garbage collection if needed.
 * lists: If given, used to mark in-use nodes during garbage collection.
 * maxbits: Size of lists.
 * pool: Memory pool to get free node from.
 */
unsafe fn get_free_node(lists: *const *mut *mut Node, maxbits: i32, pool: *mut NodePool) -> *mut Node {
    loop {
        if (*pool).next >= (*pool).nodes.offset((*pool).size as isize) {
            // Garbage collection
            for i in 0..(*pool).size {
                (*(*pool).nodes.offset(i as isize)).in_use = false;
            }
            if !lists.is_null() {
                for i in 0..maxbits*2 {
                    let mut node = *(*lists.offset((i / 2) as isize)).offset((i % 2) as isize);
                    while !node.is_null() {
                        (*node).in_use = true;
                        node = (*node).tail;
                    }
                }
            }
            (*pool).next = (*pool).nodes;
        }
        if !(*(*pool).next).in_use {
            break;
        }
        (*pool).next = (*pool).next.offset(1);
    }
    let free_node = (*pool).next;
    (*pool).next = (*pool).next.offset(1);
    free_node
}

/**
 * Performs a Boundary Package-Merge step. Puts a new chain in the given list. The
 * new chain is, depending on the weights, a leaf or a combination of two chains
 * from the previous list.
 * lists: The lists of chains.
 * maxbits: Number of lists.
 * leaves: The leaves, one per symbol.
 * numsymbols: Number of leaves.
 * pool: the node memory pool.
 * index: The index of the list in which a new chain or leaf is required.
 * final: Whether this is the last time this function is called. If it is then it
 *   is no more needed to recursively call self.
 */
unsafe fn boundary_pm(lists: *const *mut *mut Node, maxbits: i32, leaves: *const Node, numsymbols: i32, pool: *mut NodePool,
                      index: i32, is_final: bool) {
    let lastcount = (**(*lists.offset(index as isize)).offset(1)).count; // Count of last chain of list.

    if index == 0 && lastcount >= numsymbols {
        return;
    }

    let newchain = get_free_node(lists, maxbits, pool);
    let oldchain = *(*lists.offset(index as isize)).offset(1);

    // These are set up before the recursive calls below, so that there is a list
    // pointing to the new node, to let the garbage collection know it's in use.
    *(*lists.offset(index as isize)).offset(0) = oldchain;
    *(*lists.offset(index as isize)).offset(1) = newchain;

    if index == 0 {
        //New leaf node in list 0.
        init_node((*leaves.offset(lastcount as isize)).weight, lastcount + 1, null_mut(), newchain);
    } else {
        let sum = (**(*lists.offset((index - 1) as isize)).offset(0)).weight + (**(*lists.offset((index - 1) as isize)).offset(0)).weight;
        if lastcount < numsymbols && sum > (*leaves.offset(lastcount as isize)).weight {
            // New leaf inserted in list, so count is incremented.
            init_node((*leaves.offset(lastcount as isize)).weight, lastcount + 1, (*oldchain).tail, newchain);
        } else {
            init_node(sum, lastcount, *(*lists.offset((index - 1) as isize)).offset(1), newchain);
            if !is_final {
                // Two lookahead chains of previous list used up, create new ones.
                boundary_pm(lists, maxbits, leaves, numsymbols, pool, index - 1, false);
                boundary_pm(lists, maxbits, leaves, numsymbols, pool, index - 1, false);
            }
        }
    }
}

/// Initializes each list with as lookahead chains the two leaves with lowest weights.
unsafe fn init_lists(pool: *mut NodePool, leaves: *const Node, maxbits: i32, lists: *const *mut *mut Node) {
    let node0 = get_free_node(null(), maxbits, pool);
    let node1 = get_free_node(null(), maxbits, pool);
    init_node((*leaves.offset(0)).weight, 1, null_mut(), node0);
    init_node((*leaves.offset(1)).weight, 2, null_mut(), node1);
    for i in 0..maxbits {
        *(*lists.offset(i as isize)).offset(0) = node0;
        *(*lists.offset(i as isize)).offset(1) = node1;
    }
}

/**
 * Converts result of boundary package-merge to the bitlengths. The result in the
 * last chain of the last list contains the amount of active leaves in each list.
 * chain: Chain to extract the bit length from (last chain from last list).
 */
unsafe fn extract_bit_lengths(chain: *const Node, leaves: *const Node, bitlengths: *mut u32) {
    let mut node = chain;
    while node != null() {
        for i in 0..(*node).count {
            *bitlengths.offset((*leaves.offset(i as isize)).count as isize) = *bitlengths.offset((*leaves.offset(i as isize)).count as isize) + 1;
        }
        node = (*node).tail;
    }
}

/// Comparator for sorting the leaves. Has the function signature for qsort.
extern fn leaf_comparator(a: *const c_void, b: *const c_void) -> i32 {
    unsafe {
        let node_a: *const Node = a as *const Node;
        let node_b: *const Node = b as *const Node;
        (*node_a).weight as i32 - (*node_b).weight as i32
    }
}

#[link(name="c")]
extern {
    fn qsort(base: *mut c_void, nmemb: size_t, size: size_t, compar: extern fn(*const c_void, *const c_void) -> i32);
}

pub unsafe fn length_limited_code_lengths(frequencies: *const usize, n: i32, maxbits: i32, bitlengths: *mut u32) -> bool {
    // Amount of symbols with frequency > 0.
    let mut numsymbols = 0;

    // One leaf per symbol. Only numsymbols leaves will be used.
    let leaves: *mut Node = malloc((n as usize * mem::size_of::<Node>()) as size_t) as *mut Node;

    // Initialize all bitlengths at 0.
    for i in 0..n {
        *bitlengths.offset(i as isize) = 0;
    }

    // Count used symbols and place them in the leaves.
    for i in 0..n {
        if *frequencies.offset(i as isize) != 0 {
            (*leaves.offset(numsymbols)).weight = *frequencies.offset(i as isize);
            (*leaves.offset(numsymbols)).count = i;
            numsymbols = numsymbols + 1;
        }
    }

    // Check special cases and error conditions.
    if (1 << maxbits) < numsymbols {
        free(leaves as *mut c_void);
        return true; // Error, too few maxbits to represent symbols.
    }
    if numsymbols == 0 {
        free(leaves as *mut c_void);
        return false; // No symbols at all. OK.
    }
    if numsymbols == 1 {
        *bitlengths.offset((*leaves.offset(0)).count as isize) = 1;
        free(leaves as *mut c_void);
        return false; // Only one symbol, give it bitlength 1, not 0. OK.
    }

    // Sort the leaves from lightest to heaviest.
    qsort(leaves as *mut c_void, numsymbols as size_t, mem::size_of::<Node>() as size_t, leaf_comparator);

    // Initialize node memory pool.
    let pool_size = 2 * maxbits * (maxbits + 1);
    let pool_nodes = malloc((pool_size as usize * mem::size_of::<Node>()) as size_t) as *mut Node;
    let mut pool = NodePool { size: pool_size as u32, nodes: pool_nodes, next: pool_nodes };
    for i in 0..pool.size {
        (*pool.nodes.offset(i as isize)).in_use = false;
    }

    // Array of lists of chains. Each list requires only two lookahead chains at
    // a time, so each list is a array of two Node*'s.
    let lists: *const *mut *mut Node = malloc((maxbits as usize * mem::size_of::<*const *const Node>()) as size_t) as *const *mut *mut Node;
    init_lists(&mut pool, leaves, maxbits, lists);

    // In the last list, 2 * numsymbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each BoundaryPM run creates one.
    let num_boundary_pm_runs = 2 * numsymbols - 4;
    for i in 0..num_boundary_pm_runs {
        let is_final = i == num_boundary_pm_runs - 1;
        boundary_pm(lists, maxbits, leaves, numsymbols as i32, &mut pool, maxbits - 1, is_final);
    }

    extract_bit_lengths(*(*lists.offset(maxbits as isize - 1)).offset(1), leaves, bitlengths);

    free(lists as *mut c_void);
    free(leaves as *mut c_void);
    free(pool.nodes as *mut c_void);
    false
}

#[cfg(test)]
mod tests {
    use std::mem::size_of;
    use libc::{c_void, size_t};

    use super::qsort;

    extern fn comparator(a: *const c_void, b: *const c_void) -> i32 {
        unsafe {
            *(a as *const i32) - *(b as *const i32)
        }
    }

    #[test]
    fn test_qsort() {
        unsafe {
            let mut a: [i32; 5] = [3, 2, 4, 1, 5];
            qsort(a.as_mut_ptr() as *mut c_void, a.len() as size_t, size_of::<i32>() as size_t, comparator);
            assert_eq!(a, [1, 2, 3, 4, 5]);
        }
    }
}
<commit_msg>cleanup<commit_after>/*!
 * Bounded package merge algorithm, based on the paper
 * "A Fast and Space-Economical Algorithm for Length-Limited Coding
 * Jyrki Katajainen, Alistair Moffat, Andrew Turpin".
 */

use std::mem;
use std::ptr::{null, null_mut};

use libc::{c_void, size_t};
use libc::funcs::c95::stdlib::{free, malloc};

/// Nodes forming chains. Also used to represent leaves.
struct Node {
    /// Total weight (symbol count) of this chain.
    weight: usize,
    /// Previous node(s) of this chain, or 0 if none.
    tail: *mut Node,
    /// Leaf symbol index, or number of leaves before this chain.
    count: i32,
    /// Tracking for garbage collection.
    in_use: bool,
}

/// Memory pool for nodes.
struct NodePool {
    /// The pool.
    nodes: *mut Node,
    /// Pointer to a possibly free node in the pool.
    next: *mut Node,
    /// Size of the memory pool.
    size: u32,
}

/// Initializes a chain node with the given values and marks it as in use.
unsafe fn init_node(weight: usize, count: i32, tail: *mut Node, node: *mut Node) {
    (*node).weight = weight;
    (*node).count = count;
    (*node).tail = tail;
    (*node).in_use = true;
}

/**
 * Finds a free location in the memory pool. Performs garbage collection if needed.
 * lists: If given, used to mark in-use nodes during garbage collection.
 * maxbits: Size of lists.
 * pool: Memory pool to get free node from.
 */
unsafe fn get_free_node(lists: *const *mut *mut Node, maxbits: i32, pool: *mut NodePool) -> *mut Node {
    loop {
        if (*pool).next >= (*pool).nodes.offset((*pool).size as isize) {
            // Garbage collection
            for i in 0..(*pool).size {
                (*(*pool).nodes.offset(i as isize)).in_use = false;
            }
            if !lists.is_null() {
                for i in 0..maxbits*2 {
                    let mut node = *(*lists.offset((i / 2) as isize)).offset((i % 2) as isize);
                    while !node.is_null() {
                        (*node).in_use = true;
                        node = (*node).tail;
                    }
                }
            }
            (*pool).next = (*pool).nodes;
        }
        if !(*(*pool).next).in_use {
            break;
        }
        (*pool).next = (*pool).next.offset(1);
    }
    let free_node = (*pool).next;
    (*pool).next = (*pool).next.offset(1);
    free_node
}

/**
 * Performs a Boundary Package-Merge step. Puts a new chain in the given list. The
 * new chain is, depending on the weights, a leaf or a combination of two chains
 * from the previous list.
 * lists: The lists of chains.
 * maxbits: Number of lists.
 * leaves: The leaves, one per symbol.
 * numsymbols: Number of leaves.
 * pool: the node memory pool.
 * index: The index of the list in which a new chain or leaf is required.
 * final: Whether this is the last time this function is called. If it is then it
 *   is no more needed to recursively call self.
 */
unsafe fn boundary_pm(lists: *const *mut *mut Node, maxbits: i32, leaves: *const Node, numsymbols: i32, pool: *mut NodePool,
                      index: i32, is_final: bool) {
    let lastcount = (**(*lists.offset(index as isize)).offset(1)).count; // Count of last chain of list.

    if index == 0 && lastcount >= numsymbols {
        return;
    }

    let newchain = get_free_node(lists, maxbits, pool);
    let oldchain = *(*lists.offset(index as isize)).offset(1);

    // These are set up before the recursive calls below, so that there is a list
    // pointing to the new node, to let the garbage collection know it's in use.
    *(*lists.offset(index as isize)).offset(0) = oldchain;
    *(*lists.offset(index as isize)).offset(1) = newchain;

    if index == 0 {
        //New leaf node in list 0.
        init_node((*leaves.offset(lastcount as isize)).weight, lastcount + 1, null_mut(), newchain);
    } else {
        let sum = (**(*lists.offset((index - 1) as isize)).offset(0)).weight + (**(*lists.offset((index - 1) as isize)).offset(0)).weight;
        if lastcount < numsymbols && sum > (*leaves.offset(lastcount as isize)).weight {
            // New leaf inserted in list, so count is incremented.
            init_node((*leaves.offset(lastcount as isize)).weight, lastcount + 1, (*oldchain).tail, newchain);
        } else {
            init_node(sum, lastcount, *(*lists.offset((index - 1) as isize)).offset(1), newchain);
            if !is_final {
                // Two lookahead chains of previous list used up, create new ones.
                boundary_pm(lists, maxbits, leaves, numsymbols, pool, index - 1, false);
                boundary_pm(lists, maxbits, leaves, numsymbols, pool, index - 1, false);
            }
        }
    }
}

/// Initializes each list with as lookahead chains the two leaves with lowest weights.
unsafe fn init_lists(pool: *mut NodePool, leaves: *const Node, maxbits: i32, lists: *const *mut *mut Node) {
    let node0 = get_free_node(null(), maxbits, pool);
    let node1 = get_free_node(null(), maxbits, pool);
    init_node((*leaves.offset(0)).weight, 1, null_mut(), node0);
    init_node((*leaves.offset(1)).weight, 2, null_mut(), node1);
    for i in 0..maxbits {
        *(*lists.offset(i as isize)).offset(0) = node0;
        *(*lists.offset(i as isize)).offset(1) = node1;
    }
}

/**
 * Converts result of boundary package-merge to the bitlengths. The result in the
 * last chain of the last list contains the amount of active leaves in each list.
 * chain: Chain to extract the bit length from (last chain from last list).
 */
unsafe fn extract_bit_lengths(chain: *const Node, leaves: *const Node, bitlengths: *mut u32) {
    let mut node = chain;
    while node != null() {
        for i in 0..(*node).count {
            *bitlengths.offset((*leaves.offset(i as isize)).count as isize) = *bitlengths.offset((*leaves.offset(i as isize)).count as isize) + 1;
        }
        node = (*node).tail;
    }
}

/// Comparator for sorting the leaves. Has the function signature for qsort.
extern fn leaf_comparator(a: *const c_void, b: *const c_void) -> i32 {
    unsafe {
        let node_a: *const Node = a as *const Node;
        let node_b: *const Node = b as *const Node;
        (*node_a).weight as i32 - (*node_b).weight as i32
    }
}

#[link(name="c")]
extern {
    fn qsort(base: *mut c_void, nmemb: size_t, size: size_t, compar: extern fn(*const c_void, *const c_void) -> i32);
}

pub unsafe fn length_limited_code_lengths(frequencies: *const usize, n: i32, maxbits: i32, bitlengths: *mut u32) -> bool {
    // Amount of symbols with frequency > 0.
    let mut numsymbols = 0;

    // One leaf per symbol. Only numsymbols leaves will be used.
    let leaves: *mut Node = malloc((n as usize * mem::size_of::<Node>()) as size_t) as *mut Node;

    // Initialize all bitlengths at 0.
    for i in 0..n {
        *bitlengths.offset(i as isize) = 0;
    }

    // Count used symbols and place them in the leaves.
    for i in 0..n {
        if *frequencies.offset(i as isize) != 0 {
            (*leaves.offset(numsymbols)).weight = *frequencies.offset(i as isize);
            (*leaves.offset(numsymbols)).count = i; // Index of symbol this leaf represents.
            numsymbols += 1;
        }
    }

    // Check special cases and error conditions.
    if (1 << maxbits) < numsymbols {
        free(leaves as *mut c_void);
        return true; // Error, too few maxbits to represent symbols.
    }
    if numsymbols == 0 {
        free(leaves as *mut c_void);
        return false; // No symbols at all. OK.
    }
    if numsymbols == 1 {
        *bitlengths.offset((*leaves.offset(0)).count as isize) = 1;
        free(leaves as *mut c_void);
        return false; // Only one symbol, give it bitlength 1, not 0. OK.
    }

    // Sort the leaves from lightest to heaviest.
    qsort(leaves as *mut c_void, numsymbols as size_t, mem::size_of::<Node>() as size_t, leaf_comparator);

    // Initialize node memory pool.
    let pool_size = 2 * maxbits * (maxbits + 1);
    let pool_nodes = malloc((pool_size as usize * mem::size_of::<Node>()) as size_t) as *mut Node;
    let mut pool = NodePool { size: pool_size as u32, nodes: pool_nodes, next: pool_nodes };
    for i in 0..pool.size {
        (*pool.nodes.offset(i as isize)).in_use = false;
    }

    // Array of lists of chains. Each list requires only two lookahead chains at
    // a time, so each list is a array of two Node*'s.
    let lists: *const *mut *mut Node = malloc((maxbits as usize * mem::size_of::<*const *const Node>()) as size_t) as *const *mut *mut Node;
    init_lists(&mut pool, leaves, maxbits, lists);

    // In the last list, 2 * numsymbols - 2 active chains need to be created. Two
    // are already created in the initialization. Each BoundaryPM run creates one.
    let num_boundary_pm_runs = 2 * numsymbols - 4;
    for i in 0..num_boundary_pm_runs {
        let is_final = i == num_boundary_pm_runs - 1;
        boundary_pm(lists, maxbits, leaves, numsymbols as i32, &mut pool, maxbits - 1, is_final);
    }

    extract_bit_lengths(*(*lists.offset(maxbits as isize - 1)).offset(1), leaves, bitlengths);

    free(lists as *mut c_void);
    free(leaves as *mut c_void);
    free(pool.nodes as *mut c_void);
    false
}

#[cfg(test)]
mod tests {
    use std::mem::size_of;
    use libc::{c_void, size_t};

    use super::qsort;

    extern fn comparator(a: *const c_void, b: *const c_void) -> i32 {
        unsafe {
            *(a as *const i32) - *(b as *const i32)
        }
    }

    #[test]
    fn test_qsort() {
        unsafe {
            let mut a: [i32; 5] = [3, 2, 4, 1, 5];
            qsort(a.as_mut_ptr() as *mut c_void, a.len() as size_t, size_of::<i32>() as size_t, comparator);
            assert_eq!(a, [1, 2, 3, 4, 5]);
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2017-2019 int08h LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#[cfg(feature = "gcpkms")]
pub mod inner {
    extern crate base64;
    extern crate google_cloudkms1 as cloudkms1;
    extern crate hyper;
    extern crate hyper_rustls;
    extern crate yup_oauth2 as oauth2;

    use std::default::Default;
    use std::env;
    use std::path::Path;
    use std::result::Result;

    use self::cloudkms1::CloudKMS;
    use self::cloudkms1::{DecryptRequest, EncryptRequest};
    use self::hyper::net::HttpsConnector;
    use self::hyper::status::StatusCode;
    use self::hyper_rustls::TlsClient;
    use self::oauth2::{ServiceAccountAccess, ServiceAccountKey};

    use crate::kms::{EncryptedDEK, KmsError, KmsProvider, PlaintextDEK, AD};

    const GOOGLE_APP_CREDS: &str = &"GOOGLE_APPLICATION_CREDENTIALS";

    /// Google Cloud Key Management Service
    /// https://cloud.google.com/kms/
    pub struct GcpKms {
        key_resource_id: String,
        service_account: ServiceAccountKey,
    }

    impl GcpKms {
        ///
        /// Create a new GcpKms from a Google Cloud KMS key resource ID of the form
        /// `projects/*/locations/*/keyRings/*/cryptoKeys/*`
        ///
        pub fn from_resource_id(resource_id: &str) -> Result<Self, KmsError> {
            let svc_acct = load_gcp_credential()?;

            Ok(GcpKms {
                key_resource_id: resource_id.to_string(),
                service_account: svc_acct,
            })
        }

        fn new_hub(&self) -> CloudKMS<hyper::Client, ServiceAccountAccess<hyper::Client>> {
            let client1 = hyper::Client::with_connector(HttpsConnector::new(TlsClient::new()));
            let access = oauth2::ServiceAccountAccess::new(self.service_account.clone(), client1);

            let client2 = hyper::Client::with_connector(HttpsConnector::new(TlsClient::new()));
            CloudKMS::new(client2, access)
        }

        fn pretty_http_error(&self, resp: &hyper::client::Response) -> KmsError {
            let code = resp.status;
            let url = &resp.url;

            KmsError::OperationFailed(format!("Response {} from {}", code, url))
        }
    }

    impl KmsProvider for GcpKms {
        fn encrypt_dek(&self, plaintext_dek: &PlaintextDEK) -> Result<EncryptedDEK, KmsError> {
            let mut request = EncryptRequest::default();
            request.plaintext = Some(base64::encode(plaintext_dek));
            request.additional_authenticated_data = Some(base64::encode(AD));

            let hub = self.new_hub();
            let result = hub
                .projects()
                .locations_key_rings_crypto_keys_encrypt(request, &self.key_resource_id)
                .doit();

            match result {
                Ok((http_resp, enc_resp)) => {
                    if http_resp.status == StatusCode::Ok {
                        let ciphertext = enc_resp.ciphertext.unwrap();
                        let ct = base64::decode(&ciphertext)?;
                        Ok(ct)
                    } else {
                        Err(self.pretty_http_error(&http_resp))
                    }
                }
                Err(e) => Err(KmsError::OperationFailed(format!("encrypt_dek() {:?}", e))),
            }
        }

        fn decrypt_dek(&self, encrypted_dek: &EncryptedDEK) -> Result<PlaintextDEK, KmsError> {
            let mut request = DecryptRequest::default();
            request.ciphertext = Some(base64::encode(encrypted_dek));
            request.additional_authenticated_data = Some(base64::encode(AD));

            let hub = self.new_hub();
            let result = hub
                .projects()
                .locations_key_rings_crypto_keys_decrypt(request, &self.key_resource_id)
                .doit();

            match result {
                Ok((http_resp, enc_resp)) => {
                    if http_resp.status == StatusCode::Ok {
                        let plaintext = enc_resp.plaintext.unwrap();
                        let ct = base64::decode(&plaintext)?;
                        Ok(ct)
                    } else {
                        Err(self.pretty_http_error(&http_resp))
                    }
                }
                Err(e) => Err(KmsError::OperationFailed(format!("decrypt_dek() {:?}", e))),
            }
        }
    }

    /// Minimal implementation of Application Default Credentials.
    /// https://cloud.google.com/docs/authentication/production
    ///
    ///   1. Look for GOOGLE_APPLICATION_CREDENTIALS and load service account
    ///      credentials if found.
    ///   2. If not, error
    ///
    /// TODO attempt to load GCE default credentials from metadata server.
    /// This will be a bearer token instead of service account credential.

    fn load_gcp_credential() -> Result<ServiceAccountKey, KmsError> {
        if let Ok(gac) = env::var(GOOGLE_APP_CREDS.to_string()) {
            if Path::new(&gac).exists() {
                match oauth2::service_account_key_from_file(&gac) {
                    Ok(svc_acct_key) => return Ok(svc_acct_key),
                    Err(e) => {
                        return Err(KmsError::InvalidConfiguration(format!(
                            "Can't load service account credential '{}': {:?}",
                            gac, e
                        )))
                    }
                }
            } else {
                return Err(KmsError::InvalidConfiguration(format!(
                    "{} ='{}' does not exist",
                    GOOGLE_APP_CREDS, gac
                )));
            }
        }

        // TODO: call to GCE metadata service to get default credential from
        // http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/token

        panic!(
            "Failed to load service account credential. Is {} set?",
            GOOGLE_APP_CREDS
        );
    }
}
<commit_msg>Refactor: lift return out of if block<commit_after>// Copyright 2017-2019 int08h LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#[cfg(feature = "gcpkms")]
pub mod inner {
    extern crate base64;
    extern crate google_cloudkms1 as cloudkms1;
    extern crate hyper;
    extern crate hyper_rustls;
    extern crate yup_oauth2 as oauth2;

    use std::default::Default;
    use std::env;
    use std::path::Path;
    use std::result::Result;

    use self::cloudkms1::CloudKMS;
    use self::cloudkms1::{DecryptRequest, EncryptRequest};
    use self::hyper::net::HttpsConnector;
    use self::hyper::status::StatusCode;
    use self::hyper_rustls::TlsClient;
    use self::oauth2::{ServiceAccountAccess, ServiceAccountKey};

    use crate::kms::{EncryptedDEK, KmsError, KmsProvider, PlaintextDEK, AD};

    const GOOGLE_APP_CREDS: &str = &"GOOGLE_APPLICATION_CREDENTIALS";

    /// Google Cloud Key Management Service
    /// https://cloud.google.com/kms/
    pub struct GcpKms {
        key_resource_id: String,
        service_account: ServiceAccountKey,
    }

    impl GcpKms {
        ///
        /// Create a new GcpKms from a Google Cloud KMS key resource ID of the form
        /// `projects/*/locations/*/keyRings/*/cryptoKeys/*`
        ///
        pub fn from_resource_id(resource_id: &str) -> Result<Self, KmsError> {
            let svc_acct = load_gcp_credential()?;

            Ok(GcpKms {
                key_resource_id: resource_id.to_string(),
                service_account: svc_acct,
            })
        }

        fn new_hub(&self) -> CloudKMS<hyper::Client, ServiceAccountAccess<hyper::Client>> {
            let client1 = hyper::Client::with_connector(HttpsConnector::new(TlsClient::new()));
            let access = oauth2::ServiceAccountAccess::new(self.service_account.clone(), client1);

            let client2 = hyper::Client::with_connector(HttpsConnector::new(TlsClient::new()));
            CloudKMS::new(client2, access)
        }

        fn pretty_http_error(&self, resp: &hyper::client::Response) -> KmsError {
            let code = resp.status;
            let url = &resp.url;

            KmsError::OperationFailed(format!("Response {} from {}", code, url))
        }
    }

    impl KmsProvider for GcpKms {
        fn encrypt_dek(&self, plaintext_dek: &PlaintextDEK) -> Result<EncryptedDEK, KmsError> {
            let mut request = EncryptRequest::default();
            request.plaintext = Some(base64::encode(plaintext_dek));
            request.additional_authenticated_data = Some(base64::encode(AD));

            let hub = self.new_hub();
            let result = hub
                .projects()
                .locations_key_rings_crypto_keys_encrypt(request, &self.key_resource_id)
                .doit();

            match result {
                Ok((http_resp, enc_resp)) => {
                    if http_resp.status == StatusCode::Ok {
                        let ciphertext = enc_resp.ciphertext.unwrap();
                        let ct = base64::decode(&ciphertext)?;
                        Ok(ct)
                    } else {
                        Err(self.pretty_http_error(&http_resp))
                    }
                }
                Err(e) => Err(KmsError::OperationFailed(format!("encrypt_dek() {:?}", e))),
            }
        }

        fn decrypt_dek(&self, encrypted_dek: &EncryptedDEK) -> Result<PlaintextDEK, KmsError> {
            let mut request = DecryptRequest::default();
            request.ciphertext = Some(base64::encode(encrypted_dek));
            request.additional_authenticated_data = Some(base64::encode(AD));

            let hub = self.new_hub();
            let result = hub
                .projects()
                .locations_key_rings_crypto_keys_decrypt(request, &self.key_resource_id)
                .doit();

            match result {
                Ok((http_resp, enc_resp)) => {
                    if http_resp.status == StatusCode::Ok {
                        let plaintext = enc_resp.plaintext.unwrap();
                        let ct = base64::decode(&plaintext)?;
                        Ok(ct)
                    } else {
                        Err(self.pretty_http_error(&http_resp))
                    }
                }
                Err(e) => Err(KmsError::OperationFailed(format!("decrypt_dek() {:?}", e))),
            }
        }
    }

    /// Minimal implementation of Application Default Credentials.
    /// https://cloud.google.com/docs/authentication/production
    ///
    ///   1. Look for GOOGLE_APPLICATION_CREDENTIALS and load service account
    ///      credentials if found.
    ///   2. If not, error
    ///
    /// TODO attempt to load GCE default credentials from metadata server.
    /// This will be a bearer token instead of service account credential.

    fn load_gcp_credential() -> Result<ServiceAccountKey, KmsError> {
        if let Ok(gac) = env::var(GOOGLE_APP_CREDS.to_string()) {
            return if Path::new(&gac).exists() {
                match oauth2::service_account_key_from_file(&gac) {
                    Ok(svc_acct_key) => Ok(svc_acct_key),
                    Err(e) => {
                        Err(KmsError::InvalidConfiguration(format!(
                            "Can't load service account credential '{}': {:?}",
                            gac, e
                        )))
                    }
                }
            } else {
                Err(KmsError::InvalidConfiguration(format!(
                    "{} ='{}' does not exist",
                    GOOGLE_APP_CREDS, gac
                )))
            }
        }

        // TODO: call to GCE metadata service to get default credential from
        // http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/token

        panic!(
            "Failed to load service account credential. Is {} set?",
            GOOGLE_APP_CREDS
        );
    }
}
<|endoftext|>


<commit_before>pub mod container;
pub mod tree;

pub use self::container::{Container, ContainerType, Handle, Layout};
<commit_msg>Add module for graph_tree<commit_after>pub mod container;
mod graph_tree;
pub mod tree;

pub use self::container::{Container, ContainerType, Handle, Layout};
<|endoftext|>


<commit_before>use term;

use std::convert::AsRef;
use std::io;
use std::io::Write;
use std::fmt::Display;

/// The level of information to display to a logger
///
/// This enum should only be used as parameters for `Logger` or `Book` methods. Library
/// users should **not** do exhaustive pattern matching on the variants of the enums,
/// as it might grow variants later.
#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Copy)]
pub enum InfoLevel {
    /// Debug: the lowest level
    Debug = 0,
    /// Warning: won't be displayed by default
    Warning,
    /// Info: won't be displayed by default
    Info,
    /// Error
    Error,
    /// Quiet
    Quiet,

    /// Hint that destructuring should not be exhaustive
    #[doc(hidden)]
    __NonExhaustive,
}

use self::InfoLevel::*;

/// Abstract over eithr term outupt or (if it fails) io::stderr()
enum Output {
    Terminal(Box<term::StderrTerminal>),
    Stderr(io::Stderr),
}

impl Output {
    pub fn new() -> Output {
        if let Some(term) = term::stderr() {
            if (*term).supports_color() {
                return Output::Terminal(term)
            }
        }
        return Output::Stderr(io::stderr())
    }

    pub fn print_msg<S: Display>(&mut self, level: InfoLevel, msg: S) {
        let (colour, head_msg) = match level {
            Debug => (term::color::BRIGHT_BLUE, lformat!("Debug: ")),
            Warning => (term::color::BRIGHT_YELLOW, lformat!("Warning: ")),
            Info => (term::color::BRIGHT_GREEN, lformat!("Info: ")),
            Error => (term::color::BRIGHT_RED, lformat!("Error: ")),
            _ => unreachable!(),
        };
        match *self {
            Output::Stderr(ref mut stderr) => {
                writeln!(stderr, "{}{}",
                       head_msg,
                       msg)
                    .unwrap();
            },
            Output::Terminal(ref mut term) => {
                let _ = term.fg(colour);
                let _ = term.attr(term::Attr::Bold);
                write!(term, "{}",
                       head_msg)
                    .unwrap();
                let _ = term.reset();
                writeln!(term, "{}",
                       msg)
                    .unwrap();
            }
        }
    }
}

/// Logs info and warning message and choose whether to display them
/// according to verbosity
#[derive(Debug)]
pub struct Logger {
    verbosity: InfoLevel,
}

impl Logger {
    /// Creates a new logger with defined verbosity
    pub fn new(verbosity: InfoLevel) -> Logger {
        Logger { verbosity: verbosity }
    }

    /// Sets verbosity
    pub fn set_verbosity(&mut self, verbosity: InfoLevel) {
        self.verbosity = verbosity;
    }

    /// Prints a message
    pub fn display_msg<S: AsRef<str>>(level: InfoLevel, s: S) {
        let mut output = Output::new();
        output.print_msg(level, s.as_ref());
    }

    /// Prints a debug message
    pub fn display_debug<S: AsRef<str>>(s: S) {
        Self::display_msg(Debug, s);
    }

    /// Prints an info message
    pub fn display_info<S: AsRef<str>>(s: S) {
        Self::display_msg(Info, s);
    }

    /// Prints a warning message
    pub fn display_warning<S: AsRef<str>>(s: S) {
        Self::display_msg(Warning, s);
    }

    /// Prints an error message
    pub fn display_error<S: AsRef<str>>(s: S) {
        Self::display_msg(Error, s);
    }

    /// Prints a message if logger's verbosity <= level
    pub fn log<S: AsRef<str>>(&self, level: InfoLevel, s: S) {
        if level >= self.verbosity {
            match level {
                Debug => Self::display_debug(s),
                Info => Self::display_info(s),
                Warning => Self::display_warning(s),
                Error => Self::display_error(s),
                Quiet => unreachable!(),
                __NonExhaustive => unreachable!(),
            }
        }
    }

    /// Equivalent of log(Debug, s)
    pub fn debug<S: AsRef<str>>(&self, s: S) {
        self.log(InfoLevel::Debug, s);
    }

    /// Equivalent of log(Info, s)
    pub fn info<S: AsRef<str>>(&self, s: S) {
        self.log(InfoLevel::Info, s);
    }

    /// Equivalent of log(Warning, s)
    pub fn warning<S: AsRef<str>>(&self, s: S) {
        self.log(InfoLevel::Warning, s);
    }

    /// Equivalent of log(Error, s)
    pub fn error<S: AsRef<str>>(&self, s: S) {
        self.log(InfoLevel::Error, s);
    }
}

// Code to end shell colouring
pub const SHELL_COLOUR_OFF: &'static str = "\x1B[0m";
pub const SHELL_COLOUR_RED: &'static str = "\x1B[1m\x1B[31m";
pub const SHELL_COLOUR_BLUE: &'static str = "\x1B[1m\x1B[36m";
pub const SHELL_COLOUR_ORANGE: &'static str = "\x1B[1m\x1B[33m";
pub const SHELL_COLOUR_GREEN: &'static str = "\x1B[1m\x1B[32m";
<commit_msg>Add doc for Output<commit_after>use term;

use std::convert::AsRef;
use std::io;
use std::io::Write;
use std::fmt::Display;

/// The level of information to display to a logger
///
/// This enum should only be used as parameters for `Logger` or `Book` methods. Library
/// users should **not** do exhaustive pattern matching on the variants of the enums,
/// as it might grow variants later.
#[derive(Debug, PartialEq, Eq, PartialOrd, Clone, Copy)]
pub enum InfoLevel {
    /// Debug: the lowest level
    Debug = 0,
    /// Warning: won't be displayed by default
    Warning,
    /// Info: won't be displayed by default
    Info,
    /// Error
    Error,
    /// Quiet
    Quiet,

    /// Hint that destructuring should not be exhaustive
    #[doc(hidden)]
    __NonExhaustive,
}

use self::InfoLevel::*;

/// Abstract over either term outupt or (if it fails) io::stderr()
enum Output {
    Terminal(Box<term::StderrTerminal>),
    Stderr(io::Stderr),
}

impl Output {
    /// Creates a new Output
    pub fn new() -> Output {
        if let Some(term) = term::stderr() {
            if (*term).supports_color() {
                return Output::Terminal(term)
            }
        }
        return Output::Stderr(io::stderr())
    }

    /// Print a msg with prefixed by a coloured `level` message
    pub fn print_msg<S: Display>(&mut self, level: InfoLevel, msg: S) {
        let (colour, head_msg) = match level {
            Debug => (term::color::BRIGHT_BLUE, lformat!("Debug: ")),
            Warning => (term::color::BRIGHT_YELLOW, lformat!("Warning: ")),
            Info => (term::color::BRIGHT_GREEN, lformat!("Info: ")),
            Error => (term::color::BRIGHT_RED, lformat!("Error: ")),
            _ => unreachable!(),
        };
        match *self {
            Output::Stderr(ref mut stderr) => {
                writeln!(stderr, "{}{}",
                       head_msg,
                       msg)
                    .unwrap();
            },
            Output::Terminal(ref mut term) => {
                let _ = term.fg(colour);
                let _ = term.attr(term::Attr::Bold);
                write!(term, "{}",
                       head_msg)
                    .unwrap();
                let _ = term.reset();
                writeln!(term, "{}",
                       msg)
                    .unwrap();
            }
        }
    }
}

/// Logs info and warning message and choose whether to display them
/// according to verbosity
#[derive(Debug)]
pub struct Logger {
    verbosity: InfoLevel,
}

impl Logger {
    /// Creates a new logger with defined verbosity
    pub fn new(verbosity: InfoLevel) -> Logger {
        Logger { verbosity: verbosity }
    }

    /// Sets verbosity
    pub fn set_verbosity(&mut self, verbosity: InfoLevel) {
        self.verbosity = verbosity;
    }

    /// Prints a message
    pub fn display_msg<S: AsRef<str>>(level: InfoLevel, s: S) {
        let mut output = Output::new();
        output.print_msg(level, s.as_ref());
    }

    /// Prints a debug message
    pub fn display_debug<S: AsRef<str>>(s: S) {
        Self::display_msg(Debug, s);
    }

    /// Prints an info message
    pub fn display_info<S: AsRef<str>>(s: S) {
        Self::display_msg(Info, s);
    }

    /// Prints a warning message
    pub fn display_warning<S: AsRef<str>>(s: S) {
        Self::display_msg(Warning, s);
    }

    /// Prints an error message
    pub fn display_error<S: AsRef<str>>(s: S) {
        Self::display_msg(Error, s);
    }

    /// Prints a message if logger's verbosity <= level
    pub fn log<S: AsRef<str>>(&self, level: InfoLevel, s: S) {
        if level >= self.verbosity {
            match level {
                Debug => Self::display_debug(s),
                Info => Self::display_info(s),
                Warning => Self::display_warning(s),
                Error => Self::display_error(s),
                Quiet => unreachable!(),
                __NonExhaustive => unreachable!(),
            }
        }
    }

    /// Equivalent of log(Debug, s)
    pub fn debug<S: AsRef<str>>(&self, s: S) {
        self.log(InfoLevel::Debug, s);
    }

    /// Equivalent of log(Info, s)
    pub fn info<S: AsRef<str>>(&self, s: S) {
        self.log(InfoLevel::Info, s);
    }

    /// Equivalent of log(Warning, s)
    pub fn warning<S: AsRef<str>>(&self, s: S) {
        self.log(InfoLevel::Warning, s);
    }

    /// Equivalent of log(Error, s)
    pub fn error<S: AsRef<str>>(&self, s: S) {
        self.log(InfoLevel::Error, s);
    }
}

// Code to end shell colouring
pub const SHELL_COLOUR_OFF: &'static str = "\x1B[0m";
pub const SHELL_COLOUR_RED: &'static str = "\x1B[1m\x1B[31m";
pub const SHELL_COLOUR_BLUE: &'static str = "\x1B[1m\x1B[36m";
pub const SHELL_COLOUR_ORANGE: &'static str = "\x1B[1m\x1B[33m";
pub const SHELL_COLOUR_GREEN: &'static str = "\x1B[1m\x1B[32m";
<|endoftext|>