<commit_before>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// 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.

//! Input / output.

/// A producer of bytes.
pub trait InputStream {
    /// Reads at least `min_bytes` into `buf` unless EOF is encountered first. Returns the
    /// number of bytes read.
    fn try_read(&mut self, buf : &mut [u8], min_bytes : usize) -> ::std::io::Result<usize>;

    /// Reads at least `min_bytes` into `buf`, returning the number of bytes read. If EOF is
    /// encountered first, returns an error.
    fn read(&mut self, buf : &mut [u8], min_bytes : usize) -> ::std::io::Result<usize> {
        let n = try!(self.try_read(buf, min_bytes));
        if n < min_bytes {
            Err(::std::io::Error::new(::std::io::ErrorKind::Other, "Premature EOF", None))
        } else {
            Ok(n)
        }
    }

    /// Reads into `buf` until it is full. Returns an error if EOF is encountered first.
    fn read_exact(&mut self, buf : &mut [u8]) -> ::std::io::Result<()> {
        let min_bytes = buf.len();
        try!(self.read(buf, min_bytes));
        Ok(())
    }
}

impl <R> InputStream for R where R : ::std::io::Read {
    fn try_read(&mut self, buf : &mut [u8], min_bytes : usize) -> ::std::io::Result<usize> {
        let mut pos = 0;
        let buf_len = buf.len();
        while pos < min_bytes {
            let buf1 = &mut buf[pos .. buf_len];
            let n = try!(self.read(buf1));
            pos += n;
            if n == 0 { return Ok(pos); }
        }
        return Ok(pos);
    }
}

pub trait BufferedInputStream : InputStream {
    fn skip(&mut self, bytes : usize) -> ::std::io::Result<()>;
    unsafe fn get_read_buffer(&mut self) -> ::std::io::Result<(*const u8, *const u8)>;
}

pub struct BufferedInputStreamWrapper<R> {
    inner : R,
    buf : Vec<u8>,
    pos : usize,
    cap : usize
}

impl <R> BufferedInputStreamWrapper<R> {
    pub fn new(r : R) -> BufferedInputStreamWrapper<R> {
        let mut result = BufferedInputStreamWrapper {
            inner : r,
            buf : Vec::with_capacity(8192),
            pos : 0,
            cap : 0
        };
        unsafe {
            result.buf.set_len(8192)
        }
        return result;
    }
}

impl<R: InputStream> BufferedInputStream for BufferedInputStreamWrapper<R> {

   fn skip(&mut self, mut bytes : usize) -> ::std::io::Result<()> {
        let available = self.cap - self.pos;
        if bytes <= available {
            self.pos += bytes;
        } else {
            bytes -= available;
            if bytes <= self.buf.len() {
                //# Read the next buffer-full.
                let n = try!(self.inner.try_read(self.buf.as_mut_slice(), bytes));
                self.pos = bytes;
                self.cap = n;
            } else {
                //# Forward large skip to the underlying stream.
                panic!("TODO")
            }
        }
        Ok(())
    }

    unsafe fn get_read_buffer(&mut self) -> ::std::io::Result<(*const u8, *const u8)> {
        if self.cap - self.pos == 0 {
            let n = try!(self.inner.try_read(self.buf.as_mut_slice(), 1));
            self.cap = n;
            self.pos = 0;
        }
        Ok((self.buf.get_unchecked(self.pos) as *const u8,
            self.buf.get_unchecked(self.cap) as *const u8))
    }
}

impl<R: InputStream> InputStream for BufferedInputStreamWrapper<R> {
    fn try_read(&mut self, mut dst: &mut [u8], mut min_bytes : usize) -> ::std::io::Result<usize> {
        if min_bytes <= self.cap - self.pos {
            // Serve from the current buffer.
            let n = ::std::cmp::min(self.cap - self.pos, dst.len());
            ::std::slice::bytes::copy_memory(dst,
                                             &self.buf[self.pos .. self.pos + n]);
            self.pos += n;
            return Ok(n);
        } else {
            // Copy current available into destination.
            ::std::slice::bytes::copy_memory(dst,
                                             &self.buf[self.pos .. self.cap]);
            let from_first_buffer = self.cap - self.pos;

            let dst = &mut dst[from_first_buffer ..];
            min_bytes -= from_first_buffer;

            if dst.len() <= self.buf.len() {
                // Read the next buffer-full.
                let n = try!(self.inner.try_read(self.buf.as_mut_slice(), min_bytes));
                let from_second_buffer = ::std::cmp::min(n, dst.len());
                ::std::slice::bytes::copy_memory(dst,
                                                 &self.buf[0 .. from_second_buffer]);
                self.cap = n;
                self.pos = from_second_buffer;
                return Ok(from_first_buffer + from_second_buffer);
            } else {
                // Forward large read to the underlying stream.
                self.pos = 0;
                self.cap = 0;
                return Ok(from_first_buffer + try!(self.inner.try_read(dst, min_bytes)));
            }
        }
    }
}

pub struct ArrayInputStream<'a> {
    array : &'a [u8]
}

impl <'a> ArrayInputStream<'a> {
    pub fn new<'b>(array : &'b [u8]) -> ArrayInputStream<'b> {
        ArrayInputStream { array : array }
    }
}

impl <'a> InputStream for ArrayInputStream<'a> {
    fn try_read(&mut self, dst: &mut [u8], _min_bytes : usize) -> ::std::io::Result<usize> {
        let n = ::std::cmp::min(dst.len(), self.array.len());
        ::std::slice::bytes::copy_memory(dst, &self.array[0 .. n]);
        self.array = &self.array[n ..];
        Ok(n)
    }
}

impl <'a> BufferedInputStream for ArrayInputStream<'a> {
    fn skip(&mut self, bytes : usize) -> ::std::io::Result<()> {
        assert!(self.array.len() >= bytes,
                "ArrayInputStream ended prematurely.");
        self.array = &self.array[bytes ..];
        Ok(())
    }
    unsafe fn get_read_buffer(&mut self) -> ::std::io::Result<(*const u8, *const u8)> {
        let len = self.array.len();
        Ok((self.array.as_ptr() as *const u8,
           self.array.get_unchecked(len) as *const u8))
    }
}

/// A consumer of bytes.
pub trait OutputStream {
    /// Writes all of `buf`.
    fn write(&mut self, buf : &[u8]) -> ::std::io::Result<()>;
    fn flush(&mut self) -> ::std::io::Result<()> { Ok(()) }
}

impl <W> OutputStream for W where W : ::std::io::Write {
    fn write(&mut self, buf : &[u8]) -> ::std::io::Result<()> {
        self.write_all(buf)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        self.flush()
    }
}

pub trait BufferedOutputStream : OutputStream {
    unsafe fn get_write_buffer(&mut self) -> (*mut u8, *mut u8);
    unsafe fn write_ptr(&mut self, ptr: *mut u8, size: usize) -> ::std::io::Result<()>;
}

pub struct BufferedOutputStreamWrapper<'a, W:'a> {
    inner: &'a mut W,
    buf: Vec<u8>,
    pos: usize
}

impl <'a, W> BufferedOutputStreamWrapper<'a, W> {
    pub fn new<'b> (w : &'b mut W) -> BufferedOutputStreamWrapper<W> {
        let mut result = BufferedOutputStreamWrapper {
            inner: w,
            buf : Vec::with_capacity(8192),
            pos : 0
        };
        unsafe {
            result.buf.set_len(8192);
        }
        return result;
    }
}

impl<'a, W: OutputStream> BufferedOutputStream for BufferedOutputStreamWrapper<'a, W> {
    #[inline]
    unsafe fn get_write_buffer(&mut self) -> (*mut u8, *mut u8) {
        let len = self.buf.len();
        (self.buf.get_unchecked_mut(self.pos) as *mut u8,
         self.buf.get_unchecked_mut(len) as *mut u8)
    }

    #[inline]
    unsafe fn write_ptr(&mut self, ptr: *mut u8, size: usize) -> ::std::io::Result<()> {
        let easy_case = ptr == self.buf.get_unchecked_mut(self.pos) as *mut u8;
        if easy_case {
            self.pos += size;
            Ok(())
        } else {
            let buf = ::std::slice::from_raw_parts_mut::<u8>(ptr, size);
            self.write(buf)
        }
    }

}

impl<'a, W: OutputStream> OutputStream for BufferedOutputStreamWrapper<'a, W> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<()> {
        let available = self.buf.len() - self.pos;
        let mut size = buf.len();
        if size <= available {
            let dst = &mut self.buf.as_mut_slice()[self.pos ..];
            ::std::slice::bytes::copy_memory(dst, buf);
            self.pos += size;
        } else if size <= self.buf.len() {
            // Too much for this buffer, but not a full buffer's
            // worth, so we'll go ahead and copy.
            {
                let dst = &mut self.buf.as_mut_slice()[self.pos ..];
                ::std::slice::bytes::copy_memory(dst, &buf[0 .. available]);
            }
            try!(self.inner.write(self.buf.as_mut_slice()));

            size -= available;
            let src = &buf[available ..];
            let dst = &mut self.buf.as_mut_slice()[0 ..];
            ::std::slice::bytes::copy_memory(dst, src);
            self.pos = size;
        } else {
            // Writing so much data that we might as well write
            // directly to avoid a copy.
            try!(self.inner.write(&self.buf[0 .. self.pos]));
            self.pos = 0;
            try!(self.inner.write(buf));
        }
        return Ok(());
    }

    fn flush(&mut self) -> ::std::io::Result<()> {
        if self.pos > 0 {
            try!(self.inner.write(&self.buf[0 .. self.pos]));
            self.pos = 0;
        }
        self.inner.flush()
    }
}

pub struct ArrayOutputStream<'a> {
    array : &'a mut [u8],
    fill_pos : usize,
}

impl <'a> ArrayOutputStream<'a> {
    pub fn new<'b>(array : &'b mut [u8]) -> ArrayOutputStream<'b> {
        ArrayOutputStream {
            array : array,
            fill_pos : 0
        }
    }
}

