<commit_before>// Copyright 2012-2013 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.

use std::num::{One, Zero, CheckedAdd};
use std::vec::bytes::{MutableByteVector, copy_memory};

/// Write a u64 into a vector, which must be 8 bytes long. The value is written in big-endian
/// format.
pub fn write_u64_be(dst: &mut[u8], input: u64) {
    use std::cast::transmute;
    use std::unstable::intrinsics::to_be64;
    assert!(dst.len() == 8);
    unsafe {
        let x: *mut i64 = transmute(dst.unsafe_mut_ref(0));
        *x = to_be64(input as i64);
    }
}

/// Write a u32 into a vector, which must be 4 bytes long. The value is written in big-endian
/// format.
pub fn write_u32_be(dst: &mut[u8], input: u32) {
    use std::cast::transmute;
    use std::unstable::intrinsics::to_be32;
    assert!(dst.len() == 4);
    unsafe {
        let x: *mut i32 = transmute(dst.unsafe_mut_ref(0));
        *x = to_be32(input as i32);
    }
}

/// Write a u32 into a vector, which must be 4 bytes long. The value is written in little-endian
/// format.
pub fn write_u32_le(dst: &mut[u8], input: u32) {
    use std::cast::transmute;
    use std::unstable::intrinsics::to_le32;
    assert!(dst.len() == 4);
    unsafe {
        let x: *mut i32 = transmute(dst.unsafe_mut_ref(0));
        *x = to_le32(input as i32);
    }
}

/// Read a vector of bytes into a vector of u64s. The values are read in big-endian format.
pub fn read_u64v_be(dst: &mut[u64], input: &[u8]) {
    use std::cast::transmute;
    use std::unstable::intrinsics::to_be64;
    assert!(dst.len() * 8 == input.len());
    unsafe {
        let mut x: *mut i64 = transmute(dst.unsafe_mut_ref(0));
        let mut y: *i64 = transmute(input.unsafe_ref(0));
        do dst.len().times() {
            *x = to_be64(*y);
            x = x.offset(1);
            y = y.offset(1);
        }
    }
}

/// Read a vector of bytes into a vector of u32s. The values are read in big-endian format.
pub fn read_u32v_be(dst: &mut[u32], input: &[u8]) {
    use std::cast::transmute;
    use std::unstable::intrinsics::to_be32;
    assert!(dst.len() * 4 == input.len());
    unsafe {
        let mut x: *mut i32 = transmute(dst.unsafe_mut_ref(0));
        let mut y: *i32 = transmute(input.unsafe_ref(0));
        do dst.len().times() {
            *x = to_be32(*y);
            x = x.offset(1);
            y = y.offset(1);
        }
    }
}

/// Read a vector of bytes into a vector of u32s. The values are read in little-endian format.
pub fn read_u32v_le(dst: &mut[u32], input: &[u8]) {
    use std::cast::transmute;
    use std::unstable::intrinsics::to_le32;
    assert!(dst.len() * 4 == input.len());
    unsafe {
        let mut x: *mut i32 = transmute(dst.unsafe_mut_ref(0));
        let mut y: *i32 = transmute(input.unsafe_ref(0));
        do dst.len().times() {
            *x = to_le32(*y);
            x = x.offset(1);
            y = y.offset(1);
        }
    }
}

trait ToBits {
    /// Convert the value in bytes to the number of bits, a tuple where the 1st item is the
    /// high-order value and the 2nd item is the low order value.
    fn to_bits(self) -> (Self, Self);
}

impl ToBits for u64 {
    fn to_bits(self) -> (u64, u64) {
        return (self >> 61, self << 3);
    }
}

/// Adds the specified number of bytes to the bit count. fail!() if this would cause numeric
/// overflow.
pub fn add_bytes_to_bits<T: Int + CheckedAdd + ToBits>(bits: T, bytes: T) -> T {
    let (new_high_bits, new_low_bits) = bytes.to_bits();

    if new_high_bits > Zero::zero() {
        fail!("Numeric overflow occured.")
    }

    match bits.checked_add(&new_low_bits) {
        Some(x) => return x,
        None => fail!("Numeric overflow occured.")
    }
}

/// Adds the specified number of bytes to the bit count, which is a tuple where the first element is
/// the high order value. fail!() if this would cause numeric overflow.
pub fn add_bytes_to_bits_tuple
        <T: Int + Unsigned + CheckedAdd + ToBits>
        (bits: (T, T), bytes: T) -> (T, T) {
    let (new_high_bits, new_low_bits) = bytes.to_bits();
    let (hi, low) = bits;

    // Add the low order value - if there is no overflow, then add the high order values
    // If the addition of the low order values causes overflow, add one to the high order values
    // before adding them.
    match low.checked_add(&new_low_bits) {
        Some(x) => {
            if new_high_bits == Zero::zero() {
                // This is the fast path - every other alternative will rarely occur in practice
                // considering how large an input would need to be for those paths to be used.
                return (hi, x);
            } else {
                match hi.checked_add(&new_high_bits) {
                    Some(y) => return (y, x),
                    None => fail!("Numeric overflow occured.")
                }
            }
        },
        None => {
            let one: T = One::one();
            let z = match new_high_bits.checked_add(&one) {
                Some(w) => w,
                None => fail!("Numeric overflow occured.")
            };
            match hi.checked_add(&z) {
                // This re-executes the addition that was already performed earlier when overflow
                // occured, this time allowing the overflow to happen. Technically, this could be
                // avoided by using the checked add intrinsic directly, but that involves using
                // unsafe code and is not really worthwhile considering how infrequently code will
                // run in practice. This is the reason that this function requires that the type T
                // be Unsigned - overflow is not defined for Signed types. This function could be
                // implemented for signed types as well if that were needed.
                Some(y) => return (y, low + new_low_bits),
                None => fail!("Numeric overflow occured.")
            }
        }
    }
}

/// A FixedBuffer, likes its name implies, is a fixed size buffer. When the buffer becomes full, it
/// must be processed. The input() method takes care of processing and then clearing the buffer
/// automatically. However, other methods do not and require the caller to process the buffer. Any
/// method that modifies the buffer directory or provides the caller with bytes that can be modifies
/// results in those bytes being marked as used by the buffer.
pub trait FixedBuffer {
    /// Input a vector of bytes. If the buffer becomes full, proccess it with the provided
    /// function and then clear the buffer.
    fn input(&mut self, input: &[u8], func: &fn(&[u8]));

    /// Reset the buffer.
    fn reset(&mut self);

    /// Zero the buffer up until the specified index. The buffer position currently must not be
    /// greater than that index.
    fn zero_until(&mut self, idx: uint);

    /// Get a slice of the buffer of the specified size. There must be at least that many bytes
    /// remaining in the buffer.
    fn next<'s>(&'s mut self, len: uint) -> &'s mut [u8];

    /// Get the current buffer. The buffer must already be full. This clears the buffer as well.
    fn full_buffer<'s>(&'s mut self) -> &'s [u8];

    /// Get the current position of the buffer.
    fn position(&self) -> uint;

    /// Get the number of bytes remaining in the buffer until it is full.
    fn remaining(&self) -> uint;

    /// Get the size of the buffer
    fn size(&self) -> uint;
}

macro_rules! impl_fixed_buffer( ($name:ident, $size:expr) => (
    impl FixedBuffer for $name {
        fn input(&mut self, input: &[u8], func: &fn(&[u8])) {
            let mut i = 0;

            // FIXME: #6304 - This local variable shouldn't be necessary.
            let size = $size;

            // If there is already data in the buffer, copy as much as we can into it and process
            // the data if the buffer becomes full.
            if self.buffer_idx != 0 {
                let buffer_remaining = size - self.buffer_idx;
                if input.len() >= buffer_remaining {
                        copy_memory(
                            self.buffer.mut_slice(self.buffer_idx, size),
                            input.slice_to(buffer_remaining),
                            buffer_remaining);
                    self.buffer_idx = 0;
                    func(self.buffer);
                    i += buffer_remaining;
                } else {
                    copy_memory(
                        self.buffer.mut_slice(self.buffer_idx, self.buffer_idx + input.len()),
                        input,
                        input.len());
                    self.buffer_idx += input.len();
                    return;
                }
            }

            // While we have at least a full buffer size chunks's worth of data, process that data
            // without copying it into the buffer
            while input.len() - i >= size {
                func(input.slice(i, i + size));
                i += size;
            }

            // Copy any input data into the buffer. At this point in the method, the ammount of
            // data left in the input vector will be less than the buffer size and the buffer will
            // be empty.
            let input_remaining = input.len() - i;
            copy_memory(
                self.buffer.mut_slice(0, input_remaining),
                input.slice_from(i),
                input.len() - i);
            self.buffer_idx += input_remaining;
        }

        fn reset(&mut self) {
            self.buffer_idx = 0;
        }

        fn zero_until(&mut self, idx: uint) {
            assert!(idx >= self.buffer_idx);
            self.buffer.mut_slice(self.buffer_idx, idx).set_memory(0);
            self.buffer_idx = idx;
        }

        fn next<'s>(&'s mut self, len: uint) -> &'s mut [u8] {
            self.buffer_idx += len;
            return self.buffer.mut_slice(self.buffer_idx - len, self.buffer_idx);
        }

        fn full_buffer<'s>(&'s mut self) -> &'s [u8] {
            assert!(self.buffer_idx == $size);
            self.buffer_idx = 0;
            return self.buffer.slice_to($size);
        }

        fn position(&self) -> uint { self.buffer_idx }

        fn remaining(&self) -> uint { $size - self.buffer_idx }

        fn size(&self) -> uint { $size }
    }
))

/// A fixed size buffer of 64 bytes useful for cryptographic operations.
pub struct FixedBuffer64 {
    priv buffer: [u8, ..64],
    priv buffer_idx: uint,
}

