<commit_before>use std::fs::File;
use std::collections::HashMap;
use std::io::BufReader;
use std::borrow::Cow;

use encoding::{Encoding, DecoderTrap};
use encoding::all::UTF_16LE;

use errors::*;
use {ExcelReader, Range, DataType, CellErrorType};
use vba::VbaProject;
use cfb::Cfb;
use utils::{read_u16, read_u32, read_slice};

enum CfbWrap {
    Wb(Cfb, BufReader<File>),
    Vba(VbaProject),
}

/// A struct representing an old xls format file (CFB)
pub struct Xls {
    codepage: u16,
    sheets: HashMap<String, Range>,
    strings: Vec<String>,
    cfb: Option<CfbWrap>,
}

impl ExcelReader for Xls {

    fn new(r: File) -> Result<Self> {
        let len = try!(r.metadata()).len() as usize;
        let mut r = BufReader::new(r);
        let mut cfb = try!(Cfb::new(&mut r, len ));
        let wb = try!(cfb.get_stream("Workbook", &mut r));

        let mut xls = Xls {
            codepage: 1200,
            sheets: HashMap::new(),
            strings: Vec::new(),
            cfb: Some(CfbWrap::Wb(cfb, r)),
        };
        try!(xls.parse_workbook(&wb));
        Ok(xls)
    }

    fn has_vba(&mut self) -> bool {
        match self.cfb {
            Some(CfbWrap::Wb(ref cfb, _)) => cfb.has_directory("_VBA_PROJECT_CUR"),
            Some(CfbWrap::Vba(_)) => true,
            None => unreachable!(), // option is used to transfer ownership only
        }
    }

    fn vba_project(&mut self) -> Result<Cow<VbaProject>> {
        if let Some(CfbWrap::Wb(..)) = self.cfb {
            match self.cfb.take() {
                Some(CfbWrap::Wb(cfb, mut r)) => {
                    let vba = try!(VbaProject::from_cfb(&mut r, cfb));
                    self.cfb = Some(CfbWrap::Vba(vba));
                },
                _ => unreachable!(),
            }
        }

        match self.cfb {
            Some(CfbWrap::Vba(ref v)) => Ok(Cow::Borrowed(v)),
            _ => unreachable!(),
        }
    }

    /// Parses Workbook stream, no need for the relationships variable
    fn read_sheets_names(&mut self, _: &HashMap<Vec<u8>, String>)
        -> Result<HashMap<String, String>> {
        Ok(self.sheets.keys().map(|k| (k.to_string(), k.to_string())).collect())
    }

    fn read_shared_strings(&mut self) -> Result<Vec<String>> {
        Ok(self.strings.clone())
    }

    fn read_relationships(&mut self) -> Result<HashMap<Vec<u8>, String>> {
        Ok(HashMap::new())
    }

    fn read_worksheet_range(&mut self, name: &str, _: &[String]) -> Result<Range> {
        match self.sheets.get(name) {
            None => Err(format!("Sheet '{}' does not exist", name).into()),
            Some(r) => Ok(r.clone()),
        }
    }
}

impl Xls {
    fn parse_workbook(&mut self, stream: &[u8]) -> Result<()> {

        let mut sheets = Vec::new();
        {
            let mut wb = stream;
            let records = RecordIter { stream: &mut wb };
            for record in records {
                let r = try!(record);
                match r.typ {
                    0x0009 => if read_u16(&r.data[2..]) != 0x0005 {
                        return Err("Expecting Workbook BOF".into());
                    }, // BOF,
                    0x0042 => self.codepage = read_u16(&r.data), // CodePage (defines encoding)
                    0x013D => {
                        let sheet_len = r.data.len() / 2;
                        sheets.reserve(sheet_len);
                    }, // RRTabId
                    0x0085 => sheets.push(try!(parse_sheet_name(&r.data))), // BoundSheet8
                    0x00FC => self.strings = try!(parse_sst(&r.data)), // SST
                    0x000A => break, // EOF,
                    _ => (),
                }
            }
        }

        'sh: for (pos, name) in sheets.into_iter() {
            let mut sh = &stream[pos..];
            let records = RecordIter { stream: &mut sh };
            let mut range = Range::default();
            for record in records {
                let r = try!(record);
                match r.typ {
                    0x0009 => if read_u16(&r.data[2..]) != 0x0010 { continue 'sh; }, // BOF, worksheet
                    0x0200 => range = try!(parse_dimensions(&r.data)), // Dimensions
                    0x0203 => try!(parse_number(&r.data, &mut range)), // Number
                    0x0205 => try!(parse_bool_err(&r.data, &mut range)), // BoolErr
                    0x027E => try!(parse_rk(&r.data, &mut range)), // RK
                    0x00FD => try!(parse_label_sst(&r.data, &self.strings, &mut range)), // LabelSst
                    0x000A => break, // EOF,
                    _ => (),
                }
            }
            self.sheets.insert(name, range);
        }
        Ok(())
    }
}

fn parse_sheet_name(r: &[u8]) -> Result<(usize, String)> {
    let pos = read_u32(r) as usize;
    let sheet = try!(parse_short_string(&r[6..]));
    Ok((pos, sheet))
}

fn parse_sst(r: &[u8]) -> Result<Vec<String>> {
    if r.len() < 8 {
        return Err("Invalid sst length".into());
    }
    let len = read_slice::<i32>(&r[4..]) as usize;
    let mut sst = Vec::with_capacity(len);
    let mut read = &mut &r[8..];
    for _ in 0..len {
        sst.push(try!(read_rich_extended_string(read)));
    }
    Ok(sst)
}