impl <'a> OutputStream for ArrayOutputStream<'a> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<()> {
        assert!(buf.len() <= self.array.len() - self.fill_pos,
                "ArrayOutputStream's backing array was not large enough for the data written.");
        ::std::slice::bytes::copy_memory(&mut self.array[self.fill_pos ..], buf);
        self.fill_pos += buf.len();
        Ok(())
    }
}

impl <'a> BufferedOutputStream for ArrayOutputStream<'a> {
    unsafe fn get_write_buffer(&mut self) -> (*mut u8, *mut u8) {
        let len = self.array.len();
        (self.array.get_unchecked_mut(self.fill_pos) as *mut u8,
         self.array.get_unchecked_mut(len) as *mut u8)
    }
    unsafe fn write_ptr(&mut self, ptr: *mut u8, size: usize) -> ::std::io::Result<()> {
        let easy_case = ptr == self.array.get_unchecked_mut(self.fill_pos) as *mut u8;
        if easy_case {
            self.fill_pos += size;
            Ok(())
        } else {
            let buf = ::std::slice::from_raw_parts_mut::<u8>(ptr, size);
            self.write(buf)
        }
    }
}
<commit_msg>fancy slice syntax is nifty<commit_after>// Copyright (c) 2013-2015 Sandstorm Development Group, Inc. and contributors
// Licensed under the MIT License:
//
// 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.

//! Input / output.

/// A producer of bytes.
pub trait InputStream {
    /// Reads at least `min_bytes` into `buf` unless EOF is encountered first. Returns the
    /// number of bytes read.
    fn try_read(&mut self, buf : &mut [u8], min_bytes : usize) -> ::std::io::Result<usize>;

    /// Reads at least `min_bytes` into `buf`, returning the number of bytes read. If EOF is
    /// encountered first, returns an error.
    fn read(&mut self, buf : &mut [u8], min_bytes : usize) -> ::std::io::Result<usize> {
        let n = try!(self.try_read(buf, min_bytes));
        if n < min_bytes {
            Err(::std::io::Error::new(::std::io::ErrorKind::Other, "Premature EOF", None))
        } else {
            Ok(n)
        }
    }

    /// Reads into `buf` until it is full. Returns an error if EOF is encountered first.
    fn read_exact(&mut self, buf : &mut [u8]) -> ::std::io::Result<()> {
        let min_bytes = buf.len();
        try!(self.read(buf, min_bytes));
        Ok(())
    }
}

impl <R> InputStream for R where R : ::std::io::Read {
    fn try_read(&mut self, buf : &mut [u8], min_bytes : usize) -> ::std::io::Result<usize> {
        let mut pos = 0;
        while pos < min_bytes {
            let buf1 = &mut buf[pos ..];
            let n = try!(self.read(buf1));
            pos += n;
            if n == 0 { return Ok(pos); }
        }
        return Ok(pos);
    }
}

pub trait BufferedInputStream : InputStream {
    fn skip(&mut self, bytes : usize) -> ::std::io::Result<()>;
    unsafe fn get_read_buffer(&mut self) -> ::std::io::Result<(*const u8, *const u8)>;
}

pub struct BufferedInputStreamWrapper<R> {
    inner : R,
    buf : Vec<u8>,
    pos : usize,
    cap : usize
}

impl <R> BufferedInputStreamWrapper<R> {
    pub fn new(r : R) -> BufferedInputStreamWrapper<R> {
        let mut result = BufferedInputStreamWrapper {
            inner : r,
            buf : Vec::with_capacity(8192),
            pos : 0,
            cap : 0
        };
        unsafe {
            result.buf.set_len(8192)
        }
        return result;
    }
}

impl<R: InputStream> BufferedInputStream for BufferedInputStreamWrapper<R> {

   fn skip(&mut self, mut bytes : usize) -> ::std::io::Result<()> {
        let available = self.cap - self.pos;
        if bytes <= available {
            self.pos += bytes;
        } else {
            bytes -= available;
            if bytes <= self.buf.len() {
                //# Read the next buffer-full.
                let n = try!(self.inner.try_read(self.buf.as_mut_slice(), bytes));
                self.pos = bytes;
                self.cap = n;
            } else {
                //# Forward large skip to the underlying stream.
                panic!("TODO")
            }
        }
        Ok(())
    }

    unsafe fn get_read_buffer(&mut self) -> ::std::io::Result<(*const u8, *const u8)> {
        if self.cap - self.pos == 0 {
            let n = try!(self.inner.try_read(self.buf.as_mut_slice(), 1));
            self.cap = n;
            self.pos = 0;
        }
        Ok((self.buf.get_unchecked(self.pos) as *const u8,
            self.buf.get_unchecked(self.cap) as *const u8))
    }
}

impl<R: InputStream> InputStream for BufferedInputStreamWrapper<R> {
    fn try_read(&mut self, mut dst: &mut [u8], mut min_bytes : usize) -> ::std::io::Result<usize> {
        if min_bytes <= self.cap - self.pos {
            // Serve from the current buffer.
            let n = ::std::cmp::min(self.cap - self.pos, dst.len());
            ::std::slice::bytes::copy_memory(dst,
                                             &self.buf[self.pos .. self.pos + n]);
            self.pos += n;
            return Ok(n);
        } else {
            // Copy current available into destination.
            ::std::slice::bytes::copy_memory(dst,
                                             &self.buf[self.pos .. self.cap]);
            let from_first_buffer = self.cap - self.pos;

            let dst = &mut dst[from_first_buffer ..];
            min_bytes -= from_first_buffer;

            if dst.len() <= self.buf.len() {
                // Read the next buffer-full.
                let n = try!(self.inner.try_read(self.buf.as_mut_slice(), min_bytes));
                let from_second_buffer = ::std::cmp::min(n, dst.len());
                ::std::slice::bytes::copy_memory(dst,
                                                 &self.buf[0 .. from_second_buffer]);
                self.cap = n;
                self.pos = from_second_buffer;
                return Ok(from_first_buffer + from_second_buffer);
            } else {
                // Forward large read to the underlying stream.
                self.pos = 0;
                self.cap = 0;
                return Ok(from_first_buffer + try!(self.inner.try_read(dst, min_bytes)));
            }
        }
    }
}

pub struct ArrayInputStream<'a> {
    array : &'a [u8]
}

impl <'a> ArrayInputStream<'a> {
    pub fn new<'b>(array : &'b [u8]) -> ArrayInputStream<'b> {
        ArrayInputStream { array : array }
    }
}

impl <'a> InputStream for ArrayInputStream<'a> {
    fn try_read(&mut self, dst: &mut [u8], _min_bytes : usize) -> ::std::io::Result<usize> {
        let n = ::std::cmp::min(dst.len(), self.array.len());
        ::std::slice::bytes::copy_memory(dst, &self.array[0 .. n]);
        self.array = &self.array[n ..];
        Ok(n)
    }
}

impl <'a> BufferedInputStream for ArrayInputStream<'a> {
    fn skip(&mut self, bytes : usize) -> ::std::io::Result<()> {
        assert!(self.array.len() >= bytes,
                "ArrayInputStream ended prematurely.");
        self.array = &self.array[bytes ..];
        Ok(())
    }
    unsafe fn get_read_buffer(&mut self) -> ::std::io::Result<(*const u8, *const u8)> {
        let len = self.array.len();
        Ok((self.array.as_ptr() as *const u8,
           self.array.get_unchecked(len) as *const u8))
    }
}

/// A consumer of bytes.
pub trait OutputStream {
    /// Writes all of `buf`.
    fn write(&mut self, buf : &[u8]) -> ::std::io::Result<()>;
    fn flush(&mut self) -> ::std::io::Result<()> { Ok(()) }
}

impl <W> OutputStream for W where W : ::std::io::Write {
    fn write(&mut self, buf : &[u8]) -> ::std::io::Result<()> {
        self.write_all(buf)
    }
    fn flush(&mut self) -> ::std::io::Result<()> {
        self.flush()
    }
}

pub trait BufferedOutputStream : OutputStream {
    unsafe fn get_write_buffer(&mut self) -> (*mut u8, *mut u8);
    unsafe fn write_ptr(&mut self, ptr: *mut u8, size: usize) -> ::std::io::Result<()>;
}

pub struct BufferedOutputStreamWrapper<'a, W:'a> {
    inner: &'a mut W,
    buf: Vec<u8>,
    pos: usize
}

impl <'a, W> BufferedOutputStreamWrapper<'a, W> {
    pub fn new<'b> (w : &'b mut W) -> BufferedOutputStreamWrapper<W> {
        let mut result = BufferedOutputStreamWrapper {
            inner: w,
            buf : Vec::with_capacity(8192),
            pos : 0
        };
        unsafe {
            result.buf.set_len(8192);
        }
        return result;
    }
}

impl<'a, W: OutputStream> BufferedOutputStream for BufferedOutputStreamWrapper<'a, W> {
    #[inline]
    unsafe fn get_write_buffer(&mut self) -> (*mut u8, *mut u8) {
        let len = self.buf.len();
        (self.buf.get_unchecked_mut(self.pos) as *mut u8,
         self.buf.get_unchecked_mut(len) as *mut u8)
    }

    #[inline]
    unsafe fn write_ptr(&mut self, ptr: *mut u8, size: usize) -> ::std::io::Result<()> {
        let easy_case = ptr == self.buf.get_unchecked_mut(self.pos) as *mut u8;
        if easy_case {
            self.pos += size;
            Ok(())
        } else {
            let buf = ::std::slice::from_raw_parts_mut::<u8>(ptr, size);
            self.write(buf)
        }
    }

}