impl FixedBuffer64 {
    /// Create a new buffer
    pub fn new() -> FixedBuffer64 {
        return FixedBuffer64 {
            buffer: [0u8, ..64],
            buffer_idx: 0
        };
    }
}

impl_fixed_buffer!(FixedBuffer64, 64)

/// A fixed size buffer of 128 bytes useful for cryptographic operations.
pub struct FixedBuffer128 {
    priv buffer: [u8, ..128],
    priv buffer_idx: uint,
}

impl FixedBuffer128 {
    /// Create a new buffer
    pub fn new() -> FixedBuffer128 {
        return FixedBuffer128 {
            buffer: [0u8, ..128],
            buffer_idx: 0
        };
    }
}

impl_fixed_buffer!(FixedBuffer128, 128)

/// The StandardPadding trait adds a method useful for various hash algorithms to a FixedBuffer
/// struct.
pub trait StandardPadding {
    /// Add standard padding to the buffer. The buffer must not be full when this method is called
    /// and is guaranteed to have exactly rem remaining bytes when it returns. If there are not at
    /// least rem bytes available, the buffer will be zero padded, processed, cleared, and then
    /// filled with zeros again until only rem bytes are remaining.
    fn standard_padding(&mut self, rem: uint, func: &fn(&[u8]));
}

impl <T: FixedBuffer> StandardPadding for T {
    fn standard_padding(&mut self, rem: uint, func: &fn(&[u8])) {
        let size = self.size();

        self.next(1)[0] = 128;

        if self.remaining() < rem {
            self.zero_until(size);
            func(self.full_buffer());
        }

        self.zero_until(size - rem);
    }
}

#[cfg(test)]
mod test {
    use std::rand::IsaacRng;
    use std::rand::RngUtil;
    use std::vec;

    use digest::Digest;

    /// Feed 1,000,000 'a's into the digest with varying input sizes and check that the result is
    /// correct.
    pub fn test_digest_1million_random<D: Digest>(digest: &mut D, blocksize: uint, expected: &str) {
        let total_size = 1000000;
        let buffer = vec::from_elem(blocksize * 2, 'a' as u8);
        let mut rng = IsaacRng::new_unseeded();
        let mut count = 0;

        digest.reset();

        while count < total_size {
            let next: uint = rng.gen_uint_range(0, 2 * blocksize + 1);
            let remaining = total_size - count;
            let size = if next > remaining { remaining } else { next };
            digest.input(buffer.slice_to(size));
            count += size;
        }

        let result_str = digest.result_str();

        assert!(expected == result_str);
    }
}
<commit_msg>Crypto: Add tests for add_bytes_to_bits functions.<commit_after>// Copyright 2012-2013 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.

use std::num::{One, Zero, CheckedAdd};
use std::vec::bytes::{MutableByteVector, copy_memory};

/// Write a u64 into a vector, which must be 8 bytes long. The value is written in big-endian
/// format.
pub fn write_u64_be(dst: &mut[u8], input: u64) {
    use std::cast::transmute;
    use std::unstable::intrinsics::to_be64;
    assert!(dst.len() == 8);
    unsafe {
        let x: *mut i64 = transmute(dst.unsafe_mut_ref(0));
        *x = to_be64(input as i64);
    }
}

/// Write a u32 into a vector, which must be 4 bytes long. The value is written in big-endian
/// format.
pub fn write_u32_be(dst: &mut[u8], input: u32) {
    use std::cast::transmute;
    use std::unstable::intrinsics::to_be32;
    assert!(dst.len() == 4);
    unsafe {
        let x: *mut i32 = transmute(dst.unsafe_mut_ref(0));
        *x = to_be32(input as i32);
    }
}

/// Write a u32 into a vector, which must be 4 bytes long. The value is written in little-endian
/// format.
pub fn write_u32_le(dst: &mut[u8], input: u32) {
    use std::cast::transmute;
    use std::unstable::intrinsics::to_le32;
    assert!(dst.len() == 4);
    unsafe {
        let x: *mut i32 = transmute(dst.unsafe_mut_ref(0));
        *x = to_le32(input as i32);
    }
}

/// Read a vector of bytes into a vector of u64s. The values are read in big-endian format.
pub fn read_u64v_be(dst: &mut[u64], input: &[u8]) {
    use std::cast::transmute;
    use std::unstable::intrinsics::to_be64;
    assert!(dst.len() * 8 == input.len());
    unsafe {
        let mut x: *mut i64 = transmute(dst.unsafe_mut_ref(0));
        let mut y: *i64 = transmute(input.unsafe_ref(0));
        do dst.len().times() {
            *x = to_be64(*y);
            x = x.offset(1);
            y = y.offset(1);
        }
    }
}

/// Read a vector of bytes into a vector of u32s. The values are read in big-endian format.
pub fn read_u32v_be(dst: &mut[u32], input: &[u8]) {
    use std::cast::transmute;
    use std::unstable::intrinsics::to_be32;
    assert!(dst.len() * 4 == input.len());
    unsafe {
        let mut x: *mut i32 = transmute(dst.unsafe_mut_ref(0));
        let mut y: *i32 = transmute(input.unsafe_ref(0));
        do dst.len().times() {
            *x = to_be32(*y);
            x = x.offset(1);
            y = y.offset(1);
        }
    }
}

/// Read a vector of bytes into a vector of u32s. The values are read in little-endian format.
pub fn read_u32v_le(dst: &mut[u32], input: &[u8]) {
    use std::cast::transmute;
    use std::unstable::intrinsics::to_le32;
    assert!(dst.len() * 4 == input.len());
    unsafe {
        let mut x: *mut i32 = transmute(dst.unsafe_mut_ref(0));
        let mut y: *i32 = transmute(input.unsafe_ref(0));
        do dst.len().times() {
            *x = to_le32(*y);
            x = x.offset(1);
            y = y.offset(1);
        }
    }
}

trait ToBits {
    /// Convert the value in bytes to the number of bits, a tuple where the 1st item is the
    /// high-order value and the 2nd item is the low order value.
    fn to_bits(self) -> (Self, Self);
}

impl ToBits for u64 {
    fn to_bits(self) -> (u64, u64) {
        return (self >> 61, self << 3);
    }
}

/// Adds the specified number of bytes to the bit count. fail!() if this would cause numeric
/// overflow.
pub fn add_bytes_to_bits<T: Int + CheckedAdd + ToBits>(bits: T, bytes: T) -> T {
    let (new_high_bits, new_low_bits) = bytes.to_bits();

    if new_high_bits > Zero::zero() {
        fail!("Numeric overflow occured.")
    }

    match bits.checked_add(&new_low_bits) {
        Some(x) => return x,
        None => fail!("Numeric overflow occured.")
    }
}

/// Adds the specified number of bytes to the bit count, which is a tuple where the first element is
/// the high order value. fail!() if this would cause numeric overflow.
pub fn add_bytes_to_bits_tuple
        <T: Int + Unsigned + CheckedAdd + ToBits>
        (bits: (T, T), bytes: T) -> (T, T) {
    let (new_high_bits, new_low_bits) = bytes.to_bits();
    let (hi, low) = bits;

    // Add the low order value - if there is no overflow, then add the high order values
    // If the addition of the low order values causes overflow, add one to the high order values
    // before adding them.
    match low.checked_add(&new_low_bits) {
        Some(x) => {
            if new_high_bits == Zero::zero() {
                // This is the fast path - every other alternative will rarely occur in practice
                // considering how large an input would need to be for those paths to be used.
                return (hi, x);
            } else {
                match hi.checked_add(&new_high_bits) {
                    Some(y) => return (y, x),
                    None => fail!("Numeric overflow occured.")
                }
            }
        },
        None => {
            let one: T = One::one();
            let z = match new_high_bits.checked_add(&one) {
                Some(w) => w,
                None => fail!("Numeric overflow occured.")
            };
            match hi.checked_add(&z) {
                // This re-executes the addition that was already performed earlier when overflow
                // occured, this time allowing the overflow to happen. Technically, this could be
                // avoided by using the checked add intrinsic directly, but that involves using
                // unsafe code and is not really worthwhile considering how infrequently code will
                // run in practice. This is the reason that this function requires that the type T
                // be Unsigned - overflow is not defined for Signed types. This function could be
                // implemented for signed types as well if that were needed.
                Some(y) => return (y, low + new_low_bits),
                None => fail!("Numeric overflow occured.")
            }
        }
    }
}

/// A FixedBuffer, likes its name implies, is a fixed size buffer. When the buffer becomes full, it
/// must be processed. The input() method takes care of processing and then clearing the buffer
/// automatically. However, other methods do not and require the caller to process the buffer. Any
/// method that modifies the buffer directory or provides the caller with bytes that can be modifies
/// results in those bytes being marked as used by the buffer.
pub trait FixedBuffer {
    /// Input a vector of bytes. If the buffer becomes full, proccess it with the provided
    /// function and then clear the buffer.
    fn input(&mut self, input: &[u8], func: &fn(&[u8]));

    /// Reset the buffer.
    fn reset(&mut self);

    /// Zero the buffer up until the specified index. The buffer position currently must not be
    /// greater than that index.
    fn zero_until(&mut self, idx: uint);

    /// Get a slice of the buffer of the specified size. There must be at least that many bytes
    /// remaining in the buffer.
    fn next<'s>(&'s mut self, len: uint) -> &'s mut [u8];

    /// Get the current buffer. The buffer must already be full. This clears the buffer as well.
    fn full_buffer<'s>(&'s mut self) -> &'s [u8];

    /// Get the current position of the buffer.
    fn position(&self) -> uint;

    /// Get the number of bytes remaining in the buffer until it is full.
    fn remaining(&self) -> uint;

    /// Get the size of the buffer
    fn size(&self) -> uint;
}

