<commit_before>// Copyright 2015 The tiny-http Contributors
// Copyright 2015 The rust-chunked-transfer 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::io::Result as IoResult;
use std::io::Read;
use std::io::Error as IoError;
use std::io::ErrorKind;
use std::fmt;
use std::error::Error;

/// Reads HTTP chunks and sends back real data.
///
/// # Example
///
/// ```
/// use chunked_transfer::Decoder;
/// use std::io::Read;
///
/// let encoded = b"3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n";
/// let mut decoded = String::new();
///
/// let mut decoder = Decoder::new(encoded as &[u8]);
/// decoder.read_to_string(&mut decoded);
///
/// assert_eq!(decoded, "hello world!!!");
/// ```
pub struct Decoder<R> {
    // where the chunks come from
    source: R,

    // remaining size of the chunk being read
    // none if we are not in a chunk
    remaining_chunks_size: Option<usize>,
}

impl<R> Decoder<R> where R: Read {
    pub fn new(source: R) -> Decoder<R> {
        Decoder {
            source: source,
            remaining_chunks_size: None,
        }
    }
}

impl<R> Read for Decoder<R> where R: Read {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        // first possibility: we are not in a chunk
        if self.remaining_chunks_size.is_none() {
            // trying the read the chunk size
            let mut chunk_size = Vec::new();

            loop {
                let byte = match self.source.by_ref().bytes().next() {
                    Some(b) => try!(b),
                    None => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                };

                if byte == b'\r' {
                    break;
                }

                chunk_size.push(byte);
            }

            match self.source.by_ref().bytes().next() {
                Some(Ok(b'\n')) => (),
                _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
            }

            let chunk_size = match String::from_utf8(chunk_size) {
                Ok(c) => c,
                Err(_) => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError))
            };

            let chunk_size = match usize::from_str_radix(&chunk_size, 16) {
                Ok(c) => c,
                Err(_) => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError))
            };

            // if the chunk size is 0, we are at EOF
            if chunk_size == 0 {
                match self.source.by_ref().bytes().next() {
                    Some(Ok(b'\r')) => (),
                    _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                }

                match self.source.by_ref().bytes().next() {
                    Some(Ok(b'\n')) => (),
                    _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                }

                return Ok(0);
            }

            // now that we now the current chunk size, calling ourselves recursively
            self.remaining_chunks_size = Some(chunk_size);
            return self.read(buf);
        }

        assert!(self.remaining_chunks_size.is_some());

        // second possibility: we continue reading from a chunk
        if buf.len() < *self.remaining_chunks_size.as_ref().unwrap() {
            let read = try!(self.source.read(buf));
            *self.remaining_chunks_size.as_mut().unwrap() -= read;
            return Ok(read);
        }

        // third possibility: the read request goes further than the current chunk
        // we simply read until the end of the chunk and return
        assert!(buf.len() >= *self.remaining_chunks_size.as_ref().unwrap());

        let remaining_chunks_size = *self.remaining_chunks_size.as_ref().unwrap();

        let buf = &mut buf[.. remaining_chunks_size];
        let read = try!(self.source.read(buf));

        self.remaining_chunks_size = if read == remaining_chunks_size {
            match self.source.by_ref().bytes().next() {
                Some(Ok(b'\r')) => (),
                _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
            }

            match self.source.by_ref().bytes().next() {
                Some(Ok(b'\n')) => (),
                _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
            }

            None
        } else {
            Some(remaining_chunks_size - read)
        };

        return Ok(read);
    }
}

#[derive(Debug, Copy, Clone)]
struct DecoderError;

impl fmt::Display for DecoderError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "Error while decoding chunks")
    }
}

impl Error for DecoderError {
    fn description(&self) -> &str {
        "Error while decoding chunks"
    }
}

#[cfg(test)]
mod test {
    use super::Decoder;
    use std::io;
    use std::io::Read;

    #[test]
    fn test_valid_chunk_decode() {
        let source = io::Cursor::new("3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n".to_string().into_bytes());
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).unwrap();

        assert_eq!(string, "hello world!!!");
    }

    #[test]
    fn invalid_input1() {
        let source = io::Cursor::new("2\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).is_err();
    }

    #[test]
    fn invalid_input2() {
        let source = io::Cursor::new("3\rhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).is_err();
    }
}
<commit_msg>Remove all `unwrap`s in read<commit_after>// Copyright 2015 The tiny-http Contributors
// Copyright 2015 The rust-chunked-transfer 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::io::Result as IoResult;
use std::io::Read;
use std::io::Error as IoError;
use std::io::ErrorKind;
use std::fmt;
use std::error::Error;

/// Reads HTTP chunks and sends back real data.
///
/// # Example
///
/// ```
/// use chunked_transfer::Decoder;
/// use std::io::Read;
///
/// let encoded = b"3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n";
/// let mut decoded = String::new();
///
/// let mut decoder = Decoder::new(encoded as &[u8]);
/// decoder.read_to_string(&mut decoded);
///
/// assert_eq!(decoded, "hello world!!!");
/// ```
pub struct Decoder<R> {
    // where the chunks come from
    source: R,

    // remaining size of the chunk being read
    // none if we are not in a chunk
    remaining_chunks_size: Option<usize>,
}

impl<R> Decoder<R> where R: Read {
    pub fn new(source: R) -> Decoder<R> {
        Decoder {
            source: source,
            remaining_chunks_size: None,
        }
    }
}

impl<R> Read for Decoder<R> where R: Read {
    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
        let remaining_chunks_size = match self.remaining_chunks_size {
            Some(c) => c,
            None => {
                // first possibility: we are not in a chunk, so we'll attempt to determine
                // the chunks size
                let mut chunk_size = Vec::new();

                loop {
                    let byte = match self.source.by_ref().bytes().next() {
                        Some(b) => try!(b),
                        None => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                    };

                    if byte == b'\r' {
                        break;
                    }

                    chunk_size.push(byte);
                }

                match self.source.by_ref().bytes().next() {
                    Some(Ok(b'\n')) => (),
                    _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                }

                let chunk_size = match String::from_utf8(chunk_size) {
                    Ok(c) => c,
                    Err(_) => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError))
                };

                let chunk_size = match usize::from_str_radix(&chunk_size, 16) {
                    Ok(c) => c,
                    Err(_) => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError))
                };

                // if the chunk size is 0, we are at EOF
                if chunk_size == 0 {
                    match self.source.by_ref().bytes().next() {
                        Some(Ok(b'\r')) => (),
                        _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                    }

                    match self.source.by_ref().bytes().next() {
                        Some(Ok(b'\n')) => (),
                        _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
                    }

                    return Ok(0);
                }

                // now that we now the current chunk size, calling ourselves recursively
                self.remaining_chunks_size = Some(chunk_size);
                return self.read(buf);
            }
        };

        // second possibility: we continue reading from a chunk
        if buf.len() < remaining_chunks_size {
            let read = try!(self.source.read(buf));
            self.remaining_chunks_size = Some(remaining_chunks_size - read);
            return Ok(read);
        }

        // third possibility: the read request goes further than the current chunk
        // we simply read until the end of the chunk and return
        assert!(buf.len() >= remaining_chunks_size);

        let buf = &mut buf[.. remaining_chunks_size];
        let read = try!(self.source.read(buf));

        self.remaining_chunks_size = if read == remaining_chunks_size {
            match self.source.by_ref().bytes().next() {
                Some(Ok(b'\r')) => (),
                _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
            }

            match self.source.by_ref().bytes().next() {
                Some(Ok(b'\n')) => (),
                _ => return Err(IoError::new(ErrorKind::InvalidInput, DecoderError)),
            }

            None
        } else {
            Some(remaining_chunks_size - read)
        };

        return Ok(read);
    }
}

#[derive(Debug, Copy, Clone)]
struct DecoderError;

impl fmt::Display for DecoderError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(fmt, "Error while decoding chunks")
    }
}

impl Error for DecoderError {
    fn description(&self) -> &str {
        "Error while decoding chunks"
    }
}

#[cfg(test)]
mod test {
    use super::Decoder;
    use std::io;
    use std::io::Read;

    #[test]
    fn test_valid_chunk_decode() {
        let source = io::Cursor::new("3\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n\r\n".to_string().into_bytes());
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).unwrap();