impl<'a, W: OutputStream> OutputStream for BufferedOutputStreamWrapper<'a, W> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<()> {
        let available = self.buf.len() - self.pos;
        let mut size = buf.len();
        if size <= available {
            let dst = &mut self.buf.as_mut_slice()[self.pos ..];
            ::std::slice::bytes::copy_memory(dst, buf);
            self.pos += size;
        } else if size <= self.buf.len() {
            // Too much for this buffer, but not a full buffer's
            // worth, so we'll go ahead and copy.
            {
                let dst = &mut self.buf.as_mut_slice()[self.pos ..];
                ::std::slice::bytes::copy_memory(dst, &buf[0 .. available]);
            }
            try!(self.inner.write(self.buf.as_mut_slice()));

            size -= available;
            let src = &buf[available ..];
            let dst = &mut self.buf.as_mut_slice()[0 ..];
            ::std::slice::bytes::copy_memory(dst, src);
            self.pos = size;
        } else {
            // Writing so much data that we might as well write
            // directly to avoid a copy.
            try!(self.inner.write(&self.buf[0 .. self.pos]));
            self.pos = 0;
            try!(self.inner.write(buf));
        }
        return Ok(());
    }

    fn flush(&mut self) -> ::std::io::Result<()> {
        if self.pos > 0 {
            try!(self.inner.write(&self.buf[0 .. self.pos]));
            self.pos = 0;
        }
        self.inner.flush()
    }
}

pub struct ArrayOutputStream<'a> {
    array : &'a mut [u8],
    fill_pos : usize,
}

impl <'a> ArrayOutputStream<'a> {
    pub fn new<'b>(array : &'b mut [u8]) -> ArrayOutputStream<'b> {
        ArrayOutputStream {
            array : array,
            fill_pos : 0
        }
    }
}

impl <'a> OutputStream for ArrayOutputStream<'a> {
    fn write(&mut self, buf: &[u8]) -> ::std::io::Result<()> {
        assert!(buf.len() <= self.array.len() - self.fill_pos,
                "ArrayOutputStream's backing array was not large enough for the data written.");
        ::std::slice::bytes::copy_memory(&mut self.array[self.fill_pos ..], buf);
        self.fill_pos += buf.len();
        Ok(())
    }
}

impl <'a> BufferedOutputStream for ArrayOutputStream<'a> {
    unsafe fn get_write_buffer(&mut self) -> (*mut u8, *mut u8) {
        let len = self.array.len();
        (self.array.get_unchecked_mut(self.fill_pos) as *mut u8,
         self.array.get_unchecked_mut(len) as *mut u8)
    }
    unsafe fn write_ptr(&mut self, ptr: *mut u8, size: usize) -> ::std::io::Result<()> {
        let easy_case = ptr == self.array.get_unchecked_mut(self.fill_pos) as *mut u8;
        if easy_case {
            self.fill_pos += size;
            Ok(())
        } else {
            let buf = ::std::slice::from_raw_parts_mut::<u8>(ptr, size);
            self.write(buf)
        }
    }
}
<|endoftext|>


<commit_before>/// A three-part version consisting of major, minor, and sub minor components.
///
/// This can be used to represent versions of the format `J.M.N`, where `J` is the major version,
/// `M` is the minor version, and `N` is the sub minor version. A version is constructed by
/// providing the fields in the same order to the tuple. For example:
///
/// ```ignore
/// libusb::Version(0, 2, 1);
/// ```
///
/// represents the version 0.2.1.
///
/// The intended use case of `Version` is to extract meaning from the version fields in USB
/// descriptors, such as `bcdUSB` and `bcdDevice` in device descriptors.
#[derive(Debug,PartialEq)]
pub struct Version(pub u8, pub u8, pub u8);

impl Version {
  /// Extracts a version from a binary coded decimal (BCD) field. BCD fields exist in USB
  /// descriptors as 16-bit integers encoding a version as `0xJJMN`, where `JJ` is the major
  /// version, `M` is the minor version, and `N` is the sub minor version. For example, 2.0 is
  /// endoded as `0x0200` and 1.1 is encoded as `0x0110`.
  pub fn from_bcd(mut raw: u16) -> Self {
    let sub_minor: u8 = (raw & 0x000F) as u8;
    raw >>= 4;

    let minor: u8 = (raw & 0x000F) as u8;
    raw >>= 4;

    let mut major: u8 = (raw & 0x000F) as u8;
    raw >>= 4;

    major += (10 * raw) as u8;

    Version(major, minor, sub_minor)
  }

  /// Returns the major version.
  pub fn major(&self) -> u8 {
    let Version(major, _, _) = *self;
    major
  }

  /// Returns the minor version.
  pub fn minor(&self) -> u8 {
    let Version(_, minor, _) = *self;
    minor
  }

  /// Returns the sub minor version.
  pub fn sub_minor(&self) -> u8 {
    let Version(_, _, sub_minor) = *self;
    sub_minor
  }
}

#[cfg(test)]
mod version_test {
  use ::Version;

  #[test]
  fn it_returns_major_version() {
    assert_eq!(1, Version(1, 0, 0).major());
    assert_eq!(2, Version(2, 0, 0).major());
  }

  #[test]
  fn it_returns_minor_version() {
    assert_eq!(1, Version(0, 1, 0).minor());
    assert_eq!(2, Version(0, 2, 0).minor());
  }

  #[test]
  fn it_returns_sub_minor_version() {
    assert_eq!(1, Version(0, 0, 1).sub_minor());
    assert_eq!(2, Version(0, 0, 2).sub_minor());
  }

  #[test]
  fn it_parses_major_version() {
    assert_eq!(3, Version::from_bcd(0x0300).major());
  }

  #[test]
  fn it_parses_long_major_version() {
    assert_eq!(12, Version::from_bcd(0x1200).major());
  }

  #[test]
  fn it_parses_minor_version() {
    assert_eq!(1, Version::from_bcd(0x0010).minor());
    assert_eq!(2, Version::from_bcd(0x0020).minor());
  }

  #[test]
  fn it_parses_sub_minor_version() {
    assert_eq!(1, Version::from_bcd(0x0001).sub_minor());
    assert_eq!(2, Version::from_bcd(0x0002).sub_minor());
  }

  #[test]
  fn it_parses_full_version() {
    assert_eq!(Version(12, 3, 4), Version::from_bcd(0x1234));
  }
}
<commit_msg>enables compilation of doc tests<commit_after>/// A three-part version consisting of major, minor, and sub minor components.
///
/// This can be used to represent versions of the format `J.M.N`, where `J` is the major version,
/// `M` is the minor version, and `N` is the sub minor version. A version is constructed by
/// providing the fields in the same order to the tuple. For example:
///
/// ```
/// libusb::Version(0, 2, 1);
/// ```
///
/// represents the version 0.2.1.
///
/// The intended use case of `Version` is to extract meaning from the version fields in USB
/// descriptors, such as `bcdUSB` and `bcdDevice` in device descriptors.
#[derive(Debug,PartialEq)]
pub struct Version(pub u8, pub u8, pub u8);

impl Version {
  /// Extracts a version from a binary coded decimal (BCD) field. BCD fields exist in USB
  /// descriptors as 16-bit integers encoding a version as `0xJJMN`, where `JJ` is the major
  /// version, `M` is the minor version, and `N` is the sub minor version. For example, 2.0 is
  /// endoded as `0x0200` and 1.1 is encoded as `0x0110`.
  pub fn from_bcd(mut raw: u16) -> Self {
    let sub_minor: u8 = (raw & 0x000F) as u8;
    raw >>= 4;

    let minor: u8 = (raw & 0x000F) as u8;
    raw >>= 4;

    let mut major: u8 = (raw & 0x000F) as u8;
    raw >>= 4;

    major += (10 * raw) as u8;

    Version(major, minor, sub_minor)
  }

  /// Returns the major version.
  pub fn major(&self) -> u8 {
    let Version(major, _, _) = *self;
    major
  }

  /// Returns the minor version.
  pub fn minor(&self) -> u8 {
    let Version(_, minor, _) = *self;
    minor
  }

  /// Returns the sub minor version.
  pub fn sub_minor(&self) -> u8 {
    let Version(_, _, sub_minor) = *self;
    sub_minor
  }
}

#[cfg(test)]
mod version_test {
  use ::Version;

  #[test]
  fn it_returns_major_version() {
    assert_eq!(1, Version(1, 0, 0).major());
    assert_eq!(2, Version(2, 0, 0).major());
  }

  #[test]
  fn it_returns_minor_version() {
    assert_eq!(1, Version(0, 1, 0).minor());
    assert_eq!(2, Version(0, 2, 0).minor());
  }

  #[test]
  fn it_returns_sub_minor_version() {
    assert_eq!(1, Version(0, 0, 1).sub_minor());
    assert_eq!(2, Version(0, 0, 2).sub_minor());
  }

  #[test]
  fn it_parses_major_version() {
    assert_eq!(3, Version::from_bcd(0x0300).major());
  }

  #[test]
  fn it_parses_long_major_version() {
    assert_eq!(12, Version::from_bcd(0x1200).major());
  }

  #[test]
  fn it_parses_minor_version() {
    assert_eq!(1, Version::from_bcd(0x0010).minor());
    assert_eq!(2, Version::from_bcd(0x0020).minor());
  }

  #[test]
  fn it_parses_sub_minor_version() {
    assert_eq!(1, Version::from_bcd(0x0001).sub_minor());
    assert_eq!(2, Version::from_bcd(0x0002).sub_minor());
  }

  #[test]
  fn it_parses_full_version() {
    assert_eq!(Version(12, 3, 4), Version::from_bcd(0x1234));
  }
}
<|endoftext|>


<commit_before>pub struct KeyBuilder {
    key: Vec<u8>,
}

impl KeyBuilder {
    pub fn new() -> KeyBuilder {
        KeyBuilder {
            key: Vec::new(),
        }
    }

    pub fn with_capacity(size: usize) -> KeyBuilder {
        KeyBuilder {
            key: Vec::with_capacity(size),
        }
    }

    pub fn doc_key_mapping(key: &[u8]) -> KeyBuilder {
        let mut kb = KeyBuilder::with_capacity(1 + key.len());
        kb.push_char(b'k');
        kb.push_string(key);
        kb
    }

    pub fn term_dict_mapping(term: &[u8]) -> KeyBuilder {
        let mut kb = KeyBuilder::with_capacity(1 + term.len());
        kb.push_char(b't');
        kb.push_string(term);
        kb
    }