macro_rules! impl_fixed_buffer( ($name:ident, $size:expr) => (
    impl FixedBuffer for $name {
        fn input(&mut self, input: &[u8], func: &fn(&[u8])) {
            let mut i = 0;

            // FIXME: #6304 - This local variable shouldn't be necessary.
            let size = $size;

            // If there is already data in the buffer, copy as much as we can into it and process
            // the data if the buffer becomes full.
            if self.buffer_idx != 0 {
                let buffer_remaining = size - self.buffer_idx;
                if input.len() >= buffer_remaining {
                        copy_memory(
                            self.buffer.mut_slice(self.buffer_idx, size),
                            input.slice_to(buffer_remaining),
                            buffer_remaining);
                    self.buffer_idx = 0;
                    func(self.buffer);
                    i += buffer_remaining;
                } else {
                    copy_memory(
                        self.buffer.mut_slice(self.buffer_idx, self.buffer_idx + input.len()),
                        input,
                        input.len());
                    self.buffer_idx += input.len();
                    return;
                }
            }

            // While we have at least a full buffer size chunks's worth of data, process that data
            // without copying it into the buffer
            while input.len() - i >= size {
                func(input.slice(i, i + size));
                i += size;
            }

            // Copy any input data into the buffer. At this point in the method, the ammount of
            // data left in the input vector will be less than the buffer size and the buffer will
            // be empty.
            let input_remaining = input.len() - i;
            copy_memory(
                self.buffer.mut_slice(0, input_remaining),
                input.slice_from(i),
                input.len() - i);
            self.buffer_idx += input_remaining;
        }

        fn reset(&mut self) {
            self.buffer_idx = 0;
        }

        fn zero_until(&mut self, idx: uint) {
            assert!(idx >= self.buffer_idx);
            self.buffer.mut_slice(self.buffer_idx, idx).set_memory(0);
            self.buffer_idx = idx;
        }

        fn next<'s>(&'s mut self, len: uint) -> &'s mut [u8] {
            self.buffer_idx += len;
            return self.buffer.mut_slice(self.buffer_idx - len, self.buffer_idx);
        }

        fn full_buffer<'s>(&'s mut self) -> &'s [u8] {
            assert!(self.buffer_idx == $size);
            self.buffer_idx = 0;
            return self.buffer.slice_to($size);
        }

        fn position(&self) -> uint { self.buffer_idx }

        fn remaining(&self) -> uint { $size - self.buffer_idx }

        fn size(&self) -> uint { $size }
    }
))

/// A fixed size buffer of 64 bytes useful for cryptographic operations.
pub struct FixedBuffer64 {
    priv buffer: [u8, ..64],
    priv buffer_idx: uint,
}

impl FixedBuffer64 {
    /// Create a new buffer
    pub fn new() -> FixedBuffer64 {
        return FixedBuffer64 {
            buffer: [0u8, ..64],
            buffer_idx: 0
        };
    }
}

impl_fixed_buffer!(FixedBuffer64, 64)

/// A fixed size buffer of 128 bytes useful for cryptographic operations.
pub struct FixedBuffer128 {
    priv buffer: [u8, ..128],
    priv buffer_idx: uint,
}

impl FixedBuffer128 {
    /// Create a new buffer
    pub fn new() -> FixedBuffer128 {
        return FixedBuffer128 {
            buffer: [0u8, ..128],
            buffer_idx: 0
        };
    }
}

impl_fixed_buffer!(FixedBuffer128, 128)

/// The StandardPadding trait adds a method useful for various hash algorithms to a FixedBuffer
/// struct.
pub trait StandardPadding {
    /// Add standard padding to the buffer. The buffer must not be full when this method is called
    /// and is guaranteed to have exactly rem remaining bytes when it returns. If there are not at
    /// least rem bytes available, the buffer will be zero padded, processed, cleared, and then
    /// filled with zeros again until only rem bytes are remaining.
    fn standard_padding(&mut self, rem: uint, func: &fn(&[u8]));
}

impl <T: FixedBuffer> StandardPadding for T {
    fn standard_padding(&mut self, rem: uint, func: &fn(&[u8])) {
        let size = self.size();

        self.next(1)[0] = 128;

        if self.remaining() < rem {
            self.zero_until(size);
            func(self.full_buffer());
        }

        self.zero_until(size - rem);
    }
}

#[cfg(test)]
mod test {
    use std::rand::IsaacRng;
    use std::rand::RngUtil;
    use std::vec;

    use cryptoutil::{add_bytes_to_bits, add_bytes_to_bits_tuple};
    use digest::Digest;

    /// Feed 1,000,000 'a's into the digest with varying input sizes and check that the result is
    /// correct.
    pub fn test_digest_1million_random<D: Digest>(digest: &mut D, blocksize: uint, expected: &str) {
        let total_size = 1000000;
        let buffer = vec::from_elem(blocksize * 2, 'a' as u8);
        let mut rng = IsaacRng::new_unseeded();
        let mut count = 0;

        digest.reset();

        while count < total_size {
            let next: uint = rng.gen_uint_range(0, 2 * blocksize + 1);
            let remaining = total_size - count;
            let size = if next > remaining { remaining } else { next };
            digest.input(buffer.slice_to(size));
            count += size;
        }

        let result_str = digest.result_str();

        assert!(expected == result_str);
    }

    // A normal addition - no overflow occurs
    #[test]
    fn test_add_bytes_to_bits_ok() {
        assert!(add_bytes_to_bits::<u64>(100, 10) == 180);
    }

    // A simple failure case - adding 1 to the max value
    #[test]
    #[should_fail]
    fn test_add_bytes_to_bits_overflow() {
        add_bytes_to_bits::<u64>(Bounded::max_value(), 1);
    }

    // A normal addition - no overflow occurs (fast path)
    #[test]
    fn test_add_bytes_to_bits_tuple_ok() {
        assert!(add_bytes_to_bits_tuple::<u64>((5, 100), 10) == (5, 180));
    }

    // The low order value overflows into the high order value
    #[test]
    fn test_add_bytes_to_bits_tuple_ok2() {
        assert!(add_bytes_to_bits_tuple::<u64>((5, Bounded::max_value()), 1) == (6, 7));
    }

    // The value to add is too large to be converted into bits without overflowing its type
    #[test]
    fn test_add_bytes_to_bits_tuple_ok3() {
        assert!(add_bytes_to_bits_tuple::<u64>((5, 0), 0x4000000000000001) == (7, 8));
    }

    // A simple failure case - adding 1 to the max value
    #[test]
    #[should_fail]
    fn test_add_bytes_to_bits_tuple_overflow() {
        add_bytes_to_bits_tuple::<u64>((Bounded::max_value(), Bounded::max_value()), 1);
    }