        assert_eq!(string, "hello world!!!");
    }

    #[test]
    fn invalid_input1() {
        let source = io::Cursor::new("2\r\nhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).is_err();
    }

    #[test]
    fn invalid_input2() {
        let source = io::Cursor::new("3\rhel\r\nb\r\nlo world!!!\r\n0\r\n".to_string().into_bytes());
        let mut decoded = Decoder::new(source);

        let mut string = String::new();
        decoded.read_to_string(&mut string).is_err();
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2015, Sam Payson
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
// associated documentation files (the "Software"), to deal in the Software without restriction,
// including without limitation the rights to use, copy, modify, merge, publish, distribute,
// sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all copies or
// substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
// NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

// Allow dead code to silence warnings until things stabilize.
#![allow(dead_code)]

use std::error::FromError;
use std::io;
use std::mem;
use std::result;
use std::io::ReadExt;

use iter::{RexIterExt, ResultIterExt};

/// `Error` is used to report errors that occur within the decoder.
pub enum Error {
    /// `EOF` indicates that record source ended before it could be fully decoded.
    EOF,

    /// `IoError` is used to pass through std::io errors.
    IoError(io::Error),
}

impl FromError<io::Error> for Error {
    fn from_error(err: io::Error) -> Error {
        Error::IoError(err)
    }
}

/// `read_uvarint` reads a varint encoded `u64` from `r`.
fn read_uvarint<R>(r: &mut R) -> Result<u64, Error>
    where R: io::Read {

    let itr = r.bytes()
        .take_while_incl(|x| match x { &Ok(x) => x >= 0x80, _ => true });

    // The Ok(try!(..)) just lets us convert from io::Error to Error
    result::fold(itr, 0, |sum, x| (sum << 7) + (x & 0x7F) as u64)
        .map_err(FromError::from_error)
}

/// `read_varint` reads a zig-zag varint encoded `i64` from `r`.
fn read_varint<R>(r: &mut R) -> Result<i64, Error>
    where R: io::Read {

    let ux = try!(read_uvarint(r));

    Ok( if ux & 1 != 0 {
        let x = (!ux >> 1) as i64;
        -x
    } else {
        let x = (ux >> 1) as i64;
        x
    })
}

/// `read_u8` reads a single byte from `r`.
fn read_u8<R>(r: &mut R) -> Result<u8, Error>
    where R: io::Read {

    r.bytes()
        .take_or_err(1, Error::EOF)
        .next()
        // take_or_err guarantees that we'll get an Error at worst.
        .unwrap()
}

/// `read_le_u16` reads 2 bytes as a little endian `u16` from 'r'
fn read_le_u16<R>(r: &mut R) -> Result<u16, Error>
    where R: io::Read {

    let itr = r.bytes().take_or_err(2, Error::EOF);

    result::fold(itr, 0, |so_far, next| (so_far << 8) + next as u16)
}

/// `read_le_u32` reads 4 bytes as a little endian `u32` from 'r'
fn read_le_u32<R>(r: &mut R) -> Result<u32, Error>
    where R: io::Read {

    let itr = r.bytes().take_or_err(4, Error::EOF);

    result::fold(itr, 0, |so_far, next| (so_far << 8) + next as u32)
}

/// `read_le_u64` reads 8 bytes as a little endian `u64` from 'r'
fn read_le_u64<R>(r: &mut R) -> Result<u64, Error>
    where R: io::Read {

    let itr = r.bytes().take_or_err(8, Error::EOF);

    result::fold(itr, 0, |so_far, next| (so_far << 8) + next as u64)
}

/// `read_i8` reads a single byte from `r`, as a 2's complement `i8`
fn read_i8<R>(r: &mut R) -> Result<i8, Error>
    where R: io::Read {

    Ok(try!(read_u8(r)) as i8)
}

/// `read_le_i16` reads 2 bytes as a little endian 2's complement `i16` from `r`.
fn read_le_i16<R>(r: &mut R) -> Result<i16, Error>
    where R: io::Read {

    Ok(try!(read_le_u16(r)) as i16)
}

/// `read_le_i32` reads 4 bytes as a little endian 2's complement `i32` from `r`.
fn read_le_i32<R>(r: &mut R) -> Result<i32, Error>
    where R: io::Read {

    Ok(try!(read_le_u32(r)) as i32)
}

/// `read_le_i64` reads 8 bytes as a little endian 2's complement `i64` from `r`.
fn read_le_i64<R>(r: &mut R) -> Result<i64, Error>
    where R: io::Read {

    Ok(try!(read_le_u64(r)) as i64)
}

/// `read_le_f32` reads 8 bytes as a little endian ieee-754 binary32 encoded `f32` from `r`.
fn read_le_f32<R>(r: &mut R) -> Result<f32, Error>
    where R: io::Read {

    let u = try!(read_le_u32(r));

    Ok(unsafe { mem::transmute(u) })
}

/// `read_le_f64` reads 8 bytes as a little endian ieee-754 binary64 encoded `f64` from `r`.
fn read_le_f64<R>(r: &mut R) -> Result<f64, Error>
    where R: io::Read {

    let u = try!(read_le_u64(r));

    Ok(unsafe { mem::transmute(u) })
}
<commit_msg>Fleshed out decoder.rs<commit_after>// Copyright (c) 2015, Sam Payson
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
// associated documentation files (the "Software"), to deal in the Software without restriction,
// including without limitation the rights to use, copy, modify, merge, publish, distribute,
// sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all copies or
// substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
// NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

// Allow dead code to silence warnings until things stabilize.
#![allow(dead_code)]

use std::error::FromError;
use std::io;
use std::io::ReadExt;
use std::mem;
use std::result;
use std::string::FromUtf8Error;

use encoding::{CompleteEncoding, RecordEncoding, FieldEncoding, Type, FieldID};
use primitive::Primitive;

use iter::{RexIterExt, ResultIterExt};

/// `Error` is used to report errors that occur during the decoding process.
pub enum Error {
    /// `EOF` indicates that record source ended before it could be fully decoded.
    EOF,

    /// `FieldTypeMismatch` indicates a disagreement between the `Decoder` and the `Decodable`
    /// about what the type of a field is. This probably indicates one of two things:
    ///
    /// 1. The wrong `CompleteEncoding` is being used to decode this type.
    /// 2. The `Decodable` implementation is not consistent with the record definition.
    FieldTypeMismatch,

    /// `EncodingInvalid` indicates that there is an inconsistency in the `CompleteEncoding` being
    /// used. A probable fix is to regenerate the encoding (or inspect the encoding compiler for
    /// errors).
    EncodingInvalid,

    /// `BadBool` indicates that a bool was read off the wire as a value other than `0xFF` or
    /// `0x00`, indicating that the decoder is incorrect in expecting it to be a bool.
    BadBool,

    /// `Utf8Error` is used to pass through `std::str::FromUtf8Error`s.
    Utf8Error(FromUtf8Error),

    /// `IoError` is used to pass through `std::io` errors.
    IoError(io::Error),
}

impl FromError<io::Error> for Error {
    fn from_error(err: io::Error) -> Error {
        Error::IoError(err)
    }
}

impl FromError<FromUtf8Error> for Error {
    fn from_error(err: FromUtf8Error) -> Error {
        Error::Utf8Error(err)
    }
}

/// The `Decodable` trait allows an object to be decoded from a rex record.
pub trait Decodable {
    /// `set_primitive` sets the value of a single element of a field with primitive type.
    fn set_primitive(&mut self, id: FieldID, idx: usize, prim: Primitive) -> Result<(), Error>;

    /// `decode_record` is a request to call `d.decode` on a record field.
    fn decode_record<'x, R>(&mut self, d: Decoder<'x, R>, id: FieldID, idx: usize) -> Result<(), Error>
        where R: io::Read + 'x;

    /// `alloc_field` requests that the receiver allocate space for an optional or repeated field.
    /// The boolean return value can be `false` to indicate that the receiver is uninterested in
    /// this field.
    fn alloc_field(&mut self, id: FieldID, count: usize) -> Result<bool, Error>;
}

/// A `Decoder` is a struct that knows how to decode a particular record field. `Decoder`s should
/// only be used where they are passed to the `decode_record` method of a `Decodable`.
pub struct Decoder<'x, R: io::Read + 'x> {
    // `r` is the reader that the record is being read from.
    r: &'x mut R,

    // `rec` is the record type that this `Decoder` knows how to decode.
    rec: &'x RecordEncoding,

    // `deps` is a (full) slice of the `depends` field of the `CompleteEncoding` that `rec` is a
    // member of.
    deps: &'x [RecordEncoding],
}

impl<'x, R> Decoder<'x, R>
    where R: io::Read + 'x {

    /// `decode` decodes the next record on the wire into `d`.
    pub fn decode<D>(&mut self, d: &mut D) -> Result<(), Error>
        where D: Decodable {

        use encoding::Quantifier::*;

        for req_field in self.rec.req_fields.iter() {
            try!(self.decode_required(d, req_field));
        }

        let mut opt_rep_itr = self.rec.opt_rep_fields.iter();
        let mut next_field  = opt_rep_itr.next();
        let mut next_id     = FieldID(try!(read_uvarint(self.r)));

        while next_id != FieldID(0) {
            match next_field {
                Some(field) => if field.id < next_id {
                    next_field = opt_rep_itr.next();
                } else if field.id > next_id {
                    try!(self.skip_field());
                    next_id = FieldID(try!(read_uvarint(self.r)));
                } else {
                    match field.quant {
                        Required => return Err(Error::EncodingInvalid),
                        Optional => try!(self.decode_optional(d, field)),
                        Repeated => try!(self.decode_repeated(d, field)),
                    }
                    next_field = opt_rep_itr.next();
                    next_id    = FieldID(try!(read_uvarint(self.r)));
                },
                None => next_id = FieldID(try!(read_uvarint(self.r))),
            }
        }

        Ok(())
    }

    fn skip_field(&mut self) -> Result<(), Error> {
        let len = try!(read_uvarint(self.r)) as usize;
        result::fold(self.r.bytes().take_or_err(len, Error::EOF), (), |(), _| ())
    }

    fn child(&mut self, index: usize) -> Result<Decoder<R>, Error> {
        if index < self.deps.len() {
            Ok( Decoder {
                r:      self.r,
                rec:    &self.deps[index],
                deps:   self.deps,
            })
        } else {
            Err(Error::EncodingInvalid)
        }
    }

    fn decode_required<D>(&mut self, d: &mut D, f: &FieldEncoding) -> Result<(), Error>
        where D: Decodable {

        self.decode_array(d, f, 0)
    }

    fn decode_optional<D>(&mut self, d: &mut D, f: &FieldEncoding) -> Result<(), Error>
        where D: Decodable {

        // Discard the byte-size prefix. We don't need it.
        try!(read_uvarint(self.r));

        self.decode_array(d, f, 0)
    }

    fn decode_repeated<D>(&mut self, d: &mut D, f: &FieldEncoding) -> Result<(), Error>
        where D: Decodable {

        // Discard the byte-size prefix. We don't need it.
        try!(read_uvarint(self.r));

        let len = try!(read_uvarint(self.r)) as usize;

        for idx in 0..len {
            try!(self.decode_array(d, f, idx));
        }

        Ok(())
    }

    fn decode_array<D>(&mut self, d: &mut D, f: &FieldEncoding, idx: usize) -> Result<(), Error>
        where D: Decodable {

        if let Some(max) = f.bounds {
            for arr_index in 0..max {
                try!(self.decode_field(d, f, idx*max + arr_index));
            }
        } else {
            try!(self.decode_field(d, f, idx));
        }

        Ok(())
    }

    fn decode_field<D>(&mut self, d: &mut D, f: &FieldEncoding, idx: usize) -> Result<(), Error>
        where D: Decodable {

        let prim = match f.typ {
            Type::UInt8  => Primitive::UInt8(try!(read_u8(self.r))),
            Type::UInt16 => Primitive::UInt16(try!(read_le_u16(self.r))),
            Type::UInt32 => Primitive::UInt32(try!(read_uvarint(self.r)) as u32),
            Type::UInt64 => Primitive::UInt64(try!(read_uvarint(self.r))),

            Type::Int8  => Primitive::Int8(try!(read_i8(self.r))),
            Type::Int16 => Primitive::Int16(try!(read_le_i16(self.r))),
            Type::Int32 => Primitive::Int32(try!(read_varint(self.r)) as i32),
            Type::Int64 => Primitive::Int64(try!(read_varint(self.r))),

            Type::Fixed32 => Primitive::Fixed32(try!(read_le_u32(self.r))),
            Type::Fixed64 => Primitive::Fixed64(try!(read_le_u64(self.r))),

            Type::Float32 => Primitive::Float32(try!(read_le_f32(self.r))),
            Type::Float64 => Primitive::Float64(try!(read_le_f64(self.r))),

            Type::Bool => Primitive::Bool(match try!(read_u8(self.r)) {
                0xFF => true,
                0x00 => false,
                _    => return Err(Error::BadBool),
            }),

            Type::Bytes => Primitive::Bytes({
                let len = try!(read_uvarint(self.r)) as usize;
                try!(result::fold(
                    self.r.bytes().take_or_err(len, Error::EOF),
                    Vec::with_capacity(len),
                    |mut v, elem| {
                        v.push(elem);
                        v
                }))
            }),

            Type::String => Primitive::String({
                let len = try!(read_uvarint(self.r)) as usize;
                let utf8 = try!(result::fold(
                    self.r.bytes().take_or_err(len, Error::EOF),
                    Vec::with_capacity(len),
                    |mut v, elem| {
                        v.push(elem);
                        v
                }));

                try!(String::from_utf8(utf8))
            }),

            Type::Enum => Primitive::Enum(try!(read_varint(self.r))),

            // Records work a little differently. Create a child decoder and have the `Decodable`
            // run it on its own record field.
            Type::Record{index: dep_index} => {
                let child = try!(self.child(dep_index));
                return d.decode_record(child, f.id, idx);
            },
        };

        d.set_primitive(f.id, idx, prim)
    }
}