    pub fn chunk_dir_list(chunk: u32, field_ord: u32, term_ord: u32) -> KeyBuilder {
        let mut kb = KeyBuilder::new();
        kb.push_char(b'd');
        kb.push_string(field_ord.to_string().as_bytes());
        kb.separator();
        kb.push_string(term_ord.to_string().as_bytes());
        kb.separator();
        kb.push_string(chunk.to_string().as_bytes());
        kb
    }

    pub fn chunk_stat(chunk: u32, name: &[u8]) -> KeyBuilder {
        let mut kb = KeyBuilder::new();
        kb.push_char(b's');
        kb.push_string(name);
        kb.separator();
        kb.push_string(chunk.to_string().as_bytes());
        kb
    }

    pub fn chunk_del_list(chunk: u32) -> KeyBuilder {
        let mut kb = KeyBuilder::new();
        kb.push_char(b'x');
        kb.push_string(chunk.to_string().as_bytes());
        kb
    }

    #[inline]
    pub fn key(&self) -> &[u8] {
        &self.key[..]
    }

    #[inline]
    pub fn push_char(&mut self, c: u8) {
        if c == b'/' || c == b'\\' {
            self.key.push(b'\\');
        }
        self.key.push(c);
    }

    pub fn push_string(&mut self, s: &[u8]) {
        for c in s {
            self.push_char(*c);
        }
    }

    pub fn separator(&mut self) {
        self.key.push(b'/');
    }
}<commit_msg>Change key format<commit_after>pub struct KeyBuilder {
    key: Vec<u8>,
}

impl KeyBuilder {
    pub fn new() -> KeyBuilder {
        KeyBuilder {
            key: Vec::new(),
        }
    }

    pub fn with_capacity(size: usize) -> KeyBuilder {
        KeyBuilder {
            key: Vec::with_capacity(size),
        }
    }

    pub fn doc_key_mapping(key: &[u8]) -> KeyBuilder {
        let mut kb = KeyBuilder::with_capacity(1 + key.len());
        kb.push_char(b'k');
        kb.push_string(key);
        kb
    }

    pub fn term_dict_mapping(term: &[u8]) -> KeyBuilder {
        let mut kb = KeyBuilder::with_capacity(1 + term.len());
        kb.push_char(b't');
        kb.push_string(term);
        kb
    }

    pub fn chunk_dir_list(chunk: u32, field_ord: u32, term_ord: u32) -> KeyBuilder {
        let mut kb = KeyBuilder::new();
        kb.chunk(chunk);
        kb.push_char(b'd');
        kb.push_string(field_ord.to_string().as_bytes());
        kb.separator();
        kb.push_string(term_ord.to_string().as_bytes());
        kb
    }

    pub fn chunk_stat(chunk: u32, name: &[u8]) -> KeyBuilder {
        let mut kb = KeyBuilder::new();
        kb.chunk(chunk);
        kb.push_char(b's');
        kb.push_string(name);
        kb
    }

    pub fn chunk_del_list(chunk: u32) -> KeyBuilder {
        let mut kb = KeyBuilder::new();
        kb.chunk(chunk);
        kb.push_char(b'x');
        kb
    }

    #[inline]
    pub fn key(&self) -> &[u8] {
        &self.key[..]
    }

    #[inline]
    pub fn push_char(&mut self, c: u8) {
        if c == b'/' || c == b'\\' {
            self.key.push(b'\\');
        }
        self.key.push(c);
    }

    pub fn push_string(&mut self, s: &[u8]) {
        for c in s {
            self.push_char(*c);
        }
    }

    pub fn separator(&mut self) {
        self.key.push(b'/');
    }

    pub fn chunk(&mut self, chunk: u32) {
        self.key.push(b'c');
        self.push_string(chunk.to_string().as_bytes());
    }
}<|endoftext|>


<commit_before>// Copyright 2013-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.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.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(all(unix, not(target_os = "ios")))]
mod imp {
    extern crate libc;

    use self::OsRngInner::*;

    use std::io;
    use std::fs::File;
    use Rng;
    use read::ReadRng;
    use std::mem;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64",
                  target_arch = "x86",
                  target_arch = "arm",
                  target_arch = "aarch64",
                  target_arch = "powerpc")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        extern "C" {
            fn syscall(number: libc::c_long, ...) -> libc::c_long;
        }

        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
        const NR_GETRANDOM: libc::c_long = 384;
        #[cfg(target_arch = "powerpc")]
        const NR_GETRANDOM: libc::c_long = 384;

        unsafe {
            syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64",
                      target_arch = "x86",
                      target_arch = "arm",
                      target_arch = "aarch64",
                      target_arch = "powerpc"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        let len = v.len();
        while read < len {
            let result = getrandom(&mut v[read..]);
            if result == -1 {
                let err = io::Error::last_os_error();
                if err.kind() == io::ErrorKind::Interrupted {
                    continue
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as usize;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8; 4] = [0u8; 4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8; 4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8; 8] = [0u8; 8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8; 8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64",
                  target_arch = "x86",
                  target_arch = "arm",
                  target_arch = "aarch64",
                  target_arch = "powerpc")))]
    fn is_getrandom_available() -> bool {
        use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering};
        use std::sync::{Once, ONCE_INIT};

        static CHECKER: Once = ONCE_INIT;
        static AVAILABLE: AtomicBool = ATOMIC_BOOL_INIT;

        CHECKER.call_once(|| {
            let mut buf: [u8; 0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = io::Error::last_os_error().raw_os_error();
                err != Some(libc::ENOSYS)
            } else {
                true
            };
            AVAILABLE.store(available, Ordering::Relaxed);
        });

        AVAILABLE.load(Ordering::Relaxed)
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64",
                      target_arch = "x86",
                      target_arch = "arm",
                      target_arch = "aarch64",
                      target_arch = "powerpc"))))]
    fn is_getrandom_available() -> bool { false }

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReadRng(ReadRng<File>),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open("/dev/urandom"));
            let reader_rng = ReadRng::new(reader);

            Ok(OsRng { inner: OsReadRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReadRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReadRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReadRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use std::io;
    use std::mem;
    use Rng;
    use self::libc::{c_int, size_t};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        // dummy field to ensure that this struct cannot be constructed outside of this module
        _dummy: (),
    }

    #[repr(C)]
    struct SecRandom;

    #[allow(non_upper_case_globals)]
    const kSecRandomDefault: *const SecRandom = 0 as *const SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            Ok(OsRng { _dummy: () })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}", io::Error::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate winapi;
    extern crate advapi32;

    use std::io;
    use std::mem;
    use std::ptr;
    use Rng;
    use self::winapi::{CRYPT_SILENT, CRYPT_VERIFYCONTEXT, DWORD, HCRYPTPROV, PROV_RSA_FULL};
    use self::advapi32::{CryptAcquireContextA, CryptGenRandom, CryptReleaseContext};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                CryptAcquireContextA(&mut hcp, ptr::null(), ptr::null(),
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            if ret == 0 {
                Err(io::Error::last_os_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}",
                       io::Error::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}",
                       io::Error::last_os_error());
            }
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::mpsc::channel;
    use Rng;
    use OsRng;
    use std::thread;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8; 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in 0..20 {
            let (tx, rx) = channel();
            txs.push(tx);

            thread::spawn(move|| {
                // wait until all the tasks are ready to go.
                rx.recv().unwrap();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                thread::yield_now();
                let mut v = [0u8; 1000];

                for _ in 0..100 {
                    r.next_u32();
                    thread::yield_now();
                    r.next_u64();
                    thread::yield_now();
                    r.fill_bytes(&mut v);
                    thread::yield_now();
                }
            });
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(()).unwrap();
        }
    }
}
<commit_msg>Fix getrandom syscall number for aarch64-unknown-linux-gnu<commit_after>// Copyright 2013-2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.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.

//! Interfaces to the operating system provided random number
//! generators.

pub use self::imp::OsRng;

#[cfg(all(unix, not(target_os = "ios")))]
mod imp {
    extern crate libc;

    use self::OsRngInner::*;