    // The value to add is too large to convert to bytes without overflowing its type, but the high
    // order value from this conversion overflows when added to the existing high order value
    #[test]
    #[should_fail]
    fn test_add_bytes_to_bits_tuple_overflow2() {
        add_bytes_to_bits_tuple::<u64>((Bounded::max_value::<u64>() - 1, 0), 0x8000000000000000);
    }
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#![cfg_attr(not(feature = "clippy"), allow(unknown_lints))]
#![allow(doc_markdown)]

extern crate devicemapper;
extern crate libstratis;
#[macro_use]
extern crate log;
extern crate clap;
#[cfg(feature = "dbus_enabled")]
extern crate dbus;
extern crate env_logger;
extern crate libc;
extern crate libudev;
extern crate nix;

use std::cell::RefCell;
use std::env;
use std::fs::{File, OpenOptions};
use std::io::{ErrorKind, Read, Write};
use std::os::unix::io::AsRawFd;
use std::path::PathBuf;
use std::process::exit;
use std::rc::Rc;

use clap::{App, Arg, ArgMatches};
use env_logger::LogBuilder;
use libc::pid_t;
use log::{LogLevelFilter, SetLoggerError};
use nix::fcntl::{flock, FlockArg};
use nix::unistd::getpid;

#[cfg(feature = "dbus_enabled")]
use dbus::WatchEvent;

use devicemapper::Device;

use libstratis::engine::{Engine, SimEngine, StratEngine};
use libstratis::stratis::{StratisError, StratisResult, VERSION};

const STRATISD_PID_PATH: &str = "/var/run/stratisd.pid";

/// If writing a program error to stderr fails, panic.
fn print_err(err: &StratisError) -> () {
    eprintln!("{}", err);
}

/// Configure and initialize the logger.
/// If debug is true, log at debug level. Otherwise read log configuration
/// parameters from the environment if RUST_LOG is set. Otherwise, just
/// accept the default configuration.
fn initialize_log(debug: bool) -> Result<(), SetLoggerError> {
    let mut builder = LogBuilder::new();
    if debug {
        builder.filter(Some("stratisd"), LogLevelFilter::Debug);
        builder.filter(Some("libstratis"), LogLevelFilter::Debug);
    } else if let Ok(s) = env::var("RUST_LOG") {
        builder.parse(&s);
    };

    builder.init()
}

/// Given a udev event check to see if it's an add and if it is return the device node and
/// devicemapper::Device.
fn handle_udev_add(event: &libudev::Event) -> Option<(Device, PathBuf)> {
    if event.event_type() == libudev::EventType::Add {
        let device = event.device();
        return device.devnode().and_then(|devnode| {
            device
                .devnum()
                .and_then(|devnum| Some((Device::from(devnum), PathBuf::from(devnode))))
        });
    }
    None
}

/// To ensure only one instance of stratisd runs at a time, acquire an
/// exclusive lock. Return an error if lock attempt fails.
fn trylock_pid_file() -> StratisResult<File> {
    let mut f = match OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .open(STRATISD_PID_PATH)
    {
        Ok(f) => f,
        Err(e) => {
            if e.kind() == ErrorKind::PermissionDenied {
                return Err(StratisError::Error(
                    "Must be running as root in order to start daemon.".to_string(),
                ));
            }
            return Err(e.into());
        }
    };
    match flock(f.as_raw_fd(), FlockArg::LockExclusiveNonblock) {
        Ok(_) => {
            f.write_all(format!("{}\n", getpid()).as_bytes())?;
            Ok(f)
        }
        Err(_) => {
            let mut buf = String::new();
            f.read_to_string(&mut buf)?;
            // pidfile is supposed to contain pid of holder. But you never
            // know so be paranoid.
            let pid_str = buf.split_whitespace()
                .next()
                .and_then(|s| s.parse::<pid_t>().ok())
                .map(|pid| format!("{}", pid))
                .unwrap_or_else(|| "<unknown>".into());
            Err(StratisError::Error(format!(
                "Daemon already running with pid: {}",
                pid_str
            )))
        }
    }
}

fn run(matches: &ArgMatches) -> StratisResult<()> {
    // Setup a udev listener before initializing the engine. A device may
    // appear after the engine has read the /dev directory but before it has
    // completed initialization. Unless the udev event has been recorded, the
    // engine will miss the device.
    // This is especially important since stratisd must run during early boot.
    let context = libudev::Context::new()?;
    let mut monitor = libudev::Monitor::new(&context)?;
    monitor.match_subsystem_devtype("block", "disk")?;
    let mut udev = monitor.listen()?;

    let engine: Rc<RefCell<Engine>> = {
        if matches.is_present("sim") {
            info!("Using SimEngine");
            Rc::new(RefCell::new(SimEngine::default()))
        } else {
            info!("Using StratEngine");
            Rc::new(RefCell::new(StratEngine::initialize()?))
        }
    };

    /*
    The file descriptor array indexes are laid out in the following:

    0   == Always udev fd index
    1   == engine index if eventable
    1/2 == Start of dbus client file descriptor(s), 1 if engine is not eventable, else 2
    */
    const FD_INDEX_UDEV: usize = 0;
    const FD_INDEX_ENGINE: usize = 1;

    let mut fds = Vec::new();

    fds.push(libc::pollfd {
        fd: udev.as_raw_fd(),
        revents: 0,
        events: libc::POLLIN,
    });

    let eventable = engine.borrow().get_eventable();

    // The variable _dbus_client_index_start is only used when dbus support is compiled in, thus
    // we denote the value as not needed to compile when dbus support is not included.
    let (poll_timeout, _dbus_client_index_start) = match eventable {
        Some(ref evt) => {
            fds.push(libc::pollfd {
                fd: evt.get_pollable_fd(),
                revents: 0,
                events: libc::POLLIN,
            });

            (-1, FD_INDEX_ENGINE + 1)
        }
        None => (10000, FD_INDEX_ENGINE),
    };

    #[cfg(feature = "dbus_enabled")]
    let (dbus_conn, mut tree, base_object_path, dbus_context) =
        libstratis::dbus_api::connect(Rc::clone(&engine))?;

    #[cfg(feature = "dbus_enabled")]
    for (_, pool_uuid, pool) in engine.borrow().pools() {
        libstratis::dbus_api::register_pool(
            &dbus_conn,
            &dbus_context,
            &mut tree,
            pool_uuid,
            pool,
            &base_object_path,
        )?;
    }

    loop {
        // Process any udev block events
        if fds[FD_INDEX_UDEV].revents != 0 {
            while let Some(event) = udev.receive_event() {
                if let Some((device, devnode)) = handle_udev_add(&event) {
                    // If block evaluate returns an error we are going to ignore it as
                    // there is nothing we can do for a device we are getting errors with.
                    let pool_uuid = engine
                        .borrow_mut()
                        .block_evaluate(device, devnode)
                        .unwrap_or(None);

                    // We need to pretend that we aren't using the variable _pool_uuid so
                    // that we can conditionally compile out the register_pool when dbus
                    // is not enabled.
                    if let Some(_pool_uuid) = pool_uuid {
                        #[cfg(feature = "dbus_enabled")]
                        libstratis::dbus_api::register_pool(
                            &dbus_conn,
                            &dbus_context,
                            &mut tree,
                            _pool_uuid,
                            engine
                                .borrow()
                                .get_pool(_pool_uuid)
                                .expect(
                                    "block_evaluate() returned a pool UUID, pool must be available",
                                )
                                .1,
                            &base_object_path,
                        )?;
                    }
                }
            }
        }

        // Handle engine events, if the engine is eventable
        match eventable {
            Some(ref evt) => {
                if fds[FD_INDEX_ENGINE].revents != 0 {
                    fds[FD_INDEX_ENGINE].revents = 0;
                    evt.clear_event()?;
                    engine.borrow_mut().evented()?;
                }
            }
            None => {
                // Unconditionally call evented() if engine has no eventable.
                // This looks like a bad idea, but the only engine that has
                // no eventable is the sim engine, and for that engine,
                // evented() is essentially a no-op.
                engine.borrow_mut().evented()?;
            }
        }

        // Iterate through D-Bus file descriptors (if enabled)
        #[cfg(feature = "dbus_enabled")]
        {
            for pfd in fds[_dbus_client_index_start..]
                .iter()
                .filter(|pfd| pfd.revents != 0)
            {
                for item in dbus_conn.watch_handle(pfd.fd, WatchEvent::from_revents(pfd.revents)) {
                    if let Err(r) =
                        libstratis::dbus_api::handle(&dbus_conn, &item, &mut tree, &dbus_context)
                    {
                        print_err(&From::from(r));
                    }
                }
            }

            // Refresh list of dbus fds to poll for every time. This can change as
            // D-Bus clients come and go.
            fds.truncate(_dbus_client_index_start);

            fds.extend(dbus_conn.watch_fds().iter().map(|w| w.to_pollfd()));
        }

        let r = unsafe { libc::poll(fds.as_mut_ptr(), fds.len() as libc::c_ulong, poll_timeout) };
        assert!(r >= 0);
    }
}

fn main() {
    let matches = App::new("stratis")
        .version(VERSION)
        .about("Stratis storage management")
        .arg(
            Arg::with_name("debug")
                .long("debug")
                .help("Print additional output for debugging"),
        )
        .arg(
            Arg::with_name("sim")
                .long("sim")
                .help("Use simulator engine"),
        )
        .get_matches();

    let result = trylock_pid_file()
        .and_then(|_pidfile| {
            Ok(initialize_log(matches.is_present("debug"))
                .expect("This is the first and only invocation of this method; it must succeed."))
        })
        .and_then(|_| run(&matches));
    if let Err(err) = result {
        print_err(&err);
        exit(1);
    } else {
        exit(0);
    }
}
<commit_msg>Do not set revents to 0 when handling engine events<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#![cfg_attr(not(feature = "clippy"), allow(unknown_lints))]
#![allow(doc_markdown)]

extern crate devicemapper;
extern crate libstratis;
#[macro_use]
extern crate log;
extern crate clap;
#[cfg(feature = "dbus_enabled")]
extern crate dbus;
extern crate env_logger;
extern crate libc;
extern crate libudev;
extern crate nix;

use std::cell::RefCell;
use std::env;
use std::fs::{File, OpenOptions};
use std::io::{ErrorKind, Read, Write};
use std::os::unix::io::AsRawFd;
use std::path::PathBuf;
use std::process::exit;
use std::rc::Rc;

use clap::{App, Arg, ArgMatches};
use env_logger::LogBuilder;
use libc::pid_t;
use log::{LogLevelFilter, SetLoggerError};
use nix::fcntl::{flock, FlockArg};
use nix::unistd::getpid;

#[cfg(feature = "dbus_enabled")]
use dbus::WatchEvent;

use devicemapper::Device;

use libstratis::engine::{Engine, SimEngine, StratEngine};
use libstratis::stratis::{StratisError, StratisResult, VERSION};

const STRATISD_PID_PATH: &str = "/var/run/stratisd.pid";

/// If writing a program error to stderr fails, panic.
fn print_err(err: &StratisError) -> () {
    eprintln!("{}", err);
}

/// Configure and initialize the logger.
/// If debug is true, log at debug level. Otherwise read log configuration
/// parameters from the environment if RUST_LOG is set. Otherwise, just
/// accept the default configuration.
fn initialize_log(debug: bool) -> Result<(), SetLoggerError> {
    let mut builder = LogBuilder::new();
    if debug {
        builder.filter(Some("stratisd"), LogLevelFilter::Debug);
        builder.filter(Some("libstratis"), LogLevelFilter::Debug);
    } else if let Ok(s) = env::var("RUST_LOG") {
        builder.parse(&s);
    };

    builder.init()
}

/// Given a udev event check to see if it's an add and if it is return the device node and
/// devicemapper::Device.
fn handle_udev_add(event: &libudev::Event) -> Option<(Device, PathBuf)> {
    if event.event_type() == libudev::EventType::Add {
        let device = event.device();
        return device.devnode().and_then(|devnode| {
            device
                .devnum()
                .and_then(|devnum| Some((Device::from(devnum), PathBuf::from(devnode))))
        });
    }
    None
}

/// To ensure only one instance of stratisd runs at a time, acquire an
/// exclusive lock. Return an error if lock attempt fails.
fn trylock_pid_file() -> StratisResult<File> {
    let mut f = match OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .open(STRATISD_PID_PATH)
    {
        Ok(f) => f,
        Err(e) => {
            if e.kind() == ErrorKind::PermissionDenied {
                return Err(StratisError::Error(
                    "Must be running as root in order to start daemon.".to_string(),
                ));
            }
            return Err(e.into());
        }
    };
    match flock(f.as_raw_fd(), FlockArg::LockExclusiveNonblock) {
        Ok(_) => {
            f.write_all(format!("{}\n", getpid()).as_bytes())?;
            Ok(f)
        }
        Err(_) => {
            let mut buf = String::new();
            f.read_to_string(&mut buf)?;
            // pidfile is supposed to contain pid of holder. But you never
            // know so be paranoid.
            let pid_str = buf.split_whitespace()
                .next()
                .and_then(|s| s.parse::<pid_t>().ok())
                .map(|pid| format!("{}", pid))
                .unwrap_or_else(|| "<unknown>".into());
            Err(StratisError::Error(format!(
                "Daemon already running with pid: {}",
                pid_str
            )))
        }
    }
}

fn run(matches: &ArgMatches) -> StratisResult<()> {
    // Setup a udev listener before initializing the engine. A device may
    // appear after the engine has read the /dev directory but before it has
    // completed initialization. Unless the udev event has been recorded, the
    // engine will miss the device.
    // This is especially important since stratisd must run during early boot.
    let context = libudev::Context::new()?;
    let mut monitor = libudev::Monitor::new(&context)?;
    monitor.match_subsystem_devtype("block", "disk")?;
    let mut udev = monitor.listen()?;

    let engine: Rc<RefCell<Engine>> = {
        if matches.is_present("sim") {
            info!("Using SimEngine");
            Rc::new(RefCell::new(SimEngine::default()))
        } else {
            info!("Using StratEngine");
            Rc::new(RefCell::new(StratEngine::initialize()?))
        }
    };

    /*
    The file descriptor array indexes are laid out in the following:

    0   == Always udev fd index
    1   == engine index if eventable
    1/2 == Start of dbus client file descriptor(s), 1 if engine is not eventable, else 2
    */
    const FD_INDEX_UDEV: usize = 0;
    const FD_INDEX_ENGINE: usize = 1;

    let mut fds = Vec::new();

    fds.push(libc::pollfd {
        fd: udev.as_raw_fd(),
        revents: 0,
        events: libc::POLLIN,
    });

    let eventable = engine.borrow().get_eventable();

    // The variable _dbus_client_index_start is only used when dbus support is compiled in, thus
    // we denote the value as not needed to compile when dbus support is not included.
    let (poll_timeout, _dbus_client_index_start) = match eventable {
        Some(ref evt) => {
            fds.push(libc::pollfd {
                fd: evt.get_pollable_fd(),
                revents: 0,
                events: libc::POLLIN,
            });

            (-1, FD_INDEX_ENGINE + 1)
        }
        None => (10000, FD_INDEX_ENGINE),
    };

    #[cfg(feature = "dbus_enabled")]
    let (dbus_conn, mut tree, base_object_path, dbus_context) =
        libstratis::dbus_api::connect(Rc::clone(&engine))?;

    #[cfg(feature = "dbus_enabled")]
    for (_, pool_uuid, pool) in engine.borrow().pools() {
        libstratis::dbus_api::register_pool(
            &dbus_conn,
            &dbus_context,
            &mut tree,
            pool_uuid,
            pool,
            &base_object_path,
        )?;
    }

    loop {
        // Process any udev block events
        if fds[FD_INDEX_UDEV].revents != 0 {
            while let Some(event) = udev.receive_event() {
                if let Some((device, devnode)) = handle_udev_add(&event) {
                    // If block evaluate returns an error we are going to ignore it as
                    // there is nothing we can do for a device we are getting errors with.
                    let pool_uuid = engine
                        .borrow_mut()
                        .block_evaluate(device, devnode)
                        .unwrap_or(None);

                    // We need to pretend that we aren't using the variable _pool_uuid so
                    // that we can conditionally compile out the register_pool when dbus
                    // is not enabled.
                    if let Some(_pool_uuid) = pool_uuid {
                        #[cfg(feature = "dbus_enabled")]
                        libstratis::dbus_api::register_pool(
                            &dbus_conn,
                            &dbus_context,
                            &mut tree,
                            _pool_uuid,
                            engine
                                .borrow()
                                .get_pool(_pool_uuid)
                                .expect(
                                    "block_evaluate() returned a pool UUID, pool must be available",
                                )
                                .1,
                            &base_object_path,
                        )?;
                    }
                }
            }
        }

        // Handle engine events, if the engine is eventable
        match eventable {
            Some(ref evt) => {
                if fds[FD_INDEX_ENGINE].revents != 0 {
                    evt.clear_event()?;
                    engine.borrow_mut().evented()?;
                }
            }
            None => {
                // Unconditionally call evented() if engine has no eventable.
                // This looks like a bad idea, but the only engine that has
                // no eventable is the sim engine, and for that engine,
                // evented() is essentially a no-op.
                engine.borrow_mut().evented()?;
            }
        }

        // Iterate through D-Bus file descriptors (if enabled)
        #[cfg(feature = "dbus_enabled")]
        {
            for pfd in fds[_dbus_client_index_start..]
                .iter()
                .filter(|pfd| pfd.revents != 0)
            {
                for item in dbus_conn.watch_handle(pfd.fd, WatchEvent::from_revents(pfd.revents)) {
                    if let Err(r) =
                        libstratis::dbus_api::handle(&dbus_conn, &item, &mut tree, &dbus_context)
                    {
                        print_err(&From::from(r));
                    }
                }
            }

            // Refresh list of dbus fds to poll for every time. This can change as
            // D-Bus clients come and go.
            fds.truncate(_dbus_client_index_start);

            fds.extend(dbus_conn.watch_fds().iter().map(|w| w.to_pollfd()));
        }

        let r = unsafe { libc::poll(fds.as_mut_ptr(), fds.len() as libc::c_ulong, poll_timeout) };
        assert!(r >= 0);
    }
}

fn main() {
    let matches = App::new("stratis")
        .version(VERSION)
        .about("Stratis storage management")
        .arg(
            Arg::with_name("debug")
                .long("debug")
                .help("Print additional output for debugging"),
        )
        .arg(
            Arg::with_name("sim")
                .long("sim")
                .help("Use simulator engine"),
        )
        .get_matches();

    let result = trylock_pid_file()
        .and_then(|_pidfile| {
            Ok(initialize_log(matches.is_present("debug"))
                .expect("This is the first and only invocation of this method; it must succeed."))
        })
        .and_then(|_| run(&matches));
    if let Err(err) = result {
        print_err(&err);
        exit(1);
    } else {
        exit(0);
    }
}
<|endoftext|>


<commit_before>use super::block_tokenizer::BlockTokenizer;
use super::tokens::LineType;
use super::ast::{Block, Heading, HeadingLevel};
use super::input::IntoParserInput;
use super::error::ParseError;
use std::str::Lines;

macro_rules! consume_error {
    ($tokenizer:expr) => {
        match $tokenizer.consume_raw()? {
            Err(err) => return Some(Err(err)),
            // If peek() returns an error, we know that
            // consume_raw() must return an error too.
            Ok(..) => unreachable!(),
        }
    }
}

#[derive(Debug)]
pub struct BlockParser<'a, S, I>
where
    S: IntoParserInput<'a>,
    I: Iterator<Item = S>,
{
    tokenizer: BlockTokenizer<'a, S, I>,
}

#[derive(Debug)]
pub struct TextAccumulator {
    buffer: String,
}

impl<'a> BlockParser<'a, &'a str, Lines<'a>> {
    pub fn from_string(input: &'a str) -> Self {
        BlockParser { tokenizer: BlockTokenizer::from_string(input) }
    }
}

impl<'a, S, I> BlockParser<'a, S, I>
where
    S: IntoParserInput<'a>,
    I: Iterator<Item = S>,
{
    pub fn new(input: I) -> Self {
        BlockParser { tokenizer: BlockTokenizer::new(input) }
    }

    fn parse_text(&mut self) -> Option<Result<Block, ParseError>> {
        let mut accumulator = TextAccumulator::new();

        loop {
            match self.tokenizer.peek() {
                None => break,
                Some(Err(..)) => consume_error!(self.tokenizer),
                Some(Ok(LineType::Text)) => {
                    // unwrapping here is safe
                    let line = self.tokenizer.consume(LineType::Text).unwrap().unwrap();

                    accumulator.add(&line.value().unwrap());
                }
                Some(Ok(_)) => break,
            }
        }

        Some(Ok(Block::Text { content: accumulator.consume() }))
    }

    fn parse_quote(&mut self) -> Option<Result<Block, ParseError>> {
        let mut accumulator = TextAccumulator::new();

        loop {
            match self.tokenizer.peek() {
                None => break,
                Some(Err(..)) => consume_error!(self.tokenizer),
                Some(Ok(LineType::Quote)) => {
                    // unwrapping here is safe
                    let line = self.tokenizer.consume(LineType::Quote).unwrap().unwrap();

                    accumulator.add(&line.value().unwrap());
                }
                Some(Ok(_)) => break,
            }
        }

        Some(Ok(Block::Text { content: accumulator.consume() }))
    }

    fn parse_heading(&mut self, line_type: LineType) -> Option<Result<Block, ParseError>> {
        let level = match line_type {
            LineType::Heading1 => HeadingLevel::Level1,
            LineType::Heading2 => HeadingLevel::Level2,
            LineType::Heading3 => HeadingLevel::Level3,
            _ => return None,
        };

        match self.tokenizer.consume(line_type)? {
            Err(err) => Some(Err(err)),
            Ok(line) => Some(Ok(Block::from_inner(
                Heading::new(level, line.value()?.trim().into()),
            ))),
        }
    }
}

impl<'a, S, I> Iterator for BlockParser<'a, S, I>
where
    S: IntoParserInput<'a>,
    I: Iterator<Item = S>,
{
    type Item = Result<Block, ParseError>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.tokenizer.peek()? {
                Err(..) => consume_error!(self.tokenizer),
                Ok(LineType::Blank) => {
                    self.tokenizer.consume_raw();
                    continue;
                }
                Ok(line_type) => {
                    return match line_type {
                        LineType::Text => self.parse_text(),
                        LineType::Quote => self.parse_quote(),
                        LineType::Heading1 => self.parse_heading(LineType::Heading1),
                        LineType::Heading2 => self.parse_heading(LineType::Heading2),
                        LineType::Heading3 => self.parse_heading(LineType::Heading3),
                        _ => unimplemented!(),
                    };
                }
            };
        }
    }
}