/// `read_uvarint` reads a varint encoded `u64` from `r`.
fn read_uvarint<R>(r: &mut R) -> Result<u64, Error>
    where R: io::Read {

    let itr = r.bytes()
        .take_while_incl(|x| match x { &Ok(x) => x >= 0x80, _ => true });

    result::fold(itr, 0, |sum, x| (sum << 7) + (x & 0x7F) as u64)
        .map_err(FromError::from_error)
}

/// `read_varint` reads a zig-zag varint encoded `i64` from `r`.
fn read_varint<R>(r: &mut R) -> Result<i64, Error>
    where R: io::Read {

    let ux = try!(read_uvarint(r));

    Ok( if ux & 1 != 0 {
        let x = (!ux >> 1) as i64;
        -x
    } else {
        let x = (ux >> 1) as i64;
        x
    })
}

/// `read_u8` reads a single byte from `r`.
fn read_u8<R>(r: &mut R) -> Result<u8, Error>
    where R: io::Read {

    r.bytes()
        .take_or_err(1, Error::EOF)
        .next()
        // take_or_err guarantees that we'll get an Error at worst.
        .unwrap()
}

/// `read_le_u16` reads 2 bytes as a little endian `u16` from 'r'
fn read_le_u16<R>(r: &mut R) -> Result<u16, Error>
    where R: io::Read {

    let itr = r.bytes().take_or_err(2, Error::EOF);

    result::fold(itr, 0, |so_far, next| (so_far << 8) + next as u16)
}

/// `read_le_u32` reads 4 bytes as a little endian `u32` from 'r'
fn read_le_u32<R>(r: &mut R) -> Result<u32, Error>
    where R: io::Read {

    let itr = r.bytes().take_or_err(4, Error::EOF);

    result::fold(itr, 0, |so_far, next| (so_far << 8) + next as u32)
}

/// `read_le_u64` reads 8 bytes as a little endian `u64` from 'r'
fn read_le_u64<R>(r: &mut R) -> Result<u64, Error>
    where R: io::Read {

    let itr = r.bytes().take_or_err(8, Error::EOF);

    result::fold(itr, 0, |so_far, next| (so_far << 8) + next as u64)
}

/// `read_i8` reads a single byte from `r`, as a 2's complement `i8`
fn read_i8<R>(r: &mut R) -> Result<i8, Error>
    where R: io::Read {

    Ok(try!(read_u8(r)) as i8)
}

/// `read_le_i16` reads 2 bytes as a little endian 2's complement `i16` from `r`.
fn read_le_i16<R>(r: &mut R) -> Result<i16, Error>
    where R: io::Read {

    Ok(try!(read_le_u16(r)) as i16)
}

/// `read_le_i32` reads 4 bytes as a little endian 2's complement `i32` from `r`.
fn read_le_i32<R>(r: &mut R) -> Result<i32, Error>
    where R: io::Read {

    Ok(try!(read_le_u32(r)) as i32)
}

/// `read_le_i64` reads 8 bytes as a little endian 2's complement `i64` from `r`.
fn read_le_i64<R>(r: &mut R) -> Result<i64, Error>
    where R: io::Read {

    Ok(try!(read_le_u64(r)) as i64)
}

/// `read_le_f32` reads 8 bytes as a little endian ieee-754 binary32 encoded `f32` from `r`.
fn read_le_f32<R>(r: &mut R) -> Result<f32, Error>
    where R: io::Read {

    let u = try!(read_le_u32(r));

    Ok(unsafe { mem::transmute(u) })
}

/// `read_le_f64` reads 8 bytes as a little endian ieee-754 binary64 encoded `f64` from `r`.
fn read_le_f64<R>(r: &mut R) -> Result<f64, Error>
    where R: io::Read {

    let u = try!(read_le_u64(r));

    Ok(unsafe { mem::transmute(u) })
}
<|endoftext|>


<commit_before>use bigint::BigUint;

pub struct AsciiEncoder;
pub struct Utf8Encoder;

macro_rules! encode {
    ($alpha:ident, $input:ident) => ({
        if $input.len() == 0 {
            return String::new();
        }

        let base = $alpha.len() as u32;

        let mut big = BigUint::from($input);
        let mut out = Vec::with_capacity($input.len());

        loop {
            out.push($alpha[big.rem_div(base) as usize]);

            if big.is_zero() {
                break;
            }
        }

        let leaders = $input
            .iter()
            .take($input.len() - 1)
            .take_while(|i| **i == 0)
            .map(|_| $alpha[0]);

        out.extend(leaders);

        out
    })
}

impl AsciiEncoder {
    #[inline(always)]
    pub fn encode(alphabet: &[u8], input: &[u8]) -> String {
        let mut out = encode!(alphabet, input);

        out.reverse();

        String::from_utf8(out).expect("Alphabet must be ASCII")
    }
}

impl Utf8Encoder {
    #[inline(always)]
    pub fn encode(alphabet: &[char], input: &[u8]) -> String {
        let out = encode!(alphabet, input);

        out.iter().rev().map(|char| *char).collect()
    }
}
<commit_msg>IT'S ALIVE<commit_after>use bigint::BigUint;

pub struct AsciiEncoder;
pub struct Utf8Encoder;

macro_rules! encode {
    ($alpha:ident, $input:ident) => ({
        if $input.len() == 0 {
            return String::new();
        }

        let base = $alpha.len() as u32;
        let big_pow = 32 / (32 - base.leading_zeros());
        let big_base = base.pow(big_pow);

        let mut big = BigUint::from($input);
        let mut out = Vec::with_capacity($input.len());

        loop {
            let mut big_rem = big.rem_div(big_base);

            if big.is_zero() {
                loop {
                    out.push($alpha[(big_rem % base) as usize]);
                    big_rem /= base;

                    if big_rem == 0 {
                        break;
                    }
                }
                break;
            } else {
                for _ in 0..big_pow {
                    out.push($alpha[(big_rem % base) as usize]);
                    big_rem /= base;
                }
            }
        }

        let leaders = $input
            .iter()
            .take($input.len() - 1)
            .take_while(|i| **i == 0)
            .map(|_| $alpha[0]);

        out.extend(leaders);

        out
    })
}

impl AsciiEncoder {
    #[inline(always)]
    pub fn encode(alphabet: &[u8], input: &[u8]) -> String {
        let mut out = encode!(alphabet, input);

        out.reverse();

        String::from_utf8(out).expect("Alphabet must be ASCII")
    }
}

impl Utf8Encoder {
    #[inline(always)]
    pub fn encode(alphabet: &[char], input: &[u8]) -> String {
        let out = encode!(alphabet, input);

        out.iter().rev().map(|char| *char).collect()
    }
}
<|endoftext|>


<commit_before>//! Rust high level bindings for [IUP][1] -- A Portable User Interface Toolkit
//!
//! [1]: http://www.tecgraf.puc-rio.br/iup/
//!
//! The functions in this crate are a thin wrapper around the IUP functions unless noted
//! otherwise. Refer to the IUP website for their full documentation.
//!

extern crate libc;
extern crate "iup-sys" as sys;

use std::ffi::CString;
use std::mem;
use std::ptr;
use std::slice::SliceExt;
use libc::{c_char, c_int};

unsafe fn vec_to_c_array(v: Vec<Ihandle>) -> *mut *mut sys::Ihandle {
    let mut raw_v = Vec::with_capacity(v.len());
    for ih in v {
        raw_v.push(ih.ptr);
    }
    let null : *const sys::Ihandle = ptr::null();
    raw_v.push(mem::transmute(null));
    raw_v.as_mut_ptr()
}

#[allow(missing_copy_implementations)]
pub struct Ihandle {
    ptr: *mut sys::Ihandle,
}

impl Ihandle {
    fn from_ptr(ih: *mut sys::Ihandle) -> Ihandle {
        if ih.is_null() {
            panic!("Failed to create Ihandle.")
        } else {
            Ihandle { ptr: ih }
        }
    }
}

pub type IupResult = Result<(), String>;

/// Initializes the IUP toolkit. Must be called before any other IUP function.
pub fn open() -> IupResult {
    let p1 : *const c_int = ptr::null();
    let p2 : *const *const *const c_char = ptr::null();

    match unsafe { sys::IupOpen(p1, p2) } {
        sys::IUP_NOERROR => Ok(()),
        sys::IUP_OPENED => Err("IUP_OPENED: iup::open called while already open.".to_string()),
        sys::IUP_ERROR => Err("IUP_ERROR: X-Windows is not initialized".to_string()),
        _ => unreachable!(),
    }
}

pub fn close() {
    unsafe { sys::IupClose(); }
}

pub fn main_loop() {
    let ok = unsafe { sys::IupMainLoop() };
    assert_eq!(ok, sys::IUP_NOERROR);
}

pub fn show(ih: &mut Ihandle) -> IupResult {
    match unsafe { sys::IupShow(ih.ptr) } {
        sys::IUP_NOERROR => Ok(()),
        sys::IUP_ERROR => Err("IUP_ERROR: unknown error".to_string()),
        _ => unreachable!(),
    }
}

pub fn set_attribute(ih: &mut Ihandle, name: &str, value: &str) {
    let name_c = CString::from_slice(name.as_bytes());
    let value_c = CString::from_slice(value.as_bytes());
    unsafe { sys::IupSetAttribute(ih.ptr, name_c.as_ptr(), value_c.as_ptr()); }
}

pub fn hboxv(elements: Vec<Ihandle>) -> Ihandle {
    unsafe { Ihandle::from_ptr(sys::IupHboxv(vec_to_c_array(elements))) }
}

// Elements

pub fn label(text: &str) -> Ihandle {
    let text_c = CString::from_slice(text.as_bytes());
    unsafe { Ihandle::from_ptr(sys::IupLabel(text_c.as_ptr())) }
}

pub fn dialog(ih: Ihandle) -> Ihandle {
    unsafe { Ihandle::from_ptr(sys::IupDialog(ih.ptr)) }
}

pub fn button(text: &str) -> Ihandle {
    let text_c = CString::from_slice(text.as_bytes());
    let action: *const c_char = ptr::null();
    unsafe { Ihandle::from_ptr(sys::IupButton(text_c.as_ptr(), action)) }
}

pub fn text() -> Ihandle {
    let action: *const c_char = ptr::null();
    unsafe { Ihandle::from_ptr(sys::IupText(action)) }
}
<commit_msg>replace set_attribute with set_str_attribute<commit_after>//! Rust high level bindings for [IUP][1] -- A Portable User Interface Toolkit
//!
//! [1]: http://www.tecgraf.puc-rio.br/iup/
//!
//! The functions in this crate are a thin wrapper around the IUP functions unless noted
//! otherwise. Refer to the IUP website for their full documentation.
//!

extern crate libc;
extern crate "iup-sys" as sys;

use std::ffi::CString;
use std::mem;
use std::ptr;
use std::slice::SliceExt;
use libc::{c_char, c_int};