fn parse_number(r: &[u8], range: &mut Range) -> Result<()> {
    if r.len() < 14 {
        return Err("Invalid number length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);
    let v = read_slice::<f64>(&r[6..]);
    range.set_value((row as u32, col as u32), DataType::Float(v))
}

fn parse_bool_err(r: &[u8], range: &mut Range) -> Result<()> {
    if r.len() < 8 {
        return Err("Invalid BoolErr length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);
    let v = match r[7] {
        0x00 => DataType::Bool(r[6] != 0),
        0x01 => match r[6] {
            0x00 => DataType::Error(CellErrorType::Null),
            0x07 => DataType::Error(CellErrorType::Div0),
            0x0F => DataType::Error(CellErrorType::Value),
            0x17 => DataType::Error(CellErrorType::Ref),
            0x1D => DataType::Error(CellErrorType::Name),
            0x24 => DataType::Error(CellErrorType::Num),
            0x2A => DataType::Error(CellErrorType::NA),
            0x2B => DataType::Error(CellErrorType::GettingData),
            e => return Err(format!("Unrecognized error {:x}", e).into()),
        },
        e => return Err(format!("Unrecognized fError {:x}", e).into()),
    };
    range.set_value((row as u32, col as u32), v)
}

fn parse_rk(r: &[u8], range: &mut Range) -> Result<()> {
    if r.len() < 10 {
        return Err("Invalid rk length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);

    let d100 = (r[6] & 1) != 0;
    let is_int = (r[6] & 2) != 0;

    let mut v = [0u8; 8];
    v[4..].copy_from_slice(&r[6..10]);
    v[0] &= 0xFC;
    let v = if is_int {
        let v = (read_slice::<i32>(&v[4..]) >> 2) as i64;
        DataType::Int( if d100 { v/100 } else { v })
    } else {
        let v = read_slice(&v);
        DataType::Float( if d100 { v/100.0 } else { v })
    };

    range.set_value((row as u32, col as u32), v)
}

fn parse_short_string(r: &[u8]) -> Result<String> {
    if r.len() < 2 {
        return Err("Invalid short string length".into());
    }
    let mut len = r[0] as usize;
    let high_byte = r[1];
    if high_byte != 0 { len *= 2; }
    if r.len() < 2 + len {
        return Err(format!("Invalid short string length: {} < {}", r.len(), 2 + len).into());
    }
    let bytes = if high_byte == 0 {
        // add 0x00 high bytes to unicodes
        let mut bytes = vec![0; len * 2];
        for (i, sce) in r[2..2 + len].iter().enumerate() {
            bytes[2 * i] = *sce;
        }
        Cow::Owned(bytes)
    } else {
        len *= 2;
        Cow::Borrowed(&r[2..2 + len])
    };
    UTF_16LE.decode(&bytes, DecoderTrap::Ignore).map_err(|e| e.to_string().into())
}

fn parse_label_sst(r: &[u8], strings: &[String], range: &mut Range) -> Result<()> {
    if r.len() < 10 {
        return Err("Invalid short string length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);
    let i = read_u32(&r[6..]) as usize;
    range.set_value((row as u32, col as u32), DataType::String(strings[i].clone()))
}

fn parse_dimensions(r: &[u8]) -> Result<Range> {
    if r.len() != 14 {
        return Err("Invalid dimensions lengths".into());
    }
    let rw_first = read_u32(&r[0..4]);
    let rw_last = read_u32(&r[4..8]);
    let col_first = read_u16(&r[8..10]);
    let col_last = read_u16(&r[10..12]);

    let start = (rw_first, col_first as u32);
    let size = ((rw_last - rw_first) as usize, (col_last - col_first) as usize);

    Ok(Range::new(start, size))
}

fn read_rich_extended_string(r: &mut &[u8]) -> Result<String> {
    if r.len() < 3 {
        return Err("Invalid rich extended string length".into());
    }

    let mut len = read_u16(r) as usize;
    let mut start = 3;
    let flags = r[2];
    let high_byte = flags & 0x1;
    let ext_st = flags & 0x4;
    let rich_st = flags & 0x8;
    let str_len = len;

    if high_byte != 0 { len += str_len; }
    if rich_st != 0 {
        len += 4 * read_u16(&r[start..]) as usize;
        start += 2;
    }
    if ext_st != 0 {
        len += read_slice::<i32>(&r[start..]) as usize;
        start += 4;
    }
    if r.len() < start + len {
        return Err(format!("Invalid rich extended string length: {} < {}",
                           r.len(), start + len).into());
    }

    let bytes = if high_byte == 0 {
        // add 0x00 high bytes to unicodes
        let mut bytes = vec![0; str_len * 2];
        for (i, sce) in r[start..start + str_len].iter().enumerate() {
            bytes[2 * i] = *sce;
        }
        Cow::Owned(bytes)
    } else {
        Cow::Borrowed(&r[start..start + 2 * str_len])
    };

    *r = &r[start + len..];
    UTF_16LE.decode(&bytes, DecoderTrap::Ignore).map_err(|e| e.to_string().into())
}

struct Record<'a> {
    typ: u16,
    data: Cow<'a, [u8]>,
}

struct RecordIter<'a> {
    stream: &'a[u8],
}

impl<'a> Iterator for RecordIter<'a> {
    type Item=Result<Record<'a>>;
    fn next(&mut self) -> Option<Self::Item> {
        if self.stream.len() < 4 {
            return if self.stream.is_empty() {
                None
            } else {
                Some(Err("Expecting record type and length, found end of stream".into()))
            }
        }
        let t = read_u16(self.stream);
        let mut len = read_u16(&self.stream[2..]) as usize;
        if self.stream.len() < len + 4 {
            return Some(Err("Expecting record length, found end of stream".into()));
        }
        let (mut data, mut next) = self.stream.split_at(len + 4);
        self.stream = next;

        // Append next record data if it is a Continue record
        let cow = if next.len() > 4 && read_u16(next) == 0x003C {
            let mut c = data[4..].to_vec();
            while next.len() > 4 && read_u16(next) == 0x003C {
                len = read_u16(&self.stream[2..]) as usize;
                if self.stream.len() < len + 4 {
                    return Some(Err("Expecting continue record length, found end of stream".into()));
                }
                let sp = self.stream.split_at(len + 4);
                data = sp.0;
                next = sp.1;
                c.extend_from_slice(&data[4..]);
                self.stream = next;
            }
            Cow::Owned(c)
        } else {
            Cow::Borrowed(&data[4..])
        };

        Some(Ok(Record { typ: t, data: cow }))
    }
}
<commit_msg>first byte or continue record is a flag for dbcs<commit_after>use std::fs::File;
use std::collections::HashMap;
use std::io::BufReader;
use std::borrow::Cow;
use std::cmp::min;

use encoding::{Encoding, DecoderTrap};
use encoding::all::UTF_16LE;

use errors::*;
use {ExcelReader, Range, DataType, CellErrorType};
use vba::VbaProject;
use cfb::Cfb;
use utils::{read_u16, read_u32, read_slice};

enum CfbWrap {
    Wb(Cfb, BufReader<File>),
    Vba(VbaProject),
}

/// A struct representing an old xls format file (CFB)
pub struct Xls {
    codepage: u16,
    sheets: HashMap<String, Range>,
    strings: Vec<String>,
    cfb: Option<CfbWrap>,
}

impl ExcelReader for Xls {

    fn new(r: File) -> Result<Self> {
        let len = try!(r.metadata()).len() as usize;
        let mut r = BufReader::new(r);
        let mut cfb = try!(Cfb::new(&mut r, len ));
        let wb = try!(cfb.get_stream("Workbook", &mut r));

        let mut xls = Xls {
            codepage: 1200,
            sheets: HashMap::new(),
            strings: Vec::new(),
            cfb: Some(CfbWrap::Wb(cfb, r)),
        };
        try!(xls.parse_workbook(&wb));
        Ok(xls)
    }

    fn has_vba(&mut self) -> bool {
        match self.cfb {
            Some(CfbWrap::Wb(ref cfb, _)) => cfb.has_directory("_VBA_PROJECT_CUR"),
            Some(CfbWrap::Vba(_)) => true,
            None => unreachable!(), // option is used to transfer ownership only
        }
    }

    fn vba_project(&mut self) -> Result<Cow<VbaProject>> {
        if let Some(CfbWrap::Wb(..)) = self.cfb {
            match self.cfb.take() {
                Some(CfbWrap::Wb(cfb, mut r)) => {
                    let vba = try!(VbaProject::from_cfb(&mut r, cfb));
                    self.cfb = Some(CfbWrap::Vba(vba));
                },
                _ => unreachable!(),
            }
        }

        match self.cfb {
            Some(CfbWrap::Vba(ref v)) => Ok(Cow::Borrowed(v)),
            _ => unreachable!(),
        }
    }

    /// Parses Workbook stream, no need for the relationships variable
    fn read_sheets_names(&mut self, _: &HashMap<Vec<u8>, String>)
        -> Result<HashMap<String, String>> {
        Ok(self.sheets.keys().map(|k| (k.to_string(), k.to_string())).collect())
    }

    fn read_shared_strings(&mut self) -> Result<Vec<String>> {
        Ok(self.strings.clone())
    }

    fn read_relationships(&mut self) -> Result<HashMap<Vec<u8>, String>> {
        Ok(HashMap::new())
    }

    fn read_worksheet_range(&mut self, name: &str, _: &[String]) -> Result<Range> {
        match self.sheets.get(name) {
            None => Err(format!("Sheet '{}' does not exist", name).into()),
            Some(r) => Ok(r.clone()),
        }
    }
}

impl Xls {
    fn parse_workbook(&mut self, stream: &[u8]) -> Result<()> {

        let mut sheets = Vec::new();
        {
            let mut wb = stream;
            let records = RecordIter { stream: &mut wb };
            for record in records {
                let mut r = try!(record);
                match r.typ {
                    0x0009 => if read_u16(&r.data[2..]) != 0x0005 {
                        return Err("Expecting Workbook BOF".into());
                    }, // BOF,
                    0x0042 => self.codepage = read_u16(&r.data), // CodePage (defines encoding)
                    0x013D => {
                        let sheet_len = r.data.len() / 2;
                        sheets.reserve(sheet_len);
                    }, // RRTabId
                    0x0085 => sheets.push(try!(parse_sheet_name(&r.data))), // BoundSheet8
                    0x00FC => self.strings = try!(parse_sst(&mut r)), // SST
                    0x000A => break, // EOF,
                    _ => (),
                }
            }
        }

        'sh: for (pos, name) in sheets.into_iter() {
            let mut sh = &stream[pos..];
            let records = RecordIter { stream: &mut sh };
            let mut range = Range::default();
            for record in records {
                let r = try!(record);
                match r.typ {
                    0x0009 => if read_u16(&r.data[2..]) != 0x0010 { continue 'sh; }, // BOF, worksheet
                    0x0200 => range = try!(parse_dimensions(&r.data)), // Dimensions
                    0x0203 => try!(parse_number(&r.data, &mut range)), // Number
                    0x0205 => try!(parse_bool_err(&r.data, &mut range)), // BoolErr
                    0x027E => try!(parse_rk(&r.data, &mut range)), // RK
                    0x00FD => try!(parse_label_sst(&r.data, &self.strings, &mut range)), // LabelSst
                    0x000A => break, // EOF,
                    _ => (),
                }
            }
            self.sheets.insert(name, range);
        }
        Ok(())
    }
}

fn parse_sheet_name(r: &[u8]) -> Result<(usize, String)> {
    let pos = read_u32(r) as usize;
    let sheet = try!(parse_short_string(&r[6..]));
    Ok((pos, sheet))
}

fn parse_number(r: &[u8], range: &mut Range) -> Result<()> {
    if r.len() < 14 {
        return Err("Invalid number length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);
    let v = read_slice::<f64>(&r[6..]);
    range.set_value((row as u32, col as u32), DataType::Float(v))
}

fn parse_bool_err(r: &[u8], range: &mut Range) -> Result<()> {
    if r.len() < 8 {
        return Err("Invalid BoolErr length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);
    let v = match r[7] {
        0x00 => DataType::Bool(r[6] != 0),
        0x01 => match r[6] {
            0x00 => DataType::Error(CellErrorType::Null),
            0x07 => DataType::Error(CellErrorType::Div0),
            0x0F => DataType::Error(CellErrorType::Value),
            0x17 => DataType::Error(CellErrorType::Ref),
            0x1D => DataType::Error(CellErrorType::Name),
            0x24 => DataType::Error(CellErrorType::Num),
            0x2A => DataType::Error(CellErrorType::NA),
            0x2B => DataType::Error(CellErrorType::GettingData),
            e => return Err(format!("Unrecognized error {:x}", e).into()),
        },
        e => return Err(format!("Unrecognized fError {:x}", e).into()),
    };
    range.set_value((row as u32, col as u32), v)
}

fn parse_rk(r: &[u8], range: &mut Range) -> Result<()> {
    if r.len() < 10 {
        return Err("Invalid rk length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);

    let d100 = (r[6] & 1) != 0;
    let is_int = (r[6] & 2) != 0;

    let mut v = [0u8; 8];
    v[4..].copy_from_slice(&r[6..10]);
    v[0] &= 0xFC;
    let v = if is_int {
        let v = (read_slice::<i32>(&v[4..]) >> 2) as i64;
        DataType::Int( if d100 { v/100 } else { v })
    } else {
        let v = read_slice(&v);
        DataType::Float( if d100 { v/100.0 } else { v })
    };

    range.set_value((row as u32, col as u32), v)
}

fn parse_short_string(r: &[u8]) -> Result<String> {
    if r.len() < 2 {
        return Err("Invalid short string length".into());
    }
    let mut len = r[0] as usize;
    let high_byte = r[1];
    if high_byte != 0 { len *= 2; }
    if r.len() < 2 + len {
        return Err(format!("Invalid short string length: {} < {}", r.len(), 2 + len).into());
    }
    let bytes = if high_byte == 0 {
        // add 0x00 high bytes to unicodes
        let mut bytes = vec![0; len * 2];
        for (i, sce) in r[2..2 + len].iter().enumerate() {
            bytes[2 * i] = *sce;
        }
        Cow::Owned(bytes)
    } else {
        len *= 2;
        Cow::Borrowed(&r[2..2 + len])
    };
    UTF_16LE.decode(&bytes, DecoderTrap::Ignore).map_err(|e| e.to_string().into())
}

fn parse_label_sst(r: &[u8], strings: &[String], range: &mut Range) -> Result<()> {
    if r.len() < 10 {
        return Err("Invalid short string length".into());
    }
    let row = read_u16(r);
    let col = read_u16(&r[2..]);
    let i = read_u32(&r[6..]) as usize;
    range.set_value((row as u32, col as u32), DataType::String(strings[i].clone()))
}

fn parse_dimensions(r: &[u8]) -> Result<Range> {
    if r.len() != 14 {
        return Err("Invalid dimensions lengths".into());
    }
    let rw_first = read_u32(&r[0..4]);
    let rw_last = read_u32(&r[4..8]);
    let col_first = read_u16(&r[8..10]);
    let col_last = read_u16(&r[10..12]);

    let start = (rw_first, col_first as u32);
    let size = ((rw_last - rw_first) as usize, (col_last - col_first) as usize);

    Ok(Range::new(start, size))
}

fn parse_sst(r: &mut Record) -> Result<Vec<String>> {
    if r.data.len() < 8 {
        return Err("Invalid sst length".into());
    }
    let len = read_slice::<i32>(&r.data[4..]) as usize;
    let mut sst = Vec::with_capacity(len);
    r.data = &r.data[8..];
    for i in 0..len {
        sst.push(try!(read_rich_extended_string(r)));
    }
    Ok(sst)
}

fn read_rich_extended_string(r: &mut Record) -> Result<String> {
    if r.data.is_empty() && !r.continue_record() || r.data.len() < 3 {
        return Err("Invalid rich extended string length".into());
    }

    let str_len = read_u16(r.data) as usize;
    let flags = r.data[2];
    r.data = &r.data[3..];
    let high_byte = flags & 0x1;
    let ext_st = flags & 0x4;
    let rich_st = flags & 0x8;

    let mut unused_len = if rich_st != 0 {
        let l = 4 * read_u16(&r.data) as usize;
        r.data = &r.data[2..];
        l
    } else {
        0
    };
    if ext_st != 0 {
        unused_len += read_slice::<i32>(&r.data) as usize;
        r.data = &r.data[4..];
    };

    let s = try!(read_dbcs(high_byte, str_len, r));

    while unused_len > 0 {
        if r.data.is_empty() && !r.continue_record() {
            return Err("continued record too short while reading extended string".into());
        }
        let l = min(unused_len, r.data.len());
        let (_, next) = r.data.split_at(l);
        r.data = next;
        unused_len -= l;
    }

    Ok(s)
}

fn read_dbcs<'a>(mut high_byte: u8, mut len: usize, r: &mut Record) -> Result<String>
{
    let mut s = String::with_capacity(len);
    while len > 0 {
        let (l, bytes) = if high_byte == 0 {
            let l = min(r.data.len(), len);
            let (data, next) = r.data.split_at(l);
            r.data = next;

            // add 0x00 high bytes to unicodes
            let mut bytes = vec![0; l * 2];
            for (i, sce) in data.iter().enumerate() {
                bytes[2 * i] = *sce;
            }
            (l, Cow::Owned(bytes))
        } else {
            let l = min(r.data.len() / 2, len);
            let (data, next) = r.data.split_at(2 * l);
            r.data = next;
            (l, Cow::Borrowed(data))
        };

        let s_rec: Result<String> = UTF_16LE.decode(&bytes, DecoderTrap::Ignore)
            .map_err(|e| e.to_string().into());
        s.push_str(&try!(s_rec));

        len -= l;
        if len > 0 {
           if r.continue_record() {
               high_byte = r.data[0] & 0x1;
               r.data = &r.data[1..];
           } else {
               return Err("Cannot decode entire dbcs stream".into());
           }
        }
    }
    Ok(s)
}

struct Record<'a> {
    typ: u16,
    data: &'a[u8],
    cont: Option<Vec<&'a[u8]>>,
}

impl<'a> Record<'a> {
    fn continue_record(&mut self) -> bool {
        match self.cont {
            None => false,
            Some(ref mut v) => {
                if v.is_empty() {
                    false
                } else {
                    self.data = v.remove(0);
                    true
                }
            }
        }
    }
}

struct RecordIter<'a> {
    stream: &'a[u8],
}

impl<'a> Iterator for RecordIter<'a> {
    type Item=Result<Record<'a>>;
    fn next(&mut self) -> Option<Self::Item> {
        if self.stream.len() < 4 {
            return if self.stream.is_empty() {
                None
            } else {
                Some(Err("Expecting record type and length, found end of stream".into()))
            }
        }
        let t = read_u16(self.stream);
        let mut len = read_u16(&self.stream[2..]) as usize;
        if self.stream.len() < len + 4 {
            return Some(Err("Expecting record length, found end of stream".into()));
        }
        let (data, next) = self.stream.split_at(len + 4);
        self.stream = next;
        let d = &data[4..];

        // Append next record data if it is a Continue record
        let cont = if next.len() > 4 && read_u16(next) == 0x003C {
            let mut cont = Vec::new();
            while self.stream.len() > 4 && read_u16(self.stream) == 0x003C {
                len = read_u16(&self.stream[2..]) as usize;
                if self.stream.len() < len + 4 {
                    return Some(Err("Expecting continue record length, found end of stream".into()));
                }
                let sp = self.stream.split_at(len + 4);
                cont.push(&sp.0[4..]);
                self.stream = sp.1;
            }
            Some(cont)
        } else {
            None
        };

        Some(Ok(Record { typ: t, data: d, cont: cont }))
    }
}
<|endoftext|>


<commit_before>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

use std::fmt::{self, Formatter, Debug};
use std::mem::{size_of, zeroed};
use std::path::{Path, PathBuf};
use std::str;

use libc::{c_uint, c_void, memcpy};

use Pid;
use ProcessExt;

use winapi::shared::minwindef::{DWORD, FALSE, FILETIME, MAX_PATH/*, TRUE, USHORT*/};
use winapi::um::handleapi::CloseHandle;
use winapi::um::winnt::{
    HANDLE, ULARGE_INTEGER, /*THREAD_GET_CONTEXT, THREAD_QUERY_INFORMATION, THREAD_SUSPEND_RESUME,*/
    /*, PWSTR*/ PROCESS_QUERY_INFORMATION, PROCESS_TERMINATE, PROCESS_VM_READ,
};
use winapi::um::processthreadsapi::{GetProcessTimes, OpenProcess, TerminateProcess};
use winapi::um::psapi::{
    GetProcessMemoryInfo, PROCESS_MEMORY_COUNTERS, PROCESS_MEMORY_COUNTERS_EX,
    EnumProcessModulesEx, GetModuleBaseNameW, LIST_MODULES_ALL,
};
use winapi::um::sysinfoapi::GetSystemTimeAsFileTime;
use winapi::um::tlhelp32::{
    CreateToolhelp32Snapshot, Process32First, Process32Next, PROCESSENTRY32, TH32CS_SNAPPROCESS,
};

/// Enum describing the different status of a process.
#[derive(Clone, Copy, Debug)]
pub enum ProcessStatus {
    /// Currently runnable.
    Run,
}

impl ProcessStatus {
    /// Used to display `ProcessStatus`.
    pub fn to_string(&self) -> &str {
        match *self {
            ProcessStatus::Run => "Runnable",
        }
    }
}

impl fmt::Display for ProcessStatus {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.to_string())
    }
}

fn get_process_handler(pid: Pid) -> Option<HANDLE> {
    if pid == 0 {
        return None;
    }
    let options = PROCESS_QUERY_INFORMATION | PROCESS_VM_READ | PROCESS_TERMINATE;
    let process_handler = unsafe { OpenProcess(options, FALSE, pid as DWORD) };
    if process_handler.is_null() {
        let options = PROCESS_QUERY_INFORMATION | PROCESS_VM_READ;
        let process_handler = unsafe { OpenProcess(options, FALSE, pid as DWORD) };
        if process_handler.is_null() {
            None
        } else {
            Some(process_handler)
        }
    } else {
        Some(process_handler)
    }
}

/// Struct containing a process' information.
#[derive(Clone)]
pub struct Process {
    name: String,
    cmd: Vec<String>,
    exe: PathBuf,
    pid: Pid,
    environ: Vec<String>,
    cwd: PathBuf,
    root: PathBuf,
    memory: u64,
    parent: Option<Pid>,
    status: ProcessStatus,
    handle: HANDLE,
    old_cpu: u64,
    old_sys_cpu: u64,
    old_user_cpu: u64,
    start_time: u64,
    cpu_usage: f32,
}

// TODO: it's possible to get environment variables like it's done in
// https://github.com/processhacker/processhacker
//
// They have a very nice function called PhGetProcessEnvironment. Just very complicated as it
// seems...
impl ProcessExt for Process {
    fn new(pid: Pid, parent: Option<Pid>, _: u64) -> Process {
        if let Some(process_handler) = get_process_handler(pid) {
            let mut h_mod = ::std::ptr::null_mut();
            let mut process_name = [0u16; MAX_PATH + 1];
            let mut cb_needed = 0;

            unsafe {
                if EnumProcessModulesEx(process_handler,
                                        &mut h_mod,
                                        ::std::mem::size_of::<DWORD>() as DWORD,
                                        &mut cb_needed,
                                        LIST_MODULES_ALL) != 0 {
                    GetModuleBaseNameW(process_handler,
                                       h_mod,
                                       process_name.as_mut_ptr(),
                                       MAX_PATH as DWORD + 1);
                }
                let mut pos = 0;
                for x in process_name.iter() {
                    if *x == 0 {
                        break
                    }
                    pos += 1;
                }
                let name = String::from_utf16_lossy(&process_name[..pos]);
                let environ = get_proc_env(process_handler, pid as u32, &name);
                Process {
                    handle: process_handler,
                    name: name,
                    pid: pid,
                    parent: parent,
                    cmd: get_cmd_line(pid),
                    environ: environ,
                    exe: get_executable_path(pid),
                    cwd: PathBuf::new(),
                    root: PathBuf::new(),
                    status: ProcessStatus::Run,
                    memory: 0,
                    cpu_usage: 0.,
                    old_cpu: 0,
                    old_sys_cpu: 0,
                    old_user_cpu: 0,
                    start_time: get_start_time(process_handler),
                }
            }
        } else {
            Process {
                handle: ::std::ptr::null_mut(),
                name: String::new(),
                pid: pid,
                parent: parent,
                cmd: get_cmd_line(pid),
                environ: Vec::new(),
                exe: get_executable_path(pid),
                cwd: PathBuf::new(),
                root: PathBuf::new(),
                status: ProcessStatus::Run,
                memory: 0,
                cpu_usage: 0.,
                old_cpu: 0,
                old_sys_cpu: 0,
                old_user_cpu: 0,
                start_time: 0,
            }
        }
    }

    fn kill(&self, signal: ::Signal) -> bool {
        let x = unsafe { TerminateProcess(self.handle, signal as c_uint) };
        x != 0
    }

    fn name(&self) -> &str {
        &self.name
    }

    fn cmd(&self) -> &[String] {
        &self.cmd
    }

    fn exe(&self) -> &Path {
        self.exe.as_path()
    }

    fn pid(&self) -> Pid {
        self.pid
    }

    fn environ(&self) -> &[String] {
        &self.environ
    }

    fn cwd(&self) -> &Path {
        self.cwd.as_path()
    }

    fn root(&self) -> &Path {
        self.root.as_path()
    }

    fn memory(&self) -> u64 {
        self.memory
    }

    fn parent(&self) -> Option<Pid> {
        self.parent
    }

    fn status(&self) -> ProcessStatus {
        self.status
    }

    fn start_time(&self) -> u64 {
        self.start_time
    }

    fn cpu_usage(&self) -> f32 {
        self.cpu_usage
    }
}

impl Drop for Process {
    fn drop(&mut self) {
        unsafe {
            if self.handle.is_null() {
                return
            }
            CloseHandle(self.handle);
        }
    }
}

#[allow(unused_must_use)]
impl Debug for Process {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        writeln!(f, "pid: {}", self.pid);
        writeln!(f, "name: {}", self.name);
        writeln!(f, "environment:");
        for var in self.environ.iter() {
        if var.len() > 0 {
                writeln!(f, "\t{}", var);
            }
        }
        writeln!(f, "command:");
        for arg in &self.cmd {
            writeln!(f, "\t{}", arg);
        }
        writeln!(f, "executable path: {:?}", self.exe);
        writeln!(f, "current working directory: {:?}", self.cwd);
        writeln!(f, "memory usage: {} kB", self.memory);
        writeln!(f, "cpu usage: {}", self.cpu_usage);
        writeln!(f, "root path: {:?}", self.root)
    }
}

unsafe fn get_start_time(handle: HANDLE) -> u64 {
    let mut fstart: FILETIME = zeroed();
    let mut x = zeroed();

    GetProcessTimes(handle,
                    &mut fstart as *mut FILETIME,
                    &mut x as *mut FILETIME,
                    &mut x as *mut FILETIME,
                    &mut x as *mut FILETIME);
    let tmp = (fstart.dwHighDateTime as u64) << 32 | (fstart.dwLowDateTime as u64);
    tmp / 10_000_000 - 11_644_473_600
}

pub unsafe fn get_parent_process_id(pid: Pid) -> Option<Pid> {
    let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    let mut entry: PROCESSENTRY32 = zeroed();
    entry.dwSize = size_of::<PROCESSENTRY32>() as u32;
    let mut not_the_end = Process32First(snapshot, &mut entry);
    while not_the_end != 0 {
        if pid == entry.th32ProcessID as usize {
            // TODO: if some day I have the motivation to add threads:
            // ListProcessThreads(entry.th32ProcessID);
            CloseHandle(snapshot);
            return Some(entry.th32ParentProcessID as usize);
        }
        not_the_end = Process32Next(snapshot, &mut entry);
    }
    CloseHandle(snapshot);
    None
}

/*fn run_wmi(args: &[&str]) -> Option<String> {
    use std::process::Command;

    if let Ok(out) = Command::new("wmic")
                             .args(args)
                             .output() {
        if out.status.success() {
            return Some(String::from_utf8_lossy(&out.stdout).into_owned());
        }
    }
    None
}*/

fn get_cmd_line(_pid: Pid) -> Vec<String> {
    /*let where_req = format!("ProcessId={}", pid);

    if let Some(ret) = run_wmi(&["process", "where", &where_req, "get", "CommandLine"]) {
        for line in ret.lines() {
            if line.is_empty() || line == "CommandLine" {
                continue
            }
            return vec![line.to_owned()];
        }
    }*/
    Vec::new()
}

unsafe fn get_proc_env(_handle: HANDLE, _pid: u32, _name: &str) -> Vec<String> {
    let ret = Vec::new();
    /*
    println!("current pid: {}", kernel32::GetCurrentProcessId());
    if kernel32::GetCurrentProcessId() == pid {
        println!("current proc!");
        for (key, value) in env::vars() {
            ret.push(format!("{}={}", key, value));
        }
        return ret;
    }
    println!("1");
    let snapshot_handle = kernel32::CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
    if !snapshot_handle.is_null() {
        println!("2");
        let mut target_thread: THREADENTRY32 = zeroed();
        target_thread.dwSize = size_of::<THREADENTRY32>() as DWORD;
        if kernel32::Thread32First(snapshot_handle, &mut target_thread) == TRUE {
            println!("3");
            loop {
                if target_thread.th32OwnerProcessID == pid {
                    println!("4");
                    let thread_handle = kernel32::OpenThread(THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION | THREAD_GET_CONTEXT,
                                                             FALSE,
                                                             target_thread.th32ThreadID);
                    if !thread_handle.is_null() {
                        println!("5 -> {}", pid);
                        if kernel32::SuspendThread(thread_handle) != DWORD::max_value() {
                            println!("6");
                            let mut context = zeroed();
                            if kernel32::GetThreadContext(thread_handle, &mut context) != 0 {
                                println!("7 --> {:?}", context);
                                let mut x = vec![0u8; 10];
                                if kernel32::ReadProcessMemory(handle,
                                                               context.MxCsr as usize as *mut winapi::c_void,
                                                               x.as_mut_ptr() as *mut winapi::c_void,
                                                               x.len() as u64,
                                                               ::std::ptr::null_mut()) != 0 {
                                    for y in x {
                                        print!("{}", y as char);
                                    }
                                    println!("");
                                } else {
                                    println!("failure... {:?}", kernel32::GetLastError());
                                }
                            } else {
                                println!("-> {:?}", kernel32::GetLastError());
                            }
                            kernel32::ResumeThread(thread_handle);
                        }
                        kernel32::CloseHandle(thread_handle);
                    }
                    break;
                }
                if kernel32::Thread32Next(snapshot_handle, &mut target_thread) != TRUE {
                    break;
                }
            }
        }
        kernel32::CloseHandle(snapshot_handle);
    }*/
    ret
}

pub fn get_executable_path(_pid: Pid) -> PathBuf {
    /*let where_req = format!("ProcessId={}", pid);

    if let Some(ret) = run_wmi(&["process", "where", &where_req, "get", "ExecutablePath"]) {
        for line in ret.lines() {
            if line.is_empty() || line == "ExecutablePath" {
                continue
            }
            return line.to_owned();
        }
    }*/
    PathBuf::new()
}

pub fn compute_cpu_usage(p: &mut Process, nb_processors: u64) {
    unsafe {
        let mut now: ULARGE_INTEGER = ::std::mem::zeroed();
        let mut sys: ULARGE_INTEGER = ::std::mem::zeroed();
        let mut user: ULARGE_INTEGER = ::std::mem::zeroed();
        let mut ftime: FILETIME = zeroed();
        let mut fsys: FILETIME = zeroed();
        let mut fuser: FILETIME = zeroed();

        GetSystemTimeAsFileTime(&mut ftime);
        memcpy(&mut now as *mut ULARGE_INTEGER as *mut c_void,
               &mut ftime as *mut FILETIME as *mut c_void,
               size_of::<FILETIME>());

        GetProcessTimes(p.handle,
                        &mut ftime as *mut FILETIME,
                        &mut ftime as *mut FILETIME,
                        &mut fsys as *mut FILETIME,
                        &mut fuser as *mut FILETIME);
        memcpy(&mut sys as *mut ULARGE_INTEGER as *mut c_void,
               &mut fsys as *mut FILETIME as *mut c_void,
               size_of::<FILETIME>());
        memcpy(&mut user as *mut ULARGE_INTEGER as *mut c_void,
               &mut fuser as *mut FILETIME as *mut c_void,
               size_of::<FILETIME>());
        p.cpu_usage = ((*sys.QuadPart() - p.old_sys_cpu) as f32 +
            (*user.QuadPart() - p.old_user_cpu) as f32)
            / (*now.QuadPart() - p.old_cpu) as f32
            / nb_processors as f32 * 100.;
        p.old_cpu = *now.QuadPart();
        p.old_user_cpu = *user.QuadPart();
        p.old_sys_cpu = *sys.QuadPart();
    }
}

pub fn get_handle(p: &Process) -> HANDLE {
    p.handle
}

pub fn update_proc_info(p: &mut Process) {
    update_memory(p);
}

pub fn update_memory(p: &mut Process) {
    unsafe {
        let mut pmc: PROCESS_MEMORY_COUNTERS_EX = zeroed();
        if GetProcessMemoryInfo(p.handle,
                                &mut pmc as *mut PROCESS_MEMORY_COUNTERS_EX as *mut c_void as *mut PROCESS_MEMORY_COUNTERS,
                                size_of::<PROCESS_MEMORY_COUNTERS_EX>() as DWORD) != 0 {
            p.memory = (pmc.PrivateUsage as u64) >> 10u64; // / 1024;
        }
    }
}
<commit_msg>Windows: get root directory of process<commit_after>//
// Sysinfo
//
// Copyright (c) 2018 Guillaume Gomez
//

use std::fmt::{self, Formatter, Debug};
use std::mem::{size_of, zeroed};
use std::path::{Path, PathBuf};
use std::str;

use libc::{c_uint, c_void, memcpy};

use Pid;
use ProcessExt;

use winapi::shared::minwindef::{DWORD, FALSE, FILETIME, MAX_PATH/*, TRUE, USHORT*/};
use winapi::um::handleapi::CloseHandle;
use winapi::um::winnt::{
    HANDLE, ULARGE_INTEGER, /*THREAD_GET_CONTEXT, THREAD_QUERY_INFORMATION, THREAD_SUSPEND_RESUME,*/
    /*, PWSTR*/ PROCESS_QUERY_INFORMATION, PROCESS_TERMINATE, PROCESS_VM_READ,
};
use winapi::um::processthreadsapi::{GetProcessTimes, OpenProcess, TerminateProcess};
use winapi::um::psapi::{
    GetProcessMemoryInfo, PROCESS_MEMORY_COUNTERS, PROCESS_MEMORY_COUNTERS_EX,
    EnumProcessModulesEx, GetModuleBaseNameW, GetModuleFileNameExW, LIST_MODULES_ALL,
};
use winapi::um::sysinfoapi::GetSystemTimeAsFileTime;
use winapi::um::tlhelp32::{
    CreateToolhelp32Snapshot, Process32First, Process32Next, PROCESSENTRY32, TH32CS_SNAPPROCESS,
};

/// Enum describing the different status of a process.
#[derive(Clone, Copy, Debug)]
pub enum ProcessStatus {
    /// Currently runnable.
    Run,
}

impl ProcessStatus {
    /// Used to display `ProcessStatus`.
    pub fn to_string(&self) -> &str {
        match *self {
            ProcessStatus::Run => "Runnable",
        }
    }
}

impl fmt::Display for ProcessStatus {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.to_string())
    }
}

fn get_process_handler(pid: Pid) -> Option<HANDLE> {
    if pid == 0 {
        return None;
    }
    let options = PROCESS_QUERY_INFORMATION | PROCESS_VM_READ | PROCESS_TERMINATE;
    let process_handler = unsafe { OpenProcess(options, FALSE, pid as DWORD) };
    if process_handler.is_null() {
        let options = PROCESS_QUERY_INFORMATION | PROCESS_VM_READ;
        let process_handler = unsafe { OpenProcess(options, FALSE, pid as DWORD) };
        if process_handler.is_null() {
            None
        } else {
            Some(process_handler)
        }
    } else {
        Some(process_handler)
    }
}

/// Struct containing a process' information.
#[derive(Clone)]
pub struct Process {
    name: String,
    cmd: Vec<String>,
    exe: PathBuf,
    pid: Pid,
    environ: Vec<String>,
    cwd: PathBuf,
    root: PathBuf,
    memory: u64,
    parent: Option<Pid>,
    status: ProcessStatus,
    handle: HANDLE,
    old_cpu: u64,
    old_sys_cpu: u64,
    old_user_cpu: u64,
    start_time: u64,
    cpu_usage: f32,
}

// TODO: it's possible to get environment variables like it's done in
// https://github.com/processhacker/processhacker
//
// They have a very nice function called PhGetProcessEnvironment. Just very complicated as it
// seems...
impl ProcessExt for Process {
    fn new(pid: Pid, parent: Option<Pid>, _: u64) -> Process {
        if let Some(process_handler) = get_process_handler(pid) {
            let mut h_mod = ::std::ptr::null_mut();
            let mut process_name = [0u16; MAX_PATH + 1];
            let mut cb_needed = 0;

            unsafe {
                if EnumProcessModulesEx(process_handler,
                                        &mut h_mod,
                                        ::std::mem::size_of::<DWORD>() as DWORD,
                                        &mut cb_needed,
                                        LIST_MODULES_ALL) != 0 {
                    GetModuleBaseNameW(process_handler,
                                       h_mod,
                                       process_name.as_mut_ptr(),
                                       MAX_PATH as DWORD + 1);
                }
                let mut pos = 0;
                for x in process_name.iter() {
                    if *x == 0 {
                        break
                    }
                    pos += 1;
                }
                let name = String::from_utf16_lossy(&process_name[..pos]);
                let environ = get_proc_env(process_handler, pid as u32, &name);

                let mut root_buf = [0u16; MAX_PATH + 1];
                GetModuleFileNameExW(process_handler,
                                     h_mod,
                                     root_buf.as_mut_ptr(),
                                     MAX_PATH as DWORD + 1);

                pos = 0;
                for x in root_buf.iter() {
                    if *x == 0 {
                        break
                    }
                    pos += 1;
                }

                let root = String::from_utf16_lossy(&root_buf[..pos]);
                Process {
                    handle: process_handler,
                    name: name,
                    pid: pid,
                    parent: parent,
                    cmd: get_cmd_line(pid),
                    environ: environ,
                    exe: get_executable_path(pid),
                    cwd: PathBuf::new(),
                    root: PathBuf::from(root),
                    status: ProcessStatus::Run,
                    memory: 0,
                    cpu_usage: 0.,
                    old_cpu: 0,
                    old_sys_cpu: 0,
                    old_user_cpu: 0,
                    start_time: get_start_time(process_handler),
                }
            }
        } else {
            Process {
                handle: ::std::ptr::null_mut(),
                name: String::new(),
                pid: pid,
                parent: parent,
                cmd: get_cmd_line(pid),
                environ: Vec::new(),
                exe: get_executable_path(pid),
                cwd: PathBuf::new(),
                root: PathBuf::new(),
                status: ProcessStatus::Run,
                memory: 0,
                cpu_usage: 0.,
                old_cpu: 0,
                old_sys_cpu: 0,
                old_user_cpu: 0,
                start_time: 0,
            }
        }
    }

    fn kill(&self, signal: ::Signal) -> bool {
        let x = unsafe { TerminateProcess(self.handle, signal as c_uint) };
        x != 0
    }

    fn name(&self) -> &str {
        &self.name
    }

    fn cmd(&self) -> &[String] {
        &self.cmd
    }

    fn exe(&self) -> &Path {
        self.exe.as_path()
    }

    fn pid(&self) -> Pid {
        self.pid
    }

    fn environ(&self) -> &[String] {
        &self.environ
    }

    fn cwd(&self) -> &Path {
        self.cwd.as_path()
    }

    fn root(&self) -> &Path {
        self.root.as_path()
    }

    fn memory(&self) -> u64 {
        self.memory
    }

    fn parent(&self) -> Option<Pid> {
        self.parent
    }

    fn status(&self) -> ProcessStatus {
        self.status
    }

    fn start_time(&self) -> u64 {
        self.start_time
    }

    fn cpu_usage(&self) -> f32 {
        self.cpu_usage
    }
}

impl Drop for Process {
    fn drop(&mut self) {
        unsafe {
            if self.handle.is_null() {
                return
            }
            CloseHandle(self.handle);
        }
    }
}

#[allow(unused_must_use)]
impl Debug for Process {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        writeln!(f, "pid: {}", self.pid);
        writeln!(f, "name: {}", self.name);
        writeln!(f, "environment:");
        for var in self.environ.iter() {
        if var.len() > 0 {
                writeln!(f, "\t{}", var);
            }
        }
        writeln!(f, "command:");
        for arg in &self.cmd {
            writeln!(f, "\t{}", arg);
        }
        writeln!(f, "executable path: {:?}", self.exe);
        writeln!(f, "current working directory: {:?}", self.cwd);
        writeln!(f, "memory usage: {} kB", self.memory);
        writeln!(f, "cpu usage: {}", self.cpu_usage);
        writeln!(f, "root path: {:?}", self.root)
    }
}

unsafe fn get_start_time(handle: HANDLE) -> u64 {
    let mut fstart: FILETIME = zeroed();
    let mut x = zeroed();

    GetProcessTimes(handle,
                    &mut fstart as *mut FILETIME,
                    &mut x as *mut FILETIME,
                    &mut x as *mut FILETIME,
                    &mut x as *mut FILETIME);
    let tmp = (fstart.dwHighDateTime as u64) << 32 | (fstart.dwLowDateTime as u64);
    tmp / 10_000_000 - 11_644_473_600
}

pub unsafe fn get_parent_process_id(pid: Pid) -> Option<Pid> {
    let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    let mut entry: PROCESSENTRY32 = zeroed();
    entry.dwSize = size_of::<PROCESSENTRY32>() as u32;
    let mut not_the_end = Process32First(snapshot, &mut entry);
    while not_the_end != 0 {
        if pid == entry.th32ProcessID as usize {
            // TODO: if some day I have the motivation to add threads:
            // ListProcessThreads(entry.th32ProcessID);
            CloseHandle(snapshot);
            return Some(entry.th32ParentProcessID as usize);
        }
        not_the_end = Process32Next(snapshot, &mut entry);
    }
    CloseHandle(snapshot);
    None
}

/*fn run_wmi(args: &[&str]) -> Option<String> {
    use std::process::Command;

    if let Ok(out) = Command::new("wmic")
                             .args(args)
                             .output() {
        if out.status.success() {
            return Some(String::from_utf8_lossy(&out.stdout).into_owned());
        }
    }
    None
}*/

fn get_cmd_line(_pid: Pid) -> Vec<String> {
    /*let where_req = format!("ProcessId={}", pid);

    if let Some(ret) = run_wmi(&["process", "where", &where_req, "get", "CommandLine"]) {
        for line in ret.lines() {
            if line.is_empty() || line == "CommandLine" {
                continue
            }
            return vec![line.to_owned()];
        }
    }*/
    Vec::new()
}

unsafe fn get_proc_env(_handle: HANDLE, _pid: u32, _name: &str) -> Vec<String> {
    let ret = Vec::new();
    /*
    println!("current pid: {}", kernel32::GetCurrentProcessId());
    if kernel32::GetCurrentProcessId() == pid {
        println!("current proc!");
        for (key, value) in env::vars() {
            ret.push(format!("{}={}", key, value));
        }
        return ret;
    }
    println!("1");
    let snapshot_handle = kernel32::CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
    if !snapshot_handle.is_null() {
        println!("2");
        let mut target_thread: THREADENTRY32 = zeroed();
        target_thread.dwSize = size_of::<THREADENTRY32>() as DWORD;
        if kernel32::Thread32First(snapshot_handle, &mut target_thread) == TRUE {
            println!("3");
            loop {
                if target_thread.th32OwnerProcessID == pid {
                    println!("4");
                    let thread_handle = kernel32::OpenThread(THREAD_SUSPEND_RESUME | THREAD_QUERY_INFORMATION | THREAD_GET_CONTEXT,
                                                             FALSE,
                                                             target_thread.th32ThreadID);
                    if !thread_handle.is_null() {
                        println!("5 -> {}", pid);
                        if kernel32::SuspendThread(thread_handle) != DWORD::max_value() {
                            println!("6");
                            let mut context = zeroed();
                            if kernel32::GetThreadContext(thread_handle, &mut context) != 0 {
                                println!("7 --> {:?}", context);
                                let mut x = vec![0u8; 10];
                                if kernel32::ReadProcessMemory(handle,
                                                               context.MxCsr as usize as *mut winapi::c_void,
                                                               x.as_mut_ptr() as *mut winapi::c_void,
                                                               x.len() as u64,
                                                               ::std::ptr::null_mut()) != 0 {
                                    for y in x {
                                        print!("{}", y as char);
                                    }
                                    println!("");
                                } else {
                                    println!("failure... {:?}", kernel32::GetLastError());
                                }
                            } else {
                                println!("-> {:?}", kernel32::GetLastError());
                            }
                            kernel32::ResumeThread(thread_handle);
                        }
                        kernel32::CloseHandle(thread_handle);
                    }
                    break;
                }
                if kernel32::Thread32Next(snapshot_handle, &mut target_thread) != TRUE {
                    break;
                }
            }
        }
        kernel32::CloseHandle(snapshot_handle);
    }*/
    ret
}

pub fn get_executable_path(_pid: Pid) -> PathBuf {
    /*let where_req = format!("ProcessId={}", pid);

    if let Some(ret) = run_wmi(&["process", "where", &where_req, "get", "ExecutablePath"]) {
        for line in ret.lines() {
            if line.is_empty() || line == "ExecutablePath" {
                continue
            }
            return line.to_owned();
        }
    }*/
    PathBuf::new()
}

pub fn compute_cpu_usage(p: &mut Process, nb_processors: u64) {
    unsafe {
        let mut now: ULARGE_INTEGER = ::std::mem::zeroed();
        let mut sys: ULARGE_INTEGER = ::std::mem::zeroed();
        let mut user: ULARGE_INTEGER = ::std::mem::zeroed();
        let mut ftime: FILETIME = zeroed();
        let mut fsys: FILETIME = zeroed();
        let mut fuser: FILETIME = zeroed();

        GetSystemTimeAsFileTime(&mut ftime);
        memcpy(&mut now as *mut ULARGE_INTEGER as *mut c_void,
               &mut ftime as *mut FILETIME as *mut c_void,
               size_of::<FILETIME>());

        GetProcessTimes(p.handle,
                        &mut ftime as *mut FILETIME,
                        &mut ftime as *mut FILETIME,
                        &mut fsys as *mut FILETIME,
                        &mut fuser as *mut FILETIME);
        memcpy(&mut sys as *mut ULARGE_INTEGER as *mut c_void,
               &mut fsys as *mut FILETIME as *mut c_void,
               size_of::<FILETIME>());
        memcpy(&mut user as *mut ULARGE_INTEGER as *mut c_void,
               &mut fuser as *mut FILETIME as *mut c_void,
               size_of::<FILETIME>());
        p.cpu_usage = ((*sys.QuadPart() - p.old_sys_cpu) as f32 +
            (*user.QuadPart() - p.old_user_cpu) as f32)
            / (*now.QuadPart() - p.old_cpu) as f32
            / nb_processors as f32 * 100.;
        p.old_cpu = *now.QuadPart();
        p.old_user_cpu = *user.QuadPart();
        p.old_sys_cpu = *sys.QuadPart();
    }
}

pub fn get_handle(p: &Process) -> HANDLE {
    p.handle
}

pub fn update_proc_info(p: &mut Process) {
    update_memory(p);
}

pub fn update_memory(p: &mut Process) {
    unsafe {
        let mut pmc: PROCESS_MEMORY_COUNTERS_EX = zeroed();
        if GetProcessMemoryInfo(p.handle,
                                &mut pmc as *mut PROCESS_MEMORY_COUNTERS_EX as *mut c_void as *mut PROCESS_MEMORY_COUNTERS,
                                size_of::<PROCESS_MEMORY_COUNTERS_EX>() as DWORD) != 0 {
            p.memory = (pmc.PrivateUsage as u64) >> 10u64; // / 1024;
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// 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::collections::HashMap;
use std::error;
use std::path::{Path, PathBuf};
use std::result;
use std::str::{self, FromStr};

use libarchive::writer;
use libarchive::reader::{self, Reader};
use libarchive::archive::{Entry, ReadFilter, ReadFormat, ExtractOption, ExtractOptions};
use regex::Regex;

use super::{Identifiable, PackageIdent, PackageTarget};
use super::metadata::MetaFile;
use error::{Error, Result};
use crypto::{artifact, hash};

lazy_static! {
    static ref METAFILE_REGXS: HashMap<MetaFile, Regex> = {
        let mut map = HashMap::new();
        map.insert(MetaFile::CFlags,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::CFlags)).unwrap());
        map.insert(MetaFile::Config,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Config)).unwrap());
        map.insert(MetaFile::Deps,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Deps)).unwrap());
        map.insert(MetaFile::TDeps,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::TDeps)).unwrap());
        map.insert(MetaFile::Exposes,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Exposes)).unwrap());
        map.insert(MetaFile::Ident,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Ident)).unwrap());
        map.insert(MetaFile::LdRunPath,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::LdRunPath)).unwrap());
        map.insert(MetaFile::LdFlags,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::LdFlags)).unwrap());
        map.insert(MetaFile::Manifest,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Manifest)).unwrap());
        map.insert(MetaFile::Path,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Path)).unwrap());
        map.insert(MetaFile::Target,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Target)).unwrap());
        map
    };
}