impl TextAccumulator {
    pub fn new() -> Self {
        TextAccumulator { buffer: String::new() }
    }

    ///
    /// Adds a new line to the current accumulated text.
    /// Todo: this should also take care of newlines with two spaces
    ///
    pub fn add(&mut self, line: &str) {
        if self.buffer.len() > 0 {
            self.buffer.push_str(" ");
        }

        self.buffer.push_str(line);
    }

    pub fn consume(self) -> String {
        self.buffer
    }
}

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

    macro_rules! unwrap {
        ($value:expr) => {
            $value.unwrap().unwrap()
        }
    }

    #[test]
    fn it_works() {
        let mut parser = BlockParser::from_string("Lorem ipsum\ndolor sit amet");

        assert_eq!(
            Block::Text { content: "Lorem ipsum dolor sit amet".to_string() },
            unwrap!(parser.next())
        );
    }

    #[test]
    fn parsing_headings_works() {
        let mut parser = BlockParser::from_string("# hello world\n##    level 2\n### three");

        assert_eq!(
            Block::from_inner(Heading::new(HeadingLevel::Level1, "hello world".into())),
            unwrap!(parser.next())
        );

        assert_eq!(
            Block::from_inner(Heading::new(HeadingLevel::Level2, "level 2".into())),
            unwrap!(parser.next())
        );
    }
}<commit_msg>Add tests for quote and text blocks<commit_after>use super::block_tokenizer::BlockTokenizer;
use super::tokens::LineType;
use super::ast::{Block, Heading, HeadingLevel};
use super::input::IntoParserInput;
use super::error::ParseError;
use std::str::Lines;