unsafe fn vec_to_c_array(v: Vec<Ihandle>) -> *mut *mut sys::Ihandle {
    let mut raw_v = Vec::with_capacity(v.len());
    for ih in v {
        raw_v.push(ih.ptr);
    }
    let null : *const sys::Ihandle = ptr::null();
    raw_v.push(mem::transmute(null));
    raw_v.as_mut_ptr()
}

#[allow(missing_copy_implementations)]
pub struct Ihandle {
    ptr: *mut sys::Ihandle,
}

impl Ihandle {
    fn from_ptr(ih: *mut sys::Ihandle) -> Ihandle {
        if ih.is_null() {
            panic!("Failed to create Ihandle.")
        } else {
            Ihandle { ptr: ih }
        }
    }
}

pub type IupResult = Result<(), String>;

/// Initializes the IUP toolkit. Must be called before any other IUP function.
pub fn open() -> IupResult {
    let p1 : *const c_int = ptr::null();
    let p2 : *const *const *const c_char = ptr::null();

    match unsafe { sys::IupOpen(p1, p2) } {
        sys::IUP_NOERROR => Ok(()),
        sys::IUP_OPENED => Err("IUP_OPENED: iup::open called while already open.".to_string()),
        sys::IUP_ERROR => Err("IUP_ERROR: X-Windows is not initialized".to_string()),
        _ => unreachable!(),
    }
}

pub fn close() {
    unsafe { sys::IupClose(); }
}

pub fn main_loop() {
    let ok = unsafe { sys::IupMainLoop() };
    assert_eq!(ok, sys::IUP_NOERROR);
}

pub fn show(ih: &mut Ihandle) -> IupResult {
    match unsafe { sys::IupShow(ih.ptr) } {
        sys::IUP_NOERROR => Ok(()),
        sys::IUP_ERROR => Err("IUP_ERROR: unknown error".to_string()),
        _ => unreachable!(),
    }
}

pub fn set_str_attribute(ih: &mut Ihandle, name: &str, value: &str) {
    let name_c = CString::from_slice(name.as_bytes());
    let value_c = CString::from_slice(value.as_bytes());
    unsafe { sys::IupSetStrAttribute(ih.ptr, name_c.as_ptr(), value_c.as_ptr()); }
}

pub fn hboxv(elements: Vec<Ihandle>) -> Ihandle {
    unsafe { Ihandle::from_ptr(sys::IupHboxv(vec_to_c_array(elements))) }
}

// Elements

pub fn label(text: &str) -> Ihandle {
    let text_c = CString::from_slice(text.as_bytes());
    unsafe { Ihandle::from_ptr(sys::IupLabel(text_c.as_ptr())) }
}

pub fn dialog(ih: Ihandle) -> Ihandle {
    unsafe { Ihandle::from_ptr(sys::IupDialog(ih.ptr)) }
}

pub fn button(text: &str) -> Ihandle {
    let text_c = CString::from_slice(text.as_bytes());
    let action: *const c_char = ptr::null();
    unsafe { Ihandle::from_ptr(sys::IupButton(text_c.as_ptr(), action)) }
}

pub fn text() -> Ihandle {
    let action: *const c_char = ptr::null();
    unsafe { Ihandle::from_ptr(sys::IupText(action)) }
}
<|endoftext|>


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

//!
//! Protect the server's long-term key with envelope encryption and a key management system.
//!
//! Note: KMS support must be enabled at compile time, see the Roughenough's [documentation
//! on optional features](https://github.com/int08h/roughenough/blob/master/doc/OPTIONAL-FEATURES.md#key-management-system-kms-support)
//! for instructions.
//!
//! ## Motivation
//!
//! The seed for the server's [long-term key](../key/struct.LongTermKey.html) is subject to
//! contradictory requirements:
//!
//!   1. The seed must be kept secret, but
//!   2. The seed must be available at server start-up to create the
//!      [delegated on-line key](../key/struct.OnlineKey.html)
//!
//! ## Plaintext seed
//!
//! The default option is to store the seed in plaintext as part of the server's configuration.
//! This usually means the seed is present in the clear: on disk, in a repository, or otherwise
//! durably persisted where it can be compromised (accidentally or maliciously).
//!
//! ## Encrypting the seed
//!
//! Envelope encryption protects the seed by encrypting it with a locally generated 256-bit
//! Data Encryption Key (DEK). The DEK itself is then encrypted using a cloud key management
//! system (KMS). The resulting opaque encrypted "blob" (encrypted seed + encrypted DEK) is
//! stored in the Roughenough configuration.
//!
//! At server start-up the KMS is used to decrypt the DEK, which is then used to (in memory)
//! decrypt the seed. The seed is used to generate the
//! [delegated on-line key](../key/struct.OnlineKey.html) after which the seed and DEK are erased
//! from memory.
//!
//! See
//!   * [`EnvelopeEncryption`](struct.EnvelopeEncryption.html) for Roughenough's implementation.
//!   * [Google](https://cloud.google.com/kms/docs/envelope-encryption) or
//!     [Amazon](https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#enveloping)
//!     for more in-depth explanations of envelope encryption.
//!

mod envelope;

use base64;
use ring;
use std;

use crate::config::ServerConfig;
use crate::error;
use crate::key::KmsProtection;

pub use self::envelope::EnvelopeEncryption;

/// Errors generated by KMS operations
#[derive(Debug, PartialEq, Eq, PartialOrd, Hash, Clone)]
pub enum KmsError {
    OperationFailed(String),
    InvalidConfiguration(String),
    InvalidData(String),
    InvalidKey(String),
}

impl From<std::io::Error> for KmsError {
    fn from(error: std::io::Error) -> Self {
        KmsError::OperationFailed(format!("{:?}", error))
    }
}

impl From<ring::error::Unspecified> for KmsError {
    fn from(_: ring::error::Unspecified) -> Self {
        KmsError::OperationFailed("unspecified ring cryptographic failure".to_string())
    }
}

impl From<base64::DecodeError> for KmsError {
    fn from(error: base64::DecodeError) -> Self {
        KmsError::OperationFailed(format!("base64: {}", error))
    }
}

// Size of the AEAD nonce in bytes.
const NONCE_SIZE_BYTES: usize = 12;

// Size of the AEAD authentication tag in bytes.
const TAG_SIZE_BYTES: usize = 16;

// Size of the 256-bit Data Encryption Key (DEK) in bytes.
const DEK_SIZE_BYTES: usize = 32;

// Trivial domain separation to guard against KMS key reuse
const AD: &str = "roughenough";

/// An unencrypted (plaintext) 256-bit Data Encryption Key (DEK).
pub type PlaintextDEK = Vec<u8>;

/// A Data Encryption Key (DEK) that has been encrypted (wrapped) by a Key Management System (KMS).
///
/// This is an opaque, implementation-specific value. AEAD tag size, nonce size,
/// provider metadata, and so on will vary between [`KmsProvider`](trait.KmsProvider.html)
/// implementations.
pub type EncryptedDEK = Vec<u8>;

///
/// A key management system that wraps/unwraps a data encryption key (DEK).
///
pub trait KmsProvider {
    /// Make a blocking request to encrypt (wrap) the provided plaintext data encryption key.
    fn encrypt_dek(&self, plaintext_dek: &PlaintextDEK) -> Result<EncryptedDEK, KmsError>;

    /// Make a blocking request to decrypt (unwrap) a previously encrypted data encryption key.
    fn decrypt_dek(&self, encrypted_dek: &EncryptedDEK) -> Result<PlaintextDEK, KmsError>;
}

#[cfg(feature = "awskms")]
mod awskms;

#[cfg(feature = "awskms")]
pub use crate::kms::awskms::inner::AwsKms;

/// Load the seed value for the long-term key.
///
/// Loading behavior depends on the value of `config.kms_protection()`:
///
///  * If `config.kms_protection() == Plaintext` then the value returned from `config.seed()`
///    is used as-is and assumed to be a 32-byte hexadecimal value.
///
///  * Otherwise `config.seed()` is assumed to be an encrypted opaque blob generated from
///    a prior `EnvelopeEncryption::encrypt_seed` call. The value of `config.kms_protection()`
///    is parsed as a KMS key id and `EnvelopeEncryption::decrypt_seed` is called to obtain
///    the plaintext seed value.
///
#[cfg(feature = "awskms")]
pub fn load_seed(config: &Box<ServerConfig>) -> Result<Vec<u8>, error::Error> {
    match config.kms_protection() {
        KmsProtection::Plaintext => Ok(config.seed()),
        KmsProtection::AwsKmsEnvelope(key_id) => {
            info!("Unwrapping seed via AWS KMS key '{}'", key_id);
            let kms = AwsKms::from_arn(key_id)?;
            let seed = EnvelopeEncryption::decrypt_seed(&kms, &config.seed())?;
            Ok(seed)
        }
        _ => Err(error::Error::InvalidConfiguration(
            "Google KMS not supported".to_string(),
        )),
    }
}

#[cfg(feature = "gcpkms")]
mod gcpkms;

#[cfg(feature = "gcpkms")]
pub use crate::kms::gcpkms::inner::GcpKms;

/// Load the seed value for the long-term key.
///
/// Loading behavior depends on the value of `config.kms_protection()`:
///
///  * If `config.kms_protection() == Plaintext` then the value returned from `config.seed()`
///    is used as-is and assumed to be a 32-byte hexadecimal value.
///
///  * Otherwise `config.seed()` is assumed to be an encrypted opaque blob generated from
///    a prior `EnvelopeEncryption::encrypt_seed` call. The value of `config.kms_protection()`
///    is parsed as a KMS key id and `EnvelopeEncryption::decrypt_seed` is called to obtain
///    the plaintext seed value.
///
#[cfg(feature = "gcpkms")]
pub fn load_seed(config: &Box<ServerConfig>) -> Result<Vec<u8>, error::Error> {
    match config.kms_protection() {
        KmsProtection::Plaintext => Ok(config.seed()),
        KmsProtection::GoogleKmsEnvelope(resource_id) => {
            info!("Unwrapping seed via Google KMS key '{}'", resource_id);
            let kms = GcpKms::from_resource_id(resource_id)?;
            let seed = EnvelopeEncryption::decrypt_seed(&kms, &config.seed())?;
            Ok(seed)
        }
        _ => Err(error::Error::InvalidConfiguration(
            "AWS KMS not supported".to_string(),
        )),
    }
}