type Metadata = HashMap<MetaFile, String>;

#[derive(Debug)]
pub struct PackageArchive {
    pub path: PathBuf,
    metadata: Option<Metadata>,
}

impl PackageArchive {
    pub fn new<P: Into<PathBuf>>(path: P) -> Self {
        PackageArchive {
            path: path.into(),
            metadata: None,
        }
    }

    /// Calculate and return the checksum of the package archive in base64 format.
    ///
    /// # Failures
    ///
    /// * If the archive cannot be read
    pub fn checksum(&self) -> Result<String> {
        hash::hash_file(&self.path)
    }

    pub fn cflags(&mut self) -> Result<Option<String>> {
        match self.read_metadata(MetaFile::CFlags) {
            Ok(data) => Ok(data.cloned()),
            Err(e) => Err(e),
        }
    }

    pub fn config(&mut self) -> Result<Option<String>> {
        match self.read_metadata(MetaFile::Config) {
            Ok(data) => Ok(data.cloned()),
            Err(e) => Err(e),
        }
    }

    /// Returns a list of package identifiers representing the runtime package dependencies for
    /// this archive.
    ///
    /// # Failures
    ///
    /// * If a `DEPS` metafile is not found in the archive
    /// * If the archive cannot be read
    /// * If the archive cannot be verified
    pub fn deps(&mut self) -> Result<Vec<PackageIdent>> {
        self.read_deps(MetaFile::Deps)
    }