macro_rules! consume_error {
    ($tokenizer:expr) => {
        match $tokenizer.consume_raw()? {
            Err(err) => return Some(Err(err)),
            // If peek() returns an error, we know that
            // consume_raw() must return an error too.
            Ok(..) => unreachable!(),
        }
    }
}

#[derive(Debug)]
pub struct BlockParser<'a, S, I>
where
    S: IntoParserInput<'a>,
    I: Iterator<Item = S>,
{
    tokenizer: BlockTokenizer<'a, S, I>,
}

#[derive(Debug)]
pub struct TextAccumulator {
    buffer: String,
}

impl<'a> BlockParser<'a, &'a str, Lines<'a>> {
    pub fn from_string(input: &'a str) -> Self {
        BlockParser { tokenizer: BlockTokenizer::from_string(input) }
    }
}

impl<'a, S, I> BlockParser<'a, S, I>
where
    S: IntoParserInput<'a>,
    I: Iterator<Item = S>,
{
    pub fn new(input: I) -> Self {
        BlockParser { tokenizer: BlockTokenizer::new(input) }
    }

    fn parse_text(&mut self) -> Option<Result<Block, ParseError>> {
        let mut accumulator = TextAccumulator::new();

        loop {
            match self.tokenizer.peek() {
                None => break,
                Some(Err(..)) => consume_error!(self.tokenizer),
                Some(Ok(LineType::Text)) => {
                    // unwrapping here is safe
                    let line = self.tokenizer.consume(LineType::Text).unwrap().unwrap();

                    accumulator.add(&line.value().unwrap());
                }
                Some(Ok(_)) => break,
            }
        }

        Some(Ok(Block::Text { content: accumulator.consume() }))
    }

    fn parse_quote(&mut self) -> Option<Result<Block, ParseError>> {
        let mut accumulator = TextAccumulator::new();

        loop {
            match self.tokenizer.peek() {
                None => break,
                Some(Err(..)) => consume_error!(self.tokenizer),
                Some(Ok(LineType::Quote)) => {
                    // unwrapping here is safe
                    let line = self.tokenizer.consume(LineType::Quote).unwrap().unwrap();

                    accumulator.add(&line.value().unwrap());
                }
                Some(Ok(_)) => break,
            }
        }

        Some(Ok(Block::Quote { content: accumulator.consume() }))
    }

    fn parse_heading(&mut self, line_type: LineType) -> Option<Result<Block, ParseError>> {
        let level = match line_type {
            LineType::Heading1 => HeadingLevel::Level1,
            LineType::Heading2 => HeadingLevel::Level2,
            LineType::Heading3 => HeadingLevel::Level3,
            _ => return None,
        };

        match self.tokenizer.consume(line_type)? {
            Err(err) => Some(Err(err)),
            Ok(line) => Some(Ok(Block::from_inner(
                Heading::new(level, line.value()?.trim().into()),
            ))),
        }
    }
}

impl<'a, S, I> Iterator for BlockParser<'a, S, I>
where
    S: IntoParserInput<'a>,
    I: Iterator<Item = S>,
{
    type Item = Result<Block, ParseError>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.tokenizer.peek()? {
                Err(..) => consume_error!(self.tokenizer),
                Ok(LineType::Blank) => {
                    self.tokenizer.consume_raw();
                    continue;
                }
                Ok(line_type) => {
                    return match line_type {
                        LineType::Text => self.parse_text(),
                        LineType::Quote => self.parse_quote(),
                        LineType::Heading1 => self.parse_heading(LineType::Heading1),
                        LineType::Heading2 => self.parse_heading(LineType::Heading2),
                        LineType::Heading3 => self.parse_heading(LineType::Heading3),
                        _ => unimplemented!(),
                    };
                }
            };
        }
    }
}

impl TextAccumulator {
    pub fn new() -> Self {
        TextAccumulator { buffer: String::new() }
    }

    ///
    /// Adds a new line to the current accumulated text.
    /// Todo: this should also take care of newlines with two spaces
    ///
    pub fn add(&mut self, line: &str) {
        if self.buffer.len() > 0 {
            self.buffer.push_str(" ");
        }

        self.buffer.push_str(line.trim());
    }

    pub fn consume(self) -> String {
        self.buffer
    }
}

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

    macro_rules! unwrap {
        ($value:expr) => {
            $value.unwrap().unwrap()
        }
    }

    #[test]
    fn it_works() {
        let mut parser = BlockParser::from_string("Lorem ipsum\ndolor sit amet");

        assert_eq!(
            Block::Text { content: "Lorem ipsum dolor sit amet".to_string() },
            unwrap!(parser.next())
        );
    }

    #[test]
    fn parsing_headings_works() {
        let mut parser = BlockParser::from_string("# hello world\n##    level 2\n### three");

        assert_eq!(
            Block::from_inner(Heading::new(HeadingLevel::Level1, "hello world".into())),
            unwrap!(parser.next())
        );

        assert_eq!(
            Block::from_inner(Heading::new(HeadingLevel::Level2, "level 2".into())),
            unwrap!(parser.next())
        );
    }

    #[test]
    fn parsing_quote_works() {
        let mut parser = BlockParser::from_string("> Foo\n> bar baz");

        assert_eq!(
            Block::Quote { content: "Foo bar baz".into() },
            unwrap!(parser.next())
        );
    }

    #[test]
    fn parsing_text_works() {
        let mut parser = BlockParser::from_string("Foo\n    \tbar baz");

        assert_eq!(
            Block::Text { content: "Foo bar baz".into() },
            unwrap!(parser.next())
        );
    }
}<|endoftext|>


<commit_before>#[macro_use] extern crate html5ever_atoms;
extern crate html5ever;
extern crate orbclient;
extern crate orbfont;
extern crate tendril;

use std::{cmp, env};
use std::iter::repeat;
use std::default::Default;
use std::string::String;

use html5ever::parse_document;
use html5ever::rcdom::{Document, Doctype, Text, Comment, Element, RcDom, Handle};
use orbclient::{Color, Window, EventOption, K_ESC, K_DOWN, K_UP};
use orbfont::Font;
use tendril::TendrilSink;

// This is not proper HTML serialization, of course.

struct Block<'a> {
    x: i32,
    y: i32,
    w: i32,
    h: i32,
    color: Color,
    text: orbfont::Text<'a>
}

impl<'a> Block<'a> {
    fn draw(&self, window: &mut Window, offset: i32) {
        let x = self.x;
        let y = self.y - offset;
        if x + self.w > 0 && x < window.width() as i32 && y + self.h > 0 && y < window.height() as i32 {
            self.text.draw(window, x, y, self.color);
        }
    }
}