/// Load the seed value for the long-term key.
///
/// ## This build has KMS disabled
///
/// *The KMS feature is disabled in this build of Roughenough*.
///
/// The only supported `kms_protection` value in this build is `plaintext`. Any
/// other value will cause a runtime error.
///
/// ## Background
///
/// Loading behavior depends on the value of `config.kms_protection()`:
///
///  * If `config.kms_protection() == Plaintext` then the value returned from `config.seed()`
///    is used as-is and assumed to be a 32-byte hexadecimal value.
///
///  * Otherwise `config.seed()` is assumed to be an encrypted opaque blob generated from
///    a prior `EnvelopeEncryption::encrypt_seed` call. The value of `config.kms_protection()`
///    is parsed as a KMS key id and `EnvelopeEncryption::decrypt_seed` is called to obtain
///    the plaintext seed value.
///
#[cfg(not(any(feature = "awskms", feature = "gcpkms")))]
pub fn load_seed(config: &Box<ServerConfig>) -> Result<Vec<u8>, error::Error> {
    match config.kms_protection() {
        KmsProtection::Plaintext => Ok(config.seed()),
        v => Err(error::Error::InvalidConfiguration(format!(
            "kms_protection '{}' requires KMS, but server was not compiled with KMS support",
            v
        ))),
    }
}
<commit_msg>Update module docs when KMS feature is disabled<commit_after>// Copyright 2017-2019 int08h LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//!
//! Protect the server's long-term key with envelope encryption and a key management system.
//!
//! Note: KMS support must be enabled at compile time, see the Roughenough's [documentation
//! on optional features](https://github.com/int08h/roughenough/blob/master/doc/OPTIONAL-FEATURES.md#key-management-system-kms-support)
//! for instructions.
//!
//! ## Motivation
//!
//! The seed for the server's [long-term key](../key/struct.LongTermKey.html) is subject to
//! contradictory requirements:
//!
//!   1. The seed must be kept secret, but
//!   2. The seed must be available at server start-up to create the
//!      [delegated on-line key](../key/struct.OnlineKey.html)
//!
//! ## Plaintext seed
//!
//! The default option is to store the seed in plaintext as part of the server's configuration.
//! This usually means the seed is present in the clear: on disk, in a repository, or otherwise
//! durably persisted where it can be compromised (accidentally or maliciously).
//!
//! ## Encrypting the seed
//!
//! Envelope encryption protects the seed by encrypting it with a locally generated 256-bit
//! Data Encryption Key (DEK). The DEK itself is then encrypted using a cloud key management
//! system (KMS). The resulting opaque encrypted "blob" (encrypted seed + encrypted DEK) is
//! stored in the Roughenough configuration.
//!
//! At server start-up the KMS is used to decrypt the DEK, which is then used to (in memory)
//! decrypt the seed. The seed is used to generate the
//! [delegated on-line key](../key/struct.OnlineKey.html) after which the seed and DEK are erased
//! from memory.
//!
//! See
//!   * [`EnvelopeEncryption`](struct.EnvelopeEncryption.html) for Roughenough's implementation.
//!   * [Google](https://cloud.google.com/kms/docs/envelope-encryption) or
//!     [Amazon](https://docs.aws.amazon.com/kms/latest/developerguide/concepts.html#enveloping)
//!     for more in-depth explanations of envelope encryption.
//!

mod envelope;

use base64;
use ring;
use std;

use crate::config::ServerConfig;
use crate::error;
use crate::key::KmsProtection;

pub use self::envelope::EnvelopeEncryption;

/// Errors generated by KMS operations
#[derive(Debug, PartialEq, Eq, PartialOrd, Hash, Clone)]
pub enum KmsError {
    OperationFailed(String),
    InvalidConfiguration(String),
    InvalidData(String),
    InvalidKey(String),
}

impl From<std::io::Error> for KmsError {
    fn from(error: std::io::Error) -> Self {
        KmsError::OperationFailed(format!("{:?}", error))
    }
}

impl From<ring::error::Unspecified> for KmsError {
    fn from(_: ring::error::Unspecified) -> Self {
        KmsError::OperationFailed("unspecified ring cryptographic failure".to_string())
    }
}

impl From<base64::DecodeError> for KmsError {
    fn from(error: base64::DecodeError) -> Self {
        KmsError::OperationFailed(format!("base64: {}", error))
    }
}

// Size of the AEAD nonce in bytes.
const NONCE_SIZE_BYTES: usize = 12;

// Size of the AEAD authentication tag in bytes.
const TAG_SIZE_BYTES: usize = 16;

// Size of the 256-bit Data Encryption Key (DEK) in bytes.
const DEK_SIZE_BYTES: usize = 32;

// Trivial domain separation to guard against KMS key reuse
const AD: &str = "roughenough";

/// An unencrypted (plaintext) 256-bit Data Encryption Key (DEK).
pub type PlaintextDEK = Vec<u8>;

/// A Data Encryption Key (DEK) that has been encrypted (wrapped) by a Key Management System (KMS).
///
/// This is an opaque, implementation-specific value. AEAD tag size, nonce size,
/// provider metadata, and so on will vary between [`KmsProvider`](trait.KmsProvider.html)
/// implementations.
pub type EncryptedDEK = Vec<u8>;

///
/// A key management system that wraps/unwraps a data encryption key (DEK).
///
pub trait KmsProvider {
    /// Make a blocking request to encrypt (wrap) the provided plaintext data encryption key.
    fn encrypt_dek(&self, plaintext_dek: &PlaintextDEK) -> Result<EncryptedDEK, KmsError>;

    /// Make a blocking request to decrypt (unwrap) a previously encrypted data encryption key.
    fn decrypt_dek(&self, encrypted_dek: &EncryptedDEK) -> Result<PlaintextDEK, KmsError>;
}

#[cfg(feature = "awskms")]
mod awskms;

#[cfg(feature = "awskms")]
pub use crate::kms::awskms::inner::AwsKms;

/// Load the seed value for the long-term key.
///
/// Loading behavior depends on the value of `config.kms_protection()`:
///
///  * If `config.kms_protection() == Plaintext` then the value returned from `config.seed()`
///    is used as-is and assumed to be a 32-byte hexadecimal value.
///
///  * Otherwise `config.seed()` is assumed to be an encrypted opaque blob generated from
///    a prior `EnvelopeEncryption::encrypt_seed` call. The value of `config.kms_protection()`
///    is parsed as a KMS key id and `EnvelopeEncryption::decrypt_seed` is called to obtain
///    the plaintext seed value.
///
#[cfg(feature = "awskms")]
pub fn load_seed(config: &Box<ServerConfig>) -> Result<Vec<u8>, error::Error> {
    match config.kms_protection() {
        KmsProtection::Plaintext => Ok(config.seed()),
        KmsProtection::AwsKmsEnvelope(key_id) => {
            info!("Unwrapping seed via AWS KMS key '{}'", key_id);
            let kms = AwsKms::from_arn(key_id)?;
            let seed = EnvelopeEncryption::decrypt_seed(&kms, &config.seed())?;
            Ok(seed)
        }
        _ => Err(error::Error::InvalidConfiguration(
            "Google KMS not supported".to_string(),
        )),
    }
}

#[cfg(feature = "gcpkms")]
mod gcpkms;

#[cfg(feature = "gcpkms")]
pub use crate::kms::gcpkms::inner::GcpKms;

/// Load the seed value for the long-term key.
///
/// Loading behavior depends on the value of `config.kms_protection()`:
///
///  * If `config.kms_protection() == Plaintext` then the value returned from `config.seed()`
///    is used as-is and assumed to be a 32-byte hexadecimal value.
///
///  * Otherwise `config.seed()` is assumed to be an encrypted opaque blob generated from
///    a prior `EnvelopeEncryption::encrypt_seed` call. The value of `config.kms_protection()`
///    is parsed as a KMS key id and `EnvelopeEncryption::decrypt_seed` is called to obtain
///    the plaintext seed value.
///
#[cfg(feature = "gcpkms")]
pub fn load_seed(config: &Box<ServerConfig>) -> Result<Vec<u8>, error::Error> {
    match config.kms_protection() {
        KmsProtection::Plaintext => Ok(config.seed()),
        KmsProtection::GoogleKmsEnvelope(resource_id) => {
            info!("Unwrapping seed via Google KMS key '{}'", resource_id);
            let kms = GcpKms::from_resource_id(resource_id)?;
            let seed = EnvelopeEncryption::decrypt_seed(&kms, &config.seed())?;
            Ok(seed)
        }
        _ => Err(error::Error::InvalidConfiguration(
            "AWS KMS not supported".to_string(),
        )),
    }
}

/// Load the seed value for the long-term key.
///
/// ## This build has KMS disabled
///
/// *The KMS feature is disabled in this build of Roughenough*.
///
/// The only supported `kms_protection` value in this build is `plaintext`. Any
/// other value will cause a runtime error.
///
///  * `config.seed()` is used as-is and assumed to be a 32-byte hexadecimal value
///
#[cfg(not(any(feature = "awskms", feature = "gcpkms")))]
pub fn load_seed(config: &Box<ServerConfig>) -> Result<Vec<u8>, error::Error> {
    match config.kms_protection() {
        KmsProtection::Plaintext => Ok(config.seed()),
        v => Err(error::Error::InvalidConfiguration(format!(
            "kms_protection '{}' requires KMS, but server was not compiled with KMS support",
            v
        ))),
    }
}
<|endoftext|>


<commit_before>// Cymbalum, an extensible molecular simulation engine
// Copyright (C) 2015-2016 G. Fraux — BSD license

//! Logging configuration
//!
//! This module provides a `Logger` struct which can be used for logging
//! messages from Cymbalum. One and only one of the initilialization function
//! should be called, and it should only be called once. Any call to
//! initializing function after the first one will be ignored.
//!
//! Other [log](https://github.com/rust-lang/log) compatible crates can be used
//! for logging if you whant.

use std::io;
use std::io::prelude::*;
use std::fs::OpenOptions;
use std::path::Path;
use std::sync::{Arc, Mutex};

use log::{Log, LogRecord, LogMetadata, set_logger};
pub use log::LogLevel;

/// Logger with capacity to write to the standard output stream, the standard
/// error stream or a file.
 pub struct Logger {
     level: LogLevel,
     writer: Arc<Mutex<Box<Write + Send + Sync>>>,
 }

 impl Logger{
     fn new<T>(level: LogLevel, handle: T) -> Logger where T: Write + Send + Sync + 'static {
         ::chemfiles::Logger::get().log_callback(|level, message| {
             match level {
                 ::chemfiles::LogLevel::Error => error!("{}", message),
                 ::chemfiles::LogLevel::Warning => warn!("{}", message),
                 ::chemfiles::LogLevel::Info => info!("{}", message),
                 ::chemfiles::LogLevel::Debug => debug!("{}", message),
             }
         }).expect("Could not set the chemfiles logging callback");
         Logger{
             level: level,
             writer: Arc::new(Mutex::new(Box::new(handle))),
         }
     }

     fn init(logger: Logger) {
         let res = set_logger(|max_log_level| {
             max_log_level.set(logger.level.to_log_level_filter());
             Box::new(logger)
         });

         if let Err(_) = res {
             warn!("Logger was initialized more than once! This call is ignored.");
         }
     }

     /// Initialize the global logger to use the standard output stream.
     pub fn stdout() {
         let logger = Logger::new(levels::DEFAULT_LEVEL, io::stdout());
         Logger::init(logger);
     }

     /// Initialize the global logger to use the standard output stream, with
     /// the maximum log level of `level`
     pub fn stdout_at_level(level: LogLevel) {
         let logger = Logger::new(level, io::stdout());
         Logger::init(logger);
     }

     /// Initialize the global logger to use the standard error stream.
     pub fn stderr() {
         let logger = Logger::new(levels::DEFAULT_LEVEL, io::stderr());
         Logger::init(logger);
     }

     /// Initialize the global logger to use the standard error stream, with
     /// the maximum log level of `level`
     pub fn stderr_at_level(level: LogLevel) {
         let logger = Logger::new(level, io::stderr());
         Logger::init(logger);
     }

     /// Initialize the global logger to write to the file at `path`, creating
     /// the file if needed. If the file already exists, it is not removed.
     pub fn file<P: AsRef<Path>>(path: P) -> Result<(), io::Error> {
         let file = try!(OpenOptions::new().write(true).append(true).open(path));
         let logger = Logger::new(levels::DEFAULT_LEVEL, file);
         Logger::init(logger);
         Ok(())
     }

     /// Initialize the global logger to write to the file at `path`, creating
     /// the file if needed. If the file already exists, it is not removed. The
     /// maximum log level is set to `level`.
     pub fn file_at_level<P: AsRef<Path>>(path: P, level: LogLevel) -> Result<(), io::Error> {
         let file = try!(OpenOptions::new().write(true).append(true).open(path));
         let logger = Logger::new(level, file);
         Logger::init(logger);
         Ok(())
     }
 }