    /// Returns a list of package identifiers representing the transitive runtime package
    /// dependencies for this archive.
    ///
    /// # Failures
    ///
    /// * If a `TDEPS` metafile is not found in the archive
    /// * If the archive cannot be read
    /// * If the archive cannot be verified
    pub fn tdeps(&mut self) -> Result<Vec<PackageIdent>> {
        self.read_deps(MetaFile::TDeps)
    }

    pub fn exposes(&mut self) -> Result<Vec<u16>> {
        match self.read_metadata(MetaFile::Exposes) {
            Ok(Some(data)) => {
                let ports: Vec<u16> = data.split(" ")
                    .filter_map(|port| port.parse::<u16>().ok())
                    .collect();
                Ok(ports)
            }
            Ok(None) => Ok(vec![]),
            Err(e) => Err(e),
        }
    }

    pub fn ident(&mut self) -> Result<PackageIdent> {
        match self.read_metadata(MetaFile::Ident) {
            Ok(None) => Err(Error::MetaFileNotFound(MetaFile::Ident)),
            Ok(Some(data)) => PackageIdent::from_str(&data),
            Err(e) => Err(e),
        }
    }

    pub fn ld_run_path(&mut self) -> Result<Option<String>> {
        match self.read_metadata(MetaFile::LdRunPath) {
            Ok(data) => Ok(data.cloned()),
            Err(e) => Err(e),
        }
    }