    use std::io;
    use std::fs::File;
    use Rng;
    use read::ReadRng;
    use std::mem;

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64",
                  target_arch = "x86",
                  target_arch = "arm",
                  target_arch = "aarch64",
                  target_arch = "powerpc")))]
    fn getrandom(buf: &mut [u8]) -> libc::c_long {
        extern "C" {
            fn syscall(number: libc::c_long, ...) -> libc::c_long;
        }

        #[cfg(target_arch = "x86_64")]
        const NR_GETRANDOM: libc::c_long = 318;
        #[cfg(target_arch = "x86")]
        const NR_GETRANDOM: libc::c_long = 355;
        #[cfg(target_arch = "arm")]
        const NR_GETRANDOM: libc::c_long = 384;
        #[cfg(target_arch = "aarch64")]
        const NR_GETRANDOM: libc::c_long = 278;
        #[cfg(target_arch = "powerpc")]
        const NR_GETRANDOM: libc::c_long = 384;

        unsafe {
            syscall(NR_GETRANDOM, buf.as_mut_ptr(), buf.len(), 0)
        }
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64",
                      target_arch = "x86",
                      target_arch = "arm",
                      target_arch = "aarch64",
                      target_arch = "powerpc"))))]
    fn getrandom(_buf: &mut [u8]) -> libc::c_long { -1 }

    fn getrandom_fill_bytes(v: &mut [u8]) {
        let mut read = 0;
        let len = v.len();
        while read < len {
            let result = getrandom(&mut v[read..]);
            if result == -1 {
                let err = io::Error::last_os_error();
                if err.kind() == io::ErrorKind::Interrupted {
                    continue
                } else {
                    panic!("unexpected getrandom error: {}", err);
                }
            } else {
                read += result as usize;
            }
        }
    }

    fn getrandom_next_u32() -> u32 {
        let mut buf: [u8; 4] = [0u8; 4];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8; 4], u32>(buf) }
    }

    fn getrandom_next_u64() -> u64 {
        let mut buf: [u8; 8] = [0u8; 8];
        getrandom_fill_bytes(&mut buf);
        unsafe { mem::transmute::<[u8; 8], u64>(buf) }
    }

    #[cfg(all(target_os = "linux",
              any(target_arch = "x86_64",
                  target_arch = "x86",
                  target_arch = "arm",
                  target_arch = "aarch64",
                  target_arch = "powerpc")))]
    fn is_getrandom_available() -> bool {
        use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering};
        use std::sync::{Once, ONCE_INIT};

        static CHECKER: Once = ONCE_INIT;
        static AVAILABLE: AtomicBool = ATOMIC_BOOL_INIT;

        CHECKER.call_once(|| {
            let mut buf: [u8; 0] = [];
            let result = getrandom(&mut buf);
            let available = if result == -1 {
                let err = io::Error::last_os_error().raw_os_error();
                err != Some(libc::ENOSYS)
            } else {
                true
            };
            AVAILABLE.store(available, Ordering::Relaxed);
        });

        AVAILABLE.load(Ordering::Relaxed)
    }

    #[cfg(not(all(target_os = "linux",
                  any(target_arch = "x86_64",
                      target_arch = "x86",
                      target_arch = "arm",
                      target_arch = "aarch64",
                      target_arch = "powerpc"))))]
    fn is_getrandom_available() -> bool { false }

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        inner: OsRngInner,
    }

    enum OsRngInner {
        OsGetrandomRng,
        OsReadRng(ReadRng<File>),
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            if is_getrandom_available() {
                return Ok(OsRng { inner: OsGetrandomRng });
            }

            let reader = try!(File::open("/dev/urandom"));
            let reader_rng = ReadRng::new(reader);

            Ok(OsRng { inner: OsReadRng(reader_rng) })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u32(),
                OsReadRng(ref mut rng) => rng.next_u32(),
            }
        }
        fn next_u64(&mut self) -> u64 {
            match self.inner {
                OsGetrandomRng => getrandom_next_u64(),
                OsReadRng(ref mut rng) => rng.next_u64(),
            }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            match self.inner {
                OsGetrandomRng => getrandom_fill_bytes(v),
                OsReadRng(ref mut rng) => rng.fill_bytes(v)
            }
        }
    }
}

#[cfg(target_os = "ios")]
mod imp {
    extern crate libc;

    use std::io;
    use std::mem;
    use Rng;
    use self::libc::{c_int, size_t};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        // dummy field to ensure that this struct cannot be constructed outside of this module
        _dummy: (),
    }

    #[repr(C)]
    struct SecRandom;

    #[allow(non_upper_case_globals)]
    const kSecRandomDefault: *const SecRandom = 0 as *const SecRandom;

    #[link(name = "Security", kind = "framework")]
    extern "C" {
        fn SecRandomCopyBytes(rnd: *const SecRandom,
                              count: size_t, bytes: *mut u8) -> c_int;
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            Ok(OsRng { _dummy: () })
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                SecRandomCopyBytes(kSecRandomDefault, v.len() as size_t, v.as_mut_ptr())
            };
            if ret == -1 {
                panic!("couldn't generate random bytes: {}", io::Error::last_os_error());
            }
        }
    }
}

#[cfg(windows)]
mod imp {
    extern crate winapi;
    extern crate advapi32;

    use std::io;
    use std::mem;
    use std::ptr;
    use Rng;
    use self::winapi::{CRYPT_SILENT, CRYPT_VERIFYCONTEXT, DWORD, HCRYPTPROV, PROV_RSA_FULL};
    use self::advapi32::{CryptAcquireContextA, CryptGenRandom, CryptReleaseContext};

    /// A random number generator that retrieves randomness straight from
    /// the operating system. Platform sources:
    ///
    /// - Unix-like systems (Linux, Android, Mac OSX): read directly from
    ///   `/dev/urandom`, or from `getrandom(2)` system call if available.
    /// - Windows: calls `CryptGenRandom`, using the default cryptographic
    ///   service provider with the `PROV_RSA_FULL` type.
    /// - iOS: calls SecRandomCopyBytes as /dev/(u)random is sandboxed.
    ///
    /// This does not block.
    pub struct OsRng {
        hcryptprov: HCRYPTPROV
    }

    impl OsRng {
        /// Create a new `OsRng`.
        pub fn new() -> io::Result<OsRng> {
            let mut hcp = 0;
            let ret = unsafe {
                CryptAcquireContextA(&mut hcp, ptr::null(), ptr::null(),
                                     PROV_RSA_FULL,
                                     CRYPT_VERIFYCONTEXT | CRYPT_SILENT)
            };

            if ret == 0 {
                Err(io::Error::last_os_error())
            } else {
                Ok(OsRng { hcryptprov: hcp })
            }
        }
    }

    impl Rng for OsRng {
        fn next_u32(&mut self) -> u32 {
            let mut v = [0u8; 4];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn next_u64(&mut self) -> u64 {
            let mut v = [0u8; 8];
            self.fill_bytes(&mut v);
            unsafe { mem::transmute(v) }
        }
        fn fill_bytes(&mut self, v: &mut [u8]) {
            let ret = unsafe {
                CryptGenRandom(self.hcryptprov, v.len() as DWORD,
                               v.as_mut_ptr())
            };
            if ret == 0 {
                panic!("couldn't generate random bytes: {}",
                       io::Error::last_os_error());
            }
        }
    }

    impl Drop for OsRng {
        fn drop(&mut self) {
            let ret = unsafe {
                CryptReleaseContext(self.hcryptprov, 0)
            };
            if ret == 0 {
                panic!("couldn't release context: {}",
                       io::Error::last_os_error());
            }
        }
    }
}

#[cfg(test)]
mod test {
    use std::sync::mpsc::channel;
    use Rng;
    use OsRng;
    use std::thread;

    #[test]
    fn test_os_rng() {
        let mut r = OsRng::new().unwrap();

        r.next_u32();
        r.next_u64();

        let mut v = [0u8; 1000];
        r.fill_bytes(&mut v);
    }

    #[test]
    fn test_os_rng_tasks() {

        let mut txs = vec!();
        for _ in 0..20 {
            let (tx, rx) = channel();
            txs.push(tx);

            thread::spawn(move|| {
                // wait until all the tasks are ready to go.
                rx.recv().unwrap();

                // deschedule to attempt to interleave things as much
                // as possible (XXX: is this a good test?)
                let mut r = OsRng::new().unwrap();
                thread::yield_now();
                let mut v = [0u8; 1000];

                for _ in 0..100 {
                    r.next_u32();
                    thread::yield_now();
                    r.next_u64();
                    thread::yield_now();
                    r.fill_bytes(&mut v);
                    thread::yield_now();
                }
            });
        }

        // start all the tasks
        for tx in txs.iter() {
            tx.send(()).unwrap();
        }
    }
}
<|endoftext|>


<commit_before>use std::iter::repeat;
use std::io::{self, Write};
use time::{self, Timespec, Duration};
use tty::{Width, terminal_size};

macro_rules! printfl {
    ($($tt:tt)*) => {{
        print!($($tt)*);
        io::stdout().flush().ok().expect("flush() fail");
    }}
}

macro_rules! kb_fmt {
    ($n: ident) => {{
        let kb = 1024f64;
        match $n {
            $n if $n >= kb.powf(4_f64) => format!("{:.*} TB", 2, $n / kb.powf(4_f64)),
            $n if $n >= kb.powf(3_f64) => format!("{:.*} GB", 2, $n / kb.powf(3_f64)),
            $n if $n >= kb.powf(2_f64) => format!("{:.*} MB", 2, $n / kb.powf(2_f64)),
            $n if $n >= kb => format!("{:.*} KB", 2, $n / kb),
            _ => format!("{:.*} B", 0, $n)
        }
    }}
}

macro_rules! repeat {
    ($s: expr, $n: expr) => {{
        &repeat($s).take($n).collect::<String>()
    }}
}

static FORMAT: &'static str = "[=>-]";

// Output type format, indicate which format wil be used in
// the speed box.
#[derive(Debug)]
pub enum Units {
    Default,
    Bytes,
}

#[derive(Debug)]
pub struct ProgressBar {
    start_time: Timespec,
    units: Units,
    total: usize,
    current: usize,
    is_finish: bool,
    show_bar: bool,
    show_speed: bool,
    show_percent: bool,
    show_counter: bool,
    show_time_left: bool,
    bar_start: String,
    bar_current: String,
    bar_current_n: String,
    bar_remain: String,
    bar_end: String,
}

impl ProgressBar {
    /// Create a new ProgressBar with default configuration.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::thread;
    /// use pbr::{ProgressBar, Units};
    ///
    /// let count = 1000;
    /// let mut pb = ProgressBar::new(count);
    /// pb.set_units(Units::Bytes);
    ///
    /// for _ in 0..count {
    ///    pb.inc();
    ///    thread::sleep_ms(100);
    /// }
    /// ```
    pub fn new(total: usize) -> ProgressBar {
        let mut pb = ProgressBar {
            total: total,
            current: 0,
            start_time: time::get_time(),
            units: Units::Default,
            is_finish: false,
            show_bar: true,
            show_speed: false,
            show_percent: true,
            show_counter: true,
            show_time_left: true,
            bar_start: String::new(),
            bar_current: String::new(),
            bar_current_n: String::new(),
            bar_remain: String::new(),
            bar_end: String::new(),
        };
        pb.format(FORMAT);
        pb
    }

    /// set units, default is simple numbers
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pbr::{ProgressBar, Units};
    ///
    /// let mut pb = ProgressBar::new(n_bytes);
    /// pb.set_units(Units::Bytes);
    /// ```
    pub fn set_units(&mut self, u: Units) {
        self.units = u;
    }

    /// Set custom format to the drawing bar, default is `[=>-]`
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut pb = ProgressBar::new(...);
    /// pb.format("[=>_]");
    /// ```
    pub fn format(&mut self, fmt: &str) {
        if fmt.len() >= 5 {
            let v: Vec<&str> = fmt.split("").collect();
            self.bar_start = v[1].to_string();
            self.bar_current = v[2].to_string();
            self.bar_current_n = v[3].to_string();
            self.bar_remain = v[4].to_string();
            self.bar_end = v[5].to_string();
        }
    }