#[cfg(not(debug_assertions))]
 mod levels {
     use super::LogLevel;
     /// Default log level
     pub const DEFAULT_LEVEL: LogLevel = LogLevel::Warn;
 }

#[cfg(debug_assertions)]
 mod levels {
     use super::LogLevel;
     /// Default log level
     pub const DEFAULT_LEVEL: LogLevel = LogLevel::Debug;
 }

/// Implements Log trait for Logger
impl Log for Logger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= self.level
    }

    fn log(&self, record: &LogRecord) {
        let mut out = self.writer.lock().expect("Could not lock the logger.");
        let write_res = match record.level() {
            LogLevel::Info => write!(&mut out, "{}\n", record.args()),
            LogLevel::Warn => write!(
                   &mut out, "Warning: {} from {}:{}\n",
                   record.args(),
                   record.location().file(), record.location().line()
            ),
            _ => write!(
                   &mut out, "{}: {} from {}:{}\n",
                   record.level(), record.args(),
                   record.location().file(), record.location().line()
            ),
        };

        if let Err(err) = write_res {
            let error_res = write!(& mut io::stderr(),
                "Error while writing log: {:?}\nMessage was: {} at {}:{} ({})\n",
                err, record.args(),
                record.location().file(), record.location().line(),
                record.level()
            );
            if let Err(err2) = error_res {
                panic!(format!(
                    "Error while logging logger error, aborting. Error chain: {:?} -> {:?} from {}:{}",
                    err, err2,
                    record.location().file(), record.location().line(),
                ))
            }
        }
    }
}
<commit_msg>Factorize Logger method, and add missing create<commit_after>// Cymbalum, an extensible molecular simulation engine
// Copyright (C) 2015-2016 G. Fraux — BSD license

//! Logging configuration
//!
//! This module provides a `Logger` struct which can be used for logging
//! messages from Cymbalum. One and only one of the initilialization function
//! should be called, and it should only be called once. Any call to
//! initializing function after the first one will be ignored.
//!
//! Other [log](https://github.com/rust-lang/log) compatible crates can be used
//! for logging if you whant.

use std::io;
use std::io::prelude::*;
use std::fs::OpenOptions;
use std::path::Path;
use std::sync::{Arc, Mutex};

use log::{Log, LogRecord, LogMetadata, set_logger};
pub use log::LogLevel;

/// Logger with capacity to write to the standard output stream, the standard
/// error stream or a file.
 pub struct Logger {
     level: LogLevel,
     writer: Arc<Mutex<Box<Write + Send + Sync>>>,
 }

 impl Logger{
     fn new<T>(level: LogLevel, handle: T) -> Logger where T: Write + Send + Sync + 'static {
         ::chemfiles::Logger::get().log_callback(|level, message| {
             match level {
                 ::chemfiles::LogLevel::Error => error!("{}", message),
                 ::chemfiles::LogLevel::Warning => warn!("{}", message),
                 ::chemfiles::LogLevel::Info => info!("{}", message),
                 ::chemfiles::LogLevel::Debug => debug!("{}", message),
             }
         }).expect("Could not set the chemfiles logging callback");
         Logger{
             level: level,
             writer: Arc::new(Mutex::new(Box::new(handle))),
         }
     }

     fn init(logger: Logger) {
         let res = set_logger(|max_log_level| {
             max_log_level.set(logger.level.to_log_level_filter());
             Box::new(logger)
         });

         if let Err(_) = res {
             warn!("Logger was initialized more than once! This call is ignored.");
         }
     }

     /// Initialize the global logger to use the standard output stream.
     pub fn stdout() {
         Logger::stdout_at_level(levels::DEFAULT_LEVEL)
     }

     /// Initialize the global logger to use the standard output stream, with
     /// the maximum log level of `level`
     pub fn stdout_at_level(level: LogLevel) {
         let logger = Logger::new(level, io::stdout());
         Logger::init(logger);
     }

     /// Initialize the global logger to use the standard error stream.
     pub fn stderr() {
         Logger::stderr_at_level(levels::DEFAULT_LEVEL)
     }

     /// Initialize the global logger to use the standard error stream, with
     /// the maximum log level of `level`
     pub fn stderr_at_level(level: LogLevel) {
         let logger = Logger::new(level, io::stderr());
         Logger::init(logger);
     }

     /// Initialize the global logger to write to the file at `path`, creating
     /// the file if needed. If the file already exists, it is not removed.
     pub fn file<P: AsRef<Path>>(path: P) -> Result<(), io::Error> {
         Logger::file_at_level(path, levels::DEFAULT_LEVEL)
     }

     /// Initialize the global logger to write to the file at `path`, creating
     /// the file if needed. If the file already exists, it is not removed. The
     /// maximum log level is set to `level`.
     pub fn file_at_level<P: AsRef<Path>>(path: P, level: LogLevel) -> Result<(), io::Error> {
         let file = try!(OpenOptions::new().write(true).create(true).append(true).open(path));
         let logger = Logger::new(level, file);
         Logger::init(logger);
         Ok(())
     }
 }

#[cfg(not(debug_assertions))]
 mod levels {
     use super::LogLevel;
     /// Default log level
     pub const DEFAULT_LEVEL: LogLevel = LogLevel::Warn;
 }

#[cfg(debug_assertions)]
 mod levels {
     use super::LogLevel;
     /// Default log level
     pub const DEFAULT_LEVEL: LogLevel = LogLevel::Debug;
 }

/// Implements Log trait for Logger
impl Log for Logger {
    fn enabled(&self, metadata: &LogMetadata) -> bool {
        metadata.level() <= self.level
    }

    fn log(&self, record: &LogRecord) {
        let mut out = self.writer.lock().expect("Could not lock the logger.");
        let write_res = match record.level() {
            LogLevel::Info => write!(&mut out, "{}\n", record.args()),
            LogLevel::Warn => write!(
                   &mut out, "Warning: {} from {}:{}\n",
                   record.args(),
                   record.location().file(), record.location().line()
            ),
            _ => write!(
                   &mut out, "{}: {} from {}:{}\n",
                   record.level(), record.args(),
                   record.location().file(), record.location().line()
            ),
        };

        if let Err(err) = write_res {
            let error_res = write!(& mut io::stderr(),
                "Error while writing log: {:?}\nMessage was: {} at {}:{} ({})\n",
                err, record.args(),
                record.location().file(), record.location().line(),
                record.level()
            );
            if let Err(err2) = error_res {
                panic!(format!(
                    "Error while logging logger error, aborting. Error chain: {:?} -> {:?} from {}:{}",
                    err, err2,
                    record.location().file(), record.location().line(),
                ))
            }
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 pdb Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use common::*;
use source::*;

mod page_list;
use self::page_list::PageList;

use std::fmt;

type PageNumber = u32;

#[derive(Debug,Copy,Clone)]
struct Header {
    page_size: usize,
    maximum_valid_page_number: PageNumber,
}

impl Header {
    fn pages_needed_to_store(&self, bytes: usize) -> usize {
        (bytes + (self.page_size - 1)) / self.page_size
    }

    fn validate_page_number(&self, page_number: u32) -> Result<PageNumber> {
        if page_number == 0 || page_number > self.maximum_valid_page_number {
            Err(Error::PageReferenceOutOfRange(page_number))
        } else {
            Ok(page_number as PageNumber)
        }
    }
}

#[derive(Debug)]
struct BigMSF<'s, S> {
    header: Header,
    source: S,
    stream_table: StreamTable<'s>,
}

/// Represents a stream table at various stages of access
#[doc(hidden)]
#[derive(Debug)]
enum StreamTable<'s> {
    /// The MSF header gives us the size of the table in bytes, and the list of pages (usually one)
    /// where we can find the list of pages that contain the stream table.
    HeaderOnly { size_in_bytes: usize, stream_table_location_location: PageList },

    /// Given the HeaderOnly information, we can do an initial read to get the actual location of
    /// the stream table as a PageList.
    TableFound { stream_table_location: PageList },

    // Given the table location, we can access the stream table itself
    Available { stream_table_view: Box<SourceView<'s>> }
}

const BIG_MSF_MAGIC: &'static [u8] = b"Microsoft C/C++ MSF 7.00\r\n\x1a\x44\x53\x00\x00\x00";

fn view<'s>(source: &mut Source<'s>, page_list: &PageList) -> Result<Box<SourceView<'s>>> {
    // view it
    let view = source.view(page_list.source_slices())?;

    // double check our Source
    // if the Source didn't return the requested bits, that's an implementation bug, so
    // assert instead of returning an error
    assert_eq!(view.as_slice().len(), page_list.len());

    // done
    Ok(view)
}

/// The PDB header as stored on disk.
/// See the Microsoft code for reference: https://github.com/Microsoft/microsoft-pdb/blob/082c5290e5aff028ae84e43affa8be717aa7af73/PDB/msf/msf.cpp#L946
#[derive(Debug, Pread)]
#[repr(C, packed)]
struct RawHeader {
    magic: [u8; 32],
    page_size: u32,
    free_page_map: u32,
    pages_used: u32,
    directory_size: u32,
    _reserved: u32,
}