    pub fn ldflags(&mut self) -> Result<Option<String>> {
        match self.read_metadata(MetaFile::LdFlags) {
            Ok(data) => Ok(data.cloned()),
            Err(e) => Err(e),
        }
    }

    pub fn manifest(&mut self) -> Result<String> {
        match self.read_metadata(MetaFile::Manifest) {
            Ok(None) => Err(Error::MetaFileNotFound(MetaFile::Manifest)),
            Ok(Some(data)) => Ok(data.clone()),
            Err(e) => Err(e),
        }
    }

    pub fn path(&mut self) -> Result<Option<String>> {
        match self.read_metadata(MetaFile::Path) {
            Ok(data) => Ok(data.cloned()),
            Err(e) => Err(e),
        }
    }

    pub fn target(&mut self) -> Result<PackageTarget> {
        match self.read_metadata(MetaFile::Target) {
            Ok(None) => Err(Error::MetaFileNotFound(MetaFile::Target)),
            Ok(Some(data)) => PackageTarget::from_str(&data),
            Err(e) => Err(e),
        }
    }

    /// A plain string representation of the archive's file name.
    pub fn file_name(&self) -> String {
        self.path
            .file_name()
            .unwrap()
            .to_string_lossy()
            .into_owned()
    }

    /// Given a package name and a path to a file as an `&str`, verify
    /// the files signature.
    ///
    /// # Failures
    ///
    /// * Fails if it cannot verify the signature for any reason
    pub fn verify<P: AsRef<Path>>(&self, cache_key_path: &P) -> Result<(String, String)> {
        artifact::verify(&self.path, cache_key_path)
    }