fn walk<'a>(handle: Handle, indent: usize, x: &mut i32, y: &mut i32, mut size: f32, mut bold: bool, mut color: Color, mut ignore: bool, whitespace: &mut bool, font: &'a Font, font_bold: &'a Font, blocks: &mut Vec<Block<'a>>) {
    let node = handle.borrow();

    let mut new_line = false;

    print!("{}", repeat(" ").take(indent).collect::<String>());
    match node.node {
        Document
            => {
                println!("#Document")
            },

        Doctype(ref name, ref public, ref system)
            => {
                println!("<!DOCTYPE {} \"{}\" \"{}\">", *name, *public, *system);
            },

        Text(ref text)
            => {
                let mut block_text = String::new();

                for c in text.chars() {
                    match c {
                        ' ' | '\n' | '\r' => if *whitespace {
                            // Ignore
                        } else {
                            // Set whitespace
                            *whitespace = true;
                            block_text.push(' ');
                        },
                        _ => {
                            if *whitespace {
                                *whitespace = false;
                            }
                            block_text.push(c);
                        }
                    }
                }

                if ! block_text.is_empty() {
                    if ignore {
                        println!("#text: ignored");
                    } else {
                        let trimmed_left = block_text.trim_left();
                        let left_margin = block_text.len() as i32 - trimmed_left.len() as i32;
                        let trimmed_right = trimmed_left.trim_right();
                        let right_margin = trimmed_left.len() as i32 - trimmed_right.len() as i32;

                        let escaped_text = escape_default(&trimmed_right);
                        println!("#text: block {} at {}, {}: '{}'", blocks.len(), *x, *y, escaped_text);

                        *x += left_margin * 8;

                        for (word_i, word) in trimmed_right.split(' ').enumerate() {
                            if word_i > 0 {
                                *x += 8;
                            }

                            let text = if bold {
                                font_bold.render(word, size)
                            } else {
                                font.render(word, size)
                            };

                            let w = text.width() as i32;
                            let h = text.height() as i32;

                            if *x + w >= 640 && *x > 0 {
                                *x = 0;
                                *y += size.ceil() as i32;
                            }

                            blocks.push(Block {
                                x: *x,
                                y: *y,
                                w: w,
                                h: h,
                                color: color,
                                text: text
                            });

                            *x += w;
                        }

                        *x += right_margin * 8;
                    }
                } else {
                    println!("#text: empty");
                }
            },

        Comment(ref text)
            => {
                println!("<!-- {} -->", escape_default(text))
            },

        Element(ref name, _, ref _attrs) => {
            assert!(name.ns == ns!(html));
            print!("<{}", name.local);
            /*
            for attr in attrs.iter() {
                assert!(attr.name.ns == ns!());
                print!(" {}=\"{}\"", attr.name.local, attr.value);
            }
            */
            println!(">");

            match &*name.local {
                "a" => {
                    color = Color::rgb(0, 0, 255);
                },
                "b" => {
                    bold = true;
                },
                "br" => {
                    ignore = true;
                    new_line = true;
                },
                "div" => {
                    new_line = true;
                },
                "h1" => {
                    size = 32.0;
                    bold = true;
                    new_line = true;
                },
                "h2" => {
                    size = 24.0;
                    bold = true;
                    new_line = true;
                },
                "h3" => {
                    size = 18.0;
                    bold = true;
                    new_line = true;
                }
                "h4" => {
                    size = 16.0;
                    bold = true;
                    new_line = true;
                }
                "h5" => {
                    size = 14.0;
                    bold = true;
                    new_line = true;
                }
                "h6" => {
                    size = 10.0;
                    bold = true;
                    new_line = true;
                },
                "hr" => {
                    new_line = true;
                },
                "li" => {
                    new_line = true;
                },
                "p" => {
                    new_line = true;
                },

                "head" => ignore = true,
                "title" => ignore = true, //TODO: Grab title
                "link" => ignore = true,
                "meta" => ignore = true,
                "script" => ignore = true,
                "style" => ignore = true,
                _ => ()
            }
        }
    }

    for child in node.children.iter() {
        walk(child.clone(), indent + 4, x, y, size, bold, color, ignore, whitespace, font, font_bold, blocks);
    }

    if new_line {
        *whitespace = true;
        *x = 0;
        *y += size.ceil() as i32;
    }
}

// FIXME: Copy of str::escape_default from std, which is currently unstable
pub fn escape_default(s: &str) -> String {
    s.chars().flat_map(|c| c.escape_default()).collect()
}

fn event_loop(window: &mut Window){
    loop {
        for event in window.events() {
            if let EventOption::Key(key_event) = event.to_option() {
                if key_event.pressed && key_event.scancode == K_ESC {
                    return;
                }
            }
            if let EventOption::Quit(_) = event.to_option() {
                return;
            }
        }
    }
}

fn error_msg(window: &mut Window, msg: &str) {
    let mut x = 0;
    for c in msg.chars() {
        window.char(x, 0, c, Color::rgb(255, 255, 255));
        x += 8;
    }
}

fn main() {
    let title = match env::args().nth(1) {
        Some(arg) => arg.clone(),
        None => "".to_string(),
    };

    match Font::find(None, None, None) {
        Ok(font) => match Font::find(None, None, Some("Bold")) {
            Ok(font_bold) => {
                let dom = parse_document(RcDom::default(), Default::default())
                    .from_utf8()
                    .from_file(env::args().nth(1).expect("browser: no file provided"))
                    .unwrap();

                let mut blocks = vec![];
                let mut x = 0;
                let mut y = 0;
                let mut whitespace = false;
                walk(dom.document, 0, &mut x, &mut y, 16.0, false, Color::rgb(0, 0, 0), false, &mut whitespace, &font, &font_bold, &mut blocks);

                if !dom.errors.is_empty() {
                    /*
                    println!("\nParse errors:");
                    for err in dom.errors.into_iter() {
                        println!("    {}", err);
                    }
                    */
                }

                let mut window = Window::new(-1,
                                             -1,
                                             640,
                                             480,
                                             &("Browser (".to_string() + &title + ")"))
                                     .unwrap();

                let mut offset = 0;
                let mut max_offset = 0;
                for block in blocks.iter() {
                    if block.y + block.h > max_offset {
                        max_offset = block.y + block.h;
                    }
                }

                let mut redraw = true;
                loop {
                    if redraw {
                        redraw = false;

                        window.set(Color::rgb(255, 255, 255));
                        for block in blocks.iter() {
                            block.draw(&mut window, offset);
                        }
                        window.sync();
                    }

                    for event in window.events() {
                        if let EventOption::Key(key_event) = event.to_option() {
                            if key_event.pressed {
                                match key_event.scancode {
                                    K_ESC => return,
                                    K_UP => {
                                        redraw = true;
                                        offset = cmp::max(0, offset - 128);
                                    },
                                    K_DOWN => {
                                        redraw = true;
                                        offset = cmp::min(max_offset, offset + 128);
                                    },
                                    _ => ()
                                }
                            }
                        }
                        if let EventOption::Quit(_) = event.to_option() {
                            return;
                        }
                    }
                }
            },
            Err(err) => {
                let mut window = Window::new(-1,
                                             -1,
                                             320,
                                             32,
                                             &("Browser (".to_string() + &title + ")"))
                                     .unwrap();
                window.set(Color::rgb(0, 0, 0));
                error_msg(&mut window, &format!("{}", err));
                window.sync();
                event_loop(&mut window);
            }
        },
        Err(err) => {
            let mut window = Window::new(-1,
                                         -1,
                                         320,
                                         32,
                                         &("Browser (".to_string() + &title + ")"))
                                 .unwrap();
            window.set(Color::rgb(0, 0, 0));
            error_msg(&mut window, &format!("{}", err));
            window.sync();
            event_loop(&mut window);
        }
    }
}
<commit_msg>Browsing links<commit_after>#[macro_use] extern crate html5ever_atoms;
extern crate html5ever;
extern crate orbclient;
extern crate orbfont;
extern crate tendril;

use std::{cmp, env, str};
use std::iter::repeat;
use std::default::Default;
use std::fs::File;
use std::io::{stderr, Read, Write};
use std::net::TcpStream;
use std::string::String;

use html5ever::parse_document;
use html5ever::rcdom::{Document, Doctype, Text, Comment, Element, RcDom, Handle};
use orbclient::{Color, Window, EventOption, K_ESC, K_DOWN, K_UP};
use orbfont::Font;
use tendril::TendrilSink;

#[derive(Clone, Debug)]
struct Url {
    scheme: String,
    host: String,
    port: u16,
    path: Vec<String>
}

impl Url {
    fn new(url: &str) -> Url {
        let (scheme, reference) = url.split_at(url.find(':').unwrap_or(0));

        let mut parts = reference.split('/').skip(2); //skip first two slashes
        let remote = parts.next().unwrap_or("");
        let mut remote_parts = remote.split(':');
        let host = remote_parts.next().unwrap_or("");
        let port = remote_parts.next().unwrap_or("").parse::<u16>().unwrap_or(80);

        let mut path = Vec::new();
        for part in parts {
            path.push(part.to_string());
        }

        Url {
            scheme: scheme.to_string(),
            host: host.to_string(),
            port: port,
            path: path
        }
    }
}

struct Block<'a> {
    x: i32,
    y: i32,
    w: i32,
    h: i32,
    color: Color,
    string: String,
    link: Option<String>,
    text: orbfont::Text<'a>
}

impl<'a> Block<'a> {
    fn contains(&self, m_x: i32, m_y: i32, offset: i32) -> bool {
        let x = self.x;
        let y = self.y - offset;

        m_x >= x && m_x < x + self.w && m_y >= y && m_y < y + self.h
    }

    fn draw(&self, window: &mut Window, offset: i32) {
        let x = self.x;
        let y = self.y - offset;
        if x + self.w > 0 && x < window.width() as i32 && y + self.h > 0 && y < window.height() as i32 {
            self.text.draw(window, x, y, self.color);
        }
    }
}