impl<'s, S: Source<'s>> BigMSF<'s, S> {
    pub fn new(source: S, header_view: Box<SourceView>) -> Result<BigMSF<'s, S>> {
        let mut buf = ParseBuffer::from(header_view.as_slice());
        let header: RawHeader = buf.parse()?;

        if header.magic != BIG_MSF_MAGIC {
            return Err(Error::UnrecognizedFileFormat);
        }

        if header.page_size.count_ones() != 1 || header.page_size < 0x100
            || header.page_size > (128 * 0x10000) {
                return Err(Error::InvalidPageSize(header.page_size));
            }

        let header_object = Header{
            page_size: header.page_size as usize,
            maximum_valid_page_number: header.pages_used,
        };

        // calculate how many pages are needed to store the stream table
        let size_of_stream_table_in_pages = header_object.pages_needed_to_store(header.directory_size as usize);

        // now: how many pages are needed to store the list of pages that store the stream table?
        // each page entry is a u32, so multiply by four
        let size_of_stream_table_page_list_in_pages = header_object.pages_needed_to_store(size_of_stream_table_in_pages * 4);

        // read the list of stream table page list pages, which immediately follow the header
        // yes, this is a stupid level of indirection
        let mut stream_table_page_list_page_list = PageList::new(header_object.page_size);
        for _ in 0..size_of_stream_table_page_list_in_pages {
            let n = buf.parse_u32()?;
            stream_table_page_list_page_list.push(header_object.validate_page_number(n)?);
        }

        // truncate the stream table location location to the correct size
        stream_table_page_list_page_list.truncate(size_of_stream_table_in_pages * 4);

        Ok(BigMSF{
            header: header_object,
            source: source,
            stream_table: StreamTable::HeaderOnly {
                size_in_bytes: header.directory_size as usize,
                stream_table_location_location: stream_table_page_list_page_list,
            },
        })
    }

    fn find_stream_table(&mut self) -> Result<()> {
        let mut new_stream_table: Option<StreamTable> = None;

        if let StreamTable::HeaderOnly { size_in_bytes, ref stream_table_location_location } = self.stream_table {
            // the header indicated we need to read size_in_pages page numbers from the
            // specified PageList.

            // ask to view the location location
            let location_location = view(&mut self.source, stream_table_location_location)?;

            // build a PageList
            let mut page_list = PageList::new(self.header.page_size);
            let mut buf = ParseBuffer::from(location_location.as_slice());
            while buf.len() > 0 {
                let n = buf.parse_u32()?;
                page_list.push(self.header.validate_page_number(n)?);
            }

            page_list.truncate(size_in_bytes);

            // remember what we learned
            new_stream_table = Some(StreamTable::TableFound {
                stream_table_location: page_list,
            });
        }

        if let Some(st) = new_stream_table {
            self.stream_table = st;
        }

        Ok(())
    }

    fn make_stream_table_available(&mut self) -> Result<()> {
        // do the initial read if we must
        if let StreamTable::HeaderOnly { .. } = self.stream_table {
            self.find_stream_table()?;
        }

        // do we need to map the stream table itself?
        let mut new_stream_table: Option<StreamTable> = None;
        if let StreamTable::TableFound { ref stream_table_location } = self.stream_table {
            // yep
            // ask the source to view it
            let stream_table_view = view(&mut self.source, stream_table_location)?;

            // done
            new_stream_table = Some(StreamTable::Available {
                stream_table_view: stream_table_view,
            });
        }

        if let Some(st) = new_stream_table {
            self.stream_table = st;
        }

        // stream table is available
        assert!(match &self.stream_table {
            &StreamTable::Available { .. } => true,
            _ => false
        });

        Ok(())
    }

    fn look_up_stream(&mut self, stream_number: u32) -> Result<PageList> {
        // ensure the stream table is available
        self.make_stream_table_available()?;

        let header = self.header;

        // declare the things we're going to find
        let bytes_in_stream: u32;
        let page_list: PageList;

        if let StreamTable::Available { ref stream_table_view } = self.stream_table {
            let stream_table_slice = stream_table_view.as_slice();
            let mut stream_table = ParseBuffer::from(stream_table_slice);

            // the stream table is structured as:
            // stream_count
            // 0..stream_count: size of stream in bytes (0xffffffff indicating "stream does not exist")
            // stream 0: PageNumber
            // stream 1: PageNumber, PageNumber
            // stream 2: PageNumber, PageNumber, PageNumber, PageNumber, PageNumber
            // stream 3: PageNumber, PageNumber, PageNumber, PageNumber
            // (number of pages determined by number of bytes)

            let stream_count = stream_table.parse_u32()?;

            // check if we've already outworn our welcome
            if stream_number >= stream_count {
                return Err(Error::StreamNotFound(stream_number))
            }

            // we now have {stream_count} u32s describing the length of each stream

            // walk over the streams before the requested stream
            // we need to pay attention to how big each one is, since their page numbers come
            // before our page numbers in the stream table
            let mut page_numbers_to_skip: usize = 0;
            for _ in 0..stream_number {
                let bytes = stream_table.parse_u32()?;
                if bytes == 0xffffffff {
                    // stream is not present, ergo nothing to skip
                } else {
                    page_numbers_to_skip += header.pages_needed_to_store(bytes as usize);
                }
            }

            // read our stream's size
            bytes_in_stream = stream_table.parse_u32()?;
            if bytes_in_stream == 0xffffffff {
                return Err(Error::StreamNotFound(stream_number))
            }
            let pages_in_stream = header.pages_needed_to_store(bytes_in_stream as usize);

            // skip the remaining streams' byte counts
            let _ = stream_table.take((stream_count - stream_number - 1) as usize * 4)?;

            // skip the preceding streams' page numbers
            let _ = stream_table.take((page_numbers_to_skip as usize) * 4)?;

            // we're now at the list of pages for our stream
            // accumulate them into a PageList
            let mut list = PageList::new(header.page_size);
            for _ in 0..pages_in_stream {
                let page_number = stream_table.parse_u32()?;
                list.push(self.header.validate_page_number(page_number)?);
            }

            // truncate to the size of the stream
            list.truncate(bytes_in_stream as usize);

            page_list = list;
        } else {
            unreachable!();
        }

        // done!
        Ok(page_list)
    }
}

impl<'s, S: Source<'s>> MSF<'s, S> for BigMSF<'s, S> {
    fn get(&mut self, stream_number: u32, limit: Option<usize>) -> Result<Stream<'s>> {
        // look up the stream
        let mut page_list = self.look_up_stream(stream_number)?;

        // apply any limits we have
        if let Some(limit) = limit {
            page_list.truncate(limit);
        }

        // now that we know where this stream lives, we can view it
        let view = view(&mut self.source, &page_list)?;

        // pack it into a Stream
        let stream = Stream {
            source_view: view,
        };

        Ok(stream)
    }
}

const SMALL_MSF_MAGIC: &'static [u8] = b"Microsoft C/C++ program database 2.00\r\n\x1a\x4a\x47";

// TODO: implement SmallMSF

/// Represents a single Stream within the multi-stream file.
#[derive(Debug)]
pub struct Stream<'s> {
    source_view: Box<SourceView<'s>>,
}

impl<'s> Stream<'s> {
    #[inline]
    pub fn parse_buffer(&self) -> ParseBuffer {
        let slice = self.source_view.as_slice();
        ParseBuffer::from(slice)
    }
}

/// Provides access to a "multi-stream file", which is the container format used by PDBs.
pub trait MSF<'s, S> : fmt::Debug {
    /// Accesses a stream by stream number, optionally restricted by a byte limit.
    fn get(&mut self, stream_number: u32, limit: Option<usize>) -> Result<Stream<'s>>;
}

fn header_matches(actual: &[u8], expected: &[u8]) -> bool {
    actual.len() >= expected.len() && &actual[0..expected.len()] == expected
}

pub fn open_msf<'s, S: Source<'s> + 's>(mut source: S) -> Result<Box<MSF<'s, S> + 's>> {
    // map the header
    let mut header_location = PageList::new(4096);
    header_location.push(0);
    let header_view = view(&mut source, &header_location)?;

    // see if it's a BigMSF
    if header_matches(header_view.as_slice(), BIG_MSF_MAGIC) {
        // claimed!
        let bigmsf = BigMSF::new(source, header_view)?;
        return Ok(Box::new(bigmsf))
    }

    if header_matches(header_view.as_slice(), SMALL_MSF_MAGIC) {
        // sorry
        return Err(Error::UnimplementedFeature("small MSF file format"));
    }

    Err(Error::UnrecognizedFileFormat)
}

#[cfg(test)]
mod tests {

    mod header {
        use msf::Header;
        use common::Error;

        #[test]
        fn test_pages_needed_to_store() {
            let h = Header{
                page_size: 4096,
                maximum_valid_page_number: 15,
            };
            assert_eq!(h.pages_needed_to_store(0), 0);
            assert_eq!(h.pages_needed_to_store(1), 1);
            assert_eq!(h.pages_needed_to_store(1024), 1);
            assert_eq!(h.pages_needed_to_store(2048), 1);
            assert_eq!(h.pages_needed_to_store(4095), 1);
            assert_eq!(h.pages_needed_to_store(4096), 1);
            assert_eq!(h.pages_needed_to_store(4097), 2);
        }

        #[test]
        fn test_validate_page_number() {
            let h = Header{
                page_size: 4096,
                maximum_valid_page_number: 15,
            };
            assert!(match h.validate_page_number(0) { Err(Error::PageReferenceOutOfRange(0)) => true, _ => false });
            assert!(match h.validate_page_number(1) { Ok(1) => true, _ => false });
            assert!(match h.validate_page_number(2) { Ok(2) => true, _ => false });
            assert!(match h.validate_page_number(14) { Ok(14) => true, _ => false });
            assert!(match h.validate_page_number(15) { Ok(15) => true, _ => false });
            assert!(match h.validate_page_number(16) { Err(Error::PageReferenceOutOfRange(16)) => true, _ => false });
            assert!(match h.validate_page_number(17) { Err(Error::PageReferenceOutOfRange(17)) => true, _ => false });
        }
    }
}
<commit_msg>Move `BigMSF` and its related bits into a `big` submodule, and add a `small` submodule currently containing only the small MSF `MAGIC` constant.<commit_after>// Copyright 2017 pdb Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use common::*;
use source::*;

mod page_list;
use self::page_list::PageList;

use std::fmt;

type PageNumber = u32;

#[derive(Debug,Copy,Clone)]
struct Header {
    page_size: usize,
    maximum_valid_page_number: PageNumber,
}

impl Header {
    fn pages_needed_to_store(&self, bytes: usize) -> usize {
        (bytes + (self.page_size - 1)) / self.page_size
    }

    fn validate_page_number(&self, page_number: u32) -> Result<PageNumber> {
        if page_number == 0 || page_number > self.maximum_valid_page_number {
            Err(Error::PageReferenceOutOfRange(page_number))
        } else {
            Ok(page_number as PageNumber)
        }
    }
}

/// Represents a stream table at various stages of access
#[doc(hidden)]
#[derive(Debug)]
enum StreamTable<'s> {
    /// The MSF header gives us the size of the table in bytes, and the list of pages (usually one)
    /// where we can find the list of pages that contain the stream table.
    HeaderOnly { size_in_bytes: usize, stream_table_location_location: PageList },

    /// Given the HeaderOnly information, we can do an initial read to get the actual location of
    /// the stream table as a PageList.
    TableFound { stream_table_location: PageList },

    // Given the table location, we can access the stream table itself
    Available { stream_table_view: Box<SourceView<'s>> }
}

fn view<'s>(source: &mut Source<'s>, page_list: &PageList) -> Result<Box<SourceView<'s>>> {
    // view it
    let view = source.view(page_list.source_slices())?;

    // double check our Source
    // if the Source didn't return the requested bits, that's an implementation bug, so
    // assert instead of returning an error
    assert_eq!(view.as_slice().len(), page_list.len());

    // done
    Ok(view)
}