    /// Given a package name and a path to a file as an `&str`, unpack
    /// the package.
    ///
    /// # Failures
    ///
    /// * If the package cannot be unpacked
    pub fn unpack(&self, fs_root_path: Option<&Path>) -> Result<()> {
        let root = fs_root_path.unwrap_or(Path::new("/"));
        let tar_reader = artifact::get_archive_reader(&self.path)?;
        let mut builder = reader::Builder::new();
        builder.support_format(ReadFormat::Gnutar)?;
        builder.support_filter(ReadFilter::Xz)?;
        let mut reader = builder.open_stream(tar_reader)?;
        let writer = writer::Disk::new();
        let mut extract_options = ExtractOptions::new();
        extract_options.add(ExtractOption::Time);
        extract_options.add(ExtractOption::Permissions);
        writer.set_options(&extract_options)?;
        writer.set_standard_lookup()?;
        writer.write(
            &mut reader,
            Some(root.to_string_lossy().as_ref()),
        )?;
        writer.close()?;
        Ok(())
    }

    fn read_deps(&mut self, file: MetaFile) -> Result<Vec<PackageIdent>> {
        let mut deps: Vec<PackageIdent> = vec![];
        match self.read_metadata(file) {
            Ok(Some(body)) => {
                let ids: Vec<String> = body.lines().map(|d| d.to_string()).collect();
                for id in &ids {
                    let package = PackageIdent::from_str(id)?;
                    if !package.fully_qualified() {
                        // JW TODO: use a more appropriate error to describe the invalid
                        // user input here. (user because a package was generated by a user
                        // and read into program)
                        return Err(Error::InvalidPackageIdent(package.to_string()));
                    }
                    deps.push(package);
                }
                Ok(deps)
            }
            Ok(None) => Ok(vec![]),
            Err(Error::MetaFileNotFound(_)) => Ok(deps),
            Err(e) => Err(e),
        }
    }

    fn read_metadata(&mut self, file: MetaFile) -> Result<Option<&String>> {
        if let Some(ref files) = self.metadata {
            return Ok(files.get(&file));
        }
        let mut metadata = Metadata::new();
        let mut matched_count = 0u8;
        let tar_reader = artifact::get_archive_reader(&self.path)?;
        let mut builder = reader::Builder::new();
        builder.support_format(ReadFormat::Gnutar)?;
        builder.support_filter(ReadFilter::Xz)?;
        let mut reader = builder.open_stream(tar_reader)?;
        loop {
            let mut matched_type: Option<MetaFile> = None;
            if let Some(entry) = reader.next_header() {
                for (matched, regx) in METAFILE_REGXS.iter() {
                    if regx.is_match(entry.pathname()) {
                        matched_type = Some((*matched).clone());
                        matched_count += 1;
                        break;
                    }
                }
            } else {
                break;
            }

            if matched_type.is_none() {
                continue;
            }

            let mut buf = String::new();
            loop {
                match reader.read_block() {
                    Ok(Some(bytes)) => {
                        match str::from_utf8(bytes) {
                            Ok(content) => {
                                // You used to trim. Now you don't, because you were trimming
                                // in the wrong place. Sometimes a buffer ends (or starts!) with
                                // a newline.
                                buf.push_str(content);
                            }
                            Err(_) => return Err(Error::MetaFileMalformed(matched_type.unwrap())),
                        }
                    }
                    Ok(None) => {
                        // Hey, before you go - we are trimming whitespace for you. This
                        // is handy, because later on, you just want the string you want.
                        metadata.insert(matched_type.unwrap(), String::from(buf.trim()));
                        break;
                    }
                    Err(_) => return Err(Error::MetaFileMalformed(matched_type.unwrap())),
                }
            } //inner loop

            if matched_count == METAFILE_REGXS.len() as u8 {
                break;
            }
        }
        self.metadata = Some(metadata);
        Ok(self.metadata.as_ref().unwrap().get(&file))
    }
}

pub trait FromArchive: Sized {
    type Error: error::Error;

    fn from_archive(archive: &mut PackageArchive) -> result::Result<Self, Self::Error>;
}

#[cfg(test)]
mod test {
    use std::path::PathBuf;
    use os::system::{Architecture, Platform};
    use super::*;

    #[test]
    fn reading_artifact_metadata() {
        let mut hart = PackageArchive::new(fixtures().join(
            "happyhumans-possums-8.1.4-20160427165340-x86_64-linux.hart",
        ));
        let ident = hart.ident().unwrap();
        assert_eq!(ident.origin, "happyhumans");
        assert_eq!(ident.name, "possums");
        assert_eq!(ident.version, Some("8.1.4".to_string()));
        assert_eq!(ident.release, Some("20160427165340".to_string()));
    }

    pub fn root() -> PathBuf {
        PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests")
    }

    pub fn fixtures() -> PathBuf {
        root().join("fixtures")
    }

    #[test]
    fn reading_artifact_deps() {
        let mut hart = PackageArchive::new(fixtures().join(
            "happyhumans-possums-8.1.4-20160427165340-x86_64-linux.hart",
        ));
        let _ = hart.deps().unwrap();
        let _ = hart.tdeps().unwrap();
    }

    #[test]
    fn reading_artifact_large_tdeps() {
        let mut hart = PackageArchive::new(fixtures().join(
            "unhappyhumans-possums-8.1.4-20160427165340-x86_64-linux.hart",
        ));
        let tdeps = hart.tdeps().unwrap();
        assert_eq!(1024, tdeps.len());
    }

    #[test]
    fn reading_artifact_target() {
        let mut hart = PackageArchive::new(fixtures().join(
            "unhappyhumans-possums-8.1.4-20160427165340-x86_64-linux.hart",
        ));
        let target = hart.target().unwrap();
        assert_eq!(target.platform, Platform::Linux);
        assert_eq!(target.architecture, Architecture::X86_64);
    }
}
<commit_msg>Add is_a_service to API responses for packages<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// 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::collections::HashMap;
use std::error;
use std::path::{Path, PathBuf};
use std::result;
use std::str::{self, FromStr};

use libarchive::writer;
use libarchive::reader::{self, Reader};
use libarchive::archive::{Entry, ReadFilter, ReadFormat, ExtractOption, ExtractOptions};
use regex::Regex;

use super::{Identifiable, PackageIdent, PackageTarget};
use super::metadata::MetaFile;
use error::{Error, Result};
use crypto::{artifact, hash};