    /// Add to current value
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut pb = ProgressBar::new(10);
    /// pb.add(5);
    /// pb.finish();
    /// ```
    pub fn add(&mut self, i: usize) -> usize {
        self.current += i;
        if self.current <= self.total {
            self.draw()
        };
        self.current
    }

    /// Increment current value
    pub fn inc(&mut self) -> usize {
        return self.add(1);
    }

    fn draw(&self) {
        let tty_size = terminal_size();
        let width = if let Some((Width(w), _)) = tty_size {
            w as usize
        } else {
            80
        };
        let mut base = String::new();
        let mut suffix = String::new();
        let mut prefix = String::new();
        let mut out;
        // precent box
        if self.show_percent {
            let percent = self.current as f64 / (self.total as f64 / 100f64);
            suffix = suffix + &format!(" {:.*} % ", 2, percent);
        }
        // speed box
        if !self.show_speed {
            let from_start = (time::get_time() - self.start_time).num_nanoseconds().unwrap() as f64;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as f64;
            let speed = self.current as f64 / (from_start / sec_nano);
            suffix = match self.units {
                Units::Default => suffix + &format!("{:.*}/s ", 2, speed),
                Units::Bytes => suffix + &format!("{}/s ", kb_fmt!(speed)),
            };
        }
        // time left box
        if self.show_time_left {
            let from_start = time::get_time() - self.start_time;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as i32;
            let per_entry = from_start / self.current as i32;
            let mut left = per_entry * (self.total - self.current) as i32;
            left = (left / sec_nano) * sec_nano;
            if left.num_seconds() > 0 {
                if left.num_seconds() < Duration::minutes(1).num_seconds() {
                    suffix = suffix + &format!("{}s", left.num_seconds());
                } else {
                    suffix = suffix + &format!("{}m", left.num_minutes());
                }
            }
        }
        // counter box
        if self.show_counter {
            let (c, t) = (self.current as f64, self.total as f64);
            prefix = match self.units {
                Units::Default => format!("{} / {} ", c, t),
                Units::Bytes => format!("{} / {} ", kb_fmt!(c), kb_fmt!(t)),
            };
        }
        // bar box
        if self.show_bar {
            let size = width - (prefix.len() + suffix.len() + 3);
            if size > 0 {
                let curr_count =
                    ((self.current as f64 / self.total as f64) * size as f64).ceil() as usize;
                let rema_count = size - curr_count;
                base = self.bar_start.to_string();
                if rema_count > 0 {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count - 1) +
                           &self.bar_current_n;
                } else {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count);
                }
                base = base + repeat!(self.bar_remain.as_ref(), rema_count) + &self.bar_end;
            }
        }
        out = prefix + &base + &suffix;
        // pad
        if out.len() < width {
            let gap = width - out.len();
            out = out + repeat!(" ", gap);
        }
        // print
        printfl!("\r{}", out);
    }

    /// Calling finish manually will set current to total and draw
    /// the last time
    pub fn finish(&mut self) {
        if self.current < self.total {
            self.current = self.total;
            self.draw();
        }
        println!("");
        self.is_finish = true;
    }
}

// Implement io::Writer
impl Write for ProgressBar {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let n = buf.len();
        self.add(n);
        Ok(n)
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use pb::ProgressBar;

    #[test]
    fn add() {
        let mut pb = ProgressBar::new(10);
        pb.add(2);
        assert!(pb.current == 2,
                "should add the given `n` to current");
        assert!(pb.add(2) == pb.current,
                "add should return the current value");
    }

    #[test]
    fn inc() {
        let mut pb = ProgressBar::new(10);
        pb.inc();
        assert!(pb.current == 1, "should increment current by 1");
    }

    #[test]
    fn format() {
        let fmt = "[~> ]";
        let mut pb = ProgressBar::new(1);
        pb.format(fmt);
        assert!(pb.bar_start + &pb.bar_current + &pb.bar_current_n + &pb.bar_remain +
                &pb.bar_end == fmt);
    }

    #[test]
    fn finish() {
        let mut pb = ProgressBar::new(10);
        pb.finish();
        assert!(pb.current == pb.total,
                "should set current to total");
        assert!(pb.is_finish, "should set is_finish to true");
    }

    #[test]
    fn kb_fmt() {
        let kb = 1024f64;
        let mb = kb.powf(2f64);
        let gb = kb.powf(3f64);
        let tb = kb.powf(4f64);
        assert_eq!(kb_fmt!(kb), "1.00 KB");
        assert_eq!(kb_fmt!(mb), "1.00 MB");
        assert_eq!(kb_fmt!(gb), "1.00 GB");
        assert_eq!(kb_fmt!(tb), "1.00 TB");
    }
}
<commit_msg>fix(pb): struct visibility<commit_after>use std::iter::repeat;
use std::io::{self, Write};
use time::{self, Timespec, Duration};
use tty::{Width, terminal_size};

macro_rules! printfl {
    ($($tt:tt)*) => {{
        print!($($tt)*);
        io::stdout().flush().ok().expect("flush() fail");
    }}
}

macro_rules! kb_fmt {
    ($n: ident) => {{
        let kb = 1024f64;
        match $n {
            $n if $n >= kb.powf(4_f64) => format!("{:.*} TB", 2, $n / kb.powf(4_f64)),
            $n if $n >= kb.powf(3_f64) => format!("{:.*} GB", 2, $n / kb.powf(3_f64)),
            $n if $n >= kb.powf(2_f64) => format!("{:.*} MB", 2, $n / kb.powf(2_f64)),
            $n if $n >= kb => format!("{:.*} KB", 2, $n / kb),
            _ => format!("{:.*} B", 0, $n)
        }
    }}
}

macro_rules! repeat {
    ($s: expr, $n: expr) => {{
        &repeat($s).take($n).collect::<String>()
    }}
}

static FORMAT: &'static str = "[=>-]";

// Output type format, indicate which format wil be used in
// the speed box.
#[derive(Debug)]
pub enum Units {
    Default,
    Bytes,
}

#[derive(Debug)]
pub struct ProgressBar {
    start_time: Timespec,
    units: Units,
    total: usize,
    current: usize,
    bar_start: String,
    bar_current: String,
    bar_current_n: String,
    bar_remain: String,
    bar_end: String,
    pub is_finish: bool,
    pub show_bar: bool,
    pub show_speed: bool,
    pub show_percent: bool,
    pub show_counter: bool,
    pub show_time_left: bool,
}

impl ProgressBar {
    /// Create a new ProgressBar with default configuration.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use std::thread;
    /// use pbr::{ProgressBar, Units};
    ///
    /// let count = 1000;
    /// let mut pb = ProgressBar::new(count);
    /// pb.set_units(Units::Bytes);
    ///
    /// for _ in 0..count {
    ///    pb.inc();
    ///    thread::sleep_ms(100);
    /// }
    /// ```
    pub fn new(total: usize) -> ProgressBar {
        let mut pb = ProgressBar {
            total: total,
            current: 0,
            start_time: time::get_time(),
            units: Units::Default,
            is_finish: false,
            show_bar: true,
            show_speed: true,
            show_percent: true,
            show_counter: true,
            show_time_left: true,
            bar_start: String::new(),
            bar_current: String::new(),
            bar_current_n: String::new(),
            bar_remain: String::new(),
            bar_end: String::new(),
        };
        pb.format(FORMAT);
        pb
    }

    /// set units, default is simple numbers
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pbr::{ProgressBar, Units};
    ///
    /// let mut pb = ProgressBar::new(n_bytes);
    /// pb.set_units(Units::Bytes);
    /// ```
    pub fn set_units(&mut self, u: Units) {
        self.units = u;
    }

    /// Set custom format to the drawing bar, default is `[=>-]`
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut pb = ProgressBar::new(...);
    /// pb.format("[=>_]");
    /// ```
    pub fn format(&mut self, fmt: &str) {
        if fmt.len() >= 5 {
            let v: Vec<&str> = fmt.split("").collect();
            self.bar_start = v[1].to_string();
            self.bar_current = v[2].to_string();
            self.bar_current_n = v[3].to_string();
            self.bar_remain = v[4].to_string();
            self.bar_end = v[5].to_string();
        }
    }

    /// Add to current value
    ///
    /// # Examples
    ///
    /// ```no_run
    /// let mut pb = ProgressBar::new(10);
    /// pb.add(5);
    /// pb.finish();
    /// ```
    pub fn add(&mut self, i: usize) -> usize {
        self.current += i;
        if self.current <= self.total {
            self.draw()
        };
        self.current
    }

    /// Increment current value
    pub fn inc(&mut self) -> usize {
        return self.add(1);
    }