mod big {
    use super::*;

    pub const MAGIC: &'static [u8] = b"Microsoft C/C++ MSF 7.00\r\n\x1a\x44\x53\x00\x00\x00";

    /// The PDB header as stored on disk.
    /// See the Microsoft code for reference: https://github.com/Microsoft/microsoft-pdb/blob/082c5290e5aff028ae84e43affa8be717aa7af73/PDB/msf/msf.cpp#L946
    #[derive(Debug, Pread)]
    #[repr(C, packed)]
    struct RawHeader {
        magic: [u8; 32],
        page_size: u32,
        free_page_map: u32,
        pages_used: u32,
        directory_size: u32,
        _reserved: u32,
    }

    #[derive(Debug)]
    pub struct BigMSF<'s, S> {
        header: Header,
        source: S,
        stream_table: StreamTable<'s>,
    }

    impl<'s, S: Source<'s>> BigMSF<'s, S> {
        pub fn new(source: S, header_view: Box<SourceView>) -> Result<BigMSF<'s, S>> {
            let mut buf = ParseBuffer::from(header_view.as_slice());
            let header: RawHeader = buf.parse()?;

            if header.magic != MAGIC {
                return Err(Error::UnrecognizedFileFormat);
            }

            if header.page_size.count_ones() != 1 || header.page_size < 0x100
                || header.page_size > (128 * 0x10000) {
                    return Err(Error::InvalidPageSize(header.page_size));
                }

            let header_object = Header{
                page_size: header.page_size as usize,
                maximum_valid_page_number: header.pages_used,
            };

            // calculate how many pages are needed to store the stream table
            let size_of_stream_table_in_pages = header_object.pages_needed_to_store(header.directory_size as usize);

            // now: how many pages are needed to store the list of pages that store the stream table?
            // each page entry is a u32, so multiply by four
            let size_of_stream_table_page_list_in_pages = header_object.pages_needed_to_store(size_of_stream_table_in_pages * 4);

            // read the list of stream table page list pages, which immediately follow the header
            // yes, this is a stupid level of indirection
            let mut stream_table_page_list_page_list = PageList::new(header_object.page_size);
            for _ in 0..size_of_stream_table_page_list_in_pages {
                let n = buf.parse_u32()?;
                stream_table_page_list_page_list.push(header_object.validate_page_number(n)?);
            }

            // truncate the stream table location location to the correct size
            stream_table_page_list_page_list.truncate(size_of_stream_table_in_pages * 4);

            Ok(BigMSF{
                header: header_object,
                source: source,
                stream_table: StreamTable::HeaderOnly {
                    size_in_bytes: header.directory_size as usize,
                    stream_table_location_location: stream_table_page_list_page_list,
                },
            })
        }

        fn find_stream_table(&mut self) -> Result<()> {
            let mut new_stream_table: Option<StreamTable> = None;

            if let StreamTable::HeaderOnly { size_in_bytes, ref stream_table_location_location } = self.stream_table {
                // the header indicated we need to read size_in_pages page numbers from the
                // specified PageList.

                // ask to view the location location
                let location_location = view(&mut self.source, stream_table_location_location)?;

                // build a PageList
                let mut page_list = PageList::new(self.header.page_size);
                let mut buf = ParseBuffer::from(location_location.as_slice());
                while buf.len() > 0 {
                    let n = buf.parse_u32()?;
                    page_list.push(self.header.validate_page_number(n)?);
                }

                page_list.truncate(size_in_bytes);

                // remember what we learned
                new_stream_table = Some(StreamTable::TableFound {
                    stream_table_location: page_list,
                });
            }

            if let Some(st) = new_stream_table {
                self.stream_table = st;
            }

            Ok(())
        }

        fn make_stream_table_available(&mut self) -> Result<()> {
            // do the initial read if we must
            if let StreamTable::HeaderOnly { .. } = self.stream_table {
                self.find_stream_table()?;
            }

            // do we need to map the stream table itself?
            let mut new_stream_table: Option<StreamTable> = None;
            if let StreamTable::TableFound { ref stream_table_location } = self.stream_table {
                // yep
                // ask the source to view it
                let stream_table_view = view(&mut self.source, stream_table_location)?;

                // done
                new_stream_table = Some(StreamTable::Available {
                    stream_table_view: stream_table_view,
                });
            }

            if let Some(st) = new_stream_table {
                self.stream_table = st;
            }

            // stream table is available
            assert!(match &self.stream_table {
                &StreamTable::Available { .. } => true,
                _ => false
            });

            Ok(())
        }

        fn look_up_stream(&mut self, stream_number: u32) -> Result<PageList> {
            // ensure the stream table is available
            self.make_stream_table_available()?;

            let header = self.header;

            // declare the things we're going to find
            let bytes_in_stream: u32;
            let page_list: PageList;

            if let StreamTable::Available { ref stream_table_view } = self.stream_table {
                let stream_table_slice = stream_table_view.as_slice();
                let mut stream_table = ParseBuffer::from(stream_table_slice);

                // the stream table is structured as:
                // stream_count
                // 0..stream_count: size of stream in bytes (0xffffffff indicating "stream does not exist")
                // stream 0: PageNumber
                // stream 1: PageNumber, PageNumber
                // stream 2: PageNumber, PageNumber, PageNumber, PageNumber, PageNumber
                // stream 3: PageNumber, PageNumber, PageNumber, PageNumber
                // (number of pages determined by number of bytes)

                let stream_count = stream_table.parse_u32()?;

                // check if we've already outworn our welcome
                if stream_number >= stream_count {
                    return Err(Error::StreamNotFound(stream_number))
                }

                // we now have {stream_count} u32s describing the length of each stream

                // walk over the streams before the requested stream
                // we need to pay attention to how big each one is, since their page numbers come
                // before our page numbers in the stream table
                let mut page_numbers_to_skip: usize = 0;
                for _ in 0..stream_number {
                    let bytes = stream_table.parse_u32()?;
                    if bytes == 0xffffffff {
                        // stream is not present, ergo nothing to skip
                    } else {
                        page_numbers_to_skip += header.pages_needed_to_store(bytes as usize);
                    }
                }

                // read our stream's size
                bytes_in_stream = stream_table.parse_u32()?;
                if bytes_in_stream == 0xffffffff {
                    return Err(Error::StreamNotFound(stream_number))
                }
                let pages_in_stream = header.pages_needed_to_store(bytes_in_stream as usize);

                // skip the remaining streams' byte counts
                let _ = stream_table.take((stream_count - stream_number - 1) as usize * 4)?;

                // skip the preceding streams' page numbers
                let _ = stream_table.take((page_numbers_to_skip as usize) * 4)?;

                // we're now at the list of pages for our stream
                // accumulate them into a PageList
                let mut list = PageList::new(header.page_size);
                for _ in 0..pages_in_stream {
                    let page_number = stream_table.parse_u32()?;
                    list.push(self.header.validate_page_number(page_number)?);
                }

                // truncate to the size of the stream
                list.truncate(bytes_in_stream as usize);

                page_list = list;
            } else {
                unreachable!();
            }

            // done!
            Ok(page_list)
        }
    }

    impl<'s, S: Source<'s>> MSF<'s, S> for BigMSF<'s, S> {
        fn get(&mut self, stream_number: u32, limit: Option<usize>) -> Result<Stream<'s>> {
            // look up the stream
            let mut page_list = self.look_up_stream(stream_number)?;

            // apply any limits we have
            if let Some(limit) = limit {
                page_list.truncate(limit);
            }

            // now that we know where this stream lives, we can view it
            let view = view(&mut self.source, &page_list)?;

            // pack it into a Stream
            let stream = Stream {
                source_view: view,
            };

            Ok(stream)
        }
    }
}

mod small {
    pub const MAGIC: &'static [u8] = b"Microsoft C/C++ program database 2.00\r\n\x1a\x4a\x47";
    // TODO: implement SmallMSF
}

/// Represents a single Stream within the multi-stream file.
#[derive(Debug)]
pub struct Stream<'s> {
    source_view: Box<SourceView<'s>>,
}

impl<'s> Stream<'s> {
    #[inline]
    pub fn parse_buffer(&self) -> ParseBuffer {
        let slice = self.source_view.as_slice();
        ParseBuffer::from(slice)
    }
}

/// Provides access to a "multi-stream file", which is the container format used by PDBs.
pub trait MSF<'s, S> : fmt::Debug {
    /// Accesses a stream by stream number, optionally restricted by a byte limit.
    fn get(&mut self, stream_number: u32, limit: Option<usize>) -> Result<Stream<'s>>;
}

fn header_matches(actual: &[u8], expected: &[u8]) -> bool {
    actual.len() >= expected.len() && &actual[0..expected.len()] == expected
}

pub fn open_msf<'s, S: Source<'s> + 's>(mut source: S) -> Result<Box<MSF<'s, S> + 's>> {
    // map the header
    let mut header_location = PageList::new(4096);
    header_location.push(0);
    let header_view = view(&mut source, &header_location)?;

    // see if it's a BigMSF
    if header_matches(header_view.as_slice(), big::MAGIC) {
        // claimed!
        let bigmsf = big::BigMSF::new(source, header_view)?;
        return Ok(Box::new(bigmsf))
    }

    if header_matches(header_view.as_slice(), small::MAGIC) {
        // sorry
        return Err(Error::UnimplementedFeature("small MSF file format"));
    }

    Err(Error::UnrecognizedFileFormat)
}

#[cfg(test)]
mod tests {

    mod header {
        use msf::Header;
        use common::Error;

        #[test]
        fn test_pages_needed_to_store() {
            let h = Header{
                page_size: 4096,
                maximum_valid_page_number: 15,
            };
            assert_eq!(h.pages_needed_to_store(0), 0);
            assert_eq!(h.pages_needed_to_store(1), 1);
            assert_eq!(h.pages_needed_to_store(1024), 1);
            assert_eq!(h.pages_needed_to_store(2048), 1);
            assert_eq!(h.pages_needed_to_store(4095), 1);
            assert_eq!(h.pages_needed_to_store(4096), 1);
            assert_eq!(h.pages_needed_to_store(4097), 2);
        }

        #[test]
        fn test_validate_page_number() {
            let h = Header{
                page_size: 4096,
                maximum_valid_page_number: 15,
            };
            assert!(match h.validate_page_number(0) { Err(Error::PageReferenceOutOfRange(0)) => true, _ => false });
            assert!(match h.validate_page_number(1) { Ok(1) => true, _ => false });
            assert!(match h.validate_page_number(2) { Ok(2) => true, _ => false });
            assert!(match h.validate_page_number(14) { Ok(14) => true, _ => false });
            assert!(match h.validate_page_number(15) { Ok(15) => true, _ => false });
            assert!(match h.validate_page_number(16) { Err(Error::PageReferenceOutOfRange(16)) => true, _ => false });
            assert!(match h.validate_page_number(17) { Err(Error::PageReferenceOutOfRange(17)) => true, _ => false });
        }
    }
}
<|endoftext|>