lazy_static! {
    static ref METAFILE_REGXS: HashMap<MetaFile, Regex> = {
        let mut map = HashMap::new();
        map.insert(MetaFile::CFlags,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::CFlags)).unwrap());
        map.insert(MetaFile::Config,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Config)).unwrap());
        map.insert(MetaFile::Deps,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Deps)).unwrap());
        map.insert(MetaFile::TDeps,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::TDeps)).unwrap());
        map.insert(MetaFile::Exposes,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Exposes)).unwrap());
        map.insert(MetaFile::Ident,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Ident)).unwrap());
        map.insert(MetaFile::LdRunPath,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::LdRunPath)).unwrap());
        map.insert(MetaFile::LdFlags,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::LdFlags)).unwrap());
        map.insert(MetaFile::SvcUser,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::SvcUser)).unwrap());
        map.insert(MetaFile::Manifest,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Manifest)).unwrap());
        map.insert(MetaFile::Path,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Path)).unwrap());
        map.insert(MetaFile::Target,
                   Regex::new(&format!(r"^/?hab/pkgs/([^/]+)/([^/]+)/([^/]+)/([^/]+)/{}$",
                              MetaFile::Target)).unwrap());
        map
    };
}

type Metadata = HashMap<MetaFile, String>;

#[derive(Debug)]
pub struct PackageArchive {
    pub path: PathBuf,
    metadata: Option<Metadata>,
}

impl PackageArchive {
    pub fn new<P: Into<PathBuf>>(path: P) -> Self {
        PackageArchive {
            path: path.into(),
            metadata: None,
        }
    }

    /// Calculate and return the checksum of the package archive in base64 format.
    ///
    /// # Failures
    ///
    /// * If the archive cannot be read
    pub fn checksum(&self) -> Result<String> {
        hash::hash_file(&self.path)
    }

    pub fn cflags(&mut self) -> Result<Option<String>> {
        match self.read_metadata(MetaFile::CFlags) {
            Ok(data) => Ok(data.cloned()),
            Err(e) => Err(e),
        }
    }

    pub fn config(&mut self) -> Result<Option<String>> {
        match self.read_metadata(MetaFile::Config) {
            Ok(data) => Ok(data.cloned()),
            Err(e) => Err(e),
        }
    }

    // hab-plan-build.sh only generates SVC_USER and SVC_GROUP files if it thinks a package is
    // a service. It determines that by checking for the presence of a run hook file or a
    // pkg_svc_run value. Therefore, if we can detect the presence of a SVC_USER file, we can
    // consider this archive a service.
    pub fn is_a_service(&mut self) -> bool {
        match self.svc_user() {
            Ok(_) => true,
            _ => false,
        }
    }

    /// Returns a list of package identifiers representing the runtime package dependencies for
    /// this archive.
    ///
    /// # Failures
    ///
    /// * If a `DEPS` metafile is not found in the archive
    /// * If the archive cannot be read
    /// * If the archive cannot be verified
    pub fn deps(&mut self) -> Result<Vec<PackageIdent>> {
        self.read_deps(MetaFile::Deps)
    }

    /// Returns a list of package identifiers representing the transitive runtime package
    /// dependencies for this archive.
    ///
    /// # Failures
    ///
    /// * If a `TDEPS` metafile is not found in the archive
    /// * If the archive cannot be read
    /// * If the archive cannot be verified
    pub fn tdeps(&mut self) -> Result<Vec<PackageIdent>> {
        self.read_deps(MetaFile::TDeps)
    }

    pub fn exposes(&mut self) -> Result<Vec<u16>> {
        match self.read_metadata(MetaFile::Exposes) {
            Ok(Some(data)) => {
                let ports: Vec<u16> = data.split(" ")
                    .filter_map(|port| port.parse::<u16>().ok())
                    .collect();
                Ok(ports)
            }
            Ok(None) => Ok(vec![]),
            Err(e) => Err(e),
        }
    }

    pub fn ident(&mut self) -> Result<PackageIdent> {
        match self.read_metadata(MetaFile::Ident) {
            Ok(None) => Err(Error::MetaFileNotFound(MetaFile::Ident)),
            Ok(Some(data)) => PackageIdent::from_str(&data),
            Err(e) => Err(e),
        }
    }

    pub fn ld_run_path(&mut self) -> Result<Option<String>> {
        match self.read_metadata(MetaFile::LdRunPath) {
            Ok(data) => Ok(data.cloned()),
            Err(e) => Err(e),
        }
    }

    pub fn ldflags(&mut self) -> Result<Option<String>> {
        match self.read_metadata(MetaFile::LdFlags) {
            Ok(data) => Ok(data.cloned()),
            Err(e) => Err(e),
        }
    }

    pub fn svc_user(&mut self) -> Result<String> {
        match self.read_metadata(MetaFile::SvcUser) {
            Ok(Some(data)) => Ok(data.clone()),
            Ok(None) => Err(Error::MetaFileNotFound(MetaFile::SvcUser)),
            Err(e) => Err(e),
        }
    }

    pub fn manifest(&mut self) -> Result<String> {
        match self.read_metadata(MetaFile::Manifest) {
            Ok(None) => Err(Error::MetaFileNotFound(MetaFile::Manifest)),
            Ok(Some(data)) => Ok(data.clone()),
            Err(e) => Err(e),
        }
    }

    pub fn path(&mut self) -> Result<Option<String>> {
        match self.read_metadata(MetaFile::Path) {
            Ok(data) => Ok(data.cloned()),
            Err(e) => Err(e),
        }
    }

    pub fn target(&mut self) -> Result<PackageTarget> {
        match self.read_metadata(MetaFile::Target) {
            Ok(None) => Err(Error::MetaFileNotFound(MetaFile::Target)),
            Ok(Some(data)) => PackageTarget::from_str(&data),
            Err(e) => Err(e),
        }
    }

    /// A plain string representation of the archive's file name.
    pub fn file_name(&self) -> String {
        self.path
            .file_name()
            .unwrap()
            .to_string_lossy()
            .into_owned()
    }

    /// Given a package name and a path to a file as an `&str`, verify
    /// the files signature.
    ///
    /// # Failures
    ///
    /// * Fails if it cannot verify the signature for any reason
    pub fn verify<P: AsRef<Path>>(&self, cache_key_path: &P) -> Result<(String, String)> {
        artifact::verify(&self.path, cache_key_path)
    }

    /// Given a package name and a path to a file as an `&str`, unpack
    /// the package.
    ///
    /// # Failures
    ///
    /// * If the package cannot be unpacked
    pub fn unpack(&self, fs_root_path: Option<&Path>) -> Result<()> {
        let root = fs_root_path.unwrap_or(Path::new("/"));
        let tar_reader = artifact::get_archive_reader(&self.path)?;
        let mut builder = reader::Builder::new();
        builder.support_format(ReadFormat::Gnutar)?;
        builder.support_filter(ReadFilter::Xz)?;
        let mut reader = builder.open_stream(tar_reader)?;
        let writer = writer::Disk::new();
        let mut extract_options = ExtractOptions::new();
        extract_options.add(ExtractOption::Time);
        extract_options.add(ExtractOption::Permissions);
        writer.set_options(&extract_options)?;
        writer.set_standard_lookup()?;
        writer.write(
            &mut reader,
            Some(root.to_string_lossy().as_ref()),
        )?;
        writer.close()?;
        Ok(())
    }

    fn read_deps(&mut self, file: MetaFile) -> Result<Vec<PackageIdent>> {
        let mut deps: Vec<PackageIdent> = vec![];
        match self.read_metadata(file) {
            Ok(Some(body)) => {
                let ids: Vec<String> = body.lines().map(|d| d.to_string()).collect();
                for id in &ids {
                    let package = PackageIdent::from_str(id)?;
                    if !package.fully_qualified() {
                        // JW TODO: use a more appropriate error to describe the invalid
                        // user input here. (user because a package was generated by a user
                        // and read into program)
                        return Err(Error::InvalidPackageIdent(package.to_string()));
                    }
                    deps.push(package);
                }
                Ok(deps)
            }
            Ok(None) => Ok(vec![]),
            Err(Error::MetaFileNotFound(_)) => Ok(deps),
            Err(e) => Err(e),
        }
    }

    fn read_metadata(&mut self, file: MetaFile) -> Result<Option<&String>> {
        if let Some(ref files) = self.metadata {
            return Ok(files.get(&file));
        }
        let mut metadata = Metadata::new();
        let mut matched_count = 0u8;
        let tar_reader = artifact::get_archive_reader(&self.path)?;
        let mut builder = reader::Builder::new();
        builder.support_format(ReadFormat::Gnutar)?;
        builder.support_filter(ReadFilter::Xz)?;
        let mut reader = builder.open_stream(tar_reader)?;
        loop {
            let mut matched_type: Option<MetaFile> = None;
            if let Some(entry) = reader.next_header() {
                for (matched, regx) in METAFILE_REGXS.iter() {
                    if regx.is_match(entry.pathname()) {
                        matched_type = Some((*matched).clone());
                        matched_count += 1;
                        break;
                    }
                }
            } else {
                break;
            }

            if matched_type.is_none() {
                continue;
            }

            let mut buf = String::new();
            loop {
                match reader.read_block() {
                    Ok(Some(bytes)) => {
                        match str::from_utf8(bytes) {
                            Ok(content) => {
                                // You used to trim. Now you don't, because you were trimming
                                // in the wrong place. Sometimes a buffer ends (or starts!) with
                                // a newline.
                                buf.push_str(content);
                            }
                            Err(_) => return Err(Error::MetaFileMalformed(matched_type.unwrap())),
                        }
                    }
                    Ok(None) => {
                        // Hey, before you go - we are trimming whitespace for you. This
                        // is handy, because later on, you just want the string you want.
                        metadata.insert(matched_type.unwrap(), String::from(buf.trim()));
                        break;
                    }
                    Err(_) => return Err(Error::MetaFileMalformed(matched_type.unwrap())),
                }
            } //inner loop

            if matched_count == METAFILE_REGXS.len() as u8 {
                break;
            }
        }
        self.metadata = Some(metadata);
        Ok(self.metadata.as_ref().unwrap().get(&file))
    }
}

pub trait FromArchive: Sized {
    type Error: error::Error;

    fn from_archive(archive: &mut PackageArchive) -> result::Result<Self, Self::Error>;
}

#[cfg(test)]
mod test {
    use std::path::PathBuf;
    use os::system::{Architecture, Platform};
    use super::*;

    #[test]
    fn reading_artifact_metadata() {
        let mut hart = PackageArchive::new(fixtures().join(
            "happyhumans-possums-8.1.4-20160427165340-x86_64-linux.hart",
        ));
        let ident = hart.ident().unwrap();
        assert_eq!(ident.origin, "happyhumans");
        assert_eq!(ident.name, "possums");
        assert_eq!(ident.version, Some("8.1.4".to_string()));
        assert_eq!(ident.release, Some("20160427165340".to_string()));
    }

    pub fn root() -> PathBuf {
        PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests")
    }

    pub fn fixtures() -> PathBuf {
        root().join("fixtures")
    }

    #[test]
    fn reading_artifact_deps() {
        let mut hart = PackageArchive::new(fixtures().join(
            "happyhumans-possums-8.1.4-20160427165340-x86_64-linux.hart",
        ));
        let _ = hart.deps().unwrap();
        let _ = hart.tdeps().unwrap();
    }