    fn draw(&self) {
        let tty_size = terminal_size();
        let width = if let Some((Width(w), _)) = tty_size {
            w as usize
        } else {
            80
        };
        let mut base = String::new();
        let mut suffix = String::new();
        let mut prefix = String::new();
        let mut out;
        // precent box
        if self.show_percent {
            let percent = self.current as f64 / (self.total as f64 / 100f64);
            suffix = suffix + &format!(" {:.*} % ", 2, percent);
        }
        // speed box
        if self.show_speed {
            let from_start = (time::get_time() - self.start_time).num_nanoseconds().unwrap() as f64;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as f64;
            let speed = self.current as f64 / (from_start / sec_nano);
            suffix = match self.units {
                Units::Default => suffix + &format!("{:.*}/s ", 2, speed),
                Units::Bytes => suffix + &format!("{}/s ", kb_fmt!(speed)),
            };
        }
        // time left box
        if self.show_time_left {
            let from_start = time::get_time() - self.start_time;
            let sec_nano = Duration::seconds(1).num_nanoseconds().unwrap() as i32;
            let per_entry = from_start / self.current as i32;
            let mut left = per_entry * (self.total - self.current) as i32;
            left = (left / sec_nano) * sec_nano;
            if left.num_seconds() > 0 {
                if left.num_seconds() < Duration::minutes(1).num_seconds() {
                    suffix = suffix + &format!("{}s", left.num_seconds());
                } else {
                    suffix = suffix + &format!("{}m", left.num_minutes());
                }
            }
        }
        // counter box
        if self.show_counter {
            let (c, t) = (self.current as f64, self.total as f64);
            prefix = match self.units {
                Units::Default => format!("{} / {} ", c, t),
                Units::Bytes => format!("{} / {} ", kb_fmt!(c), kb_fmt!(t)),
            };
        }
        // bar box
        if self.show_bar {
            let size = width - (prefix.len() + suffix.len() + 3);
            if size > 0 {
                let curr_count =
                    ((self.current as f64 / self.total as f64) * size as f64).ceil() as usize;
                let rema_count = size - curr_count;
                base = self.bar_start.to_string();
                if rema_count > 0 {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count - 1) +
                           &self.bar_current_n;
                } else {
                    base = base + repeat!(self.bar_current.as_ref(), curr_count);
                }
                base = base + repeat!(self.bar_remain.as_ref(), rema_count) + &self.bar_end;
            }
        }
        out = prefix + &base + &suffix;
        // pad
        if out.len() < width {
            let gap = width - out.len();
            out = out + repeat!(" ", gap);
        }
        // print
        printfl!("\r{}", out);
    }

    /// Calling finish manually will set current to total and draw
    /// the last time
    pub fn finish(&mut self) {
        if self.current < self.total {
            self.current = self.total;
            self.draw();
        }
        println!("");
        self.is_finish = true;
    }
}

// Implement io::Writer
impl Write for ProgressBar {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let n = buf.len();
        self.add(n);
        Ok(n)
    }
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use pb::ProgressBar;

    #[test]
    fn add() {
        let mut pb = ProgressBar::new(10);
        pb.add(2);
        assert!(pb.current == 2,
                "should add the given `n` to current");
        assert!(pb.add(2) == pb.current,
                "add should return the current value");
    }

    #[test]
    fn inc() {
        let mut pb = ProgressBar::new(10);
        pb.inc();
        assert!(pb.current == 1, "should increment current by 1");
    }

    #[test]
    fn format() {
        let fmt = "[~> ]";
        let mut pb = ProgressBar::new(1);
        pb.format(fmt);
        assert!(pb.bar_start + &pb.bar_current + &pb.bar_current_n + &pb.bar_remain +
                &pb.bar_end == fmt);
    }

    #[test]
    fn finish() {
        let mut pb = ProgressBar::new(10);
        pb.finish();
        assert!(pb.current == pb.total,
                "should set current to total");
        assert!(pb.is_finish, "should set is_finish to true");
    }

    #[test]
    fn kb_fmt() {
        let kb = 1024f64;
        let mb = kb.powf(2f64);
        let gb = kb.powf(3f64);
        let tb = kb.powf(4f64);
        assert_eq!(kb_fmt!(kb), "1.00 KB");
        assert_eq!(kb_fmt!(mb), "1.00 MB");
        assert_eq!(kb_fmt!(gb), "1.00 GB");
        assert_eq!(kb_fmt!(tb), "1.00 TB");
    }
}
<|endoftext|>


<commit_before>use alloc::borrow::Cow;
use alloc::vec::Vec;
use std::slice;

use goblin::pe;
use scroll::Pread;

use {
    Machine, Object, ObjectSection, ObjectSegment, SectionKind, Symbol, SymbolKind, SymbolMap,
};

/// A PE object file.
#[derive(Debug)]
pub struct PeFile<'data> {
    pe: pe::PE<'data>,
    symbol_table_offset: usize,
    data: &'data [u8],
}

/// An iterator over the loadable sections of a `PeFile`.
#[derive(Debug)]
pub struct PeSegmentIterator<'data, 'file>
where
    'data: 'file,
{
    file: &'file PeFile<'data>,
    iter: slice::Iter<'file, pe::section_table::SectionTable>,
}

/// A loadable section of a `PeFile`.
#[derive(Debug)]
pub struct PeSegment<'data, 'file>
where
    'data: 'file,
{
    file: &'file PeFile<'data>,
    section: &'file pe::section_table::SectionTable,
}

/// An iterator over the sections of a `PeFile`.
#[derive(Debug)]
pub struct PeSectionIterator<'data, 'file>
where
    'data: 'file,
{
    file: &'file PeFile<'data>,
    iter: slice::Iter<'file, pe::section_table::SectionTable>,
}

/// A section of a `PeFile`.
#[derive(Debug)]
pub struct PeSection<'data, 'file>
where
    'data: 'file,
{
    file: &'file PeFile<'data>,
    section: &'file pe::section_table::SectionTable,
}

/// An iterator over the symbols of a `PeFile`.
#[derive(Debug)]
pub struct PeSymbolIterator<'data, 'file>
where
    'data: 'file,
{
    exports: slice::Iter<'file, pe::export::Export<'data>>,
    imports: slice::Iter<'file, pe::import::Import<'data>>,
}

impl<'data> PeFile<'data> {
    /// Get the PE headers of the file.
    // TODO: this is temporary to allow access to features this crate doesn't provide yet
    #[inline]
    pub fn pe(&self) -> &pe::PE<'data> {
        &self.pe
    }

    /// Parse the raw PE file data.
    pub fn parse(data: &'data [u8]) -> Result<Self, &'static str> {
        let pe = pe::PE::parse(data).map_err(|_| "Could not parse PE header")?;
        // get offset to secret symbol table that follows the COFF symbol table
        let symbol_table_offset = (
            pe.header.coff_header.pointer_to_symbol_table +
                pe.header.coff_header.number_of_symbol_table * 18
        ) as usize;
        Ok(PeFile { pe, symbol_table_offset, data })
    }

    /// Get section name from special PE symbol table
    fn get_section_name(
        &self,
        section: &'data pe::section_table::SectionTable,
    ) -> Option<&'data str> {
        let mut name = section.name().ok()?;
        // if name start with "/" then it's an index into the secret symbol table
        let mut name_chars = name.chars();
        if name_chars.next() == Some('/') {
            let index = name_chars.as_str().parse::<usize>().ok()?;
            name = self.data.pread(self.symbol_table_offset + index).ok()?;
        };
        Some(name)
    }
}

impl<'data, 'file> Object<'data, 'file> for PeFile<'data>
where
    'data: 'file,
{
    type Segment = PeSegment<'data, 'file>;
    type SegmentIterator = PeSegmentIterator<'data, 'file>;
    type Section = PeSection<'data, 'file>;
    type SectionIterator = PeSectionIterator<'data, 'file>;
    type SymbolIterator = PeSymbolIterator<'data, 'file>;

    fn machine(&self) -> Machine {
        match self.pe.header.coff_header.machine {
            // TODO: Arm/Arm64
            pe::header::COFF_MACHINE_X86 => Machine::X86,
            pe::header::COFF_MACHINE_X86_64 => Machine::X86_64,
            _ => Machine::Other,
        }
    }

    fn segments(&'file self) -> PeSegmentIterator<'data, 'file> {
        PeSegmentIterator {
            file: self,
            iter: self.pe.sections.iter(),
        }
    }

    fn section_data_by_name(&self, section_name: &str) -> Option<Cow<'data, [u8]>> {
        for section in &self.pe.sections {
            if let Some(name) = self.get_section_name(&section) {
                if name == section_name {
                    return Some(Cow::from(
                        &self.data[section.pointer_to_raw_data as usize..]
                            [..section.virtual_size as usize],
                    ));
                }
            }
        }
        None
    }

    fn sections(&'file self) -> PeSectionIterator<'data, 'file> {
        PeSectionIterator {
            file: self,
            iter: self.pe.sections.iter(),
        }
    }

    fn symbols(&'file self) -> PeSymbolIterator<'data, 'file> {
        // TODO: return COFF symbols for object files
        PeSymbolIterator {
            exports: [].iter(),
            imports: [].iter(),
        }
    }

    fn dynamic_symbols(&'file self) -> PeSymbolIterator<'data, 'file> {
        PeSymbolIterator {
            exports: self.pe.exports.iter(),
            imports: self.pe.imports.iter(),
        }
    }

    fn symbol_map(&self) -> SymbolMap<'data> {
        // TODO: untested
        let mut symbols: Vec<_> = self.symbols().filter(SymbolMap::filter).collect();
        symbols.sort_by_key(|x| x.address);
        SymbolMap { symbols }
    }

    #[inline]
    fn is_little_endian(&self) -> bool {
        // TODO: always little endian?  The COFF header has some bits in the
        // characteristics flags, but these are obsolete.
        true
    }

    fn has_debug_symbols(&self) -> bool {
        // TODO: check if CodeView-in-PE still works
        for section in &self.pe.sections {
            if let Some(name) = self.get_section_name(&section) {
                if name == ".debug_info" {
                    return true;
                }
            }
        }
        false
    }

    fn entry(&self) -> u64 {
        self.pe.entry as u64
    }
}

impl<'data, 'file> Iterator for PeSegmentIterator<'data, 'file> {
    type Item = PeSegment<'data, 'file>;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next().map(|section| PeSegment {
            file: self.file,
            section,
        })
    }
}

impl<'data, 'file> ObjectSegment<'data> for PeSegment<'data, 'file> {
    #[inline]
    fn address(&self) -> u64 {
        u64::from(self.section.virtual_address)
    }

    #[inline]
    fn size(&self) -> u64 {
        u64::from(self.section.virtual_size)
    }

    fn data(&self) -> &'data [u8] {
        &self.file.data[self.section.pointer_to_raw_data as usize..]
            [..self.section.size_of_raw_data as usize]
    }

    #[inline]
    fn name(&self) -> Option<&str> {
        self.section.name().ok()
    }
}

impl<'data, 'file> Iterator for PeSectionIterator<'data, 'file> {
    type Item = PeSection<'data, 'file>;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next().map(|section| PeSection {
            file: self.file,
            section,
        })
    }
}