fn walk<'a>(handle: Handle, indent: usize, x: &mut i32, y: &mut i32, mut size: f32, mut bold: bool, mut color: Color, mut ignore: bool, whitespace: &mut bool, mut link: Option<String>, font: &'a Font, font_bold: &'a Font, blocks: &mut Vec<Block<'a>>) {
    let node = handle.borrow();

    let mut new_line = false;

    print!("{}", repeat(" ").take(indent).collect::<String>());
    match node.node {
        Document
            => {
                println!("#Document")
            },

        Doctype(ref name, ref public, ref system)
            => {
                println!("<!DOCTYPE {} \"{}\" \"{}\">", *name, *public, *system);
            },

        Text(ref text)
            => {
                let mut block_text = String::new();

                for c in text.chars() {
                    match c {
                        ' ' | '\n' | '\r' => if *whitespace {
                            // Ignore
                        } else {
                            // Set whitespace
                            *whitespace = true;
                            block_text.push(' ');
                        },
                        _ => {
                            if *whitespace {
                                *whitespace = false;
                            }
                            block_text.push(c);
                        }
                    }
                }

                if ! block_text.is_empty() {
                    if ignore {
                        println!("#text: ignored");
                    } else {
                        let trimmed_left = block_text.trim_left();
                        let left_margin = block_text.len() as i32 - trimmed_left.len() as i32;
                        let trimmed_right = trimmed_left.trim_right();
                        let right_margin = trimmed_left.len() as i32 - trimmed_right.len() as i32;

                        let escaped_text = escape_default(&trimmed_right);
                        println!("#text: block {} at {}, {}: '{}'", blocks.len(), *x, *y, escaped_text);

                        *x += left_margin * 8;

                        for (word_i, word) in trimmed_right.split(' ').enumerate() {
                            if word_i > 0 {
                                *x += 8;
                            }

                            let text = if bold {
                                font_bold.render(word, size)
                            } else {
                                font.render(word, size)
                            };

                            let w = text.width() as i32;
                            let h = text.height() as i32;

                            if *x + w >= 640 && *x > 0 {
                                *x = 0;
                                *y += size.ceil() as i32;
                            }

                            blocks.push(Block {
                                x: *x,
                                y: *y,
                                w: w,
                                h: h,
                                color: color,
                                string: word.to_string(),
                                link: link.clone(),
                                text: text
                            });

                            *x += w;
                        }

                        *x += right_margin * 8;
                    }
                } else {
                    println!("#text: empty");
                }
            },

        Comment(ref text)
            => {
                println!("<!-- {} -->", escape_default(text))
            },

        Element(ref name, _, ref attrs) => {
            assert!(name.ns == ns!(html));
            print!("<{}", name.local);
            for attr in attrs.iter() {
                assert!(attr.name.ns == ns!());
                print!(" {}=\"{}\"", attr.name.local, attr.value);
            }
            println!(">");

            match &*name.local {
                "a" => {
                    color = Color::rgb(0, 0, 255);
                    for attr in attrs.iter() {
                        match &*attr.name.local {
                            "href" => link = Some(attr.value.to_string()),
                            _ => ()
                        }
                    }
                },
                "b" => {
                    bold = true;
                },
                "br" => {
                    ignore = true;
                    new_line = true;
                },
                "div" => {
                    new_line = true;
                },
                "h1" => {
                    size = 32.0;
                    bold = true;
                    new_line = true;
                },
                "h2" => {
                    size = 24.0;
                    bold = true;
                    new_line = true;
                },
                "h3" => {
                    size = 18.0;
                    bold = true;
                    new_line = true;
                }
                "h4" => {
                    size = 16.0;
                    bold = true;
                    new_line = true;
                }
                "h5" => {
                    size = 14.0;
                    bold = true;
                    new_line = true;
                }
                "h6" => {
                    size = 10.0;
                    bold = true;
                    new_line = true;
                },
                "hr" => {
                    new_line = true;
                },
                "li" => {
                    new_line = true;
                },
                "p" => {
                    new_line = true;
                },

                "head" => ignore = true,
                "title" => ignore = true, //TODO: Grab title
                "link" => ignore = true,
                "meta" => ignore = true,
                "script" => ignore = true,
                "style" => ignore = true,
                _ => ()
            }
        }
    }

    for child in node.children.iter() {
        walk(child.clone(), indent + 4, x, y, size, bold, color, ignore, whitespace, link.clone(), font, font_bold, blocks);
    }

    if new_line {
        *whitespace = true;
        *x = 0;
        *y += size.ceil() as i32;
    }
}

// FIXME: Copy of str::escape_default from std, which is currently unstable
pub fn escape_default(s: &str) -> String {
    s.chars().flat_map(|c| c.escape_default()).collect()
}

fn read_blocks<'a, R: Read>(r: &mut R, font: &'a Font, font_bold: &'a Font) -> Vec<Block<'a>> {
    let mut blocks = vec![];

    let dom = parse_document(RcDom::default(), Default::default())
        .from_utf8()
        .read_from(r)
        .unwrap();

    let mut x = 0;
    let mut y = 0;
    let mut whitespace = false;
    walk(dom.document, 0, &mut x, &mut y, 16.0, false, Color::rgb(0, 0, 0), false, &mut whitespace, None, font, font_bold, &mut blocks);

    if !dom.errors.is_empty() {
        /*
        println!("\nParse errors:");
        for err in dom.errors.into_iter() {
            println!("    {}", err);
        }
        */
    }

    blocks
}

fn file_blocks<'a>(url: &Url, font: &'a Font, font_bold: &'a Font) -> Vec<Block<'a>> {
    let mut parts = url.path.iter();
    let mut path = parts.next().map_or(String::new(), |s| s.clone());
    for part in parts {
        path.push('/');
        path.push_str(part);
    }

    if let Ok(mut file) = File::open(&path) {
        read_blocks(&mut file, &font, &font_bold)
    } else {
        vec![]
    }
}

fn http_blocks<'a>(url: &Url, font: &'a Font, font_bold: &'a Font) -> Vec<Block<'a>> {
    let mut parts = url.path.iter();
    let mut path = parts.next().map_or(String::new(), |s| s.clone());
    for part in parts {
        path.push('/');
        path.push_str(part);
    }

    write!(stderr(), "* Connecting to {}:{}\n", url.host, url.port).unwrap();

    let mut stream = TcpStream::connect((url.host.as_str(), url.port)).unwrap();

    write!(stderr(), "* Requesting {}\n", path).unwrap();

    let request = format!("GET /{} HTTP/1.1\r\nHost: {}\r\nConnection: close\r\n\r\n", path, url.host);
    stream.write(request.as_bytes()).unwrap();
    stream.flush().unwrap();

    write!(stderr(), "* Waiting for response\n").unwrap();

    let mut response = Vec::new();

    loop {
        let mut buf = [0; 65536];
        let count = stream.read(&mut buf).unwrap();
        if count == 0 {
            break;
        }
        response.extend_from_slice(&buf[.. count]);
    }

    write!(stderr(), "* Received {} bytes\n", response.len()).unwrap();

    let mut header_end = 0;
    while header_end < response.len() {
        if response[header_end..].starts_with(b"\r\n\r\n") {
            break;
        }
        header_end += 1;
    }

    for line in unsafe { str::from_utf8_unchecked(&response[..header_end]) }.lines() {
        write!(stderr(), "> {}\n", line).unwrap();
    }

    read_blocks(&mut &response[header_end + 4 ..], font, font_bold)
}

fn url_blocks<'a>(url: &Url, font: &'a Font, font_bold: &'a Font) -> Vec<Block<'a>> {
    if url.scheme == "http" {
        http_blocks(url, font, font_bold)
    } else if url.scheme == "file" || url.scheme.is_empty() {
        file_blocks(url, font, font_bold)
    } else {
        vec![]
    }
}

fn main_window(arg: &str, font: &Font, font_bold: &Font) {
    let mut url = Url::new(arg);

    let mut window = Window::new(-1, -1, 640, 480,  &format!("Browser ({})", arg)).unwrap();

    let mut blocks = url_blocks(&url, &font, &font_bold);

    let mut offset = 0;
    let mut max_offset = 0;
    for block in blocks.iter() {
        if block.y + block.h > max_offset {
            max_offset = block.y + block.h;
        }
    }

    let mut mouse_down = false;

    let mut redraw = true;
    loop {
        if redraw {
            redraw = false;

            window.set(Color::rgb(255, 255, 255));

            for block in blocks.iter() {
                block.draw(&mut window, offset);
            }

            window.sync();
        }

        for event in window.events() {
            match event.to_option() {
                EventOption::Key(key_event) => if key_event.pressed {
                    match key_event.scancode {
                        K_ESC => return,
                        K_UP => {
                            redraw = true;
                            offset = cmp::max(0, offset - 128);
                        },
                        K_DOWN => {
                            redraw = true;
                            offset = cmp::min(max_offset, offset + 128);
                        },
                        _ => ()
                    }
                },
                EventOption::Mouse(mouse_event) => if mouse_event.left_button {
                    mouse_down = true;
                } else if mouse_down {
                    mouse_down = false;

                    let mut link_opt = None;
                    for block in blocks.iter() {
                        if block.contains(mouse_event.x, mouse_event.y, offset) {
                            println!("Click {}", block.string);
                            if let Some(ref link) = block.link {
                                link_opt = Some(link.clone());
                                break;
                            }
                        }
                    }

                    if let Some(link) = link_opt {
                        if link.starts_with('#') {
                            println!("Find anchor {}", link);
                        } else {
                            if link.find(':').is_some() {
                                url = Url::new(&link);
                            } else if link.starts_with('/') {
                                url.path.clear();
                                for part in link[1..].split('/') {
                                    url.path.push(part.to_string());
                                }
                            } else {
                                url.path.push(link.clone());
                            };

                            println!("Navigate {}: {:#?}", link, url);

                            blocks = url_blocks(&url, &font, &font_bold);

                            offset = 0;
                            max_offset = 0;
                            for block in blocks.iter() {
                                if block.y + block.h > max_offset {
                                    max_offset = block.y + block.h;
                                }
                            }

                            redraw = true;
                        }
                    }
                },
                EventOption::Quit(_) => return,
                _ => ()
            }
        }
    }
}

fn main() {
    let err_window = |msg: &str| {
        let mut window = Window::new(-1, -1, 320, 32, "Browser").unwrap();

        window.set(Color::rgb(0, 0, 0));

        let mut x = 0;
        for c in msg.chars() {
            window.char(x, 0, c, Color::rgb(255, 255, 255));
            x += 8;
        }

        window.sync();

        loop {
            for event in window.events() {
                if let EventOption::Key(key_event) = event.to_option() {
                    if key_event.pressed && key_event.scancode == K_ESC {
                        return;
                    }
                }
                if let EventOption::Quit(_) = event.to_option() {
                    return;
                }
            }
        }
    };

    match env::args().nth(1) {
        Some(path) => match Font::find(None, None, None) {
            Ok(font) => match Font::find(None, None, Some("Bold")) {
                Ok(font_bold) => main_window(&path, &font, &font_bold),
                Err(err) => err_window(&format!("{}", err))
            },
            Err(err) => err_window(&format!("{}", err))
        },
        None => err_window("no file argument")
    }
}
<|endoftext|>