    #[test]
    fn reading_artifact_large_tdeps() {
        let mut hart = PackageArchive::new(fixtures().join(
            "unhappyhumans-possums-8.1.4-20160427165340-x86_64-linux.hart",
        ));
        let tdeps = hart.tdeps().unwrap();
        assert_eq!(1024, tdeps.len());
    }

    #[test]
    fn reading_artifact_target() {
        let mut hart = PackageArchive::new(fixtures().join(
            "unhappyhumans-possums-8.1.4-20160427165340-x86_64-linux.hart",
        ));
        let target = hart.target().unwrap();
        assert_eq!(target.platform, Platform::Linux);
        assert_eq!(target.architecture, Architecture::X86_64);
    }
}
<|endoftext|>


<commit_before>pub use crate::inner::{
    auto::{Spawnable, HasInner, Abstract},
    application::{AApplication, ApplicationBase, ApplicationInner},
    member::{AMember, MemberBase, MemberInner},
    button::{AButton, ButtonInner, NewButtonInner},
    text::{AText, TextInner, NewTextInner},
    control::{AControl, ControlBase, ControlInner, OuterControl},
    container::{AContainer, ContainerInner},
    container_single::{ASingleContainer, SingleContainerInner},
    container_multi::{AMultiContainer, MultiContainerInner},
    tray::{ATray, TrayInner, NewTrayInner},
    has_native_id::{HasNativeIdInner, NativeId},
    drawable::{Drawable, OuterDrawable},
    closeable::CloseableInner,
    clickable::ClickableInner,
    item_clickable::ItemClickableInner,
    has_label::HasLabelInner,
    has_layout::HasLayoutInner,
    has_image::HasImageInner,
    has_progress::HasProgressInner,
    has_size::HasSizeInner,
    has_visibility::HasVisibilityInner,
    has_orientation::HasOrientationInner,
    window::{AWindow, WindowBase, WindowInner, NewWindowInner},
    message::{AMessage, MessageInner},
    image::{AImage, ImageInner, NewImageInner},
    frame::{AFrame, FrameInner, NewFrameInner},
    layout_linear::{ALinearLayout, LinearLayoutInner, NewLinearLayoutInner},
    progress_bar::{AProgressBar, ProgressBarInner, NewProgressBarInner},
    splitted::{ASplitted, SplittedInner, NewSplittedInner},
    adapted::{AAdapted, AdaptedInner, AdapterInnerCallback},
    adapter::{AdapterInner},
    list::{AList, ListInner, NewListInner},
};
<commit_msg>Revert unmanageable NewInners<commit_after>pub use crate::inner::{
    auto::{Spawnable, HasInner, Abstract},
    application::{AApplication, ApplicationBase, ApplicationInner},
    member::{AMember, MemberBase, MemberInner},
    button::{AButton, ButtonInner, NewButtonInner},
    text::{AText, TextInner, NewTextInner},
    control::{AControl, ControlBase, ControlInner, OuterControl},
    container::{AContainer, ContainerInner},
    container_single::{ASingleContainer, SingleContainerInner},
    container_multi::{AMultiContainer, MultiContainerInner},
    tray::{ATray, TrayInner},
    has_native_id::{HasNativeIdInner, NativeId},
    drawable::{Drawable, OuterDrawable},
    closeable::CloseableInner,
    clickable::ClickableInner,
    item_clickable::ItemClickableInner,
    has_label::HasLabelInner,
    has_layout::HasLayoutInner,
    has_image::HasImageInner,
    has_progress::HasProgressInner,
    has_size::HasSizeInner,
    has_visibility::HasVisibilityInner,
    has_orientation::HasOrientationInner,
    window::{AWindow, WindowBase, WindowInner},
    message::{AMessage, MessageInner},
    image::{AImage, ImageInner, NewImageInner},
    frame::{AFrame, FrameInner, NewFrameInner},
    layout_linear::{ALinearLayout, LinearLayoutInner, NewLinearLayoutInner},
    progress_bar::{AProgressBar, ProgressBarInner, NewProgressBarInner},
    splitted::{ASplitted, SplittedInner, NewSplittedInner},
    adapted::{AAdapted, AdaptedInner, AdapterInnerCallback},
    adapter::{AdapterInner},
    list::{AList, ListInner, NewListInner},
};
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use serde::{Deserialize, Serialize};
use std::{
    convert::From,
    fmt::{self, Display},
};
use warp::{http::StatusCode, reject::Reject};

use crate::U64;

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct Error {
    pub code: u16,
    pub message: String,
    /// Aptos blockchain latest onchain ledger version.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub aptos_ledger_version: Option<U64>,
}

impl Error {
    pub fn new(code: StatusCode, message: String) -> Self {
        Self {
            code: code.as_u16(),
            message,
            aptos_ledger_version: None,
        }
    }

    pub fn from_anyhow_error(code: StatusCode, err: anyhow::Error) -> Self {
        Self::new(code, err.to_string())
    }

    pub fn bad_request<S: Display>(msg: S) -> Self {
        Self::new(StatusCode::BAD_REQUEST, msg.to_string())
    }

    pub fn not_found<S: Display>(resource: &str, identifier: S, ledger_version: u64) -> Self {
        Self::new(
            StatusCode::NOT_FOUND,
            format!("{} not found by {}", resource, identifier),
        )
        .aptos_ledger_version(ledger_version)
    }

    pub fn invalid_param<S: Display>(name: &str, value: S) -> Self {
        Self::bad_request(format!("invalid parameter {}: {}", name, value))
    }

    pub fn invalid_request_body<S: Display>(msg: S) -> Self {
        Self::bad_request(format!("invalid request body: {}", msg))
    }

    pub fn insufficient_storage<S: Display>(msg: S) -> Self {
        Self::new(StatusCode::INSUFFICIENT_STORAGE, msg.to_string())
    }

    pub fn internal(err: anyhow::Error) -> Self {
        Self::from_anyhow_error(StatusCode::INTERNAL_SERVER_ERROR, err)
    }

    pub fn status_code(&self) -> StatusCode {
        StatusCode::from_u16(self.code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR)
    }

    pub fn aptos_ledger_version(mut self, ledger_version: u64) -> Self {
        self.aptos_ledger_version = Some(ledger_version.into());
        self
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}: {}", self.status_code(), &self.message)?;
        if let Some(val) = &self.aptos_ledger_version {
            write!(f, "\ndiem ledger version: {}", val)?;
        }
        Ok(())
    }
}

impl Reject for Error {}

impl From<anyhow::Error> for Error {
    fn from(e: anyhow::Error) -> Self {
        Self::internal(e)
    }
}

impl From<serde_json::error::Error> for Error {
    fn from(err: serde_json::error::Error) -> Self {
        Self::internal(err.into())
    }
}

#[cfg(test)]
mod tests {
    use crate::error::Error;
    use warp::http::StatusCode;

    #[test]
    fn test_to_string() {
        let err = Error::new(StatusCode::BAD_REQUEST, "invalid address".to_owned());
        assert_eq!(err.to_string(), "400 Bad Request: invalid address")
    }

    #[test]
    fn test_from_anyhow_error_as_internal_error() {
        let err = Error::from(anyhow::format_err!("hello"));
        assert_eq!(err.to_string(), "500 Internal Server Error: hello")
    }

    #[test]
    fn test_to_string_with_aptos_ledger_version() {
        let err = Error::new(StatusCode::BAD_REQUEST, "invalid address".to_owned())
            .aptos_ledger_version(123);
        assert_eq!(
            err.to_string(),
            "400 Bad Request: invalid address\ndiem ledger version: 123"
        )
    }

    #[test]
    fn test_internal_error() {
        let err = Error::internal(anyhow::format_err!("hello"));
        assert_eq!(err.to_string(), "500 Internal Server Error: hello")
    }
}
<commit_msg>[Cleanup] Renaming (#2179)<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use serde::{Deserialize, Serialize};
use std::{
    convert::From,
    fmt::{self, Display},
};
use warp::{http::StatusCode, reject::Reject};

use crate::U64;

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct Error {
    pub code: u16,
    pub message: String,
    /// Aptos blockchain latest onchain ledger version.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub aptos_ledger_version: Option<U64>,
}

impl Error {
    pub fn new(code: StatusCode, message: String) -> Self {
        Self {
            code: code.as_u16(),
            message,
            aptos_ledger_version: None,
        }
    }

    pub fn from_anyhow_error(code: StatusCode, err: anyhow::Error) -> Self {
        Self::new(code, err.to_string())
    }

    pub fn bad_request<S: Display>(msg: S) -> Self {
        Self::new(StatusCode::BAD_REQUEST, msg.to_string())
    }

    pub fn not_found<S: Display>(resource: &str, identifier: S, ledger_version: u64) -> Self {
        Self::new(
            StatusCode::NOT_FOUND,
            format!("{} not found by {}", resource, identifier),
        )
        .aptos_ledger_version(ledger_version)
    }

    pub fn invalid_param<S: Display>(name: &str, value: S) -> Self {
        Self::bad_request(format!("invalid parameter {}: {}", name, value))
    }

    pub fn invalid_request_body<S: Display>(msg: S) -> Self {
        Self::bad_request(format!("invalid request body: {}", msg))
    }

    pub fn insufficient_storage<S: Display>(msg: S) -> Self {
        Self::new(StatusCode::INSUFFICIENT_STORAGE, msg.to_string())
    }

    pub fn internal(err: anyhow::Error) -> Self {
        Self::from_anyhow_error(StatusCode::INTERNAL_SERVER_ERROR, err)
    }

    pub fn status_code(&self) -> StatusCode {
        StatusCode::from_u16(self.code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR)
    }

    pub fn aptos_ledger_version(mut self, ledger_version: u64) -> Self {
        self.aptos_ledger_version = Some(ledger_version.into());
        self
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}: {}", self.status_code(), &self.message)?;
        if let Some(val) = &self.aptos_ledger_version {
            write!(f, "\nAptos ledger version: {}", val)?;
        }
        Ok(())
    }
}

impl Reject for Error {}

impl From<anyhow::Error> for Error {
    fn from(e: anyhow::Error) -> Self {
        Self::internal(e)
    }
}

impl From<serde_json::error::Error> for Error {
    fn from(err: serde_json::error::Error) -> Self {
        Self::internal(err.into())
    }
}

#[cfg(test)]
mod tests {
    use crate::error::Error;
    use warp::http::StatusCode;

    #[test]
    fn test_to_string() {
        let err = Error::new(StatusCode::BAD_REQUEST, "invalid address".to_owned());
        assert_eq!(err.to_string(), "400 Bad Request: invalid address")
    }

    #[test]
    fn test_from_anyhow_error_as_internal_error() {
        let err = Error::from(anyhow::format_err!("hello"));
        assert_eq!(err.to_string(), "500 Internal Server Error: hello")
    }

    #[test]
    fn test_to_string_with_aptos_ledger_version() {
        let err = Error::new(StatusCode::BAD_REQUEST, "invalid address".to_owned())
            .aptos_ledger_version(123);
        assert_eq!(
            err.to_string(),
            "400 Bad Request: invalid address\nAptos ledger version: 123"
        )
    }

    #[test]
    fn test_internal_error() {
        let err = Error::internal(anyhow::format_err!("hello"));
        assert_eq!(err.to_string(), "500 Internal Server Error: hello")
    }
}
<|endoftext|>