impl<'data, 'file> ObjectSection<'data> for PeSection<'data, 'file> {
    #[inline]
    fn address(&self) -> u64 {
        u64::from(self.section.virtual_address)
    }

    #[inline]
    fn size(&self) -> u64 {
        u64::from(self.section.virtual_size)
    }

    fn data(&self) -> Cow<'data, [u8]> {
        Cow::from(
            &self.file.data[self.section.pointer_to_raw_data as usize..]
                [..self.section.size_of_raw_data as usize],
        )
    }

    fn name(&self) -> Option<&str> {
        self.section.name().ok()
    }

    #[inline]
    fn segment_name(&self) -> Option<&str> {
        None
    }

    #[inline]
    fn kind(&self) -> SectionKind {
        if self.section.characteristics
            & (pe::section_table::IMAGE_SCN_CNT_CODE | pe::section_table::IMAGE_SCN_MEM_EXECUTE)
            != 0
        {
            SectionKind::Text
        } else if self.section.characteristics & pe::section_table::IMAGE_SCN_CNT_INITIALIZED_DATA
            != 0
        {
            SectionKind::Data
        } else if self.section.characteristics & pe::section_table::IMAGE_SCN_CNT_UNINITIALIZED_DATA
            != 0
        {
            SectionKind::UninitializedData
        } else {
            SectionKind::Unknown
        }
    }
}

impl<'data, 'file> Iterator for PeSymbolIterator<'data, 'file> {
    type Item = Symbol<'data>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(export) = self.exports.next() {
            return Some(Symbol {
                kind: SymbolKind::Unknown,
                section_kind: Some(SectionKind::Unknown),
                global: true,
                name: export.name,
                address: export.rva as u64,
                size: 0,
            });
        }
        if let Some(import) = self.imports.next() {
            let name = match import.name {
                Cow::Borrowed(name) => Some(name),
                _ => None,
            };
            return Some(Symbol {
                kind: SymbolKind::Unknown,
                section_kind: None,
                global: true,
                name: name,
                address: 0,
                size: 0,
            });
        }
        None
    }
}
<commit_msg>Issue #40 -- Use Gimli support for DWARF-in-PE.<commit_after>use alloc::borrow::Cow;
use alloc::vec::Vec;
use std::slice;

use goblin::pe;

use {
    Machine, Object, ObjectSection, ObjectSegment, SectionKind, Symbol, SymbolKind, SymbolMap,
};

/// A PE object file.
#[derive(Debug)]
pub struct PeFile<'data> {
    pe: pe::PE<'data>,
    data: &'data [u8],
}

/// An iterator over the loadable sections of a `PeFile`.
#[derive(Debug)]
pub struct PeSegmentIterator<'data, 'file>
where
    'data: 'file,
{
    file: &'file PeFile<'data>,
    iter: slice::Iter<'file, pe::section_table::SectionTable>,
}

/// A loadable section of a `PeFile`.
#[derive(Debug)]
pub struct PeSegment<'data, 'file>
where
    'data: 'file,
{
    file: &'file PeFile<'data>,
    section: &'file pe::section_table::SectionTable,
}

/// An iterator over the sections of a `PeFile`.
#[derive(Debug)]
pub struct PeSectionIterator<'data, 'file>
where
    'data: 'file,
{
    file: &'file PeFile<'data>,
    iter: slice::Iter<'file, pe::section_table::SectionTable>,
}

/// A section of a `PeFile`.
#[derive(Debug)]
pub struct PeSection<'data, 'file>
where
    'data: 'file,
{
    file: &'file PeFile<'data>,
    section: &'file pe::section_table::SectionTable,
}

/// An iterator over the symbols of a `PeFile`.
#[derive(Debug)]
pub struct PeSymbolIterator<'data, 'file>
where
    'data: 'file,
{
    exports: slice::Iter<'file, pe::export::Export<'data>>,
    imports: slice::Iter<'file, pe::import::Import<'data>>,
}

impl<'data> PeFile<'data> {
    /// Get the PE headers of the file.
    // TODO: this is temporary to allow access to features this crate doesn't provide yet
    #[inline]
    pub fn pe(&self) -> &pe::PE<'data> {
        &self.pe
    }

    /// Parse the raw PE file data.
    pub fn parse(data: &'data [u8]) -> Result<Self, &'static str> {
        let pe = pe::PE::parse(data).map_err(|_| "Could not parse PE header")?;
        Ok(PeFile { pe, data })
    }
}

impl<'data, 'file> Object<'data, 'file> for PeFile<'data>
where
    'data: 'file,
{
    type Segment = PeSegment<'data, 'file>;
    type SegmentIterator = PeSegmentIterator<'data, 'file>;
    type Section = PeSection<'data, 'file>;
    type SectionIterator = PeSectionIterator<'data, 'file>;
    type SymbolIterator = PeSymbolIterator<'data, 'file>;

    fn machine(&self) -> Machine {
        match self.pe.header.coff_header.machine {
            // TODO: Arm/Arm64
            pe::header::COFF_MACHINE_X86 => Machine::X86,
            pe::header::COFF_MACHINE_X86_64 => Machine::X86_64,
            _ => Machine::Other,
        }
    }

    fn segments(&'file self) -> PeSegmentIterator<'data, 'file> {
        PeSegmentIterator {
            file: self,
            iter: self.pe.sections.iter(),
        }
    }

    fn section_data_by_name(&self, section_name: &str) -> Option<Cow<'data, [u8]>> {
        for section in &self.pe.sections {
            if let Ok(name) = section.name() {
                if name == section_name {
                    return Some(Cow::from(
                        &self.data[section.pointer_to_raw_data as usize..]
                            [..section.virtual_size as usize],
                    ));
                }
            }
        }
        None
    }

    fn sections(&'file self) -> PeSectionIterator<'data, 'file> {
        PeSectionIterator {
            file: self,
            iter: self.pe.sections.iter(),
        }
    }

    fn symbols(&'file self) -> PeSymbolIterator<'data, 'file> {
        // TODO: return COFF symbols for object files
        PeSymbolIterator {
            exports: [].iter(),
            imports: [].iter(),
        }
    }

    fn dynamic_symbols(&'file self) -> PeSymbolIterator<'data, 'file> {
        PeSymbolIterator {
            exports: self.pe.exports.iter(),
            imports: self.pe.imports.iter(),
        }
    }

    fn symbol_map(&self) -> SymbolMap<'data> {
        // TODO: untested
        let mut symbols: Vec<_> = self.symbols().filter(SymbolMap::filter).collect();
        symbols.sort_by_key(|x| x.address);
        SymbolMap { symbols }
    }

    #[inline]
    fn is_little_endian(&self) -> bool {
        // TODO: always little endian?  The COFF header has some bits in the
        // characteristics flags, but these are obsolete.
        true
    }

    fn has_debug_symbols(&self) -> bool {
        // TODO: check if CodeView-in-PE still works
        for section in &self.pe.sections {
            if let Ok(name) = section.name() {
                if name == ".debug_info" {
                    return true;
                }
            }
        }
        false
    }

    fn entry(&self) -> u64 {
        self.pe.entry as u64
    }
}

impl<'data, 'file> Iterator for PeSegmentIterator<'data, 'file> {
    type Item = PeSegment<'data, 'file>;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next().map(|section| PeSegment {
            file: self.file,
            section,
        })
    }
}

impl<'data, 'file> ObjectSegment<'data> for PeSegment<'data, 'file> {
    #[inline]
    fn address(&self) -> u64 {
        u64::from(self.section.virtual_address)
    }

    #[inline]
    fn size(&self) -> u64 {
        u64::from(self.section.virtual_size)
    }

    fn data(&self) -> &'data [u8] {
        &self.file.data[self.section.pointer_to_raw_data as usize..]
            [..self.section.size_of_raw_data as usize]
    }

    #[inline]
    fn name(&self) -> Option<&str> {
        self.section.name().ok()
    }
}

impl<'data, 'file> Iterator for PeSectionIterator<'data, 'file> {
    type Item = PeSection<'data, 'file>;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next().map(|section| PeSection {
            file: self.file,
            section,
        })
    }
}

impl<'data, 'file> ObjectSection<'data> for PeSection<'data, 'file> {
    #[inline]
    fn address(&self) -> u64 {
        u64::from(self.section.virtual_address)
    }

    #[inline]
    fn size(&self) -> u64 {
        u64::from(self.section.virtual_size)
    }

    fn data(&self) -> Cow<'data, [u8]> {
        Cow::from(
            &self.file.data[self.section.pointer_to_raw_data as usize..]
                [..self.section.size_of_raw_data as usize],
        )
    }

    fn name(&self) -> Option<&str> {
        self.section.name().ok()
    }

    #[inline]
    fn segment_name(&self) -> Option<&str> {
        None
    }

    #[inline]
    fn kind(&self) -> SectionKind {
        if self.section.characteristics
            & (pe::section_table::IMAGE_SCN_CNT_CODE | pe::section_table::IMAGE_SCN_MEM_EXECUTE)
            != 0
        {
            SectionKind::Text
        } else if self.section.characteristics & pe::section_table::IMAGE_SCN_CNT_INITIALIZED_DATA
            != 0
        {
            SectionKind::Data
        } else if self.section.characteristics & pe::section_table::IMAGE_SCN_CNT_UNINITIALIZED_DATA
            != 0
        {
            SectionKind::UninitializedData
        } else {
            SectionKind::Unknown
        }
    }
}

impl<'data, 'file> Iterator for PeSymbolIterator<'data, 'file> {
    type Item = Symbol<'data>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(export) = self.exports.next() {
            return Some(Symbol {
                kind: SymbolKind::Unknown,
                section_kind: Some(SectionKind::Unknown),
                global: true,
                name: export.name,
                address: export.rva as u64,
                size: 0,
            });
        }
        if let Some(import) = self.imports.next() {
            let name = match import.name {
                Cow::Borrowed(name) => Some(name),
                _ => None,
            };
            return Some(Symbol {
                kind: SymbolKind::Unknown,
                section_kind: None,
                global: true,
                name: name,
                address: 0,
                size: 0,
            });
        }
        None
    }
}
<|endoftext|>